1use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5 CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceSlice, LaunchConfig,
6 PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13const SDPA_NAIVE_SMEM_MAX: usize = 48 * 1024;
18
19const SDPA_NAIVE_GMEM_WS_MAX: usize = 1 << 30;
23
24#[cfg(debug_assertions)]
25pub(crate) fn debug_assert_tensor_stream_device<T>(
26 tensor: &CudaSlice<T>,
27 stream: &CudaStream,
28 site: &str,
29) {
30 let tensor_dev = tensor.ordinal();
31 let stream_dev = stream.context().ordinal();
32 assert_eq!(
33 tensor_dev, stream_dev,
34 "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
35 );
36}
37
38fn ensure_tensor_stream_device<T>(
39 tensor: &impl DeviceSlice<T>,
40 stream: &CudaStream,
41 site: &str,
42) -> Result<(), Box<dyn std::error::Error>> {
43 let tensor_dev = tensor.stream().context().ordinal();
44 let stream_dev = stream.context().ordinal();
45 if tensor_dev != stream_dev {
46 return Err(format!(
47 "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
48 stream on dev{stream_dev}"
49 )
50 .into());
51 }
52 Ok(())
53}
54
55pub use memra_gguf;
56pub use memra_runtime;
57
58pub mod forward;
59pub mod hybrid;
60pub mod hybrid_forward;
61pub mod model;
62pub mod sigrouter_contract;
63pub mod vision;
64pub mod vision_gemma;
65pub mod vision_pre;
66pub mod cache {
69 pub use memra_kv::*;
70}
71pub mod decode;
72pub mod decode_batch;
73pub mod dflash;
74pub mod eagle;
75pub mod gemma_spec;
76pub mod graph_update;
77pub mod mla;
81pub mod moesd;
82pub mod parallel;
83pub mod plan_backend;
84pub mod pp;
85pub mod round_stream;
86pub mod spec;
87pub mod tp;
88pub use memra_sampling as sampler;
89
90pub fn moe_f16g_mode() -> u8 {
134 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
135 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
136 Ok("0") => 0,
137 Ok("2") => 2,
138 Ok("3") => 3,
139 Ok(_) => 1,
140 Err(_) => 2,
143 })
144}
145pub fn moe_f16g_sk_params() -> (i32, i32) {
159 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
160 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
161 Ok("0") => (-1, 0),
162 Ok("32") => (0, i32::MAX),
163 Ok("128") => (0, 1),
164 _ => {
165 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
166 .ok()
167 .and_then(|v| v.parse().ok())
168 .unwrap_or(64);
169 (0, cross)
170 }
171 })
172}
173pub fn moe_f16g_direct_on(qtype: i32) -> bool {
184 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
185 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
186 Ok("0") => 0,
187 Ok("kq") => 1,
188 _ => 2,
189 });
190 match m {
191 0 => false,
192 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
193 _ => true,
194 }
195}
196pub fn moe_f16g_tail_on() -> bool {
205 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
207}
208
209pub fn moe_f16g_gemma_on() -> bool {
216 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
217 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
218}
219
220pub fn moe_fuse_actq_on() -> bool {
224 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
225 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
226}
227
228pub fn router_prefill_exact_on() -> bool {
238 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
239 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
240}
241
242pub fn router_kernel_on() -> bool {
243 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244 *ON.get_or_init(|| {
245 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
246 if !on {
247 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
248 }
249 on
250 })
251}
252
253pub const ROUTER_BATCH_MIN_T: usize = 8;
268pub fn router_batch_on() -> bool {
269 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
271}
272mod cpu_experts;
273#[cfg(memra_cutlass)]
274pub mod cutlass_ffi;
275pub mod dsv4_ffi;
276pub mod dsv4_gpu;
277pub mod f16_ffi;
278pub mod fp8_ffi;
279pub mod mmq_ffi;
280pub mod moe_cache;
281pub mod prime_graph;
282pub mod spill;
283mod spill_pread;
284
285const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
292const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
293const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
294const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
295const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
296const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
297const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
299
300fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
306 assert!(
307 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
308 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
309 );
310 match std::env::var("MEMRA_GEMM_FATBIN") {
311 Ok(path) => std::borrow::Cow::Owned(
312 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
313 ),
314 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
315 }
316}
317
318pub(crate) const fn portable_mma_gated() -> bool {
325 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
326}
327
328#[track_caller]
345pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
346 assert!(
347 !portable_mma_gated(),
348 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
349 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
350 );
351}
352
353pub(crate) const fn gdn_mma_default_on() -> bool {
362 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
363}
364
365const fn konst_eq(a: &str, b: &str) -> bool {
367 let (a, b) = (a.as_bytes(), b.as_bytes());
368 if a.len() != b.len() {
369 return false;
370 }
371 let mut i = 0;
372 while i < a.len() {
373 if a[i] != b[i] {
374 return false;
375 }
376 i += 1;
377 }
378 true
379}
380
381const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
386 (!portable_cuda || hopper_mma) && !no_gemm
387}
388
389const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
397const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
398const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
399const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
400const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
401
402pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
405
406fn flash_fatbin_bytes() -> &'static [u8] {
408 match kv_cache_formats() {
409 ("q8_0", "q5_1") => FLASH_FATBIN,
410 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
411 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
412 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
413 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
414 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
415 other => unreachable!("kv_cache_formats returned {other:?}"),
416 }
417}
418
419fn k1_launch_override() -> Option<(u32, u32, u32)> {
426 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
427 *K1.get_or_init(|| {
428 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
429 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
430 match p.as_slice() {
431 [bm, bn, w] => Some((*bm, *bn, *w)),
432 _ => None,
433 }
434 })
435}
436
437pub(crate) fn wgmma_gemm_enabled() -> bool {
444 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
445 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
446}
447
448pub const FA_VEC_MIN_TKV: usize = 96;
463pub fn fa_vec_min_tkv() -> usize {
467 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
468 *V.get_or_init(|| {
469 std::env::var("MEMRA_FA_VEC_MIN")
470 .ok()
471 .and_then(|v| v.parse().ok())
472 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
473 })
474}
475
476pub fn fa_f16pv_on() -> bool {
487 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
488 *ON.get_or_init(|| {
489 std::env::var("MEMRA_FA_F16PV")
490 .map(|v| v != "0")
491 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
492 })
493}
494
495pub fn fa512_hp_on() -> bool {
499 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
500 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
501}
502
503pub fn faw_hp_on() -> bool {
507 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
508 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
509}
510
511pub fn fa512_wide_warps() -> usize {
515 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
516 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
517 Ok("1") => 4,
518 _ => 2,
519 })
520}
521
522pub fn fa512_min_tkv() -> usize {
525 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
526 *FA512_MIN.get_or_init(|| {
527 std::env::var("MEMRA_FA512_MIN")
528 .ok()
529 .and_then(|v| v.parse().ok())
530 .unwrap_or(512)
531 })
532}
533pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
537 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
538pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
542pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
548 std::sync::atomic::AtomicBool::new(false);
549pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
556 std::sync::atomic::AtomicBool::new(true);
557pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
558 std::sync::atomic::AtomicUsize::new(16);
559pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
564pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
566pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
571pub use memra_kv::KV_FP8_FORCE;
574pub(crate) fn mmv_block() -> u32 {
579 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
580 *V.get_or_init(|| {
581 std::env::var("MEMRA_MMV_BLOCK")
582 .ok()
583 .and_then(|v| v.parse().ok())
584 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
585 .unwrap_or(128)
586 })
587}
588
589pub(crate) fn sig_expf_dev_on() -> bool {
594 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
595 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
596}
597
598pub(crate) fn topk_fast_on() -> bool {
599 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
600 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
601}
602
603pub(crate) fn rms_block() -> u32 {
604 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
605 *V.get_or_init(|| {
606 std::env::var("MEMRA_RMS_BLOCK")
607 .ok()
608 .and_then(|v| v.parse().ok())
609 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
610 })
611}
612
613pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
614 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
615 if let Some(forced) = *S.get_or_init(|| {
616 std::env::var("MEMRA_FA_SPLIT")
617 .ok()
618 .and_then(|v| v.parse().ok())
619 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
620 }) {
621 return forced;
622 }
623 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
641 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
642 {
643 return if t_kv <= 8192 {
644 16
645 } else if t_kv <= 16384 {
646 64
647 } else {
648 128
649 };
650 }
651 let big_rig = fa_sm_count() >= 128;
652 if big_rig {
653 let _ = n_head_kv;
654 if t_kv <= 2048 {
655 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
664 if let Some(sp) = *SHORT.get_or_init(|| {
665 std::env::var("MEMRA_FA_SP_SHORT")
666 .ok()
667 .and_then(|v| v.parse().ok())
668 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
669 }) {
670 return sp;
671 }
672 16
673 } else if t_kv <= 16384 {
674 64
675 } else {
676 128
677 }
678 } else if n_head_kv <= 4 {
679 if t_kv <= 512 {
700 8
701 } else if t_kv <= 16384 {
702 64
703 } else {
704 128
705 }
706 } else {
707 if t_kv <= 8192 {
708 32
709 } else if t_kv <= 16384 {
710 64
711 } else {
712 128
713 }
714 }
715}
716
717pub(crate) fn fa_sm_count() -> i32 {
720 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
721 *N.get_or_init(|| {
722 cudarc::driver::result::init().ok();
723 cudarc::driver::result::device::get(0)
724 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
725 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
726 .unwrap_or(82)
727 })
728}
729
730fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
734 match head_dim {
735 256 => Ok(""),
736 128 => Ok("_hd128"),
737 d => Err(format!(
738 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
739 callers must gate to sdpa_naive"
740 )
741 .into()),
742 }
743}
744
745pub const QT_Q8_0: i32 = 0;
747pub const QT_Q4_K: i32 = 1;
748pub const QT_Q6_K: i32 = 2;
749pub const QT_Q5_K: i32 = 3;
750pub const QT_Q3_K: i32 = 4;
751pub const QT_IQ4_XS: i32 = 5;
752pub const QT_IQ3_S: i32 = 6;
753pub const QT_NVFP4: i32 = 7;
754pub const QT_F8_E4M3: i32 = 10;
760pub const QT_NVFP4_RP: i32 = 9;
763pub const QT_F32: i32 = 8;
765pub const QT_BF16: i32 = 11;
766pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
771pub const QT_F8_E4M3_BLK: i32 = 14;
787
788pub struct Engine {
790 pub gpu: memra_runtime::Gpu,
791 module: Arc<CudaModule>,
792 hybrid: Arc<CudaModule>,
793 qmatvec: Arc<CudaModule>,
794 flash: Arc<CudaModule>,
795 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
799 gemm: Arc<CudaModule>,
800 router: Arc<CudaModule>,
801 sample: Arc<CudaModule>,
803 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
807 moe_cache_layout: Mutex<Option<Vec<usize>>>,
811 capture_keep_on: std::sync::atomic::AtomicBool,
817 verify_exact: std::sync::atomic::AtomicBool,
822 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
823 pub copy_stream: Arc<CudaStream>,
825 #[cfg(memra_cutlass)]
832 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
833 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
837 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
840 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
844 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
848 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
850 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
851 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
856 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
861 router_stage: Mutex<Option<PinnedStage>>,
865}
866
867fn fa_v2_on() -> bool {
877 std::env::var("MEMRA_FA_V2")
883 .map(|v| v != "0")
884 .unwrap_or(true)
885}
886
887pub(crate) fn fa_v3_on() -> bool {
895 std::env::var("MEMRA_FA_V3")
899 .map(|v| v != "0")
900 .unwrap_or(true)
901}
902
903fn fa_v4_mode() -> &'static str {
908 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
909 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
910}
911fn fa_v4_on() -> bool {
912 fa_v4_mode() != "0"
913} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
922 std::sync::atomic::AtomicUsize::new(1024);
923pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
924 std::sync::atomic::AtomicUsize::new(usize::MAX);
925pub fn fa_v4_at_pub(t_kv: usize) -> bool {
926 fa_v4_at(t_kv)
927}
928fn fa_v4_at(t_kv: usize) -> bool {
929 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
930 let mx = *M.get_or_init(|| {
931 std::env::var("MEMRA_FA_V4_MAX")
932 .ok()
933 .and_then(|v| v.parse().ok())
934 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
935 });
936 fa_v4_on() && t_kv < mx
937}
938pub const FA_DEEP_MIN_DEFAULT: usize = 0;
952fn fa_deep_at(t_kv: usize) -> bool {
953 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
954 return false;
955 }
956 let min = std::env::var("MEMRA_FA_DEEP_MIN")
957 .ok()
958 .and_then(|v| v.parse().ok())
959 .unwrap_or(FA_DEEP_MIN_DEFAULT);
960 t_kv >= min
961}
962pub fn fa_deep_at_pub(t_kv: usize) -> bool {
964 fa_deep_at(t_kv)
965}
966
967fn fa_v3_active(head_dim: usize) -> bool {
968 fa_v3_on()
971 && head_dim % 128 == 0
972 && kv_cache_formats() == ("q8_0", "q5_1")
973 && !Engine::kv_fp8_on()
974}
975
976pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
984 std::env::var("MEMRA_NO_FA_VEC").is_err()
985 && t_kv >= fa_vec_min_tkv()
986 && head_dim == 256
987 && fa_v4_at(t_kv)
988 && !matches!(fa_v4_mode(), "noB3" | "stage")
989 && !Engine::kv_fp8_on()
990}
991pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
993 fa_split_keys(t_kv, n_head_kv)
994}
995
996struct PinnedStage {
1001 ptr: *mut u8,
1002 cap: usize,
1003}
1004unsafe impl Send for PinnedStage {}
1005impl PinnedStage {
1006 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1007 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1008 Ok(PinnedStage { ptr, cap })
1009 }
1010}
1011impl Drop for PinnedStage {
1012 fn drop(&mut self) {
1013 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1014 }
1015}
1016
1017pub const ARGMAX_NB: usize = 256;
1020
1021pub(crate) use memra_fa3_vl as fa3_vl_raw;
1023
1024unsafe extern "C" {
1025 fn memra_fa3_prefill(
1027 q16: *const core::ffi::c_void,
1028 k16: *const core::ffi::c_void,
1029 v16: *const core::ffi::c_void,
1030 o: *mut f32,
1031 t: i32,
1032 h: i32,
1033 hkv: i32,
1034 d: i32,
1035 scale: f32,
1036 stream: *mut core::ffi::c_void,
1037 ) -> i32;
1038 pub(crate) fn memra_fa3_vl(
1040 q16s: *const *const core::ffi::c_void,
1041 k16s: *const *const core::ffi::c_void,
1042 v16s: *const *const core::ffi::c_void,
1043 os: *const *mut f32,
1044 ts: *const i32,
1045 b: i32,
1046 h: i32,
1047 hkv: i32,
1048 d: i32,
1049 scale: f32,
1050 stream: *mut core::ffi::c_void,
1051 ) -> i32;
1052}
1053
1054#[repr(C)]
1059#[derive(Clone, Copy)]
1060pub struct WPtr8(pub [u64; 8]);
1061unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1062
1063#[repr(C)]
1068#[derive(Clone, Copy, Default)]
1069pub struct GdnSeqVl {
1070 pub kb16: u64,
1071 pub gcum: u64,
1072 pub beta: u64,
1073 pub u: u64,
1074 pub wb16: u64,
1075 pub y: u64,
1076 pub ssnap: u64,
1077 pub state_in: u64,
1078 pub state_out: u64,
1079 pub q: u64,
1080 pub p: u64,
1081 pub o: u64,
1082 pub k: u64,
1083 pub v: u64,
1084 pub g: u64,
1085 pub a: u64,
1086 pub w: u64,
1087 pub t: i32,
1088 pub nc: i32,
1089}
1090unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1091#[repr(C)]
1092#[derive(Clone, Copy)]
1093pub struct GdnVl8(pub [GdnSeqVl; 8]);
1094unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1095
1096#[repr(C)]
1099#[derive(Clone, Copy, Default)]
1100pub struct GdnWVl {
1101 pub qb16: u64,
1102 pub pb16: u64,
1103}
1104unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1105#[repr(C)]
1106#[derive(Clone, Copy)]
1107pub struct GdnWVl8(pub [GdnWVl; 8]);
1108unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1109
1110#[repr(C)]
1112#[derive(Clone, Copy, Default)]
1113pub struct GdnPrepVl {
1114 pub qkv: u64,
1115 pub conv_state: u64,
1116 pub conv_out: u64,
1117 pub q_g: u64,
1118 pub k_g: u64,
1119 pub v_g: u64,
1120 pub q_l2: u64,
1121 pub k_l2: u64,
1122 pub beta_raw: u64,
1123 pub alpha: u64,
1124 pub beta: u64,
1125 pub g_log: u64,
1126 pub o: u64,
1127 pub z: u64,
1128 pub gn: u64,
1129 pub gn16: u64,
1130 pub kb16: u64,
1131 pub qb16: u64,
1132 pub t: i32,
1133 pub pad: i32,
1134}
1135unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1136#[repr(C)]
1137#[derive(Clone, Copy)]
1138pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1139unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1140
1141#[repr(C)]
1143#[derive(Clone, Copy, Default)]
1144pub struct FaSeqVl {
1145 pub q: u64,
1146 pub k16: u64,
1147 pub v16: u64,
1148 pub o: u64,
1149 pub kf: u64,
1150 pub vf: u64,
1151 pub t: i32,
1152 pub pad: i32,
1153}
1154unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1155#[repr(C)]
1156#[derive(Clone, Copy)]
1157pub struct FaVl8(pub [FaSeqVl; 8]);
1158unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1159
1160#[repr(C)]
1162#[derive(Clone, Copy, Default)]
1163pub struct AttnPreVl {
1164 pub qf: u64,
1165 pub kf: u64,
1166 pub vf: u64,
1167 pub q: u64,
1168 pub gate: u64,
1169 pub qn: u64,
1170 pub kn: u64,
1171 pub kc: u64,
1172 pub vc: u64,
1173 pub t: i32,
1174 pub pad: i32,
1175}
1176unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1177#[repr(C)]
1178#[derive(Clone, Copy)]
1179pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1180unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1181
1182pub struct GdnChunkBufs {
1185 pub gcum: CudaSlice<f32>,
1186 pub a: CudaSlice<f32>,
1187 pub p: CudaSlice<f32>,
1188 pub u: CudaSlice<f32>,
1189 pub w: CudaSlice<f32>,
1190 pub kb16: CudaSlice<u8>,
1191 pub wb16: CudaSlice<u8>,
1192 pub y16: CudaSlice<u8>,
1193 pub ssnap16: CudaSlice<u8>,
1194 pub qb16: CudaSlice<u8>,
1195 pub pb16: CudaSlice<u8>,
1196 pub o: CudaSlice<f32>,
1197 pub t: usize,
1198 pub nc: usize,
1199}
1200
1201#[repr(C)]
1203#[derive(Clone, Copy)]
1204pub struct F32x8(pub [f32; 8]);
1205unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1206
1207pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1211
1212impl Engine {
1213 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1214 let gpu = memra_runtime::Gpu::new(ordinal)?;
1215 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1219 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1220 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1221 .and_then(|d| unsafe {
1222 Ok((
1223 cudarc::driver::result::device::get_attribute(
1224 d,
1225 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1226 )?,
1227 cudarc::driver::result::device::get_attribute(
1228 d,
1229 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1230 )?,
1231 ))
1232 })
1233 .unwrap_or((0, 0));
1234 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1235 let ok = matches!(
1236 (built, maj, min),
1237 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1238 );
1239 if !ok {
1240 return Err(format!(
1241 "memra was built for sm_{built} but device {ordinal} reports compute \
1242 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1243 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1244 )
1245 .into());
1246 }
1247 }
1248 unsafe {
1253 use cudarc::driver::sys;
1254 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1255 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1256 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1257 let mut thresh: u64 = u64::MAX;
1258 let _ = sys::cuMemPoolSetAttribute(
1259 pool,
1260 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1261 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1262 );
1263 }
1264 }
1265 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1266 let hybrid = gpu
1267 .ctx
1268 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1269 let qmatvec = gpu
1270 .ctx
1271 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1272 let flash = gpu
1273 .ctx
1274 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1275 let gemm = gpu
1276 .ctx
1277 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1278 let router = gpu
1279 .ctx
1280 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1281 let sample = gpu
1282 .ctx
1283 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1284 let copy_stream = gpu.ctx.new_stream()?;
1285 if std::env::var("MEMRA_EVT")
1301 .map(|v| v == "1")
1302 .unwrap_or(false)
1303 {
1304 } else {
1306 unsafe {
1307 gpu.ctx.disable_event_tracking();
1308 }
1309 }
1310 Ok(Self {
1311 gpu,
1312 module,
1313 hybrid,
1314 qmatvec,
1315 flash,
1316 flash_g: std::sync::OnceLock::new(),
1317 gemm,
1318 router,
1319 sample,
1320 moe_cache: Mutex::new(None),
1321 moe_cache_layout: Mutex::new(None),
1322 copy_stream,
1323 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1324 verify_exact: std::sync::atomic::AtomicBool::new(false),
1325 capture_keep: Mutex::new(Vec::new()),
1326 argmax_partials: Mutex::new(None),
1327 prime_deqw_ws: Mutex::new(None),
1328 router_stage: Mutex::new(None),
1329 fp8_scratch: Mutex::new(None),
1330 fa_vf16_scratch: Mutex::new(None),
1331 fa_part_pool: Mutex::new(None),
1332 fa_part_retired: Mutex::new(Vec::new()),
1333 fn_cache: Mutex::new(Default::default()),
1334 f16_scratch: Mutex::new(None),
1335 #[cfg(memra_cutlass)]
1336 cutlass_scratch: Mutex::new(None),
1337 })
1338 }
1339
1340 pub fn ctx(&self) -> &Arc<CudaContext> {
1341 &self.gpu.ctx
1342 }
1343
1344 pub fn pool_cached_bytes(&self) -> usize {
1362 let (reserved, used) = self.pool_reserved_used();
1363 reserved.saturating_sub(used)
1364 }
1365
1366 pub fn device_graph_mem_reserved(&self) -> usize {
1376 use cudarc::driver::sys as cus;
1377 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1378 return 0;
1379 };
1380 let mut bytes: u64 = 0;
1381 let rc = unsafe {
1382 cus::cuDeviceGetGraphMemAttribute(
1383 dev,
1384 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1385 &mut bytes as *mut u64 as *mut std::ffi::c_void,
1386 )
1387 };
1388 if rc == cus::cudaError_enum::CUDA_SUCCESS {
1389 bytes as usize
1390 } else {
1391 0
1392 }
1393 }
1394
1395 pub fn pool_reserved_used(&self) -> (usize, usize) {
1402 use cudarc::driver::sys;
1403 unsafe {
1404 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1405 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1406 != sys::CUresult::CUDA_SUCCESS
1407 {
1408 return (0, 0);
1409 }
1410 let (mut reserved, mut used) = (0u64, 0u64);
1411 if sys::cuMemPoolGetAttribute(
1412 pool,
1413 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1414 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1415 ) != sys::CUresult::CUDA_SUCCESS
1416 {
1417 return (0, 0);
1418 }
1419 if sys::cuMemPoolGetAttribute(
1420 pool,
1421 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1422 &mut used as *mut u64 as *mut core::ffi::c_void,
1423 ) != sys::CUresult::CUDA_SUCCESS
1424 {
1425 return (0, 0);
1426 }
1427 (reserved as usize, used as usize)
1428 }
1429 }
1430
1431 pub fn stream(&self) -> Arc<CudaStream> {
1434 self.gpu.stream()
1435 }
1436 pub fn gkv_on() -> bool {
1439 memra_kv::gkv_on()
1440 }
1441
1442 pub fn wkv_on() -> bool {
1454 memra_kv::wkv_on()
1455 }
1456
1457 pub fn kv_fp8_on() -> bool {
1463 memra_kv::kv_fp8_on()
1464 }
1465
1466 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1469 if head_dim == 512 && Self::gkv_on() {
1470 self.func_g(name)
1471 } else {
1472 self.func(name)
1473 }
1474 }
1475
1476 fn func_g(&self, name: &str) -> CudaFunction {
1480 let m = self.flash_g.get_or_init(|| {
1481 self.gpu
1482 .ctx
1483 .load_module(cudarc::nvrtc::Ptx::from_binary(
1484 FLASH_FATBIN_KF8VF8.to_vec(),
1485 ))
1486 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1487 });
1488 let key = format!("g:{name}");
1489 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1490 return f.clone();
1491 }
1492 let f = match m.load_function(name) {
1493 Ok(f) => f,
1494 Err(_) => self.func(name),
1495 };
1496 self.fn_cache.lock().unwrap().insert(key, f.clone());
1497 f
1498 }
1499
1500 fn func(&self, name: &str) -> CudaFunction {
1501 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1504 return f.clone();
1505 }
1506 let f = self
1507 .module
1508 .load_function(name)
1509 .or_else(|_| self.hybrid.load_function(name))
1510 .or_else(|_| self.qmatvec.load_function(name))
1511 .or_else(|_| self.flash.load_function(name))
1512 .or_else(|_| self.gemm.load_function(name))
1513 .or_else(|_| self.router.load_function(name))
1514 .or_else(|_| self.sample.load_function(name))
1515 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1516 self.fn_cache
1517 .lock()
1518 .unwrap()
1519 .insert(name.to_string(), f.clone());
1520 f
1521 }
1522
1523 pub fn scatter_trim_logits(
1526 &self,
1527 src: &CudaSlice<f32>,
1528 d2t: &CudaSlice<u32>,
1529 dst: &mut CudaSlice<f32>,
1530 d_vocab: usize,
1531 n_vocab: usize,
1532 ) -> Result<(), Box<dyn std::error::Error>> {
1533 let f1 = self.func("scatter_trim_logits_f32");
1534 let f2 = self.func("scatter_trim_logits_pass2_f32");
1535 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1536 let cfg1 = LaunchConfig {
1537 grid_dim: (256, 1, 1),
1538 block_dim: (256, 1, 1),
1539 shared_mem_bytes: 0,
1540 };
1541 let __s_b1 = self.gpu.stream();
1542 let mut b1 = __s_b1.launch_builder(&f1);
1543 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1544 unsafe {
1545 b1.launch(cfg1)?;
1546 }
1547 let cfg2 = LaunchConfig {
1548 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1549 block_dim: (256, 1, 1),
1550 shared_mem_bytes: 0,
1551 };
1552 let __s_b2 = self.gpu.stream();
1553 let mut b2 = __s_b2.launch_builder(&f2);
1554 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1555 unsafe {
1556 b2.launch(cfg2)?;
1557 }
1558 Ok(())
1559 }
1560
1561 #[allow(clippy::too_many_arguments)]
1567 pub fn filter_stats(
1568 &self,
1569 x: &CudaSlice<f32>,
1570 row_stride: usize,
1571 rows: &CudaSlice<i32>,
1572 out_th: &mut CudaSlice<f32>,
1573 out_z: &mut CudaSlice<f32>,
1574 out_max: &mut CudaSlice<f32>,
1575 n: usize,
1576 nrow: usize,
1577 temp: f32,
1578 top_k: i32,
1579 top_p: f32,
1580 min_p: f32,
1581 ) -> Result<(), Box<dyn std::error::Error>> {
1582 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1597 let coop_on =
1598 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1599 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1600 if coop_on && 16 * nrow <= self.sm_count() as usize {
1601 let f = self.func("filter_stats_coop_f32");
1602 let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1603 let cfg = LaunchConfig {
1604 grid_dim: (16, nrow as u32, 1),
1605 block_dim: (512, 1, 1),
1606 shared_mem_bytes: 0,
1607 };
1608 let __s_b = self.gpu.stream();
1609 let mut b = __s_b.launch_builder(&f);
1610 b.arg(x)
1611 .arg(&rs)
1612 .arg(rows)
1613 .arg(&mut *out_th)
1614 .arg(&mut *out_z)
1615 .arg(&mut *out_max)
1616 .arg(&mut ws)
1617 .arg(&ni)
1618 .arg(&nr)
1619 .arg(&temp)
1620 .arg(&top_k)
1621 .arg(&top_p)
1622 .arg(&min_p);
1623 unsafe {
1624 b.launch_cooperative(cfg)?;
1625 }
1626 return Ok(());
1627 }
1628 let f = self.func("filter_stats_f32");
1629 let cfg = LaunchConfig {
1630 grid_dim: (nrow as u32, 1, 1),
1631 block_dim: (1024, 1, 1),
1632 shared_mem_bytes: 0,
1633 };
1634 let __s_b = self.gpu.stream();
1635 let mut b = __s_b.launch_builder(&f);
1636 b.arg(x)
1637 .arg(&rs)
1638 .arg(rows)
1639 .arg(&mut *out_th)
1640 .arg(&mut *out_z)
1641 .arg(&mut *out_max)
1642 .arg(&ni)
1643 .arg(&nr)
1644 .arg(&temp)
1645 .arg(&top_k)
1646 .arg(&top_p)
1647 .arg(&min_p);
1648 unsafe {
1649 b.launch(cfg)?;
1650 }
1651 Ok(())
1652 }
1653
1654 #[allow(clippy::too_many_arguments)]
1656 pub fn softmax_gather_filtered(
1657 &self,
1658 x: &CudaSlice<f32>,
1659 row_stride: usize,
1660 ids: &CudaSlice<u32>,
1661 rows: &CudaSlice<i32>,
1662 th: &CudaSlice<f32>,
1663 z: &CudaSlice<f32>,
1664 out: &mut CudaSlice<f32>,
1665 n: usize,
1666 npair: usize,
1667 temp: f32,
1668 ) -> Result<(), Box<dyn std::error::Error>> {
1669 let f = self.func("softmax_gather_filtered_f32");
1670 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1671 let cfg = LaunchConfig {
1672 grid_dim: (npair as u32, 1, 1),
1673 block_dim: (256, 1, 1),
1674 shared_mem_bytes: 0,
1675 };
1676 let __s_b = self.gpu.stream();
1677 let mut b = __s_b.launch_builder(&f);
1678 b.arg(x)
1679 .arg(&rs)
1680 .arg(ids)
1681 .arg(rows)
1682 .arg(th)
1683 .arg(z)
1684 .arg(&mut *out)
1685 .arg(&ni)
1686 .arg(&np)
1687 .arg(&temp);
1688 unsafe {
1689 b.launch(cfg)?;
1690 }
1691 Ok(())
1692 }
1693
1694 #[allow(clippy::too_many_arguments)]
1696 pub fn residual_sample_filtered(
1697 &self,
1698 p: &CudaSlice<f32>,
1699 q: Option<&CudaSlice<f32>>,
1700 n: usize,
1701 temp: f32,
1702 seed: u64,
1703 stream_pos: u32,
1704 p_stats: (f32, f32, f32),
1705 q_stats: (f32, f32, f32),
1706 out_tok: &mut CudaSlice<u32>,
1707 ) -> Result<(), Box<dyn std::error::Error>> {
1708 let f = self.func("residual_sample_filtered_f32");
1709 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1710 let has_q: i32 = q.is_some() as i32;
1711 let qbuf = q.unwrap_or(p);
1712 let (pm, pth, pz) = p_stats;
1713 let (qm, qth, qz) = q_stats;
1714 let cfg = LaunchConfig {
1715 grid_dim: (1, 1, 1),
1716 block_dim: (1024, 1, 1),
1717 shared_mem_bytes: 0,
1718 };
1719 let __s_b = self.gpu.stream();
1720 let mut b = __s_b.launch_builder(&f);
1721 b.arg(p)
1722 .arg(qbuf)
1723 .arg(&has_q)
1724 .arg(&ni)
1725 .arg(&temp)
1726 .arg(&slo)
1727 .arg(&shi)
1728 .arg(&stream_pos)
1729 .arg(&pm)
1730 .arg(&pth)
1731 .arg(&pz)
1732 .arg(&qm)
1733 .arg(&qth)
1734 .arg(&qz)
1735 .arg(&mut *out_tok);
1736 unsafe {
1737 b.launch(cfg)?;
1738 }
1739 Ok(())
1740 }
1741
1742 #[allow(clippy::too_many_arguments)]
1748 pub fn residual_sample_sparse_q(
1749 &self,
1750 p: &CudaSlice<f32>,
1751 cand_ids: &CudaSlice<u32>,
1752 q_probs: &CudaSlice<f32>,
1753 n_cand: usize,
1754 n: usize,
1755 temp: f32,
1756 seed: u64,
1757 stream_pos: u32,
1758 p_stats: (f32, f32, f32),
1759 out_tok: &mut CudaSlice<u32>,
1760 ) -> Result<(), Box<dyn std::error::Error>> {
1761 assert!(
1762 n_cand >= 1 && n_cand <= 32,
1763 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1764 );
1765 let f = self.func("residual_sample_sparse_q_f32");
1766 let (ni, nc) = (n as i32, n_cand as i32);
1767 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1768 let (pm, pth, pz) = p_stats;
1769 let cfg = LaunchConfig {
1770 grid_dim: (1, 1, 1),
1771 block_dim: (1024, 1, 1),
1772 shared_mem_bytes: 0,
1773 };
1774 let __s_b = self.gpu.stream();
1775 let mut b = __s_b.launch_builder(&f);
1776 b.arg(p)
1777 .arg(cand_ids)
1778 .arg(q_probs)
1779 .arg(&nc)
1780 .arg(&ni)
1781 .arg(&temp)
1782 .arg(&slo)
1783 .arg(&shi)
1784 .arg(&stream_pos)
1785 .arg(&pm)
1786 .arg(&pth)
1787 .arg(&pz)
1788 .arg(&mut *out_tok);
1789 unsafe {
1790 b.launch(cfg)?;
1791 }
1792 Ok(())
1793 }
1794
1795 #[allow(clippy::too_many_arguments)]
1797 pub fn gumbel_perturb_filtered(
1798 &self,
1799 x: &CudaSlice<f32>,
1800 y: &mut CudaSlice<f32>,
1801 n: usize,
1802 seed: u64,
1803 stream_pos: u32,
1804 temp: f32,
1805 row_max: f32,
1806 th: f32,
1807 ) -> Result<(), Box<dyn std::error::Error>> {
1808 let f = self.func("gumbel_perturb_filtered_f32");
1809 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1810 let cfg = LaunchConfig {
1811 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1812 block_dim: (256, 1, 1),
1813 shared_mem_bytes: 0,
1814 };
1815 let __s_b = self.gpu.stream();
1816 let mut b = __s_b.launch_builder(&f);
1817 b.arg(x)
1818 .arg(&mut *y)
1819 .arg(&ni)
1820 .arg(&slo)
1821 .arg(&shi)
1822 .arg(&stream_pos)
1823 .arg(&temp)
1824 .arg(&row_max)
1825 .arg(&th);
1826 unsafe {
1827 b.launch(cfg)?;
1828 }
1829 Ok(())
1830 }
1831
1832 #[allow(clippy::too_many_arguments)]
1836 pub fn penalize_logits(
1837 &self,
1838 x: &mut CudaSlice<f32>,
1839 hist: &CudaSlice<u32>,
1840 n_hist: usize,
1841 rep: f32,
1842 freq: f32,
1843 present: f32,
1844 n: usize,
1845 ) -> Result<(), Box<dyn std::error::Error>> {
1846 if n_hist == 0 {
1847 return Ok(());
1848 }
1849 let f = self.func("penalize_logits_f32");
1850 let (nh, ni) = (n_hist as i32, n as i32);
1851 let cfg = LaunchConfig {
1852 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1853 block_dim: (128, 1, 1),
1854 shared_mem_bytes: 0,
1855 };
1856 let __s_b = self.gpu.stream();
1857 let mut b = __s_b.launch_builder(&f);
1858 b.arg(&mut *x)
1859 .arg(hist)
1860 .arg(&nh)
1861 .arg(&rep)
1862 .arg(&freq)
1863 .arg(&present)
1864 .arg(&ni);
1865 unsafe {
1866 b.launch(cfg)?;
1867 }
1868 Ok(())
1869 }
1870
1871 #[allow(clippy::too_many_arguments)]
1873 pub fn penalize_logits_rows(
1874 &self,
1875 x: &mut CudaSlice<f32>,
1876 hist: &CudaSlice<u32>,
1877 n_hist: usize,
1878 rep: f32,
1879 freq: f32,
1880 present: f32,
1881 n: usize,
1882 nrow: usize,
1883 ) -> Result<(), Box<dyn std::error::Error>> {
1884 if n_hist == 0 || nrow == 0 {
1885 return Ok(());
1886 }
1887 let f = self.func("penalize_logits_rows_f32");
1888 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1889 let cfg = LaunchConfig {
1890 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1891 block_dim: (128, 1, 1),
1892 shared_mem_bytes: 0,
1893 };
1894 let __s_b = self.gpu.stream();
1895 let mut b = __s_b.launch_builder(&f);
1896 b.arg(&mut *x)
1897 .arg(hist)
1898 .arg(&nh)
1899 .arg(&rep)
1900 .arg(&freq)
1901 .arg(&present)
1902 .arg(&ni)
1903 .arg(&nr);
1904 unsafe {
1905 b.launch(cfg)?;
1906 }
1907 Ok(())
1908 }
1909
1910 #[allow(clippy::too_many_arguments)]
1916 pub fn penalize_logits_sparse_rows(
1917 &self,
1918 x: &mut CudaSlice<f32>,
1919 ids: &[u32],
1920 counts: &[u32],
1921 offsets: &[i32],
1922 rows: &[i32],
1923 reps: &[f32],
1924 freqs: &[f32],
1925 presents: &[f32],
1926 n: usize,
1927 ) -> Result<(), Box<dyn std::error::Error>> {
1928 let nrow = rows.len();
1929 if nrow == 0 {
1930 return Ok(());
1931 }
1932 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
1933 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
1934 let entry_count =
1935 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
1936 if ids.len() != counts.len()
1937 || offsets.len() != nrow + 1
1938 || reps.len() != nrow
1939 || freqs.len() != nrow
1940 || presents.len() != nrow
1941 || offsets.first().copied() != Some(0)
1942 || offsets.last().copied() != Some(entry_count)
1943 {
1944 return Err("sparse penalty row metadata shape mismatch".into());
1945 }
1946 if counts.contains(&0) {
1947 return Err("sparse penalty counts must be positive".into());
1948 }
1949 let mut max_len = 0usize;
1950 for pair in offsets.windows(2) {
1951 if pair[0] < 0 || pair[1] < pair[0] {
1952 return Err("sparse penalty offsets must be monotonic".into());
1953 }
1954 max_len = max_len.max((pair[1] - pair[0]) as usize);
1955 }
1956 if max_len == 0 {
1957 return Ok(());
1958 }
1959
1960 let mut seen = std::collections::HashSet::with_capacity(ids.len());
1961 for (r, &row) in rows.iter().enumerate() {
1962 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
1963 return Err("sparse penalty row index exceeds logits shape".into());
1964 }
1965 let begin = offsets[r] as usize;
1966 let end = offsets[r + 1] as usize;
1967 for &id in &ids[begin..end] {
1968 if id as usize >= n {
1969 return Err("sparse penalty token id exceeds logits row".into());
1970 }
1971 if !seen.insert((row, id)) {
1972 return Err("sparse penalty entries must be unique per logits row".into());
1973 }
1974 }
1975 }
1976
1977 unsafe {
1979 self.penalize_logits_sparse_rows_unchecked(
1980 x, ids, counts, offsets, rows, reps, freqs, presents, n,
1981 )
1982 }
1983 }
1984
1985 #[allow(clippy::too_many_arguments)]
1994 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
1995 &self,
1996 x: &mut CudaSlice<f32>,
1997 ids: &[u32],
1998 counts: &[u32],
1999 offsets: &[i32],
2000 rows: &[i32],
2001 reps: &[f32],
2002 freqs: &[f32],
2003 presents: &[f32],
2004 n: usize,
2005 ) -> Result<(), Box<dyn std::error::Error>> {
2006 let nrow = rows.len();
2007 if nrow == 0 {
2008 return Ok(());
2009 }
2010 let max_len = offsets
2011 .windows(2)
2012 .map(|pair| (pair[1] - pair[0]) as usize)
2013 .max()
2014 .unwrap_or(0);
2015 if max_len == 0 {
2016 return Ok(());
2017 }
2018 let ids_d = self.htod_u32_v(ids)?;
2019 let counts_d = self.htod_u32_v(counts)?;
2020 let offsets_d = self.htod_i32(offsets)?;
2021 let rows_d = self.htod_i32(rows)?;
2022 let reps_d = self.htod(reps)?;
2023 let freqs_d = self.htod(freqs)?;
2024 let presents_d = self.htod(presents)?;
2025 let f = self.func("penalize_logits_sparse_rows_f32");
2026 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2027 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2028 let cfg = LaunchConfig {
2029 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2030 block_dim: (128, 1, 1),
2031 shared_mem_bytes: 0,
2032 };
2033 let __s_b = self.gpu.stream();
2034 let mut b = __s_b.launch_builder(&f);
2035 b.arg(&mut *x)
2036 .arg(&ids_d)
2037 .arg(&counts_d)
2038 .arg(&offsets_d)
2039 .arg(&rows_d)
2040 .arg(&reps_d)
2041 .arg(&freqs_d)
2042 .arg(&presents_d)
2043 .arg(&ni)
2044 .arg(&nr);
2045 unsafe {
2046 b.launch(cfg)?;
2047 }
2048 Ok(())
2049 }
2050
2051 #[allow(clippy::too_many_arguments)]
2059 pub fn penalize_logits_rows_inc(
2060 &self,
2061 x: &mut CudaSlice<f32>,
2062 hist: &CudaSlice<u32>,
2063 n_hist0: usize,
2064 rep: f32,
2065 freq: f32,
2066 present: f32,
2067 n: usize,
2068 nrow: usize,
2069 win: usize,
2070 ) -> Result<(), Box<dyn std::error::Error>> {
2071 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2072 return Ok(());
2073 }
2074 debug_assert!(
2075 hist.len() >= n_hist0 + nrow - 1,
2076 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2077 );
2078 let f = self.func("penalize_logits_rows_inc_f32");
2079 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2080 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2081 let cfg = LaunchConfig {
2082 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2083 block_dim: (128, 1, 1),
2084 shared_mem_bytes: 0,
2085 };
2086 let __s_b = self.gpu.stream();
2087 let mut b = __s_b.launch_builder(&f);
2088 b.arg(&mut *x)
2089 .arg(hist)
2090 .arg(&nh)
2091 .arg(&rep)
2092 .arg(&freq)
2093 .arg(&present)
2094 .arg(&ni)
2095 .arg(&nr)
2096 .arg(&wi);
2097 unsafe {
2098 b.launch(cfg)?;
2099 }
2100 Ok(())
2101 }
2102
2103 pub fn wpf_level() -> u32 {
2111 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2112 *ON.get_or_init(|| {
2113 std::env::var("MEMRA_WPF")
2114 .ok()
2115 .and_then(|v| v.parse().ok())
2116 .unwrap_or(1)
2117 })
2118 }
2119
2120 pub fn set_verify_exact(&self, on: bool) {
2132 self.verify_exact
2133 .store(on, std::sync::atomic::Ordering::Relaxed);
2134 }
2135 pub(crate) fn verify_exact_on(&self) -> bool {
2136 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2137 }
2138
2139 pub fn qkv_append_on() -> bool {
2142 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2143 *ON.get_or_init(|| {
2144 std::env::var("MEMRA_QKV_APPEND")
2145 .map(|v| v != "0")
2146 .unwrap_or(true)
2147 })
2148 }
2149
2150 pub fn pdl_wb_on() -> bool {
2153 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2154 *ON.get_or_init(|| {
2155 std::env::var("MEMRA_PDL_WB")
2156 .map(|v| v != "0")
2157 .unwrap_or(true)
2158 })
2159 }
2160
2161 pub fn norm_ilp_on() -> bool {
2169 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2170 *ON.get_or_init(|| {
2171 std::env::var("MEMRA_NORM_ILP")
2172 .map(|v| v != "0")
2173 .unwrap_or(true)
2174 })
2175 }
2176
2177 pub fn tk_ffn_dual_on() -> bool {
2185 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2186 *ON.get_or_init(|| {
2187 std::env::var("MEMRA_TK_FFN_DUAL")
2188 .map(|v| v != "0")
2189 .unwrap_or(true)
2190 })
2191 }
2192
2193 pub fn pdl_mmvq_on() -> bool {
2197 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2198 *ON.get_or_init(|| {
2199 std::env::var("MEMRA_PDL_MMVQ")
2200 .map(|v| v != "0")
2201 .unwrap_or(true)
2202 })
2203 }
2204
2205 pub fn pdl_on() -> bool {
2206 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2207 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2208 }
2209
2210 pub fn pdl_nvfp4q8_on() -> bool {
2216 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2217 *ON.get_or_init(|| {
2218 std::env::var("MEMRA_PDL_NVFP4")
2219 .map(|v| v != "0")
2220 .unwrap_or(true)
2221 })
2222 }
2223
2224 fn q40_mr1_on() -> bool {
2230 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2231 match *Q40MR.get_or_init(|| {
2232 std::env::var("MEMRA_Q40_MR")
2233 .ok()
2234 .and_then(|v| v.parse().ok())
2235 }) {
2236 Some(v) => v == 1,
2237 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2238 }
2239 }
2240
2241 fn pdl_func_flash(
2246 &self,
2247 g: bool,
2248 name: &'static str,
2249 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2250 use cudarc::driver::sys as cu;
2251 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2258 std::sync::Mutex::new(None);
2259 static FNS: std::sync::Mutex<
2260 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2261 > = std::sync::Mutex::new(None);
2262 let ctx_key = self.ctx().cu_ctx() as usize;
2263 if let Some(&f) = FNS
2264 .lock()
2265 .unwrap()
2266 .get_or_insert_with(Default::default)
2267 .get(&(ctx_key, g, name))
2268 {
2269 return Ok(f as cu::CUfunction);
2270 }
2271 let module = {
2272 let mut mods = MODS.lock().unwrap();
2273 let map = mods.get_or_insert_with(Default::default);
2274 match map.get(&(ctx_key, g)) {
2275 Some(&m) => m,
2276 None => {
2277 let m = self.pdl_load_module_in_ctx(if g {
2278 FLASH_FATBIN_KF8VF8
2279 } else {
2280 FLASH_FATBIN
2281 })?;
2282 map.insert((ctx_key, g), m);
2283 m
2284 }
2285 }
2286 };
2287 let cname = std::ffi::CString::new(name)?;
2288 let mut f: cu::CUfunction = std::ptr::null_mut();
2289 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2290 if r != cu::CUresult::CUDA_SUCCESS {
2291 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2292 }
2293 FNS.lock()
2294 .unwrap()
2295 .get_or_insert_with(Default::default)
2296 .insert((ctx_key, g, name), f as usize);
2297 Ok(f)
2298 }
2299
2300 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2305 use cudarc::driver::sys as cu;
2306 let mut prev: cu::CUcontext = std::ptr::null_mut();
2307 unsafe {
2308 cu::cuCtxGetCurrent(&mut prev).result()?;
2309 }
2310 self.ctx().bind_to_thread()?;
2311 let mut m: cu::CUmodule = std::ptr::null_mut();
2312 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2313 let restore = if prev.is_null() {
2314 cu::CUresult::CUDA_SUCCESS
2315 } else {
2316 unsafe { cu::cuCtxSetCurrent(prev) }
2317 };
2318 if r != cu::CUresult::CUDA_SUCCESS {
2319 return Err(format!("pdl module load: {r:?}").into());
2320 }
2321 if restore != cu::CUresult::CUDA_SUCCESS {
2322 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2323 }
2324 Ok(m as usize)
2325 }
2326
2327 pub fn raw_kernel_function(
2330 &self,
2331 name: &'static str,
2332 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2333 self.pdl_func(name)
2334 }
2335
2336 fn pdl_func(
2337 &self,
2338 name: &'static str,
2339 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2340 use cudarc::driver::sys as cu;
2341 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2344 std::sync::Mutex::new(None);
2345 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2348 std::sync::Mutex::new(None);
2349 static FNS: std::sync::Mutex<
2350 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2351 > = std::sync::Mutex::new(None);
2352 let ctx_key = self.ctx().cu_ctx() as usize;
2353 if let Some(&f) = FNS
2354 .lock()
2355 .unwrap()
2356 .get_or_insert_with(Default::default)
2357 .get(&(ctx_key, name))
2358 {
2359 return Ok(f as cu::CUfunction);
2360 }
2361 let module = {
2362 let mut mods = MODULES.lock().unwrap();
2363 let map = mods.get_or_insert_with(Default::default);
2364 match map.get(&ctx_key) {
2365 Some(&m) => m,
2366 None => {
2367 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2368 map.insert(ctx_key, m);
2369 m
2370 }
2371 }
2372 };
2373 let cname = std::ffi::CString::new(name)?;
2374 let mut f: cu::CUfunction = std::ptr::null_mut();
2375 let mut r =
2376 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2377 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2378 let qmodule = {
2379 let mut mods = QMODULES.lock().unwrap();
2380 let map = mods.get_or_insert_with(Default::default);
2381 match map.get(&ctx_key) {
2382 Some(&m) => m,
2383 None => {
2384 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2385 map.insert(ctx_key, m);
2386 m
2387 }
2388 }
2389 };
2390 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2391 }
2392 if r != cu::CUresult::CUDA_SUCCESS {
2393 return Err(format!("pdl_func {name}: {r:?}").into());
2394 }
2395 FNS.lock()
2396 .unwrap()
2397 .get_or_insert_with(Default::default)
2398 .insert((ctx_key, name), f as usize);
2399 Ok(f)
2400 }
2401
2402 unsafe fn launch_pdl_flash(
2414 &self,
2415 g: bool,
2416 name: &'static str,
2417 grid: (u32, u32, u32),
2418 block: (u32, u32, u32),
2419 smem: u32,
2420 params: &mut [*mut std::ffi::c_void],
2421 ) -> Result<(), Box<dyn std::error::Error>> {
2422 use cudarc::driver::sys as cu;
2423 let f = self.pdl_func_flash(g, name)?;
2424 if smem > 0 {
2425 let r =
2427 unsafe {
2428 cu::cuFuncSetAttribute(f,
2429 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2430 smem as i32)
2431 };
2432 if r != cu::CUresult::CUDA_SUCCESS {
2433 return Err(format!("pdl smem attr {name}: {r:?}").into());
2434 }
2435 }
2436 let mut attr = cu::CUlaunchAttribute {
2437 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2438 pad: [0; 4],
2439 value: cu::CUlaunchAttributeValue {
2440 programmaticStreamSerializationAllowed: 1,
2441 },
2442 };
2443 let cfg = cu::CUlaunchConfig {
2444 gridDimX: grid.0,
2445 gridDimY: grid.1,
2446 gridDimZ: grid.2,
2447 blockDimX: block.0,
2448 blockDimY: block.1,
2449 blockDimZ: block.2,
2450 sharedMemBytes: smem,
2451 hStream: self.gpu.stream().cu_stream(),
2452 attrs: &mut attr,
2453 numAttrs: 1,
2454 };
2455 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2456 if r != cu::CUresult::CUDA_SUCCESS {
2457 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2458 }
2459 Ok(())
2460 }
2461
2462 unsafe fn launch_pdl(
2463 &self,
2464 name: &'static str,
2465 grid: (u32, u32, u32),
2466 block: (u32, u32, u32),
2467 params: &mut [*mut std::ffi::c_void],
2468 ) -> Result<(), Box<dyn std::error::Error>> {
2469 use cudarc::driver::sys as cu;
2470 let f = self.pdl_func(name)?;
2471 let mut attr = cu::CUlaunchAttribute {
2472 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2473 pad: [0; 4],
2474 value: cu::CUlaunchAttributeValue {
2475 programmaticStreamSerializationAllowed: 1,
2476 },
2477 };
2478 let cfg = cu::CUlaunchConfig {
2479 gridDimX: grid.0,
2480 gridDimY: grid.1,
2481 gridDimZ: grid.2,
2482 blockDimX: block.0,
2483 blockDimY: block.1,
2484 blockDimZ: block.2,
2485 sharedMemBytes: 0,
2486 hStream: self.gpu.stream().cu_stream(),
2487 attrs: &mut attr,
2488 numAttrs: 1,
2489 };
2490 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2491 if r != cu::CUresult::CUDA_SUCCESS {
2492 return Err(format!("launch_pdl {name}: {r:?}").into());
2493 }
2494 Ok(())
2495 }
2496
2497 pub fn prefetch_weight_l2(
2500 &self,
2501 w: &crate::model::GpuTensor,
2502 ) -> Result<(), Box<dyn std::error::Error>> {
2503 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2504 let p = rp4.as_ref().unwrap_or(bytes);
2505 self.prefetch_l2(p, p.len())?;
2506 }
2507 Ok(())
2508 }
2509
2510 pub fn gather_row_bf16(
2513 &self,
2514 table: &CudaSlice<u8>,
2515 tok: &CudaSlice<u32>,
2516 idx: usize,
2517 dst: &mut CudaSlice<f32>,
2518 ncols: usize,
2519 ) -> Result<(), Box<dyn std::error::Error>> {
2520 let f = self.func("gather_row_bf16_f32");
2521 let cfg = LaunchConfig {
2522 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2523 block_dim: (256, 1, 1),
2524 shared_mem_bytes: 0,
2525 };
2526 let (nc, ix) = (ncols as i32, idx as i32);
2527 let __s_b = self.gpu.stream();
2528 let mut b = __s_b.launch_builder(&f);
2529 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2530 unsafe {
2531 b.launch(cfg)?;
2532 }
2533 Ok(())
2534 }
2535
2536 #[allow(clippy::too_many_arguments)]
2542 pub fn dflash2_dynconv(
2543 &self,
2544 x: &CudaSlice<f32>,
2545 dyn_: &CudaSlice<f32>,
2546 base: &CudaSlice<f32>,
2547 out: &mut CudaSlice<f32>,
2548 rows: usize,
2549 hidden: usize,
2550 group_size: usize,
2551 ksize: usize,
2552 half: usize,
2553 ) -> Result<(), Box<dyn std::error::Error>> {
2554 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2555 let f = self.func("dflash2_dynconv_f32");
2556 let n = rows * hidden;
2557 let cfg = LaunchConfig {
2558 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2559 block_dim: (256, 1, 1),
2560 shared_mem_bytes: 0,
2561 };
2562 let (ri, hi, gi, ki, hf) = (
2563 rows as i32,
2564 hidden as i32,
2565 group_size as i32,
2566 ksize as i32,
2567 half as i32,
2568 );
2569 let __s_b = self.gpu.stream();
2570 let mut b = __s_b.launch_builder(&f);
2571 b.arg(x)
2572 .arg(dyn_)
2573 .arg(base)
2574 .arg(out)
2575 .arg(&ri)
2576 .arg(&hi)
2577 .arg(&gi)
2578 .arg(&ki)
2579 .arg(&hf);
2580 unsafe {
2581 b.launch(cfg)?;
2582 }
2583 Ok(())
2584 }
2585
2586 pub fn topk_rows(
2590 &self,
2591 logits: &CudaSlice<f32>,
2592 n_rows: usize,
2593 n_cols: usize,
2594 k: usize,
2595 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2596 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2597 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2598 let f = self.func("topk_rows_f32");
2599 let nth = 256usize;
2600 let mut vals = self.uninit(n_rows * k)?;
2601 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2602 let cfg = LaunchConfig {
2603 grid_dim: (n_rows as u32, 1, 1),
2604 block_dim: (nth as u32, 1, 1),
2605 shared_mem_bytes: (nth * k * 8) as u32,
2606 };
2607 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2608 let __s_b = self.gpu.stream();
2609 let mut b = __s_b.launch_builder(&f);
2610 b.arg(logits)
2611 .arg(&nr)
2612 .arg(&nc)
2613 .arg(&ki)
2614 .arg(&mut vals)
2615 .arg(&mut idxs);
2616 unsafe {
2617 b.launch(cfg)?;
2618 }
2619 Ok((vals, idxs))
2620 }
2621
2622 pub fn add_row_inplace(
2624 &self,
2625 logits: &mut CudaSlice<f32>,
2626 bias: &CudaSlice<f32>,
2627 n: usize,
2628 row_off: usize,
2629 ) -> Result<(), Box<dyn std::error::Error>> {
2630 let f = self.func("add_row_inplace_f32");
2631 let cfg = LaunchConfig {
2632 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2633 block_dim: (256, 1, 1),
2634 shared_mem_bytes: 0,
2635 };
2636 let (ni, off) = (n as i32, row_off as i64);
2637 let __s_b = self.gpu.stream();
2638 let mut b = __s_b.launch_builder(&f);
2639 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2640 unsafe {
2641 b.launch(cfg)?;
2642 }
2643 Ok(())
2644 }
2645
2646 pub fn prefetch_l2(
2648 &self,
2649 p: &CudaSlice<u8>,
2650 n: usize,
2651 ) -> Result<(), Box<dyn std::error::Error>> {
2652 let f = self.func("prefetch_l2_bytes");
2653 let lines = n.div_ceil(128);
2654 let ni = n as i64;
2655 let cfg = LaunchConfig {
2656 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2657 block_dim: (256, 1, 1),
2658 shared_mem_bytes: 0,
2659 };
2660 let __s_b = self.gpu.stream();
2661 let mut b = __s_b.launch_builder(&f);
2662 b.arg(p).arg(&ni);
2663 unsafe {
2664 b.launch(cfg)?;
2665 }
2666 Ok(())
2667 }
2668
2669 pub fn router_gemv(
2672 &self,
2673 w: &CudaSlice<f32>,
2674 x: &CudaSlice<f32>,
2675 n_embd: usize,
2676 n_experts: usize,
2677 t: usize,
2678 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2679 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2685 Ok("0") => false,
2686 Ok(_) => true,
2687 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2688 };
2689 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2699 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2700 }
2701
2702 pub fn router_gemv_form(
2705 &self,
2706 w: &CudaSlice<f32>,
2707 x: &CudaSlice<f32>,
2708 n_embd: usize,
2709 n_experts: usize,
2710 t: usize,
2711 w8: bool,
2712 batch: bool,
2713 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2714 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2715 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2716 let f = if batch {
2717 self.func("router_gemv_f32_w8_batch")
2718 } else if w8 {
2719 self.func("router_gemv_f32_w8")
2720 } else {
2721 self.func("router_gemv_f32")
2722 };
2723 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2724 let cfg = if batch {
2725 LaunchConfig {
2726 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2727 block_dim: (32, 8, 1),
2728 shared_mem_bytes: 0,
2729 }
2730 } else {
2731 LaunchConfig {
2732 grid_dim: (n_experts as u32, t as u32, 1),
2733 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2734 shared_mem_bytes: 0,
2735 }
2736 };
2737 let __s_b = self.gpu.stream();
2738 let mut b = __s_b.launch_builder(&f);
2739 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2740 unsafe {
2741 b.launch(cfg)?;
2742 }
2743 Ok(y)
2744 }
2745
2746 pub fn router_gemv_into(
2749 &self,
2750 w: &CudaSlice<f32>,
2751 x: &CudaSlice<f32>,
2752 y: &mut CudaSlice<f32>,
2753 n_embd: usize,
2754 n_experts: usize,
2755 t: usize,
2756 ) -> Result<(), Box<dyn std::error::Error>> {
2757 if y.len() < t * n_experts {
2758 return Err("router_gemv_into output too small".into());
2759 }
2760 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2761 Ok("0") => false,
2762 Ok(_) => true,
2763 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2764 };
2765 let f = if w8 {
2766 self.func("router_gemv_f32_w8")
2767 } else {
2768 self.func("router_gemv_f32")
2769 };
2770 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2771 let cfg = LaunchConfig {
2772 grid_dim: (n_experts as u32, t as u32, 1),
2773 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2774 shared_mem_bytes: 0,
2775 };
2776 let __s_b = self.gpu.stream();
2777 let mut b = __s_b.launch_builder(&f);
2778 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2779 unsafe {
2780 b.launch(cfg)?;
2781 }
2782 Ok(())
2783 }
2784
2785 pub fn rows_permute(
2787 &self,
2788 src: &CudaSlice<f32>,
2789 idx: &CudaSlice<i32>,
2790 nrows: usize,
2791 ncols: usize,
2792 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2793 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2794 let f = self.func("rows_permute_f32");
2795 let (nc, nr) = (ncols as i32, nrows as i32);
2796 let cfg = LaunchConfig {
2797 grid_dim: (nrows as u32, 1, 1),
2798 block_dim: (256, 1, 1),
2799 shared_mem_bytes: 0,
2800 };
2801 let __s_b = self.gpu.stream();
2802 let mut b = __s_b.launch_builder(&f);
2803 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2804 unsafe {
2805 b.launch(cfg)?;
2806 }
2807 Ok(dst)
2808 }
2809
2810 pub fn sigmoid_dot_rows(
2815 &self,
2816 x: &CudaSlice<f32>,
2817 w: &CudaSlice<f32>,
2818 n_embd: usize,
2819 t: usize,
2820 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2821 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2824 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2825 let gs = self.linear(x, w, t, n_embd, 1)?;
2826 let mut g = self.uninit(t)?;
2827 self.sigmoid(&gs, &mut g, t)?;
2828 return Ok(g);
2829 }
2830 let mut g = self.alloc_uninit::<f32>(t)?;
2836 let f = self.func("sigmoid_dot_rows_f32");
2837 let (ne, ti) = (n_embd as i32, t as i32);
2838 let cfg = LaunchConfig {
2839 grid_dim: (t as u32, 1, 1),
2840 block_dim: (32, 8, 1),
2841 shared_mem_bytes: 0,
2842 };
2843 let __s_b = self.gpu.stream();
2844 let mut b = __s_b.launch_builder(&f);
2845 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2846 unsafe {
2847 b.launch(cfg)?;
2848 }
2849 Ok(g)
2850 }
2851
2852 pub fn sigmoid_dot_rows_into(
2854 &self,
2855 x: &CudaSlice<f32>,
2856 w: &CudaSlice<f32>,
2857 g: &mut CudaSlice<f32>,
2858 n_embd: usize,
2859 t: usize,
2860 ) -> Result<(), Box<dyn std::error::Error>> {
2861 if g.len() < t {
2862 return Err("sigmoid_dot_rows_into output too small".into());
2863 }
2864 let f = self.func("sigmoid_dot_rows_f32");
2865 let (ne, ti) = (n_embd as i32, t as i32);
2866 let cfg = LaunchConfig {
2867 grid_dim: (t as u32, 1, 1),
2868 block_dim: (32, 8, 1),
2869 shared_mem_bytes: 0,
2870 };
2871 let __s_b = self.gpu.stream();
2872 let mut b = __s_b.launch_builder(&f);
2873 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2874 unsafe {
2875 b.launch(cfg)?;
2876 }
2877 Ok(())
2878 }
2879
2880 pub fn spec_rollback_stream(
2882 &self,
2883 len_ptrs: &CudaSlice<u64>,
2884 pos_start: &CudaSlice<i32>,
2885 acc: &CudaSlice<u32>,
2886 base: usize,
2887 n_rows: usize,
2888 ) -> Result<(), Box<dyn std::error::Error>> {
2889 let f = self.func("spec_rollback_stream");
2890 let (b, nr) = (base as i32, n_rows as i32);
2891 let cfg = LaunchConfig {
2892 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2893 block_dim: (64, 1, 1),
2894 shared_mem_bytes: 0,
2895 };
2896 let __s_bl = self.gpu.stream();
2897 let mut bl = __s_bl.launch_builder(&f);
2898 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2899 unsafe {
2900 bl.launch(cfg)?;
2901 }
2902 Ok(())
2903 }
2904
2905 pub fn plain_tok_ring(
2907 &self,
2908 vam: &CudaSlice<u32>,
2909 pos_start: &CudaSlice<i32>,
2910 base: usize,
2911 ring: &mut CudaSlice<u32>,
2912 ) -> Result<(), Box<dyn std::error::Error>> {
2913 let f = self.func("plain_tok_ring");
2914 let (b, cap) = (base as i32, ring.len() as i32);
2915 let cfg = LaunchConfig {
2916 grid_dim: (1, 1, 1),
2917 block_dim: (32, 1, 1),
2918 shared_mem_bytes: 0,
2919 };
2920 let __s_bl = self.gpu.stream();
2921 let mut bl = __s_bl.launch_builder(&f);
2922 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2923 unsafe {
2924 bl.launch(cfg)?;
2925 }
2926 Ok(())
2927 }
2928
2929 pub fn spec_ring_commit(
2931 &self,
2932 vtok: &CudaSlice<u32>,
2933 acc: &CudaSlice<u32>,
2934 brk: &CudaSlice<u32>,
2935 ring: &mut CudaSlice<u32>,
2936 pend: &mut CudaSlice<u32>,
2937 ) -> Result<(), Box<dyn std::error::Error>> {
2938 let f = self.func("spec_ring_commit");
2939 let cfg = LaunchConfig {
2940 grid_dim: (1, 1, 1),
2941 block_dim: (32, 1, 1),
2942 shared_mem_bytes: 0,
2943 };
2944 let __s_b = self.gpu.stream();
2945 let mut b = __s_b.launch_builder(&f);
2946 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2947 unsafe {
2948 b.launch(cfg)?;
2949 }
2950 Ok(())
2951 }
2952 pub fn i32_copy_add(
2953 &self,
2954 src: &CudaSlice<i32>,
2955 dst: &mut CudaSlice<i32>,
2956 delta: i32,
2957 ) -> Result<(), Box<dyn std::error::Error>> {
2958 let f = self.func("i32_copy_add");
2959 let cfg = LaunchConfig {
2960 grid_dim: (1, 1, 1),
2961 block_dim: (32, 1, 1),
2962 shared_mem_bytes: 0,
2963 };
2964 let __s_b = self.gpu.stream();
2965 let mut b = __s_b.launch_builder(&f);
2966 b.arg(src).arg(dst).arg(&delta);
2967 unsafe {
2968 b.launch(cfg)?;
2969 }
2970 Ok(())
2971 }
2972 pub fn u32_copy(
2973 &self,
2974 src: &CudaSlice<u32>,
2975 dst: &mut CudaSlice<u32>,
2976 ) -> Result<(), Box<dyn std::error::Error>> {
2977 let f = self.func("u32_copy");
2978 let cfg = LaunchConfig {
2979 grid_dim: (1, 1, 1),
2980 block_dim: (32, 1, 1),
2981 shared_mem_bytes: 0,
2982 };
2983 let __s_b = self.gpu.stream();
2984 let mut b = __s_b.launch_builder(&f);
2985 b.arg(src).arg(dst);
2986 unsafe {
2987 b.launch(cfg)?;
2988 }
2989 Ok(())
2990 }
2991
2992 pub fn spec_adapt_k(
2996 &self,
2997 acc: &CudaSlice<u32>,
2998 brk: &mut CudaSlice<u32>,
2999 floor: usize,
3000 cap: usize,
3001 ) -> Result<(), Box<dyn std::error::Error>> {
3002 let f = self.func("spec_adapt_k");
3003 let (fl, cp) = (floor as i32, cap as i32);
3004 let cfg = LaunchConfig {
3005 grid_dim: (1, 1, 1),
3006 block_dim: (32, 1, 1),
3007 shared_mem_bytes: 0,
3008 };
3009 let __s_b = self.gpu.stream();
3010 let mut b = __s_b.launch_builder(&f);
3011 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3012 unsafe {
3013 b.launch(cfg)?;
3014 }
3015 Ok(())
3016 }
3017
3018 pub fn spec_accept_greedy_dc(
3020 &self,
3021 preds: &CudaSlice<u32>,
3022 vtok: &CudaSlice<u32>,
3023 last_pred: &CudaSlice<u32>,
3024 brk: &CudaSlice<u32>,
3025 out: &mut CudaSlice<u32>,
3026 ) -> Result<(), Box<dyn std::error::Error>> {
3027 let f = self.func("spec_accept_greedy_dc");
3028 let cfg = LaunchConfig {
3029 grid_dim: (1, 1, 1),
3030 block_dim: (32, 1, 1),
3031 shared_mem_bytes: 0,
3032 };
3033 let __s_b = self.gpu.stream();
3034 let mut b = __s_b.launch_builder(&f);
3035 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3036 unsafe {
3037 b.launch(cfg)?;
3038 }
3039 Ok(())
3040 }
3041
3042 pub fn pos_iota(
3044 &self,
3045 pos0: &CudaSlice<i32>,
3046 out: &mut CudaSlice<i32>,
3047 t: usize,
3048 ) -> Result<(), Box<dyn std::error::Error>> {
3049 let f = self.func("pos_iota_i32");
3050 let ti = t as i32;
3051 let cfg = LaunchConfig {
3052 grid_dim: (1, 1, 1),
3053 block_dim: (t.max(1) as u32, 1, 1),
3054 shared_mem_bytes: 0,
3055 };
3056 let __s_b = self.gpu.stream();
3057 let mut b = __s_b.launch_builder(&f);
3058 b.arg(pos0).arg(out).arg(&ti);
3059 unsafe {
3060 b.launch(cfg)?;
3061 }
3062 Ok(())
3063 }
3064 #[allow(clippy::too_many_arguments)]
3065 pub fn append_kv_quantized_rows_dc(
3066 &self,
3067 k_rows: &CudaSlice<f32>,
3068 v_rows: &CudaSlice<f32>,
3069 kc: &mut CudaSlice<u8>,
3070 vc: &mut CudaSlice<u8>,
3071 t0_dev: &CudaSlice<i32>,
3072 t: usize,
3073 kv_dim_k: usize,
3074 kv_dim_v: usize,
3075 k_tok_bytes: usize,
3076 v_tok_bytes: usize,
3077 g: bool,
3078 ) -> Result<(), Box<dyn std::error::Error>> {
3079 let f = if g {
3080 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3081 } else {
3082 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3083 };
3084 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3085 let cfg = LaunchConfig {
3086 grid_dim: (nblk, t as u32, 1),
3087 block_dim: (32, 1, 1),
3088 shared_mem_bytes: 0,
3089 };
3090 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3091 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3092 let __s_b = self.gpu.stream();
3093 let mut b = __s_b.launch_builder(&f);
3094 b.arg(k_rows)
3095 .arg(v_rows)
3096 .arg(kc)
3097 .arg(vc)
3098 .arg(t0_dev)
3099 .arg(&kdk)
3100 .arg(&kdv)
3101 .arg(&ktb)
3102 .arg(&vtb);
3103 unsafe {
3104 b.launch(cfg)?;
3105 }
3106 Ok(())
3107 }
3108
3109 #[allow(clippy::too_many_arguments)]
3112 pub fn append_kv_quantized_row_dc_inc(
3113 &self,
3114 k_row: &CudaSlice<f32>,
3115 v_row: &CudaSlice<f32>,
3116 kc: &mut CudaSlice<u8>,
3117 vc: &mut CudaSlice<u8>,
3118 t0_dev: &mut CudaSlice<i32>,
3119 kv_dim_k: usize,
3120 kv_dim_v: usize,
3121 k_tok_bytes: usize,
3122 v_tok_bytes: usize,
3123 g: bool,
3124 ) -> Result<(), Box<dyn std::error::Error>> {
3125 let f = if g {
3126 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3127 } else {
3128 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3129 };
3130 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3131 let cfg = LaunchConfig {
3132 grid_dim: (1, 1, 1),
3133 block_dim: (nthreads, 1, 1),
3134 shared_mem_bytes: 0,
3135 };
3136 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3137 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3138 let __s_b = self.gpu.stream();
3139 let mut b = __s_b.launch_builder(&f);
3140 b.arg(k_row)
3141 .arg(v_row)
3142 .arg(kc)
3143 .arg(vc)
3144 .arg(t0_dev)
3145 .arg(&kdk)
3146 .arg(&kdv)
3147 .arg(&ktb)
3148 .arg(&vtb);
3149 unsafe {
3150 b.launch(cfg)?;
3151 }
3152 Ok(())
3153 }
3154
3155 pub fn pack_tok_p(
3157 &self,
3158 tok: &CudaSlice<u32>,
3159 p: &CudaSlice<f32>,
3160 out: &mut CudaSlice<u32>,
3161 slot: usize,
3162 ) -> Result<(), Box<dyn std::error::Error>> {
3163 let f = self.func("pack_tok_p");
3164 let sl = slot as i32;
3165 let cfg = LaunchConfig {
3166 grid_dim: (1, 1, 1),
3167 block_dim: (32, 1, 1),
3168 shared_mem_bytes: 0,
3169 };
3170 let __s_b = self.gpu.stream();
3171 let mut b = __s_b.launch_builder(&f);
3172 b.arg(tok).arg(p).arg(out).arg(&sl);
3173 unsafe {
3174 b.launch(cfg)?;
3175 }
3176 Ok(())
3177 }
3178 pub fn tok_map_u32(
3179 &self,
3180 tok: &mut CudaSlice<u32>,
3181 map: &CudaSlice<u32>,
3182 ) -> Result<(), Box<dyn std::error::Error>> {
3183 let f = self.func("tok_map_u32");
3184 let cfg = LaunchConfig {
3185 grid_dim: (1, 1, 1),
3186 block_dim: (32, 1, 1),
3187 shared_mem_bytes: 0,
3188 };
3189 let __s_b = self.gpu.stream();
3190 let mut b = __s_b.launch_builder(&f);
3191 b.arg(tok).arg(map);
3192 unsafe {
3193 b.launch(cfg)?;
3194 }
3195 Ok(())
3196 }
3197
3198 #[allow(clippy::too_many_arguments)]
3200 pub fn spec_assemble_verify(
3201 &self,
3202 tokp: &CudaSlice<u32>,
3203 pend: &CudaSlice<u32>,
3204 d2t: Option<&CudaSlice<u32>>,
3205 vtok: &mut CudaSlice<u32>,
3206 brk: &mut CudaSlice<u32>,
3207 p_min: f32,
3208 k: usize,
3209 pmin0: bool,
3210 ) -> Result<(), Box<dyn std::error::Error>> {
3211 let f = self.func("spec_assemble_verify");
3212 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3213 let cfg = LaunchConfig {
3214 grid_dim: (1, 1, 1),
3215 block_dim: (32, 1, 1),
3216 shared_mem_bytes: 0,
3217 };
3218 let __s_b = self.gpu.stream();
3219 let mut b = __s_b.launch_builder(&f);
3220 match d2t {
3221 Some(m) => {
3222 b.arg(tokp)
3223 .arg(pend)
3224 .arg(m)
3225 .arg(vtok)
3226 .arg(brk)
3227 .arg(&p_min)
3228 .arg(&ki)
3229 .arg(&pm);
3230 unsafe {
3231 b.launch(cfg)?;
3232 }
3233 }
3234 None => {
3235 let null: u64 = 0;
3236 b.arg(tokp)
3237 .arg(pend)
3238 .arg(&null)
3239 .arg(vtok)
3240 .arg(brk)
3241 .arg(&p_min)
3242 .arg(&ki)
3243 .arg(&pm);
3244 unsafe {
3245 b.launch(cfg)?;
3246 }
3247 }
3248 }
3249 Ok(())
3250 }
3251
3252 #[allow(clippy::too_many_arguments)]
3254 pub fn ssm_conv_ring_rebuild_dc(
3255 &self,
3256 qkv_tm: &CudaSlice<f32>,
3257 ring_old: &CudaSlice<f32>,
3258 conv_state: &mut CudaSlice<f32>,
3259 conv_dim: usize,
3260 acc: &CudaSlice<u32>,
3261 base: usize,
3262 t_v: usize,
3263 d_conv: usize,
3264 ) -> Result<(), Box<dyn std::error::Error>> {
3265 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3266 let n = conv_dim * (d_conv - 1);
3267 let cfg = LaunchConfig::for_num_elems(n as u32);
3268 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3269 let __s_b = self.gpu.stream();
3270 let mut b = __s_b.launch_builder(&f);
3271 b.arg(qkv_tm)
3272 .arg(ring_old)
3273 .arg(conv_state)
3274 .arg(&cd)
3275 .arg(acc)
3276 .arg(&b0)
3277 .arg(&tv)
3278 .arg(&dc);
3279 unsafe {
3280 b.launch(cfg)?;
3281 }
3282 Ok(())
3283 }
3284 #[allow(clippy::too_many_arguments)]
3285 pub fn gdn_scan_s128_dc(
3286 &self,
3287 q: &CudaSlice<f32>,
3288 k: &CudaSlice<f32>,
3289 v: &CudaSlice<f32>,
3290 g: &CudaSlice<f32>,
3291 beta: &CudaSlice<f32>,
3292 state_in: &CudaSlice<f32>,
3293 state_out: &mut CudaSlice<f32>,
3294 o: &mut CudaSlice<f32>,
3295 n_head: usize,
3296 acc: &CudaSlice<u32>,
3297 base: usize,
3298 t_v: usize,
3299 scale: f32,
3300 ) -> Result<(), Box<dyn std::error::Error>> {
3301 let f = self.func("gdn_scan_s128_dc");
3302 const S_V: u32 = 128;
3303 const WARP: u32 = 32;
3304 const COLS_PER_BLOCK: u32 = 4;
3305 let cfg = LaunchConfig {
3306 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3307 block_dim: (WARP, COLS_PER_BLOCK, 1),
3308 shared_mem_bytes: 0,
3309 };
3310 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3311 let __s_b = self.gpu.stream();
3312 let mut b = __s_b.launch_builder(&f);
3313 b.arg(q)
3314 .arg(k)
3315 .arg(v)
3316 .arg(g)
3317 .arg(beta)
3318 .arg(state_in)
3319 .arg(state_out)
3320 .arg(o)
3321 .arg(&h)
3322 .arg(acc)
3323 .arg(&b0)
3324 .arg(&tv)
3325 .arg(&scale);
3326 unsafe {
3327 b.launch(cfg)?;
3328 }
3329 Ok(())
3330 }
3331
3332 pub fn spec_rollback_kv(
3334 &self,
3335 len_ptrs: &CudaSlice<u64>,
3336 saved: &CudaSlice<i32>,
3337 acc: &CudaSlice<u32>,
3338 base: usize,
3339 n_layer: usize,
3340 ) -> Result<(), Box<dyn std::error::Error>> {
3341 let f = self.func("spec_rollback_kv");
3342 let (b, nl) = (base as i32, n_layer as i32);
3343 let cfg = LaunchConfig {
3344 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3345 block_dim: (64, 1, 1),
3346 shared_mem_bytes: 0,
3347 };
3348 let __s_bl = self.gpu.stream();
3349 let mut bl = __s_bl.launch_builder(&f);
3350 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3351 unsafe {
3352 bl.launch(cfg)?;
3353 }
3354 Ok(())
3355 }
3356
3357 pub fn spec_fork_valid(
3359 &self,
3360 acc: &CudaSlice<u32>,
3361 optimistic_pending: u32,
3362 valid: &mut CudaSlice<u32>,
3363 ) -> Result<(), Box<dyn std::error::Error>> {
3364 let f = self.func("spec_fork_valid");
3365 let cfg = LaunchConfig {
3366 grid_dim: (1, 1, 1),
3367 block_dim: (1, 1, 1),
3368 shared_mem_bytes: 0,
3369 };
3370 let __s_bl = self.gpu.stream();
3371 let mut bl = __s_bl.launch_builder(&f);
3372 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3373 unsafe {
3374 bl.launch(cfg)?;
3375 }
3376 Ok(())
3377 }
3378
3379 pub fn spec_fork_reconcile_kv(
3381 &self,
3382 len_ptrs: &CudaSlice<u64>,
3383 saved: &CudaSlice<i32>,
3384 acc: &CudaSlice<u32>,
3385 valid: &CudaSlice<u32>,
3386 base: usize,
3387 n_layer: usize,
3388 ) -> Result<(), Box<dyn std::error::Error>> {
3389 let f = self.func("spec_fork_reconcile_kv");
3390 let (b, nl) = (base as i32, n_layer as i32);
3391 let cfg = LaunchConfig {
3392 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3393 block_dim: (64, 1, 1),
3394 shared_mem_bytes: 0,
3395 };
3396 let __s_bl = self.gpu.stream();
3397 let mut bl = __s_bl.launch_builder(&f);
3398 bl.arg(len_ptrs)
3399 .arg(saved)
3400 .arg(acc)
3401 .arg(valid)
3402 .arg(&b)
3403 .arg(&nl);
3404 unsafe {
3405 bl.launch(cfg)?;
3406 }
3407 Ok(())
3408 }
3409
3410 pub fn spec_fork_restore_f32(
3412 &self,
3413 snapshot: &CudaSlice<f32>,
3414 state: &mut CudaSlice<f32>,
3415 valid: &CudaSlice<u32>,
3416 ) -> Result<(), Box<dyn std::error::Error>> {
3417 assert_eq!(
3418 snapshot.len(),
3419 state.len(),
3420 "fork recurrent snapshot shape mismatch"
3421 );
3422 let f = self.func("spec_fork_restore_f32");
3423 let n = state.len() as i32;
3424 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3425 let cfg = LaunchConfig {
3426 grid_dim: (blocks, 1, 1),
3427 block_dim: (256, 1, 1),
3428 shared_mem_bytes: 0,
3429 };
3430 let __s_bl = self.gpu.stream();
3431 let mut bl = __s_bl.launch_builder(&f);
3432 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3433 unsafe {
3434 bl.launch(cfg)?;
3435 }
3436 Ok(())
3437 }
3438
3439 pub fn spec_seed_gather(
3442 &self,
3443 vx: &CudaSlice<f32>,
3444 fill_prev: &CudaSlice<f32>,
3445 acc: &CudaSlice<u32>,
3446 h_seed: &mut CudaSlice<f32>,
3447 base: usize,
3448 n_embd: usize,
3449 ) -> Result<(), Box<dyn std::error::Error>> {
3450 let f = self.func("spec_seed_gather");
3451 let (b, ne) = (base as i32, n_embd as i32);
3452 let cfg = LaunchConfig {
3453 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3454 block_dim: (256, 1, 1),
3455 shared_mem_bytes: 0,
3456 };
3457 let __s_bl = self.gpu.stream();
3458 let mut bl = __s_bl.launch_builder(&f);
3459 bl.arg(vx)
3460 .arg(fill_prev)
3461 .arg(acc)
3462 .arg(h_seed)
3463 .arg(&b)
3464 .arg(&ne);
3465 unsafe {
3466 bl.launch(cfg)?;
3467 }
3468 Ok(())
3469 }
3470
3471 pub fn spec_accept_greedy(
3473 &self,
3474 preds: &CudaSlice<u32>,
3475 draft: &CudaSlice<u32>,
3476 last_pred: u32,
3477 base: usize,
3478 k_round: usize,
3479 out: &mut CudaSlice<u32>,
3480 ) -> Result<(), Box<dyn std::error::Error>> {
3481 let f = self.func("spec_accept_greedy");
3482 let (b, k) = (base as i32, k_round as i32);
3483 let cfg = LaunchConfig {
3484 grid_dim: (1, 1, 1),
3485 block_dim: (32, 1, 1),
3486 shared_mem_bytes: 0,
3487 };
3488 let __s_bl = self.gpu.stream();
3489 let mut bl = __s_bl.launch_builder(&f);
3490 bl.arg(preds)
3491 .arg(draft)
3492 .arg(&last_pred)
3493 .arg(&b)
3494 .arg(&k)
3495 .arg(out);
3496 unsafe {
3497 bl.launch(cfg)?;
3498 }
3499 Ok(())
3500 }
3501
3502 pub fn gumbel_perturb(
3509 &self,
3510 x: &CudaSlice<f32>,
3511 y: &mut CudaSlice<f32>,
3512 n: usize,
3513 seed: u64,
3514 stream_pos: u32,
3515 temp: f32,
3516 ) -> Result<(), Box<dyn std::error::Error>> {
3517 let f = self.func("gumbel_perturb_f32");
3518 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3519 let cfg = LaunchConfig {
3520 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3521 block_dim: (256, 1, 1),
3522 shared_mem_bytes: 0,
3523 };
3524 let __s_b = self.gpu.stream();
3525 let mut b = __s_b.launch_builder(&f);
3526 b.arg(x)
3527 .arg(&mut *y)
3528 .arg(&ni)
3529 .arg(&slo)
3530 .arg(&shi)
3531 .arg(&stream_pos)
3532 .arg(&temp);
3533 unsafe {
3534 b.launch(cfg)?;
3535 }
3536 Ok(())
3537 }
3538
3539 pub fn mask_logits_col(
3547 &self,
3548 logits: &mut CudaSlice<f32>,
3549 mask: &CudaSlice<u32>,
3550 col: usize,
3551 n: usize,
3552 mask_words: usize,
3553 ) -> Result<(), Box<dyn std::error::Error>> {
3554 let f = self.func("mask_logits_f32");
3555 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3556 let cfg = LaunchConfig {
3557 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3558 block_dim: (256, 1, 1),
3559 shared_mem_bytes: 0,
3560 };
3561 let __s_b = self.gpu.stream();
3562 let mut b = __s_b.launch_builder(&f);
3563 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3564 unsafe {
3565 b.launch(cfg)?;
3566 }
3567 Ok(())
3568 }
3569
3570 pub fn gumbel_perturb_col(
3577 &self,
3578 x: &CudaSlice<f32>,
3579 col: usize,
3580 y: &mut CudaSlice<f32>,
3581 n: usize,
3582 seed: u64,
3583 stream_pos: u32,
3584 temp: f32,
3585 ) -> Result<(), Box<dyn std::error::Error>> {
3586 let f = self.func("gumbel_perturb_f32");
3587 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3588 let col_view = x.slice(col * n..(col + 1) * n);
3589 let cfg = LaunchConfig {
3590 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3591 block_dim: (256, 1, 1),
3592 shared_mem_bytes: 0,
3593 };
3594 let __s_b = self.gpu.stream();
3595 let mut b = __s_b.launch_builder(&f);
3596 b.arg(&col_view)
3597 .arg(&mut *y)
3598 .arg(&ni)
3599 .arg(&slo)
3600 .arg(&shi)
3601 .arg(&stream_pos)
3602 .arg(&temp);
3603 unsafe {
3604 b.launch(cfg)?;
3605 }
3606 Ok(())
3607 }
3608
3609 #[allow(clippy::too_many_arguments)]
3615 pub fn gumbel_perturb_filtered_col(
3616 &self,
3617 x: &CudaSlice<f32>,
3618 col: usize,
3619 y: &mut CudaSlice<f32>,
3620 n: usize,
3621 seed: u64,
3622 stream_pos: u32,
3623 temp: f32,
3624 stat_max: &CudaSlice<f32>,
3625 stat_th: &CudaSlice<f32>,
3626 stat_idx: usize,
3627 ) -> Result<(), Box<dyn std::error::Error>> {
3628 let f = self.func("gumbel_perturb_filtered_col_f32");
3629 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3630 let (ci, si) = (col as i32, stat_idx as i32);
3631 let cfg = LaunchConfig {
3632 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3633 block_dim: (256, 1, 1),
3634 shared_mem_bytes: 0,
3635 };
3636 let __s_b = self.gpu.stream();
3637 let mut b = __s_b.launch_builder(&f);
3638 b.arg(x)
3639 .arg(&ci)
3640 .arg(&mut *y)
3641 .arg(&ni)
3642 .arg(&slo)
3643 .arg(&shi)
3644 .arg(&stream_pos)
3645 .arg(&temp)
3646 .arg(stat_max)
3647 .arg(stat_th)
3648 .arg(&si);
3649 unsafe {
3650 b.launch(cfg)?;
3651 }
3652 Ok(())
3653 }
3654
3655 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3660 let f = self.func("memra_sctr_inc");
3661 let cfg = LaunchConfig {
3662 grid_dim: (1, 1, 1),
3663 block_dim: (1, 1, 1),
3664 shared_mem_bytes: 0,
3665 };
3666 let __s_b = self.gpu.stream();
3667 let mut b = __s_b.launch_builder(&f);
3668 b.arg(&mut *ctr);
3669 unsafe {
3670 b.launch(cfg)?;
3671 }
3672 Ok(())
3673 }
3674
3675 pub fn gumbel_perturb_ctr(
3680 &self,
3681 x: &CudaSlice<f32>,
3682 y: &mut CudaSlice<f32>,
3683 n: usize,
3684 seed: u64,
3685 ctr: &CudaSlice<u32>,
3686 temp: f32,
3687 ) -> Result<(), Box<dyn std::error::Error>> {
3688 let f = self.func("gumbel_perturb_ctr_f32");
3689 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3690 let cfg = LaunchConfig {
3691 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3692 block_dim: (256, 1, 1),
3693 shared_mem_bytes: 0,
3694 };
3695 let __s_b = self.gpu.stream();
3696 let mut b = __s_b.launch_builder(&f);
3697 b.arg(x)
3698 .arg(&mut *y)
3699 .arg(&ni)
3700 .arg(&slo)
3701 .arg(&shi)
3702 .arg(ctr)
3703 .arg(&temp);
3704 unsafe {
3705 b.launch(cfg)?;
3706 }
3707 Ok(())
3708 }
3709
3710 pub fn softmax_gather(
3714 &self,
3715 x: &CudaSlice<f32>,
3716 row_stride: usize,
3717 ids: &CudaSlice<u32>,
3718 rows: &CudaSlice<i32>,
3719 out: &mut CudaSlice<f32>,
3720 n: usize,
3721 npair: usize,
3722 temp: f32,
3723 ) -> Result<(), Box<dyn std::error::Error>> {
3724 let f = self.func("softmax_gather_f32");
3725 let (ni, rs) = (n as i32, row_stride as i64);
3726 let np = npair as i32;
3727 let cfg = LaunchConfig {
3728 grid_dim: (npair as u32, 1, 1),
3729 block_dim: (256, 1, 1),
3730 shared_mem_bytes: 0,
3731 };
3732 let __s_b = self.gpu.stream();
3733 let mut b = __s_b.launch_builder(&f);
3734 b.arg(x)
3735 .arg(&rs)
3736 .arg(ids)
3737 .arg(rows)
3738 .arg(&mut *out)
3739 .arg(&ni)
3740 .arg(&np)
3741 .arg(&temp);
3742 unsafe {
3743 b.launch(cfg)?;
3744 }
3745 Ok(())
3746 }
3747
3748 pub fn residual_sample(
3752 &self,
3753 p: &CudaSlice<f32>,
3754 q: Option<&CudaSlice<f32>>,
3755 n: usize,
3756 temp: f32,
3757 seed: u64,
3758 stream_pos: u32,
3759 out_tok: &mut CudaSlice<u32>,
3760 ) -> Result<(), Box<dyn std::error::Error>> {
3761 let f = self.func("residual_sample_f32");
3762 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3763 let nth = 1024u32;
3764 let cfg = LaunchConfig {
3765 grid_dim: (1, 1, 1),
3766 block_dim: (nth, 1, 1),
3767 shared_mem_bytes: 0,
3768 };
3769 let has_q: i32 = q.is_some() as i32;
3770 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3772 let mut b = __s_b.launch_builder(&f);
3773 b.arg(p)
3774 .arg(qbuf)
3775 .arg(&has_q)
3776 .arg(&ni)
3777 .arg(&temp)
3778 .arg(&slo)
3779 .arg(&shi)
3780 .arg(&stream_pos)
3781 .arg(&mut *out_tok);
3782 unsafe {
3783 b.launch(cfg)?;
3784 }
3785 Ok(())
3786 }
3787
3788 pub fn with_moe_cache<R>(
3793 &self,
3794 max_block_bytes: usize,
3795 f: impl FnOnce(
3796 &mut crate::moe_cache::MoeSlotCache,
3797 &Engine,
3798 ) -> Result<R, Box<dyn std::error::Error>>,
3799 ) -> Result<R, Box<dyn std::error::Error>> {
3800 let mut guard = self.moe_cache.lock().unwrap();
3801 if guard.is_none() {
3802 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3803 }
3804 let cache = guard.as_mut().unwrap();
3805 f(cache, self)
3806 }
3807
3808 pub fn freeze_moe_cache(&self) {
3811 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3812 cache.freeze();
3813 }
3814 }
3815
3816 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3819 self.moe_cache
3820 .lock()
3821 .unwrap()
3822 .as_ref()
3823 .map(crate::moe_cache::MoeSlotCache::export_residency)
3824 }
3825
3826 pub(crate) fn moe_cache_frozen(&self) -> bool {
3827 self.moe_cache
3828 .lock()
3829 .unwrap()
3830 .as_ref()
3831 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3832 }
3833
3834 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3841 crate::cpu_experts::configured()
3842 && self.moe_cache_frozen()
3843 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3844 }
3845
3846 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3848 assert!(
3849 self.moe_cache.lock().unwrap().is_none(),
3850 "MoE cache layout configured after cache construction"
3851 );
3852 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3853 }
3854
3855 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3856 self.moe_cache_layout.lock().unwrap().clone()
3857 }
3858
3859 pub fn moe_cache_enabled() -> bool {
3861 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3862 }
3863
3864 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3867 let guard = self.moe_cache.lock().unwrap();
3868 guard
3869 .as_ref()
3870 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3871 }
3872
3873 pub fn cpu_expert_stats(
3877 &self,
3878 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3879 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3880 }
3881
3882 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3885 crate::cpu_experts::predictor_stats()
3886 }
3887
3888 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3889 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3890 }
3891
3892 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3895 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3896 }
3897
3898 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3901 let guard = self.moe_cache.lock().unwrap();
3902 guard
3903 .as_ref()
3904 .and_then(|cache| cache.pread_stats())
3905 .map(|stats| {
3906 (
3907 stats.reads,
3908 stats.bytes,
3909 stats.read_errors,
3910 stats.short_reads,
3911 stats.fallbacks,
3912 stats.buffer_waits,
3913 stats.ring_full,
3914 )
3915 })
3916 }
3917
3918 pub fn spill_config_fallbacks(&self) -> u64 {
3920 crate::spill_pread::config_fallbacks()
3921 }
3922
3923 pub fn moe_cache_reset_counters(&self) {
3925 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3926 c.reset_counters();
3927 }
3928 }
3929
3930 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3931 Ok(self.gpu.stream().clone_htod(v)?)
3932 }
3933
3934 pub fn htod_bytes_padded(
3938 &self,
3939 v: &[u8],
3940 pad: usize,
3941 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3942 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3943 {
3944 let mut view = d.slice_mut(0..v.len());
3945 self.gpu.stream().memcpy_htod(v, &mut view)?;
3946 }
3947 Ok(d)
3948 }
3949
3950 pub fn copy_into(
3952 &self,
3953 dst: &mut CudaSlice<f32>,
3954 off: usize,
3955 src: &CudaSlice<f32>,
3956 len: usize,
3957 ) -> Result<(), Box<dyn std::error::Error>> {
3958 let mut view = dst.slice_mut(off..off + len);
3959 self.gpu
3960 .stream()
3961 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3962 Ok(())
3963 }
3964
3965 pub fn copy_u8_into(
3968 &self,
3969 dst: &mut CudaSlice<u8>,
3970 off: usize,
3971 src: &CudaSlice<u8>,
3972 len: usize,
3973 ) -> Result<(), Box<dyn std::error::Error>> {
3974 let mut view = dst.slice_mut(off..off + len);
3975 self.gpu
3976 .stream()
3977 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3978 Ok(())
3979 }
3980
3981 pub fn copy_u8_range_into(
3983 &self,
3984 dst: &mut CudaSlice<u8>,
3985 dst_off: usize,
3986 src: &CudaSlice<u8>,
3987 src_off: usize,
3988 len: usize,
3989 ) -> Result<(), Box<dyn std::error::Error>> {
3990 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3991 self.gpu
3992 .stream()
3993 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3994 Ok(())
3995 }
3996
3997 pub fn prepare_kv_append(
4001 &self,
4002 kv: &mut crate::cache::KvLayer,
4003 retain_from: usize,
4004 append_rows: usize,
4005 ) -> Result<usize, Box<dyn std::error::Error>> {
4006 let Some(plan) = kv
4007 .ring
4008 .as_ref()
4009 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4010 .transpose()?
4011 else {
4012 return Ok(kv.len);
4013 };
4014 match plan {
4015 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4016 crate::cache::KvRingAppend::Rebase {
4017 src_row,
4018 keep_rows,
4019 new_base,
4020 write_row,
4021 } => {
4022 if keep_rows > 0 {
4023 let k_len = keep_rows * kv.k_tok_bytes;
4024 let v_len = keep_rows * kv.v_tok_bytes;
4025 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4026 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4027 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4028 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4029 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4030 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4031 }
4032 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4033 Ok(write_row)
4034 }
4035 }
4036 }
4037
4038 pub fn htod_u8_into(
4041 &self,
4042 dst: &mut CudaSlice<u8>,
4043 off: usize,
4044 src: &[u8],
4045 ) -> Result<(), Box<dyn std::error::Error>> {
4046 let mut view = dst.slice_mut(off..off + src.len());
4047 self.gpu.stream().memcpy_htod(src, &mut view)?;
4048 Ok(())
4049 }
4050
4051 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4052 b.slice(0..len)
4053 }
4054
4055 pub fn view_u8_range<'a>(
4058 &self,
4059 b: &'a CudaSlice<u8>,
4060 start: usize,
4061 end: usize,
4062 ) -> cudarc::driver::CudaView<'a, u8> {
4063 b.slice(start..end)
4064 }
4065 pub fn view_u8<'a>(
4066 &self,
4067 b: &'a CudaSlice<u8>,
4068 len: usize,
4069 ) -> cudarc::driver::CudaView<'a, u8> {
4070 b.slice(0..len)
4071 }
4072
4073 pub fn append_kv_quantized(
4077 &self,
4078 k_row: &CudaSlice<f32>,
4079 v_row: &CudaSlice<f32>,
4080 kc: &mut CudaSlice<u8>,
4081 vc: &mut CudaSlice<u8>,
4082 t: usize,
4083 kv_dim_k: usize,
4084 kv_dim_v: usize,
4085 k_tok_bytes: usize,
4086 v_tok_bytes: usize,
4087 g: bool,
4088 ) -> Result<(), Box<dyn std::error::Error>> {
4089 let f = if g {
4090 self.func_g("append_quantize_kv_q8_0_q5_1")
4091 } else {
4092 self.func("append_quantize_kv_q8_0_q5_1")
4093 };
4094 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4095 let cfg = LaunchConfig {
4096 grid_dim: (nblk, 1, 1),
4097 block_dim: (32, 1, 1),
4098 shared_mem_bytes: 0,
4099 };
4100 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4101 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4102 let __s_b = self.gpu.stream();
4103 let mut b = __s_b.launch_builder(&f);
4104 b.arg(k_row)
4105 .arg(v_row)
4106 .arg(kc)
4107 .arg(vc)
4108 .arg(&ti)
4109 .arg(&kdk)
4110 .arg(&kdv)
4111 .arg(&ktb)
4112 .arg(&vtb);
4113 unsafe {
4114 b.launch(cfg)?;
4115 }
4116 Ok(())
4117 }
4118
4119 pub fn append_kv_quantized_dc(
4123 &self,
4124 k_row: &CudaSlice<f32>,
4125 v_row: &CudaSlice<f32>,
4126 kc: &mut CudaSlice<u8>,
4127 vc: &mut CudaSlice<u8>,
4128 t_dev: &CudaSlice<i32>,
4129 kv_dim_k: usize,
4130 kv_dim_v: usize,
4131 k_tok_bytes: usize,
4132 v_tok_bytes: usize,
4133 g: bool,
4134 ) -> Result<(), Box<dyn std::error::Error>> {
4135 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4136 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4137 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4138 if Self::pdl_on() && Self::pdl_wb_on() {
4140 use cudarc::driver::{DevicePtr, DevicePtrMut};
4141 let s = &self.gpu.stream();
4142 let (pk, _g0) = k_row.device_ptr(s);
4143 let (pv, _g1) = v_row.device_ptr(s);
4144 let (pkc, _g2) = kc.device_ptr_mut(s);
4145 let (pvc, _g3) = vc.device_ptr_mut(s);
4146 let (pt, _g4) = t_dev.device_ptr(s);
4147 let mut ps = [
4148 &pk as *const _ as *mut std::ffi::c_void,
4149 &pv as *const _ as *mut _,
4150 &pkc as *const _ as *mut _,
4151 &pvc as *const _ as *mut _,
4152 &pt as *const _ as *mut _,
4153 &kdk as *const _ as *mut _,
4154 &kdv as *const _ as *mut _,
4155 &ktb as *const _ as *mut _,
4156 &vtb as *const _ as *mut _,
4157 ];
4158 unsafe {
4159 self.launch_pdl_flash(
4160 g,
4161 "append_quantize_kv_q8_0_q5_1_dc",
4162 (nblk, 1, 1),
4163 (32, 1, 1),
4164 0,
4165 &mut ps,
4166 )?;
4167 }
4168 return Ok(());
4169 }
4170 let f = if g {
4171 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4172 } else {
4173 self.func("append_quantize_kv_q8_0_q5_1_dc")
4174 };
4175 let cfg = LaunchConfig {
4176 grid_dim: (nblk, 1, 1),
4177 block_dim: (32, 1, 1),
4178 shared_mem_bytes: 0,
4179 };
4180 let __s_b = self.gpu.stream();
4181 let mut b = __s_b.launch_builder(&f);
4182 b.arg(k_row)
4183 .arg(v_row)
4184 .arg(kc)
4185 .arg(vc)
4186 .arg(t_dev)
4187 .arg(&kdk)
4188 .arg(&kdv)
4189 .arg(&ktb)
4190 .arg(&vtb);
4191 unsafe {
4192 b.launch(cfg)?;
4193 }
4194 Ok(())
4195 }
4196
4197 #[allow(clippy::too_many_arguments)]
4204 pub fn append_kv_quantized_rows(
4205 &self,
4206 k_rows: &CudaSlice<f32>,
4207 v_rows: &CudaSlice<f32>,
4208 kc: &mut CudaSlice<u8>,
4209 vc: &mut CudaSlice<u8>,
4210 t0: usize,
4211 t: usize,
4212 kv_dim_k: usize,
4213 kv_dim_v: usize,
4214 k_tok_bytes: usize,
4215 v_tok_bytes: usize,
4216 g: bool,
4217 ) -> Result<(), Box<dyn std::error::Error>> {
4218 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4219 for i in 0..t {
4220 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4221 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4222 self.append_kv_quantized_view(
4223 &k_row,
4224 &v_row,
4225 kc,
4226 vc,
4227 t0 + i,
4228 kv_dim_k,
4229 kv_dim_v,
4230 k_tok_bytes,
4231 v_tok_bytes,
4232 g,
4233 )?;
4234 }
4235 return Ok(());
4236 }
4237 let f = if g {
4238 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4239 } else {
4240 self.func("append_quantize_kv_q8_0_q5_1_rows")
4241 };
4242 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4243 let cfg = LaunchConfig {
4244 grid_dim: (nblk, t as u32, 1),
4245 block_dim: (32, 1, 1),
4246 shared_mem_bytes: 0,
4247 };
4248 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4249 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4250 let __s_b = self.gpu.stream();
4251 let mut b = __s_b.launch_builder(&f);
4252 b.arg(k_rows)
4253 .arg(v_rows)
4254 .arg(kc)
4255 .arg(vc)
4256 .arg(&t0i)
4257 .arg(&kdk)
4258 .arg(&kdv)
4259 .arg(&ktb)
4260 .arg(&vtb);
4261 unsafe {
4262 b.launch(cfg)?;
4263 }
4264 Ok(())
4265 }
4266
4267 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4271 let f = self.func("inc_i32");
4272 let cfg = LaunchConfig {
4273 grid_dim: (1, 1, 1),
4274 block_dim: (1, 1, 1),
4275 shared_mem_bytes: 0,
4276 };
4277 let __s_b = self.gpu.stream();
4278 let mut b = __s_b.launch_builder(&f);
4279 b.arg(p);
4280 unsafe {
4281 b.launch(cfg)?;
4282 }
4283 Ok(())
4284 }
4285
4286 pub fn append_kv_quantized_view(
4289 &self,
4290 k_row: &cudarc::driver::CudaView<f32>,
4291 v_row: &cudarc::driver::CudaView<f32>,
4292 kc: &mut CudaSlice<u8>,
4293 vc: &mut CudaSlice<u8>,
4294 t: usize,
4295 kv_dim_k: usize,
4296 kv_dim_v: usize,
4297 k_tok_bytes: usize,
4298 v_tok_bytes: usize,
4299 g: bool,
4300 ) -> Result<(), Box<dyn std::error::Error>> {
4301 let stream = self.gpu.stream();
4302 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4303 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4304 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4305 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4306 let f = if g {
4307 self.func_g("append_quantize_kv_q8_0_q5_1")
4308 } else {
4309 self.func("append_quantize_kv_q8_0_q5_1")
4310 };
4311 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4312 let cfg = LaunchConfig {
4313 grid_dim: (nblk, 1, 1),
4314 block_dim: (32, 1, 1),
4315 shared_mem_bytes: 0,
4316 };
4317 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4318 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4319 let mut b = stream.launch_builder(&f);
4320 b.arg(k_row)
4321 .arg(v_row)
4322 .arg(kc)
4323 .arg(vc)
4324 .arg(&ti)
4325 .arg(&kdk)
4326 .arg(&kdv)
4327 .arg(&ktb)
4328 .arg(&vtb);
4329 unsafe {
4330 b.launch(cfg)?;
4331 }
4332 Ok(())
4333 }
4334
4335 pub fn copy_view_into(
4338 &self,
4339 dst: &mut CudaSlice<f32>,
4340 off: usize,
4341 src: &cudarc::driver::CudaView<f32>,
4342 len: usize,
4343 ) -> Result<(), Box<dyn std::error::Error>> {
4344 let mut view = dst.slice_mut(off..off + len);
4345 self.gpu
4346 .stream()
4347 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4348 Ok(())
4349 }
4350
4351 pub fn clone_dtod(
4355 &self,
4356 src: &CudaSlice<f32>,
4357 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4358 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4359 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4360 Ok(dst)
4361 }
4362
4363 pub fn dtod_copy_view(
4366 &self,
4367 src: &cudarc::driver::CudaView<f32>,
4368 dst: &mut CudaSlice<f32>,
4369 ) -> Result<(), Box<dyn std::error::Error>> {
4370 self.gpu.stream().memcpy_dtod(src, dst)?;
4371 Ok(())
4372 }
4373
4374 pub fn dtod_copy_view_i8(
4376 &self,
4377 src: &cudarc::driver::CudaView<i8>,
4378 dst: &mut CudaSlice<i8>,
4379 ) -> Result<(), Box<dyn std::error::Error>> {
4380 self.gpu.stream().memcpy_dtod(src, dst)?;
4381 Ok(())
4382 }
4383
4384 pub fn dtod_copy_into(
4386 &self,
4387 src: &CudaSlice<f32>,
4388 dst: &mut CudaSlice<f32>,
4389 offset: usize,
4390 ) -> Result<(), Box<dyn std::error::Error>> {
4391 let n = src.len();
4392 let mut dv = dst.slice_mut(offset..offset + n);
4393 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4394 Ok(())
4395 }
4396
4397 pub fn copy_batch_uniform_f32(
4403 &self,
4404 table: &CudaSlice<u64>,
4405 n: usize,
4406 words: usize,
4407 ) -> Result<(), Box<dyn std::error::Error>> {
4408 if n == 0 || words == 0 {
4409 return Ok(());
4410 }
4411 debug_assert!(
4412 table.len() >= 2 * n,
4413 "pointer table must hold n srcs + n dsts"
4414 );
4415 let f = self.func("copy_batch_uniform_f32");
4416 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4419 let (ni, wi) = (n as i32, words as i32);
4420 let cfg = LaunchConfig {
4421 grid_dim: (chunks, n as u32, 1),
4422 block_dim: (256, 1, 1),
4423 shared_mem_bytes: 0,
4424 };
4425 let __s = self.gpu.stream();
4426 let mut b = __s.launch_builder(&f);
4427 b.arg(table).arg(&ni).arg(&wi);
4428 unsafe {
4429 b.launch(cfg)?;
4430 }
4431 Ok(())
4432 }
4433
4434 pub fn htod_u64_into(
4437 &self,
4438 v: &[u64],
4439 dst: &mut CudaSlice<u64>,
4440 ) -> Result<(), Box<dyn std::error::Error>> {
4441 let mut view = dst.slice_mut(0..v.len());
4442 self.gpu.stream().memcpy_htod(v, &mut view)?;
4443 Ok(())
4444 }
4445
4446 pub fn copy_indirect_src_f32(
4451 &self,
4452 src_entry: &cudarc::driver::CudaView<u64>,
4453 dst: &mut CudaSlice<f32>,
4454 dst_off: usize,
4455 words: usize,
4456 ) -> Result<(), Box<dyn std::error::Error>> {
4457 let f = self.func("copy_indirect_src_f32");
4458 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4459 let wi = words as i32;
4460 let cfg = LaunchConfig {
4461 grid_dim: (chunks, 1, 1),
4462 block_dim: (256, 1, 1),
4463 shared_mem_bytes: 0,
4464 };
4465 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4466 let __s = self.gpu.stream();
4467 let mut b = __s.launch_builder(&f);
4468 b.arg(src_entry).arg(&mut dv).arg(&wi);
4469 unsafe {
4470 b.launch(cfg)?;
4471 }
4472 Ok(())
4473 }
4474
4475 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4477 self.alloc_uninit::<i8>(n)
4478 }
4479
4480 pub fn qmatvec(
4482 &self,
4483 w: &CudaSlice<u8>,
4484 x: &CudaSlice<f32>,
4485 m: usize,
4486 in_f: usize,
4487 out_f: usize,
4488 qtype: i32,
4489 row_bytes: usize,
4490 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4491 let f = self.func("qmatvec_f32");
4492 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4494 grid_dim: (out_f as u32, m as u32, 1),
4495 block_dim: (256, 1, 1),
4496 shared_mem_bytes: 0,
4497 };
4498 let (inf, outf, mi, qt, rb) =
4499 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4500 let __s_b = self.gpu.stream();
4501 let mut b = __s_b.launch_builder(&f);
4502 b.arg(w)
4503 .arg(x)
4504 .arg(&mut y)
4505 .arg(&inf)
4506 .arg(&outf)
4507 .arg(&mi)
4508 .arg(&qt)
4509 .arg(&rb);
4510 unsafe {
4511 b.launch(cfg)?;
4512 }
4513 Ok(y)
4514 }
4515
4516 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4518 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4519 self.keep_if_capturing(&s);
4520 Ok(s)
4521 }
4522
4523 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4527 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4528 self.keep_if_capturing(&s);
4529 Ok(s)
4530 }
4531
4532 pub fn memset_zeros_view(
4535 &self,
4536 dst: &mut cudarc::driver::CudaViewMut<f32>,
4537 ) -> Result<(), Box<dyn std::error::Error>> {
4538 self.gpu.stream().memset_zeros(dst)?;
4539 Ok(())
4540 }
4541
4542 pub fn stage_expert(
4548 &self,
4549 host_bytes: &[u8],
4550 scratch: &mut CudaSlice<u8>,
4551 off: usize,
4552 ) -> Result<(), Box<dyn std::error::Error>> {
4553 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4556 }
4557
4558 pub fn moe_router_topk(
4564 &self,
4565 logits: &CudaSlice<f32>,
4566 t: usize,
4567 n_expert: usize,
4568 n_used: usize,
4569 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4570 let f = self.func("moe_router_topk_f32");
4571 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 {
4574 grid_dim: (t as u32, 1, 1),
4575 block_dim: (n_expert as u32, 1, 1),
4576 shared_mem_bytes: 0,
4577 };
4578 let (ne, nu) = (n_expert as i32, n_used as i32);
4579 let __s_b = self.gpu.stream();
4580 let mut b = __s_b.launch_builder(&f);
4581 b.arg(logits)
4582 .arg(&mut sel_idx)
4583 .arg(&mut sel_w)
4584 .arg(&ne)
4585 .arg(&nu);
4586 unsafe {
4587 b.launch(cfg)?;
4588 }
4589 Ok((sel_idx, sel_w))
4590 }
4591
4592 pub fn moe_router_topk_scaled(
4595 &self,
4596 logits: &CudaSlice<f32>,
4597 t: usize,
4598 n_expert: usize,
4599 n_used: usize,
4600 ex_scale: &CudaSlice<f32>,
4601 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4602 let f = self.func("moe_router_topk_scaled_f32");
4607 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4608 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4609 let cfg = LaunchConfig {
4610 grid_dim: (t as u32, 1, 1),
4611 block_dim: (n_expert as u32, 1, 1),
4612 shared_mem_bytes: 0,
4613 };
4614 let (ne, nu) = (n_expert as i32, n_used as i32);
4615 let __s_b = self.gpu.stream();
4616 let mut b = __s_b.launch_builder(&f);
4617 b.arg(logits)
4618 .arg(&mut sel_idx)
4619 .arg(&mut sel_w)
4620 .arg(&ne)
4621 .arg(&nu)
4622 .arg(ex_scale);
4623 unsafe {
4624 b.launch(cfg)?;
4625 }
4626 Ok((sel_idx, sel_w))
4627 }
4628
4629 pub fn moe_router_topk_host(
4637 &self,
4638 logits: &CudaSlice<f32>,
4639 t: usize,
4640 n_expert: usize,
4641 n_used: usize,
4642 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4643 let f = self.func("moe_router_topk_f32");
4644 let n = t * n_used;
4645 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4646 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4647 let cfg = LaunchConfig {
4648 grid_dim: (t as u32, 1, 1),
4649 block_dim: (n_expert as u32, 1, 1),
4650 shared_mem_bytes: 0,
4651 };
4652 let (ne, nu) = (n_expert as i32, n_used as i32);
4653 let __s_b = self.gpu.stream();
4654 let mut b = __s_b.launch_builder(&f);
4655 b.arg(logits)
4656 .arg(&mut sel_idx)
4657 .arg(&mut sel_w)
4658 .arg(&ne)
4659 .arg(&nu);
4660 unsafe {
4661 b.launch(cfg)?;
4662 }
4663 let bytes = n * 8;
4665 let mut guard = self.router_stage.lock().unwrap();
4666 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4667 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4668 }
4669 let stage = guard.as_mut().unwrap();
4670 let (si, sw) = unsafe {
4671 (
4672 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4673 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4674 )
4675 };
4676 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()))
4680 }
4681
4682 #[allow(clippy::too_many_arguments)]
4686 pub fn moe_router_sigmoid_topk(
4687 &self,
4688 logits: &CudaSlice<f32>,
4689 t: usize,
4690 n_expert: usize,
4691 n_used: usize,
4692 active_count: usize,
4693 correction_bias: &CudaSlice<f32>,
4694 active: &CudaSlice<u8>,
4695 scaling_factor: f32,
4696 route_norm: bool,
4697 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4698 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4699 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4700 return Err(format!(
4701 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4702 )
4703 .into());
4704 }
4705 if logits.len() < t * n_expert
4706 || correction_bias.len() != n_expert
4707 || active.len() != n_expert
4708 {
4709 return Err(format!(
4710 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4711 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4712 ).into());
4713 }
4714 let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4715 self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4718 } else if crate::sig_expf_dev_on() {
4719 self.func("moe_router_sigmoid_topk_f32_dexp")
4720 } else if crate::topk_fast_on() {
4721 self.func("moe_router_sigmoid_topk_f32_fast")
4722 } else {
4723 self.func("moe_router_sigmoid_topk_f32")
4724 };
4725 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4726 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4727 let threads = n_expert.div_ceil(32) * 32;
4728 let cfg = LaunchConfig {
4729 grid_dim: (t as u32, 1, 1),
4730 block_dim: (threads as u32, 1, 1),
4731 shared_mem_bytes: 0,
4732 };
4733 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4734 let __s_b = self.gpu.stream();
4735 let mut b = __s_b.launch_builder(&f);
4736 b.arg(logits)
4737 .arg(correction_bias)
4738 .arg(active)
4739 .arg(&mut sel_idx)
4740 .arg(&mut sel_w)
4741 .arg(&ne)
4742 .arg(&nu)
4743 .arg(&scaling_factor)
4744 .arg(&rn);
4745 unsafe {
4746 b.launch(cfg)?;
4747 }
4748 Ok((sel_idx, sel_w))
4749 }
4750
4751 #[allow(clippy::too_many_arguments)]
4754 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4758 if ptr == 0 {
4759 return Err("ring_flag_raw: unarmed flag".into());
4760 }
4761 let f = self.func("memra_ring_flag");
4762 let cfg = LaunchConfig {
4763 grid_dim: (1, 1, 1),
4764 block_dim: (32, 1, 1),
4765 shared_mem_bytes: 0,
4766 };
4767 let __s_b = self.gpu.stream();
4768 let mut b = __s_b.launch_builder(&f);
4769 b.arg(&ptr).arg(&value);
4770 unsafe {
4771 b.launch(cfg)?;
4772 }
4773 Ok(())
4774 }
4775
4776 pub fn moe_sel_w_mirror(
4779 &self,
4780 sel_src: &CudaSlice<i32>,
4781 w_src: &CudaSlice<f32>,
4782 sel_dst: &mut CudaSlice<i32>,
4783 w_dst: &mut CudaSlice<f32>,
4784 n: usize,
4785 ) -> Result<(), Box<dyn std::error::Error>> {
4786 if n == 0
4787 || n > 32
4788 || sel_src.len() < n
4789 || w_src.len() < n
4790 || sel_dst.len() < n
4791 || w_dst.len() < n
4792 {
4793 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4794 }
4795 let f = self.func("moe_sel_w_mirror");
4796 let cfg = LaunchConfig {
4797 grid_dim: (1, 1, 1),
4798 block_dim: (32, 1, 1),
4799 shared_mem_bytes: 0,
4800 };
4801 let ni = n as i32;
4802 let __s_b = self.gpu.stream();
4803 let mut b = __s_b.launch_builder(&f);
4804 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
4805 unsafe {
4806 b.launch(cfg)?;
4807 }
4808 Ok(())
4809 }
4810
4811 pub fn moe_router_sigmoid_topk_into(
4812 &self,
4813 logits: &CudaSlice<f32>,
4814 t: usize,
4815 n_expert: usize,
4816 n_used: usize,
4817 active_count: usize,
4818 correction_bias: &CudaSlice<f32>,
4819 active: &CudaSlice<u8>,
4820 scaling_factor: f32,
4821 route_norm: bool,
4822 sel_idx: &mut CudaSlice<i32>,
4823 sel_w: &mut CudaSlice<f32>,
4824 ) -> Result<(), Box<dyn std::error::Error>> {
4825 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4826 if n_expert == 0
4827 || n_expert > 1024
4828 || n_used == 0
4829 || n_used > 32 || n_used > n_expert
4831 || logits.len() < t * n_expert
4832 || correction_bias.len() != n_expert
4833 || active.len() != n_expert
4834 || sel_idx.len() < t * n_used
4835 || sel_w.len() < t * n_used
4836 {
4837 return Err("sigmoid router _into geometry mismatch".into());
4838 }
4839 let f = if crate::sig_expf_dev_on() {
4840 self.func("moe_router_sigmoid_topk_f32_dexp")
4841 } else if crate::topk_fast_on() {
4842 self.func("moe_router_sigmoid_topk_f32_fast")
4843 } else {
4844 self.func("moe_router_sigmoid_topk_f32")
4845 };
4846 let threads = n_expert.div_ceil(32) * 32;
4847 let cfg = LaunchConfig {
4848 grid_dim: (t as u32, 1, 1),
4849 block_dim: (threads as u32, 1, 1),
4850 shared_mem_bytes: 0,
4851 };
4852 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4853 let __s_b = self.gpu.stream();
4854 let mut b = __s_b.launch_builder(&f);
4855 b.arg(logits)
4856 .arg(correction_bias)
4857 .arg(active)
4858 .arg(&mut *sel_idx)
4859 .arg(&mut *sel_w)
4860 .arg(&ne)
4861 .arg(&nu)
4862 .arg(&scaling_factor)
4863 .arg(&rn);
4864 unsafe {
4865 b.launch(cfg)?;
4866 }
4867 Ok(())
4868 }
4869
4870 #[allow(clippy::too_many_arguments)]
4873 pub fn moe_router_sigmoid_topk_host(
4874 &self,
4875 logits: &CudaSlice<f32>,
4876 t: usize,
4877 n_expert: usize,
4878 n_used: usize,
4879 active_count: usize,
4880 correction_bias: &CudaSlice<f32>,
4881 active: &CudaSlice<u8>,
4882 scaling_factor: f32,
4883 route_norm: bool,
4884 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4885 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4886 logits,
4887 t,
4888 n_expert,
4889 n_used,
4890 active_count,
4891 correction_bias,
4892 active,
4893 scaling_factor,
4894 route_norm,
4895 )?;
4896 let n = t * n_used;
4897 let bytes = n * 8;
4898 let mut guard = self.router_stage.lock().unwrap();
4899 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4900 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4901 }
4902 let stage = guard.as_mut().unwrap();
4903 let (si, sw) = unsafe {
4904 (
4905 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4906 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4907 )
4908 };
4909 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4910 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4911 self.gpu.stream().synchronize()?;
4912 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4913 }
4914
4915 pub fn stage_expert_async(
4919 &self,
4920 host_bytes: &[u8],
4921 scratch: &mut CudaSlice<u8>,
4922 off: usize,
4923 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4924 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4925 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4926 Ok(self.copy_stream.record_event(None)?)
4927 }
4928
4929 pub fn compute_wait(
4931 &self,
4932 ev: &cudarc::driver::CudaEvent,
4933 ) -> Result<(), Box<dyn std::error::Error>> {
4934 self.gpu.stream().wait(ev)?;
4935 Ok(())
4936 }
4937
4938 pub fn qmatvec_view(
4943 &self,
4944 w: &CudaSlice<u8>,
4945 range: std::ops::Range<usize>,
4946 x: &cudarc::driver::CudaView<f32>,
4947 m: usize,
4948 in_f: usize,
4949 out_f: usize,
4950 qtype: i32,
4951 row_bytes: usize,
4952 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4953 let f = self.func("qmatvec_f32");
4954 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4957 grid_dim: (out_f as u32, m as u32, 1),
4958 block_dim: (256, 1, 1),
4959 shared_mem_bytes: 0,
4960 };
4961 let (inf, outf, mi, qt, rb) =
4962 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4963 let __s_b = self.gpu.stream();
4964 let mut b = __s_b.launch_builder(&f);
4965 b.arg(&wv)
4966 .arg(x)
4967 .arg(&mut y)
4968 .arg(&inf)
4969 .arg(&outf)
4970 .arg(&mi)
4971 .arg(&qt)
4972 .arg(&rb);
4973 unsafe {
4974 b.launch(cfg)?;
4975 }
4976 Ok(y)
4977 }
4978
4979 #[allow(clippy::too_many_arguments)]
4986 pub fn moe_gate_up_silu8_q8(
4990 &self,
4991 gp: WPtr8,
4992 up: WPtr8,
4993 aq: &CudaSlice<i8>,
4994 ad: &CudaSlice<f32>,
4995 in_f: usize,
4996 n_ff: usize,
4997 n_used: usize,
4998 qt_g: i32,
4999 qt_u: i32,
5000 rb_g: usize,
5001 rb_u: usize,
5002 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5003 let f = self.func("moe_gate_up_silu8_q8");
5004 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5005 let cfg = LaunchConfig {
5006 grid_dim: (n_ff as u32, n_used as u32, 1),
5007 block_dim: (32, 1, 1),
5008 shared_mem_bytes: 0,
5009 };
5010 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5011 let __s_b = self.gpu.stream();
5012 let mut b = __s_b.launch_builder(&f);
5013 b.arg(&gp)
5014 .arg(&up)
5015 .arg(aq)
5016 .arg(ad)
5017 .arg(&mut act)
5018 .arg(&inf)
5019 .arg(&nff)
5020 .arg(&qt_g)
5021 .arg(&qt_u)
5022 .arg(&rbg)
5023 .arg(&rbu);
5024 unsafe {
5025 b.launch(cfg)?;
5026 }
5027 Ok(act)
5028 }
5029
5030 #[allow(clippy::too_many_arguments)]
5031 pub fn moe_down8_fma_q8(
5032 &self,
5033 dp: WPtr8,
5034 w: F32x8,
5035 aq2: &CudaSlice<i8>,
5036 ad2: &CudaSlice<f32>,
5037 dst: &mut cudarc::driver::CudaViewMut<f32>,
5038 in_f: usize,
5039 out_f: usize,
5040 n_used: usize,
5041 qt: i32,
5042 rb: usize,
5043 ) -> Result<(), Box<dyn std::error::Error>> {
5044 let f = self.func("moe_down8_fma_q8");
5045 let cfg = LaunchConfig {
5046 grid_dim: (out_f as u32, 1, 1),
5047 block_dim: (32, 1, 1),
5048 shared_mem_bytes: 0,
5049 };
5050 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5051 let __s_b = self.gpu.stream();
5052 let mut b = __s_b.launch_builder(&f);
5053 b.arg(&dp)
5054 .arg(&w)
5055 .arg(aq2)
5056 .arg(ad2)
5057 .arg(dst)
5058 .arg(&inf)
5059 .arg(&outf)
5060 .arg(&nu)
5061 .arg(&qt)
5062 .arg(&rbi);
5063 unsafe {
5064 b.launch(cfg)?;
5065 }
5066 Ok(())
5067 }
5068
5069 pub fn qmatvec_expert_q8(
5071 &self,
5072 w: &CudaSlice<u8>,
5073 range: std::ops::Range<usize>,
5074 aq: &CudaSlice<i8>,
5075 ad: &CudaSlice<f32>,
5076 m: usize,
5077 in_f: usize,
5078 out_f: usize,
5079 qtype: i32,
5080 row_bytes: usize,
5081 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5082 let f = self.func("qmatvec_expert_q8");
5083 let wv = w.slice(range);
5084 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5085 const ROWS: u32 = 4; let cfg = LaunchConfig {
5087 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5088 block_dim: (32, ROWS, 1),
5089 shared_mem_bytes: 0,
5090 };
5091 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5092 let __s_b = self.gpu.stream();
5093 let mut b = __s_b.launch_builder(&f);
5094 b.arg(&wv)
5095 .arg(aq)
5096 .arg(ad)
5097 .arg(&mut y)
5098 .arg(&inf)
5099 .arg(&outf)
5100 .arg(&mi)
5101 .arg(&qtype)
5102 .arg(&rbi);
5103 unsafe {
5104 b.launch(cfg)?;
5105 }
5106 Ok(y)
5107 }
5108
5109 pub fn moe_gate_up_silu8(
5110 &self,
5111 gp: WPtr8,
5112 up: WPtr8,
5113 x: &cudarc::driver::CudaView<f32>,
5114 in_f: usize,
5115 n_ff: usize,
5116 n_used: usize,
5117 qt_g: i32,
5118 qt_u: i32,
5119 rb_g: usize,
5120 rb_u: usize,
5121 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5122 let f = self.func("moe_gate_up_silu8_f32");
5123 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5125 grid_dim: (n_ff as u32, n_used as u32, 1),
5126 block_dim: (256, 1, 1),
5127 shared_mem_bytes: 0,
5128 };
5129 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5130 let __s_b = self.gpu.stream();
5131 let mut b = __s_b.launch_builder(&f);
5132 b.arg(&gp)
5133 .arg(&up)
5134 .arg(x)
5135 .arg(&mut act)
5136 .arg(&inf)
5137 .arg(&nff)
5138 .arg(&qt_g)
5139 .arg(&qt_u)
5140 .arg(&rbg)
5141 .arg(&rbu);
5142 unsafe {
5143 b.launch(cfg)?;
5144 }
5145 Ok(act)
5146 }
5147
5148 #[allow(clippy::too_many_arguments)]
5154 pub fn moe_down8_fma_into(
5155 &self,
5156 dp: WPtr8,
5157 w: F32x8,
5158 act: &CudaSlice<f32>,
5159 dst: &mut cudarc::driver::CudaViewMut<f32>,
5160 in_f: usize,
5161 out_f: usize,
5162 n_used: usize,
5163 qt: i32,
5164 rb: usize,
5165 ) -> Result<(), Box<dyn std::error::Error>> {
5166 let f = self.func("moe_down8_fma_f32");
5167 let cfg = LaunchConfig {
5168 grid_dim: (out_f as u32, 1, 1),
5169 block_dim: (256, 1, 1),
5170 shared_mem_bytes: 0,
5171 };
5172 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5173 let __s_b = self.gpu.stream();
5174 let mut b = __s_b.launch_builder(&f);
5175 b.arg(&dp)
5176 .arg(&w)
5177 .arg(act)
5178 .arg(dst)
5179 .arg(&inf)
5180 .arg(&outf)
5181 .arg(&nu)
5182 .arg(&qt)
5183 .arg(&rbv);
5184 unsafe {
5185 b.launch(cfg)?;
5186 }
5187 Ok(())
5188 }
5189
5190 #[allow(clippy::too_many_arguments)]
5195 #[allow(clippy::too_many_arguments)]
5210 #[allow(clippy::too_many_arguments)]
5212 pub fn moe_pairs_matvec_q8(
5213 &self,
5214 table: &CudaSlice<u64>,
5215 proj: i32,
5216 pair_tok: &CudaSlice<i32>,
5217 pair_ex: &CudaSlice<i32>,
5218 aq: &CudaSlice<i8>,
5219 ad: &CudaSlice<f32>,
5220 in_f: usize,
5221 out_f: usize,
5222 n_expert: usize,
5223 n_pairs: usize,
5224 qtype: i32,
5225 row_bytes: usize,
5226 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5227 let f = self.func("moe_pairs_matvec_q8");
5228 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5229 const ROWS: u32 = 4;
5230 let cfg = LaunchConfig {
5231 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5232 block_dim: (32, ROWS, 1),
5233 shared_mem_bytes: 0,
5234 };
5235 let (inf, outf, ne, np, rbi) = (
5236 in_f as i32,
5237 out_f as i32,
5238 n_expert as i32,
5239 n_pairs as i32,
5240 row_bytes as i64,
5241 );
5242 let __s_b = self.gpu.stream();
5243 let mut b = __s_b.launch_builder(&f);
5244 b.arg(table)
5245 .arg(&proj)
5246 .arg(pair_tok)
5247 .arg(pair_ex)
5248 .arg(aq)
5249 .arg(ad)
5250 .arg(&mut y)
5251 .arg(&inf)
5252 .arg(&outf)
5253 .arg(&ne)
5254 .arg(&np)
5255 .arg(&qtype)
5256 .arg(&rbi);
5257 unsafe {
5258 b.launch(cfg)?;
5259 }
5260 Ok(y)
5261 }
5262
5263 #[allow(clippy::too_many_arguments)]
5265 pub fn moe_pairs_matvec_q8_em(
5266 &self,
5267 table: &CudaSlice<u64>,
5268 proj: i32,
5269 ex_ids: &CudaSlice<i32>,
5270 ex_off: &CudaSlice<i32>,
5271 ex_pairs: &CudaSlice<i32>,
5272 pair_tok: &CudaSlice<i32>,
5273 aq: &CudaSlice<i8>,
5274 ad: &CudaSlice<f32>,
5275 in_f: usize,
5276 out_f: usize,
5277 n_expert: usize,
5278 n_active: usize,
5279 n_pairs: usize,
5280 qtype: i32,
5281 row_bytes: usize,
5282 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5283 let f = self.func("moe_pairs_matvec_q8_em");
5284 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5285 const ROWS: u32 = 4;
5286 let cfg = LaunchConfig {
5287 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5288 block_dim: (32, ROWS, 1),
5289 shared_mem_bytes: 0,
5290 };
5291 let (inf, outf, ne, na, rbi) = (
5292 in_f as i32,
5293 out_f as i32,
5294 n_expert as i32,
5295 n_active as i32,
5296 row_bytes as i64,
5297 );
5298 let __s_b = self.gpu.stream();
5299 let mut b = __s_b.launch_builder(&f);
5300 b.arg(table)
5301 .arg(&proj)
5302 .arg(ex_ids)
5303 .arg(ex_off)
5304 .arg(ex_pairs)
5305 .arg(pair_tok)
5306 .arg(aq)
5307 .arg(ad)
5308 .arg(&mut y)
5309 .arg(&inf)
5310 .arg(&outf)
5311 .arg(&ne)
5312 .arg(&na)
5313 .arg(&qtype)
5314 .arg(&rbi);
5315 unsafe {
5316 b.launch(cfg)?;
5317 }
5318 Ok(y)
5319 }
5320
5321 #[allow(clippy::too_many_arguments)]
5324 pub fn moe_pairs_matvec_q8_dec(
5325 &self,
5326 table: &CudaSlice<u64>,
5327 proj: i32,
5328 ex_ids: &CudaSlice<i32>,
5329 ex_off: &CudaSlice<i32>,
5330 ex_pairs: &CudaSlice<i32>,
5331 pair_tok: &CudaSlice<i32>,
5332 aq: &CudaSlice<i8>,
5333 ad: &CudaSlice<f32>,
5334 in_f: usize,
5335 out_f: usize,
5336 n_expert: usize,
5337 n_active: usize,
5338 n_pairs: usize,
5339 qtype: i32,
5340 row_bytes: usize,
5341 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5342 let f = self.func("moe_pairs_matvec_q8_dec");
5343 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5344 const ROWS: u32 = 4;
5345 let cfg = LaunchConfig {
5346 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5347 block_dim: (32, ROWS, 1),
5348 shared_mem_bytes: 0,
5349 };
5350 let (inf, outf, ne, na, rbi) = (
5351 in_f as i32,
5352 out_f as i32,
5353 n_expert as i32,
5354 n_active as i32,
5355 row_bytes as i64,
5356 );
5357 let __s_b = self.gpu.stream();
5358 let mut b = __s_b.launch_builder(&f);
5359 b.arg(table)
5360 .arg(&proj)
5361 .arg(ex_ids)
5362 .arg(ex_off)
5363 .arg(ex_pairs)
5364 .arg(pair_tok)
5365 .arg(aq)
5366 .arg(ad)
5367 .arg(&mut y)
5368 .arg(&inf)
5369 .arg(&outf)
5370 .arg(&ne)
5371 .arg(&na)
5372 .arg(&qtype)
5373 .arg(&rbi);
5374 unsafe {
5375 b.launch(cfg)?;
5376 }
5377 Ok(y)
5378 }
5379
5380 pub fn moe_pairs_gelu_mul(
5381 &self,
5382 gate: &CudaSlice<f32>,
5383 up: &CudaSlice<f32>,
5384 n: usize,
5385 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5386 let f = self.func("moe_pairs_gelu_mul");
5387 let mut act = self.alloc_uninit::<f32>(n)?;
5388 let cfg = LaunchConfig::for_num_elems(n as u32);
5389 let nl = n as i64;
5390 let __s_b = self.gpu.stream();
5391 let mut b = __s_b.launch_builder(&f);
5392 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5393 unsafe {
5394 b.launch(cfg)?;
5395 }
5396 Ok(act)
5397 }
5398
5399 pub fn moe_pairs_silu_mul(
5400 &self,
5401 gate: &CudaSlice<f32>,
5402 up: &CudaSlice<f32>,
5403 n: usize,
5404 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5405 let f = self.func("moe_pairs_silu_mul");
5406 let mut act = self.alloc_uninit::<f32>(n)?;
5407 let cfg = LaunchConfig::for_num_elems(n as u32);
5408 let nl = n as i64;
5409 let __s_b = self.gpu.stream();
5410 let mut b = __s_b.launch_builder(&f);
5411 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5412 unsafe {
5413 b.launch(cfg)?;
5414 }
5415 Ok(act)
5416 }
5417
5418 #[allow(clippy::too_many_arguments)]
5419 pub fn moe_pairs_scatter(
5420 &self,
5421 y_down: &CudaSlice<f32>,
5422 pair_w: &CudaSlice<f32>,
5423 tok_pair_off: &CudaSlice<i32>,
5424 tok_pair_ids: &CudaSlice<i32>,
5425 moe_out: &mut CudaSlice<f32>,
5426 t: usize,
5427 n_embd: usize,
5428 ) -> Result<(), Box<dyn std::error::Error>> {
5429 let f = self.func("moe_pairs_scatter");
5430 let cfg = LaunchConfig {
5431 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5432 block_dim: (256, 1, 1),
5433 shared_mem_bytes: 0,
5434 };
5435 let ne = n_embd as i32;
5436 let __s_b = self.gpu.stream();
5437 let mut b = __s_b.launch_builder(&f);
5438 b.arg(y_down)
5439 .arg(pair_w)
5440 .arg(tok_pair_off)
5441 .arg(tok_pair_ids)
5442 .arg(moe_out)
5443 .arg(&ne);
5444 unsafe {
5445 b.launch(cfg)?;
5446 }
5447 Ok(())
5448 }
5449
5450 #[allow(clippy::too_many_arguments)]
5454 pub fn moe_gate_up_gelu8_dev_q8(
5455 &self,
5456 table: &CudaSlice<u64>,
5457 sel: &cudarc::driver::CudaView<i32>,
5458 aq: &CudaSlice<i8>,
5459 ad: &CudaSlice<f32>,
5460 in_f: usize,
5461 n_ff: usize,
5462 n_used: usize,
5463 n_expert: usize,
5464 qt_g: i32,
5465 qt_u: i32,
5466 rb_g: usize,
5467 rb_u: usize,
5468 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5469 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5470 let (inf, nff, ne, rbg, rbu) = (
5471 in_f as i32,
5472 n_ff as i32,
5473 n_expert as i32,
5474 rb_g as i64,
5475 rb_u as i64,
5476 );
5477 let f = self.func("moe_gate_up_gelu8_dev_q8");
5478 let cfg = LaunchConfig {
5479 grid_dim: (n_ff as u32, n_used as u32, 1),
5480 block_dim: (32, 1, 1),
5481 shared_mem_bytes: 0,
5482 };
5483 let __s_b = self.gpu.stream();
5484 let mut b = __s_b.launch_builder(&f);
5485 b.arg(table)
5486 .arg(sel)
5487 .arg(aq)
5488 .arg(ad)
5489 .arg(&mut act)
5490 .arg(&inf)
5491 .arg(&nff)
5492 .arg(&ne)
5493 .arg(&qt_g)
5494 .arg(&qt_u)
5495 .arg(&rbg)
5496 .arg(&rbu);
5497 unsafe {
5498 b.launch(cfg)?;
5499 }
5500 Ok(act)
5501 }
5502
5503 #[allow(clippy::too_many_arguments)]
5505 pub fn moe_gate_up_gelu8_dev_q8_rows(
5506 &self,
5507 table: &CudaSlice<u64>,
5508 sel: &CudaSlice<i32>,
5509 aq: &CudaSlice<i8>,
5510 ad: &CudaSlice<f32>,
5511 t: usize,
5512 in_f: usize,
5513 n_ff: usize,
5514 n_used: usize,
5515 n_expert: usize,
5516 qt_g: i32,
5517 qt_u: i32,
5518 rb_g: usize,
5519 rb_u: usize,
5520 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5521 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5522 let (inf, nff, ne, rbg, rbu, nu) = (
5523 in_f as i32,
5524 n_ff as i32,
5525 n_expert as i32,
5526 rb_g as i64,
5527 rb_u as i64,
5528 n_used as i32,
5529 );
5530 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5531 let cfg = LaunchConfig {
5532 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5533 block_dim: (32, 1, 1),
5534 shared_mem_bytes: 0,
5535 };
5536 let __s_b = self.gpu.stream();
5537 let mut b = __s_b.launch_builder(&f);
5538 b.arg(table)
5539 .arg(sel)
5540 .arg(aq)
5541 .arg(ad)
5542 .arg(&mut act)
5543 .arg(&inf)
5544 .arg(&nff)
5545 .arg(&ne)
5546 .arg(&qt_g)
5547 .arg(&qt_u)
5548 .arg(&rbg)
5549 .arg(&rbu)
5550 .arg(&nu);
5551 unsafe {
5552 b.launch(cfg)?;
5553 }
5554 Ok(act)
5555 }
5556
5557 #[allow(clippy::too_many_arguments)]
5559 pub fn moe_gate_up_gelu8_dev_q8_csr(
5560 &self,
5561 table: &CudaSlice<u64>,
5562 sel: &CudaSlice<i32>,
5563 aq: &CudaSlice<i8>,
5564 ad: &CudaSlice<f32>,
5565 n_pairs: usize,
5566 in_f: usize,
5567 n_ff: usize,
5568 n_used: usize,
5569 n_expert: usize,
5570 qt_g: i32,
5571 qt_u: i32,
5572 rb_g: usize,
5573 rb_u: usize,
5574 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5575 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5576 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5577 in_f as i32,
5578 n_ff as i32,
5579 n_expert as i32,
5580 rb_g as i64,
5581 rb_u as i64,
5582 n_used as i32,
5583 n_pairs as i32,
5584 );
5585 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5586 let cfg = LaunchConfig {
5587 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5588 block_dim: (32, 1, 1),
5589 shared_mem_bytes: 0,
5590 };
5591 let __s_b = self.gpu.stream();
5592 let mut b = __s_b.launch_builder(&f);
5593 b.arg(table)
5594 .arg(sel)
5595 .arg(aq)
5596 .arg(ad)
5597 .arg(&mut act)
5598 .arg(&inf)
5599 .arg(&nff)
5600 .arg(&ne)
5601 .arg(&qt_g)
5602 .arg(&qt_u)
5603 .arg(&rbg)
5604 .arg(&rbu)
5605 .arg(&nu)
5606 .arg(&npi);
5607 unsafe {
5608 b.launch(cfg)?;
5609 }
5610 Ok(act)
5611 }
5612
5613 #[allow(clippy::too_many_arguments)]
5615 pub fn moe_down8_fma_dev_q8_rows_g(
5616 &self,
5617 table: &CudaSlice<u64>,
5618 sel: &CudaSlice<i32>,
5619 w: &CudaSlice<f32>,
5620 aq2: &CudaSlice<i8>,
5621 ad2: &CudaSlice<f32>,
5622 dst: &mut CudaSlice<f32>,
5623 t: usize,
5624 in_f: usize,
5625 out_f: usize,
5626 n_used: usize,
5627 n_expert: usize,
5628 qt: i32,
5629 rb: usize,
5630 ) -> Result<(), Box<dyn std::error::Error>> {
5631 let (inf, outf, nu, ne, rbi) = (
5632 in_f as i32,
5633 out_f as i32,
5634 n_used as i32,
5635 n_expert as i32,
5636 rb as i64,
5637 );
5638 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5642 let f = self.func(if step_b1_w8 {
5643 "moe_down8_fma_dev_q8_rows_w8"
5644 } else {
5645 "moe_down8_fma_dev_q8_rows_g"
5646 });
5647 let cfg = LaunchConfig {
5648 grid_dim: (out_f as u32, 1, t as u32),
5649 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5650 shared_mem_bytes: 0,
5651 };
5652 let __s_b = self.gpu.stream();
5653 let mut b = __s_b.launch_builder(&f);
5654 b.arg(table)
5655 .arg(sel)
5656 .arg(w)
5657 .arg(aq2)
5658 .arg(ad2)
5659 .arg(dst)
5660 .arg(&inf)
5661 .arg(&outf)
5662 .arg(&nu)
5663 .arg(&ne)
5664 .arg(&qt)
5665 .arg(&rbi);
5666 unsafe {
5667 b.launch(cfg)?;
5668 }
5669 Ok(())
5670 }
5671
5672 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5676 let (out_f, in_f) = (2048usize, 2816usize);
5677 let nblk = in_f / 32;
5678 let mut seed = 0x9E3779B97F4A7C15u64;
5679 let mut rng = move || {
5680 seed = seed
5681 .wrapping_mul(6364136223846793005)
5682 .wrapping_add(1442695040888963407);
5683 (seed >> 33) as u8
5684 };
5685 let mut w = vec![0u8; out_f * nblk * 18];
5686 for b in w.iter_mut() {
5687 *b = rng();
5688 }
5689 for r in 0..out_f {
5690 for g in 0..nblk {
5691 let off = (r * nblk + g) * 18;
5692 w[off] = 0x00;
5693 w[off + 1] = 0x2C; }
5695 }
5696 let qplane = out_f * nblk * 16;
5697 let mut wrp = vec![0u8; w.len()];
5698 for r in 0..out_f {
5699 for g in 0..nblk {
5700 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5701 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5702 .copy_from_slice(&src[0..2]);
5703 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5704 }
5705 }
5706 let w_d = self.htod_bytes(&w)?;
5707 let wrp_d = self.htod_bytes(&wrp)?;
5708 let mut aq = vec![0i8; m * in_f];
5709 for v in aq.iter_mut() {
5710 *v = rng() as i8;
5711 }
5712 let aq_d = self.htod_i8(&aq)?;
5713 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5714 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5715 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5716 const RPB: u32 = 4;
5717 let cfg = LaunchConfig {
5718 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5719 block_dim: (32, RPB, 1),
5720 shared_mem_bytes: 0,
5721 };
5722 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5723 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5724 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5725 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5726 {
5727 let __s_b = self.gpu.stream();
5728 let mut b = __s_b.launch_builder(&fb);
5729 b.arg(&w_d)
5730 .arg(&aq_d)
5731 .arg(&ad_d)
5732 .arg(&mut y0)
5733 .arg(&inf)
5734 .arg(&outf)
5735 .arg(&mi)
5736 .arg(&rb);
5737 unsafe {
5738 b.launch(cfg)?;
5739 }
5740 let __s_b = self.gpu.stream();
5741 let mut b = __s_b.launch_builder(&fr);
5742 b.arg(&wrp_d)
5743 .arg(&aq_d)
5744 .arg(&ad_d)
5745 .arg(&mut y1)
5746 .arg(&inf)
5747 .arg(&outf)
5748 .arg(&mi)
5749 .arg(&qp);
5750 unsafe {
5751 b.launch(cfg)?;
5752 }
5753 }
5754 self.gpu.stream().synchronize()?;
5755 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5756 let nd = h0
5757 .iter()
5758 .zip(&h1)
5759 .filter(|(a, b)| a.to_bits() != b.to_bits())
5760 .count();
5761 if nd != 0 {
5762 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5763 }
5764 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5765 self.gpu.stream().synchronize()?;
5766 let t0 = std::time::Instant::now();
5767 for _ in 0..500 {
5768 if rp {
5769 let __s_b = self.gpu.stream();
5770 let mut b = __s_b.launch_builder(&fr);
5771 b.arg(&wrp_d)
5772 .arg(&aq_d)
5773 .arg(&ad_d)
5774 .arg(&mut y1)
5775 .arg(&inf)
5776 .arg(&outf)
5777 .arg(&mi)
5778 .arg(&qp);
5779 unsafe {
5780 b.launch(cfg)?;
5781 }
5782 } else {
5783 let __s_b = self.gpu.stream();
5784 let mut b = __s_b.launch_builder(&fb);
5785 b.arg(&w_d)
5786 .arg(&aq_d)
5787 .arg(&ad_d)
5788 .arg(&mut y0)
5789 .arg(&inf)
5790 .arg(&outf)
5791 .arg(&mi)
5792 .arg(&rb);
5793 unsafe {
5794 b.launch(cfg)?;
5795 }
5796 }
5797 }
5798 self.gpu.stream().synchronize()?;
5799 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5800 };
5801 let _ = time(false)?;
5802 let _ = time(true)?; Ok((time(false)?, time(true)?))
5804 }
5805
5806 pub fn build_q4_rp4(
5811 &self,
5812 t: &mut crate::model::GpuTensor,
5813 ) -> Result<(), Box<dyn std::error::Error>> {
5814 use crate::model::GpuTensor;
5815 let GpuTensor::Quant {
5816 bytes,
5817 qtype,
5818 row_bytes,
5819 ne,
5820 rp4,
5821 ..
5822 } = t
5823 else {
5824 return Ok(());
5825 };
5826 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5827 return Ok(());
5828 }
5829 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5830 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5831 return Ok(());
5832 }
5833 let nblk = in_f / 32;
5834 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5835 let f = self.func("q4_0_split_rp_build");
5836 let n = (out_f * nblk) as i32;
5837 let cfg = LaunchConfig {
5838 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5839 block_dim: (256, 1, 1),
5840 shared_mem_bytes: 0,
5841 };
5842 let (of, nb) = (out_f as i32, nblk as i32);
5843 let _ = n;
5844 let __s_b = self.gpu.stream();
5845 let mut b = __s_b.launch_builder(&f);
5846 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5847 unsafe {
5848 b.launch(cfg)?;
5849 }
5850 *rp4 = Some(dst);
5851 Ok(())
5852 }
5853
5854 pub fn build_q8_rp4(
5859 &self,
5860 t: &mut crate::model::GpuTensor,
5861 ) -> Result<(), Box<dyn std::error::Error>> {
5862 use crate::model::GpuTensor;
5863 let GpuTensor::Quant {
5864 bytes,
5865 qtype,
5866 row_bytes,
5867 ne,
5868 rp4,
5869 ..
5870 } = t
5871 else {
5872 return Ok(());
5873 };
5874 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5875 return Ok(());
5876 }
5877 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5878 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5879 return Ok(());
5880 }
5881 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5882 Ok(())
5883 }
5884
5885 pub fn build_q8_rp4_raw(
5888 &self,
5889 bytes: &CudaSlice<u8>,
5890 in_f: usize,
5891 out_f: usize,
5892 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5893 assert!(in_f % 32 == 0);
5894 let nblk = in_f / 32;
5895 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5896 let f = self.func("q8_0_split_rp_build");
5897 let cfg = LaunchConfig {
5898 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5899 block_dim: (256, 1, 1),
5900 shared_mem_bytes: 0,
5901 };
5902 let (of, nb) = (out_f as i32, nblk as i32);
5903 let __s_b = self.gpu.stream();
5904 let mut b = __s_b.launch_builder(&f);
5905 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5906 unsafe {
5907 b.launch(cfg)?;
5908 }
5909 Ok(dst)
5910 }
5911
5912 pub fn build_q4k_rp4(
5920 &self,
5921 t: &mut crate::model::GpuTensor,
5922 ) -> Result<(), Box<dyn std::error::Error>> {
5923 use crate::model::GpuTensor;
5924 let GpuTensor::Quant {
5925 bytes,
5926 qtype,
5927 row_bytes,
5928 ne,
5929 rp4,
5930 ..
5931 } = t
5932 else {
5933 return Ok(());
5934 };
5935 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5936 return Ok(());
5937 }
5938 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5939 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5940 return Ok(());
5941 }
5942 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5943 Ok(())
5944 }
5945
5946 pub fn build_q6k_rp4(
5947 &self,
5948 t: &mut crate::model::GpuTensor,
5949 ) -> Result<(), Box<dyn std::error::Error>> {
5950 use crate::model::GpuTensor;
5951 let GpuTensor::Quant {
5952 bytes,
5953 qtype,
5954 row_bytes,
5955 ne,
5956 rp4,
5957 ..
5958 } = t
5959 else {
5960 return Ok(());
5961 };
5962 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5963 return Ok(());
5964 }
5965 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5966 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5967 return Ok(());
5968 }
5969 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5970 Ok(())
5971 }
5972
5973 pub fn build_kq_rp4_raw(
5975 &self,
5976 bytes: &CudaSlice<u8>,
5977 in_f: usize,
5978 out_f: usize,
5979 qtype: i32,
5980 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5981 assert!(in_f % 256 == 0);
5982 let nsbk = in_f / 256;
5983 let (sb_bytes, kname) = match qtype {
5984 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5985 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5986 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5987 };
5988 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5989 let f = self.func(kname);
5990 let cfg = LaunchConfig {
5991 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5992 block_dim: (256, 1, 1),
5993 shared_mem_bytes: 0,
5994 };
5995 let (of, nb) = (out_f as i32, nsbk as i32);
5996 let __s_b = self.gpu.stream();
5997 let mut b = __s_b.launch_builder(&f);
5998 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5999 unsafe {
6000 b.launch(cfg)?;
6001 }
6002 Ok(dst)
6003 }
6004
6005 pub fn kqrp_enabled() -> bool {
6009 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6010 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6011 Ok("0") => false,
6012 Ok(_) => true,
6013 Err(_) => cfg!(memra_hopper_mma),
6014 })
6015 }
6016
6017 pub fn build_q4_rp_swap(
6023 &self,
6024 t: &mut crate::model::GpuTensor,
6025 ) -> Result<bool, Box<dyn std::error::Error>> {
6026 use crate::model::GpuTensor;
6027 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6037 return Ok(false);
6038 }
6039 self.build_q4_rp4(t)?;
6040 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6042 return Ok(false);
6043 };
6044 match rp4.take() {
6045 Some(split) => {
6046 *bytes = split; *rp = true;
6048 Ok(true)
6049 }
6050 None => Ok(false),
6051 }
6052 }
6053
6054 pub fn q4rp_enabled() -> bool {
6056 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6057 *ON.get_or_init(|| {
6058 std::env::var("MEMRA_Q4RP")
6059 .map(|v| v != "0")
6060 .unwrap_or(true)
6061 })
6062 }
6063
6064 pub fn copy_rows_strided(
6067 &self,
6068 src: &CudaSlice<f32>,
6069 dst: &mut CudaSlice<f32>,
6070 row_elems: usize,
6071 n_rows: usize,
6072 src_stride: usize,
6073 src_off: usize,
6074 ) -> Result<(), Box<dyn std::error::Error>> {
6075 let f = self.func("copy_rows_strided_f32");
6076 let cfg = LaunchConfig {
6077 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6078 block_dim: (256, 1, 1),
6079 shared_mem_bytes: 0,
6080 };
6081 let (re, nr) = (row_elems as i32, n_rows as i32);
6082 let (st, off) = (src_stride as i64, src_off as i64);
6083 let __s_b = self.gpu.stream();
6084 let mut b = __s_b.launch_builder(&f);
6085 b.arg(src)
6086 .arg(&mut *dst)
6087 .arg(&re)
6088 .arg(&nr)
6089 .arg(&st)
6090 .arg(&off);
6091 unsafe {
6092 b.launch(cfg)?;
6093 }
6094 Ok(())
6095 }
6096
6097 pub fn place_rows_strided(
6103 &self,
6104 src: &CudaSlice<f32>,
6105 dst: &mut CudaSlice<f32>,
6106 row_elems: usize,
6107 n_rows: usize,
6108 dst_stride: usize,
6109 dst_off: usize,
6110 ) -> Result<(), Box<dyn std::error::Error>> {
6111 if row_elems == 0 || n_rows == 0 {
6112 return Err("strided row placement requires nonzero rows and row width".into());
6113 }
6114 let src_len = n_rows
6115 .checked_mul(row_elems)
6116 .ok_or("strided row placement source size overflow")?;
6117 let dst_len = n_rows
6118 .checked_sub(1)
6119 .and_then(|rows| rows.checked_mul(dst_stride))
6120 .and_then(|base| base.checked_add(dst_off))
6121 .and_then(|base| base.checked_add(row_elems))
6122 .ok_or("strided row placement destination size overflow")?;
6123 let row_end = dst_off
6124 .checked_add(row_elems)
6125 .ok_or("strided row placement row size overflow")?;
6126 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6127 return Err(format!(
6128 "strided row placement geometry mismatch: src={} need_src={src_len} \
6129 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6130 dst_stride={dst_stride} dst_off={dst_off}",
6131 src.len(),
6132 dst.len(),
6133 )
6134 .into());
6135 }
6136 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6137 return Err("strided row placement exceeds CUDA kernel geometry".into());
6138 }
6139 let f = self.func("place_rows_strided_f32");
6140 let cfg = LaunchConfig {
6141 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6142 block_dim: (256, 1, 1),
6143 shared_mem_bytes: 0,
6144 };
6145 let (re, nr) = (row_elems as i32, n_rows as i32);
6146 let (st, off) = (dst_stride as i64, dst_off as i64);
6147 let __s_b = self.gpu.stream();
6148 let mut b = __s_b.launch_builder(&f);
6149 b.arg(src)
6150 .arg(&mut *dst)
6151 .arg(&re)
6152 .arg(&nr)
6153 .arg(&st)
6154 .arg(&off);
6155 unsafe {
6156 b.launch(cfg)?;
6157 }
6158 Ok(())
6159 }
6160
6161 pub fn u32_set_k(
6163 &self,
6164 dst: &mut CudaSlice<u32>,
6165 v: u32,
6166 idx: usize,
6167 ) -> Result<(), Box<dyn std::error::Error>> {
6168 let f = self.func("u32_set_k");
6169 let cfg = LaunchConfig {
6170 grid_dim: (1, 1, 1),
6171 block_dim: (1, 1, 1),
6172 shared_mem_bytes: 0,
6173 };
6174 let ii = idx as i32;
6175 let __s_b = self.gpu.stream();
6176 let mut b = __s_b.launch_builder(&f);
6177 b.arg(dst).arg(&v).arg(&ii);
6178 unsafe {
6179 b.launch(cfg)?;
6180 }
6181 Ok(())
6182 }
6183
6184 pub fn i32_add_k(
6186 &self,
6187 d: &mut CudaSlice<i32>,
6188 v: i32,
6189 ) -> Result<(), Box<dyn std::error::Error>> {
6190 let f = self.func("i32_add_k");
6191 let cfg = LaunchConfig {
6192 grid_dim: (1, 1, 1),
6193 block_dim: (32, 1, 1),
6194 shared_mem_bytes: 0,
6195 };
6196 let __s_b = self.gpu.stream();
6197 let mut b = __s_b.launch_builder(&f);
6198 b.arg(d).arg(&v);
6199 unsafe {
6200 b.launch(cfg)?;
6201 }
6202 Ok(())
6203 }
6204
6205 pub fn i32_iota_from(
6207 &self,
6208 ctr: &CudaSlice<i32>,
6209 dst: &mut CudaSlice<i32>,
6210 n: usize,
6211 ) -> Result<(), Box<dyn std::error::Error>> {
6212 let f = self.func("i32_iota_from");
6213 let cfg = LaunchConfig::for_num_elems(n as u32);
6214 let ni = n as i32;
6215 let __s_b = self.gpu.stream();
6216 let mut b = __s_b.launch_builder(&f);
6217 b.arg(ctr).arg(dst).arg(&ni);
6218 unsafe {
6219 b.launch(cfg)?;
6220 }
6221 Ok(())
6222 }
6223
6224 pub fn u32_map_k(
6226 &self,
6227 buf: &mut CudaSlice<u32>,
6228 map: &CudaSlice<u32>,
6229 idx: usize,
6230 ) -> Result<(), Box<dyn std::error::Error>> {
6231 let f = self.func("u32_map_k");
6232 let cfg = LaunchConfig {
6233 grid_dim: (1, 1, 1),
6234 block_dim: (1, 1, 1),
6235 shared_mem_bytes: 0,
6236 };
6237 let ii = idx as i32;
6238 let __s_b = self.gpu.stream();
6239 let mut b = __s_b.launch_builder(&f);
6240 b.arg(buf).arg(map).arg(&ii);
6241 unsafe {
6242 b.launch(cfg)?;
6243 }
6244 Ok(())
6245 }
6246
6247 #[allow(clippy::too_many_arguments)]
6249 pub fn u32_pack2(
6250 &self,
6251 a: &CudaSlice<u32>,
6252 off_a: usize,
6253 n1: usize,
6254 b_in: &CudaSlice<u32>,
6255 n2: usize,
6256 out: &mut CudaSlice<u32>,
6257 ) -> Result<(), Box<dyn std::error::Error>> {
6258 let f = self.func("u32_pack2");
6259 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6260 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6261 let __s_b = self.gpu.stream();
6262 let mut b = __s_b.launch_builder(&f);
6263 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6264 unsafe {
6265 b.launch(cfg)?;
6266 }
6267 Ok(())
6268 }
6269
6270 pub fn moe_w_exscale(
6272 &self,
6273 w: &mut CudaSlice<f32>,
6274 sel: &CudaSlice<i32>,
6275 s: &CudaSlice<f32>,
6276 n: usize,
6277 ) -> Result<(), Box<dyn std::error::Error>> {
6278 let f = self.func("moe_w_exscale");
6279 let cfg = LaunchConfig::for_num_elems(n as u32);
6280 let ni = n as i32;
6281 let __s_b = self.gpu.stream();
6282 let mut b = __s_b.launch_builder(&f);
6283 b.arg(w).arg(sel).arg(s).arg(&ni);
6284 unsafe {
6285 b.launch(cfg)?;
6286 }
6287 Ok(())
6288 }
6289
6290 pub fn moe_w_scale_by_expert(
6293 &self,
6294 w: &mut CudaSlice<f32>,
6295 sel: &CudaSlice<i32>,
6296 macros: &CudaSlice<f32>,
6297 n_expert: usize,
6298 n: usize,
6299 ) -> Result<(), Box<dyn std::error::Error>> {
6300 let f = self.func("moe_w_scale_by_expert");
6301 let cfg = LaunchConfig {
6302 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6303 block_dim: (64, 1, 1),
6304 shared_mem_bytes: 0,
6305 };
6306 let (ne, nn) = (n_expert as i32, n as i32);
6307 let __s_b = self.gpu.stream();
6308 let mut b = __s_b.launch_builder(&f);
6309 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6310 unsafe {
6311 b.launch(cfg)?;
6312 }
6313 Ok(())
6314 }
6315
6316 pub fn moe_gate_up_silu8_dev_q8(
6317 &self,
6318 table: &CudaSlice<u64>,
6319 sel: &cudarc::driver::CudaView<i32>,
6320 aq: &CudaSlice<i8>,
6321 ad: &CudaSlice<f32>,
6322 in_f: usize,
6323 n_ff: usize,
6324 n_used: usize,
6325 n_expert: usize,
6326 qt_g: i32,
6327 qt_u: i32,
6328 rb_g: usize,
6329 rb_u: usize,
6330 macros: &CudaSlice<f32>,
6331 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6332 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6333 let (mode, wpb) = GU.get_or_init(|| {
6334 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6335 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6336 .ok()
6337 .and_then(|v| v.parse().ok())
6338 .unwrap_or(4u32)
6339 .clamp(1, 16);
6340 (mode, wpb)
6341 });
6342 let (mode, wpb) = (mode.as_str(), *wpb);
6343 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6344 let (inf, nff, ne, rbg, rbu) = (
6345 in_f as i32,
6346 n_ff as i32,
6347 n_expert as i32,
6348 rb_g as i64,
6349 rb_u as i64,
6350 );
6351 let (f, cfg) = match mode {
6352 "1" | "2" | "4" => {
6353 let rpw: u32 = mode.parse().unwrap();
6354 let f = self.func(match rpw {
6355 1 => "moe_gate_up_silu8_dev_q8_r1",
6356 2 => "moe_gate_up_silu8_dev_q8_r2",
6357 _ => "moe_gate_up_silu8_dev_q8_r4",
6358 });
6359 let rows_per_block = (rpw * wpb) as usize;
6360 let gx = n_ff.div_ceil(rows_per_block) as u32;
6361 (
6362 f,
6363 LaunchConfig {
6364 grid_dim: (gx, n_used as u32, 1),
6365 block_dim: (32, wpb, 1),
6366 shared_mem_bytes: 0,
6367 },
6368 )
6369 }
6370 "j8" if n_used <= 32 => (
6371 self.func("moe_gate_up_silu8_dev_q8_j8"),
6372 LaunchConfig {
6373 grid_dim: (n_ff as u32, 1, 1),
6374 block_dim: (32, n_used as u32, 1),
6375 shared_mem_bytes: 0,
6376 },
6377 ),
6378 "vsm2" => {
6380 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6381 let sh = (rb_g + rb_u) as u32;
6382 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6383 f.set_attribute(
6384 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6385 sh as i32,
6386 )?;
6387 (
6388 f,
6389 LaunchConfig {
6390 grid_dim: (n_ff as u32, n_used as u32, 1),
6391 block_dim: (32, 1, 1),
6392 shared_mem_bytes: sh,
6393 },
6394 )
6395 }
6396 "vsm" => {
6397 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6398 let sh = (rb_g + rb_u) as u32;
6399 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6400 f.set_attribute(
6401 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6402 sh as i32,
6403 )?;
6404 (
6405 f,
6406 LaunchConfig {
6407 grid_dim: (n_ff as u32, n_used as u32, 1),
6408 block_dim: (32, 1, 1),
6409 shared_mem_bytes: sh,
6410 },
6411 )
6412 }
6413 "sg" => (
6414 self.func("moe_gate_up_silu8_dev_q8_sg"),
6415 LaunchConfig {
6416 grid_dim: (n_ff as u32, n_used as u32, 1),
6417 block_dim: (32, 1, 1),
6418 shared_mem_bytes: 0,
6419 },
6420 ),
6421 "j8sg" if n_used <= 32 => (
6422 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6423 LaunchConfig {
6424 grid_dim: (n_ff as u32, 1, 1),
6425 block_dim: (32, n_used as u32, 1),
6426 shared_mem_bytes: 0,
6427 },
6428 ),
6429 "u64" if in_f == 2048 => (
6430 self.func("moe_gate_up_silu8_dev_q8_u64"),
6431 LaunchConfig {
6432 grid_dim: (n_ff as u32, n_used as u32, 1),
6433 block_dim: (32, 1, 1),
6434 shared_mem_bytes: 0,
6435 },
6436 ),
6437 "gs4" if in_f == 2048 => (
6438 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6439 LaunchConfig {
6440 grid_dim: (n_ff as u32, n_used as u32, 1),
6441 block_dim: (32, 4, 1),
6442 shared_mem_bytes: 0,
6443 },
6444 ),
6445 "v" | "" => (
6447 self.func("moe_gate_up_silu8_dev_q8_v"),
6448 LaunchConfig {
6449 grid_dim: (n_ff as u32, n_used as u32, 1),
6450 block_dim: (32, 1, 1),
6451 shared_mem_bytes: 0,
6452 },
6453 ),
6454 "s2" => (
6455 self.func("moe_gate_up_silu8_dev_q8_s2"),
6456 LaunchConfig {
6457 grid_dim: (n_ff as u32, n_used as u32, 1),
6458 block_dim: (32, 2, 1),
6459 shared_mem_bytes: 0,
6460 },
6461 ),
6462 "s2z" => {
6463 let rz = wpb.min(16); (
6465 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6466 LaunchConfig {
6467 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6468 block_dim: (32, 2, rz),
6469 shared_mem_bytes: 0,
6470 },
6471 )
6472 }
6473 _ => (
6474 self.func("moe_gate_up_silu8_dev_q8"),
6475 LaunchConfig {
6476 grid_dim: (n_ff as u32, n_used as u32, 1),
6477 block_dim: (32, 1, 1),
6478 shared_mem_bytes: 0,
6479 },
6480 ),
6481 };
6482 let __s_b = self.gpu.stream();
6483 let mut b = __s_b.launch_builder(&f);
6484 b.arg(table)
6485 .arg(sel)
6486 .arg(aq)
6487 .arg(ad)
6488 .arg(&mut act)
6489 .arg(&inf)
6490 .arg(&nff)
6491 .arg(&ne)
6492 .arg(&qt_g)
6493 .arg(&qt_u)
6494 .arg(&rbg)
6495 .arg(&rbu)
6496 .arg(macros);
6497 unsafe {
6498 b.launch(cfg)?;
6499 }
6500 Ok(act)
6501 }
6502
6503 #[allow(clippy::too_many_arguments)]
6504 pub fn moe_down8_fma_dev_q8(
6505 &self,
6506 table: &CudaSlice<u64>,
6507 sel: &cudarc::driver::CudaView<i32>,
6508 w: &cudarc::driver::CudaView<f32>,
6509 aq2: &CudaSlice<i8>,
6510 ad2: &CudaSlice<f32>,
6511 dst: &mut cudarc::driver::CudaViewMut<f32>,
6512 in_f: usize,
6513 out_f: usize,
6514 n_used: usize,
6515 n_expert: usize,
6516 qt: i32,
6517 rb: usize,
6518 ) -> Result<(), Box<dyn std::error::Error>> {
6519 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6520 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6521 let (inf, outf, nu, ne, rbi) = (
6522 in_f as i32,
6523 out_f as i32,
6524 n_used as i32,
6525 n_expert as i32,
6526 rb as i64,
6527 );
6528 let (f, cfg) = match mode.as_str() {
6531 m @ ("1" | "2" | "4") if n_used <= 8 => {
6532 let rpw: usize = m.parse().unwrap();
6533 let f = self.func(match rpw {
6534 1 => "moe_down8_fma_dev_q8_w8r1",
6535 2 => "moe_down8_fma_dev_q8_w8r2",
6536 _ => "moe_down8_fma_dev_q8_w8r4",
6537 });
6538 (
6539 f,
6540 LaunchConfig {
6541 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6542 block_dim: (32, n_used as u32, 1),
6543 shared_mem_bytes: 0,
6544 },
6545 )
6546 }
6547 "h2" if in_f == 512 => (
6548 self.func("moe_down8_fma_dev_q8_h2"),
6549 LaunchConfig {
6550 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6551 block_dim: (32, 1, 1),
6552 shared_mem_bytes: 0,
6553 },
6554 ),
6555 "" if in_f == 704 && n_used <= 8 => (
6558 self.func("moe_down8_fma_dev_q8_w8r2"),
6559 LaunchConfig {
6560 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6561 block_dim: (32, n_used as u32, 1),
6562 shared_mem_bytes: 0,
6563 },
6564 ),
6565 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6569 self.func("moe_down8_fma_dev_q8_w8h2v"),
6570 LaunchConfig {
6571 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6572 block_dim: (32, n_used as u32, 1),
6573 shared_mem_bytes: 0,
6574 },
6575 ),
6576 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6577 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6578 LaunchConfig {
6579 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6580 block_dim: (32, n_used as u32, 1),
6581 shared_mem_bytes: 0,
6582 },
6583 ),
6584 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6585 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6586 LaunchConfig {
6587 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6588 block_dim: (32, n_used as u32, 1),
6589 shared_mem_bytes: 0,
6590 },
6591 ),
6592 "w8h2" if in_f == 512 && n_used <= 8 => (
6593 self.func("moe_down8_fma_dev_q8_w8h2"),
6594 LaunchConfig {
6595 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6596 block_dim: (32, n_used as u32, 1),
6597 shared_mem_bytes: 0,
6598 },
6599 ),
6600 _ => (
6601 self.func("moe_down8_fma_dev_q8"),
6602 LaunchConfig {
6603 grid_dim: (out_f as u32, 1, 1),
6604 block_dim: (32, 1, 1),
6605 shared_mem_bytes: 0,
6606 },
6607 ),
6608 };
6609 let __s_b = self.gpu.stream();
6610 let mut b = __s_b.launch_builder(&f);
6611 b.arg(table)
6612 .arg(sel)
6613 .arg(w)
6614 .arg(aq2)
6615 .arg(ad2)
6616 .arg(dst)
6617 .arg(&inf)
6618 .arg(&outf)
6619 .arg(&nu)
6620 .arg(&ne)
6621 .arg(&qt)
6622 .arg(&rbi);
6623 unsafe {
6624 b.launch(cfg)?;
6625 }
6626 Ok(())
6627 }
6628
6629 #[allow(clippy::too_many_arguments)]
6636 pub fn moe_gate_up_silu8_dev_q8_rows(
6637 &self,
6638 table: &CudaSlice<u64>,
6639 sel: &CudaSlice<i32>,
6640 aq: &CudaSlice<i8>,
6641 ad: &CudaSlice<f32>,
6642 t: usize,
6643 in_f: usize,
6644 n_ff: usize,
6645 n_used: usize,
6646 n_expert: usize,
6647 qt_g: i32,
6648 qt_u: i32,
6649 rb_g: usize,
6650 rb_u: usize,
6651 macros: &CudaSlice<f32>,
6652 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6653 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6654 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6655 let cfg = LaunchConfig {
6656 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6657 block_dim: (32, 1, 1),
6658 shared_mem_bytes: 0,
6659 };
6660 let (inf, nff, ne, nu, rbg, rbu) = (
6661 in_f as i32,
6662 n_ff as i32,
6663 n_expert as i32,
6664 n_used as i32,
6665 rb_g as i64,
6666 rb_u as i64,
6667 );
6668 let __s_b = self.gpu.stream();
6669 let mut b = __s_b.launch_builder(&f);
6670 b.arg(table)
6671 .arg(sel)
6672 .arg(aq)
6673 .arg(ad)
6674 .arg(&mut act)
6675 .arg(&inf)
6676 .arg(&nff)
6677 .arg(&ne)
6678 .arg(&qt_g)
6679 .arg(&qt_u)
6680 .arg(&rbg)
6681 .arg(&rbu)
6682 .arg(&nu)
6683 .arg(macros);
6684 unsafe {
6685 b.launch(cfg)?;
6686 }
6687 Ok(act)
6688 }
6689
6690 #[allow(clippy::too_many_arguments)]
6695 pub fn moe_down8_fma_dev_q8_rows(
6696 &self,
6697 table: &CudaSlice<u64>,
6698 sel: &CudaSlice<i32>,
6699 w: &CudaSlice<f32>,
6700 aq2: &CudaSlice<i8>,
6701 ad2: &CudaSlice<f32>,
6702 dst: &mut CudaSlice<f32>,
6703 t: usize,
6704 in_f: usize,
6705 out_f: usize,
6706 n_used: usize,
6707 n_expert: usize,
6708 qt: i32,
6709 rb: usize,
6710 ) -> Result<(), Box<dyn std::error::Error>> {
6711 assert!(
6712 in_f == 512 && n_used <= 8,
6713 "down rows twin is w8h2v shape-gated"
6714 );
6715 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6716 let cfg = LaunchConfig {
6717 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6718 block_dim: (32, n_used as u32, 1),
6719 shared_mem_bytes: 0,
6720 };
6721 let (inf, outf, nu, ne, rbi) = (
6722 in_f as i32,
6723 out_f as i32,
6724 n_used as i32,
6725 n_expert as i32,
6726 rb as i64,
6727 );
6728 let __s_b = self.gpu.stream();
6729 let mut b = __s_b.launch_builder(&f);
6730 b.arg(table)
6731 .arg(sel)
6732 .arg(w)
6733 .arg(aq2)
6734 .arg(ad2)
6735 .arg(dst)
6736 .arg(&inf)
6737 .arg(&outf)
6738 .arg(&nu)
6739 .arg(&ne)
6740 .arg(&qt)
6741 .arg(&rbi);
6742 unsafe {
6743 b.launch(cfg)?;
6744 }
6745 Ok(())
6746 }
6747
6748 #[allow(clippy::too_many_arguments)]
6752 pub fn moe_gate_up_silu8_dev_q8_csr(
6753 &self,
6754 table: &CudaSlice<u64>,
6755 sel: &CudaSlice<i32>,
6756 aq: &CudaSlice<i8>,
6757 ad: &CudaSlice<f32>,
6758 n_pairs: usize,
6759 in_f: usize,
6760 n_ff: usize,
6761 n_used: usize,
6762 n_expert: usize,
6763 qt_g: i32,
6764 qt_u: i32,
6765 rb_g: usize,
6766 rb_u: usize,
6767 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6768 let f = if qt_g == crate::QT_NVFP4 {
6771 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6772 } else {
6773 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6774 };
6775 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6776 let cfg = LaunchConfig {
6777 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6778 block_dim: (32, 1, 1),
6779 shared_mem_bytes: 0,
6780 };
6781 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6782 in_f as i32,
6783 n_ff as i32,
6784 n_expert as i32,
6785 n_used as i32,
6786 n_pairs as i32,
6787 rb_g as i64,
6788 rb_u as i64,
6789 );
6790 let __s_b = self.gpu.stream();
6791 let mut b = __s_b.launch_builder(&f);
6792 b.arg(table)
6793 .arg(sel)
6794 .arg(aq)
6795 .arg(ad)
6796 .arg(&mut act)
6797 .arg(&inf)
6798 .arg(&nff)
6799 .arg(&ne)
6800 .arg(&qt_g)
6801 .arg(&qt_u)
6802 .arg(&rbg)
6803 .arg(&rbu)
6804 .arg(&nu)
6805 .arg(&npi);
6806 unsafe {
6807 b.launch(cfg)?;
6808 }
6809 Ok(act)
6810 }
6811
6812 #[allow(clippy::too_many_arguments)]
6816 pub fn moe_down8_fma_dev_q8_variant(
6817 &self,
6818 variant: &str,
6819 table: &CudaSlice<u64>,
6820 sel: &cudarc::driver::CudaView<i32>,
6821 w: &cudarc::driver::CudaView<f32>,
6822 aq2: &CudaSlice<i8>,
6823 ad2: &CudaSlice<f32>,
6824 dst: &mut cudarc::driver::CudaViewMut<f32>,
6825 in_f: usize,
6826 out_f: usize,
6827 n_used: usize,
6828 n_expert: usize,
6829 qt: i32,
6830 rb: usize,
6831 ) -> Result<(), Box<dyn std::error::Error>> {
6832 let (inf, outf, nu, ne, rbi) = (
6833 in_f as i32,
6834 out_f as i32,
6835 n_used as i32,
6836 n_expert as i32,
6837 rb as i64,
6838 );
6839 let (f, cfg) = match variant {
6840 "w8h2" | "w8h2v" => (
6841 self.func(if variant == "w8h2" {
6842 "moe_down8_fma_dev_q8_w8h2"
6843 } else {
6844 "moe_down8_fma_dev_q8_w8h2v"
6845 }),
6846 LaunchConfig {
6847 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6848 block_dim: (32, n_used as u32, 1),
6849 shared_mem_bytes: 0,
6850 },
6851 ),
6852 "w8h2r2" | "w8h2r2v" => (
6853 self.func(if variant == "w8h2r2" {
6854 "moe_down8_fma_dev_q8_w8h2r2"
6855 } else {
6856 "moe_down8_fma_dev_q8_w8h2r2v"
6857 }),
6858 LaunchConfig {
6859 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6860 block_dim: (32, n_used as u32, 1),
6861 shared_mem_bytes: 0,
6862 },
6863 ),
6864 _ => (
6865 self.func("moe_down8_fma_dev_q8"),
6866 LaunchConfig {
6867 grid_dim: (out_f as u32, 1, 1),
6868 block_dim: (32, 1, 1),
6869 shared_mem_bytes: 0,
6870 },
6871 ),
6872 };
6873 let __s_b = self.gpu.stream();
6874 let mut b = __s_b.launch_builder(&f);
6875 b.arg(table)
6876 .arg(sel)
6877 .arg(w)
6878 .arg(aq2)
6879 .arg(ad2)
6880 .arg(dst)
6881 .arg(&inf)
6882 .arg(&outf)
6883 .arg(&nu)
6884 .arg(&ne)
6885 .arg(&qt)
6886 .arg(&rbi);
6887 unsafe {
6888 b.launch(cfg)?;
6889 }
6890 Ok(())
6891 }
6892
6893 #[allow(clippy::too_many_arguments)]
6895 pub fn moe_gate_up_silu8_dev_q8_variant(
6896 &self,
6897 variant: &str,
6898 table: &CudaSlice<u64>,
6899 sel: &cudarc::driver::CudaView<i32>,
6900 aq: &CudaSlice<i8>,
6901 ad: &CudaSlice<f32>,
6902 in_f: usize,
6903 n_ff: usize,
6904 n_used: usize,
6905 n_expert: usize,
6906 qt_g: i32,
6907 qt_u: i32,
6908 rb_g: usize,
6909 rb_u: usize,
6910 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6911 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6912 let (inf, nff, ne, rbg, rbu) = (
6913 in_f as i32,
6914 n_ff as i32,
6915 n_expert as i32,
6916 rb_g as i64,
6917 rb_u as i64,
6918 );
6919 let f = self.func(if variant == "v" {
6920 "moe_gate_up_silu8_dev_q8_v"
6921 } else {
6922 "moe_gate_up_silu8_dev_q8"
6923 });
6924 let cfg = LaunchConfig {
6925 grid_dim: (n_ff as u32, n_used as u32, 1),
6926 block_dim: (32, 1, 1),
6927 shared_mem_bytes: 0,
6928 };
6929 let __s_b = self.gpu.stream();
6930 let mut b = __s_b.launch_builder(&f);
6931 b.arg(table)
6932 .arg(sel)
6933 .arg(aq)
6934 .arg(ad)
6935 .arg(&mut act)
6936 .arg(&inf)
6937 .arg(&nff)
6938 .arg(&ne)
6939 .arg(&qt_g)
6940 .arg(&qt_u)
6941 .arg(&rbg)
6942 .arg(&rbu);
6943 unsafe {
6944 b.launch(cfg)?;
6945 }
6946 Ok(act)
6947 }
6948
6949 pub fn moe_gate_up_silu8_dev(
6950 &self,
6951 table: &CudaSlice<u64>,
6952 sel: &cudarc::driver::CudaView<i32>,
6953 x: &cudarc::driver::CudaView<f32>,
6954 in_f: usize,
6955 n_ff: usize,
6956 n_used: usize,
6957 n_expert: usize,
6958 qt_g: i32,
6959 qt_u: i32,
6960 rb_g: usize,
6961 rb_u: usize,
6962 macros: &CudaSlice<f32>,
6963 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6964 let f = self.func("moe_gate_up_silu8_dev");
6965 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6967 grid_dim: (n_ff as u32, n_used as u32, 1),
6968 block_dim: (256, 1, 1),
6969 shared_mem_bytes: 0,
6970 };
6971 let (inf, nff, ne, rbg, rbu) = (
6972 in_f as i32,
6973 n_ff as i32,
6974 n_expert as i32,
6975 rb_g as i64,
6976 rb_u as i64,
6977 );
6978 let __s_b = self.gpu.stream();
6979 let mut b = __s_b.launch_builder(&f);
6980 b.arg(table)
6981 .arg(sel)
6982 .arg(x)
6983 .arg(&mut act)
6984 .arg(&inf)
6985 .arg(&nff)
6986 .arg(&ne)
6987 .arg(&qt_g)
6988 .arg(&qt_u)
6989 .arg(&rbg)
6990 .arg(&rbu)
6991 .arg(macros);
6992 unsafe {
6993 b.launch(cfg)?;
6994 }
6995 Ok(act)
6996 }
6997
6998 #[allow(clippy::too_many_arguments)]
7001 pub fn moe_down8_fma_dev(
7002 &self,
7003 table: &CudaSlice<u64>,
7004 sel: &cudarc::driver::CudaView<i32>,
7005 w: &cudarc::driver::CudaView<f32>,
7006 act: &CudaSlice<f32>,
7007 dst: &mut cudarc::driver::CudaViewMut<f32>,
7008 in_f: usize,
7009 out_f: usize,
7010 n_used: usize,
7011 n_expert: usize,
7012 qt: i32,
7013 rb: usize,
7014 ) -> Result<(), Box<dyn std::error::Error>> {
7015 let f = self.func("moe_down8_fma_dev");
7016 let cfg = LaunchConfig {
7017 grid_dim: (out_f as u32, 1, 1),
7018 block_dim: (256, 1, 1),
7019 shared_mem_bytes: 0,
7020 };
7021 let (inf, outf, nu, ne, rbv) = (
7022 in_f as i32,
7023 out_f as i32,
7024 n_used as i32,
7025 n_expert as i32,
7026 rb as i64,
7027 );
7028 let __s_b = self.gpu.stream();
7029 let mut b = __s_b.launch_builder(&f);
7030 b.arg(table)
7031 .arg(sel)
7032 .arg(w)
7033 .arg(act)
7034 .arg(dst)
7035 .arg(&inf)
7036 .arg(&outf)
7037 .arg(&nu)
7038 .arg(&ne)
7039 .arg(&qt)
7040 .arg(&rbv);
7041 unsafe {
7042 b.launch(cfg)?;
7043 }
7044 Ok(())
7045 }
7046
7047 pub fn axpy_into(
7049 &self,
7050 src: &CudaSlice<f32>,
7051 alpha: f32,
7052 dst: &mut cudarc::driver::CudaViewMut<f32>,
7053 n: usize,
7054 ) -> Result<(), Box<dyn std::error::Error>> {
7055 let f = self.func("axpy_f32");
7056 let cfg = LaunchConfig::for_num_elems(n as u32);
7057 let (a, ni) = (alpha, n as i32);
7058 let __s_b = self.gpu.stream();
7059 let mut b = __s_b.launch_builder(&f);
7060 b.arg(src).arg(dst).arg(&a).arg(&ni);
7061 unsafe {
7062 b.launch(cfg)?;
7063 }
7064 Ok(())
7065 }
7066
7067 pub fn axpy_host_into(
7069 &self,
7070 src: &cudarc::driver::CudaView<'_, f32>,
7071 alpha: f32,
7072 dst: &mut cudarc::driver::CudaViewMut<f32>,
7073 n: usize,
7074 ) -> Result<(), Box<dyn std::error::Error>> {
7075 let f = self.func("axpy_host_f32");
7076 let cfg = LaunchConfig::for_num_elems(n as u32);
7077 let (a, ni) = (alpha, n as i32);
7078 let __s_b = self.gpu.stream();
7079 let mut b = __s_b.launch_builder(&f);
7080 b.arg(src).arg(dst).arg(&a).arg(&ni);
7081 unsafe {
7082 b.launch(cfg)?;
7083 }
7084 Ok(())
7085 }
7086
7087 pub fn add_scaled_rows(
7089 &self,
7090 src: &CudaSlice<f32>,
7091 scale: &CudaSlice<f32>,
7092 dst: &mut CudaSlice<f32>,
7093 ncols: usize,
7094 nrows: usize,
7095 ) -> Result<(), Box<dyn std::error::Error>> {
7096 let f = self.func("add_scaled_rows_f32");
7097 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7098 let (nc, nr) = (ncols as i32, nrows as i32);
7099 let __s_b = self.gpu.stream();
7100 let mut b = __s_b.launch_builder(&f);
7101 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7102 unsafe {
7103 b.launch(cfg)?;
7104 }
7105 Ok(())
7106 }
7107
7108 pub fn gather_rows(
7112 &self,
7113 src: &CudaSlice<f32>,
7114 idx: &CudaSlice<i32>,
7115 dst: &mut CudaSlice<f32>,
7116 ncols: usize,
7117 m_e: usize,
7118 ) -> Result<(), Box<dyn std::error::Error>> {
7119 let f = self.func("gather_rows_f32");
7120 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7121 let (nc, me) = (ncols as i32, m_e as i32);
7122 let __s_b = self.gpu.stream();
7123 let mut b = __s_b.launch_builder(&f);
7124 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7125 unsafe {
7126 b.launch(cfg)?;
7127 }
7128 Ok(())
7129 }
7130
7131 pub fn scatter_slot(
7136 &self,
7137 src: &CudaSlice<f32>,
7138 tok_idx: &CudaSlice<i32>,
7139 slot_idx: &CudaSlice<i32>,
7140 weight: &CudaSlice<f32>,
7141 dst: &mut CudaSlice<f32>,
7142 wbuf: &mut CudaSlice<f32>,
7143 ncols: usize,
7144 n_used: usize,
7145 m_e: usize,
7146 ) -> Result<(), Box<dyn std::error::Error>> {
7147 let f = self.func("scatter_add_slot_f32");
7148 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7149 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7150 let __s_b = self.gpu.stream();
7151 let mut b = __s_b.launch_builder(&f);
7152 b.arg(src)
7153 .arg(tok_idx)
7154 .arg(slot_idx)
7155 .arg(weight)
7156 .arg(dst)
7157 .arg(wbuf)
7158 .arg(&nc)
7159 .arg(&nu)
7160 .arg(&me);
7161 unsafe {
7162 b.launch(cfg)?;
7163 }
7164 Ok(())
7165 }
7166
7167 pub fn reduce_slots(
7171 &self,
7172 slots: &CudaSlice<f32>,
7173 wbuf: &CudaSlice<f32>,
7174 dst: &mut CudaSlice<f32>,
7175 ncols: usize,
7176 n_used: usize,
7177 t: usize,
7178 ) -> Result<(), Box<dyn std::error::Error>> {
7179 let f = self.func("reduce_slots_f32");
7180 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7181 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7182 let __s_b = self.gpu.stream();
7183 let mut b = __s_b.launch_builder(&f);
7184 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7185 unsafe {
7186 b.launch(cfg)?;
7187 }
7188 Ok(())
7189 }
7190
7191 pub fn reduce_slots_host(
7196 &self,
7197 slots: &CudaSlice<f32>,
7198 wbuf: &CudaSlice<f32>,
7199 dst: &mut CudaSlice<f32>,
7200 ncols: usize,
7201 n_used: usize,
7202 t: usize,
7203 ) -> Result<(), Box<dyn std::error::Error>> {
7204 let f = self.func("reduce_slots_host_f32");
7205 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7206 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7207 let __s_b = self.gpu.stream();
7208 let mut b = __s_b.launch_builder(&f);
7209 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7210 unsafe {
7211 b.launch(cfg)?;
7212 }
7213 Ok(())
7214 }
7215
7216 pub fn quantize_q8_1_view(
7223 &self,
7224 x: &cudarc::driver::CudaView<f32>,
7225 m: usize,
7226 in_f: usize,
7227 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7228 let f = self.func("quantize_q8_1");
7229 let nblk = in_f / 32;
7230 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7231 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7232 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7233 let (inf, mi) = (in_f as i32, m as i32);
7234 let __s_b = self.gpu.stream();
7235 let mut b = __s_b.launch_builder(&f);
7236 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7237 unsafe {
7238 b.launch(cfg)?;
7239 }
7240 Ok((q, d))
7241 }
7242
7243 pub fn quantize_q8_1(
7244 &self,
7245 x: &CudaSlice<f32>,
7246 m: usize,
7247 in_f: usize,
7248 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7249 let nblk = in_f / 32;
7250 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);
7254 let (inf, mi) = (in_f as i32, m as i32);
7255 if Self::pdl_on() && Self::pdl_wb_on() {
7256 {
7257 use cudarc::driver::{DevicePtr, DevicePtrMut};
7258 let s = &self.gpu.stream();
7259 let (px, _g0) = x.device_ptr(s);
7260 let (pq, _g1) = q.device_ptr_mut(s);
7261 let (pd, _g2) = d.device_ptr_mut(s);
7262 let mut ps = [
7263 &px as *const _ as *mut std::ffi::c_void,
7264 &pq as *const _ as *mut _,
7265 &pd as *const _ as *mut _,
7266 &inf as *const _ as *mut _,
7267 &mi as *const _ as *mut _,
7268 ];
7269 unsafe {
7270 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7271 }
7272 }
7273 return Ok((q, d));
7274 }
7275 let f = self.func("quantize_q8_1");
7276 let __s_b = self.gpu.stream();
7277 let mut b = __s_b.launch_builder(&f);
7278 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7279 unsafe {
7280 b.launch(cfg)?;
7281 }
7282 Ok((q, d))
7283 }
7284
7285 pub fn quantize_fp4_act(
7289 &self,
7290 x: &CudaSlice<f32>,
7291 m: usize,
7292 in_f: usize,
7293 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7294 let f = self.func("quantize_fp4_act");
7295 let nb16 = in_f / 16;
7296 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);
7299 let (inf, mi) = (in_f as i32, m as i32);
7300 let __s_b = self.gpu.stream();
7301 let mut b = __s_b.launch_builder(&f);
7302 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7303 unsafe {
7304 b.launch(cfg)?;
7305 }
7306 Ok((aq4, ad4))
7307 }
7308
7309 pub fn qmatvec_gemm_nvfp4_fp4(
7314 &self,
7315 bytes: &CudaSlice<u8>,
7316 x: &CudaSlice<f32>,
7317 m: usize,
7318 in_f: usize,
7319 out_f: usize,
7320 row_bytes: usize,
7321 scale: f32,
7322 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7323 assert!(
7324 in_f % 64 == 0,
7325 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7326 );
7327 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7328 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7329 if scale != 1.0 {
7330 self.scale_inplace(&mut y, scale, m * out_f)?;
7331 }
7332 Ok(y)
7333 }
7334
7335 fn fp4_gemm_launch(
7338 &self,
7339 bytes: &CudaSlice<u8>,
7340 aq4: &CudaSlice<u32>,
7341 ad4: &CudaSlice<u8>,
7342 m: usize,
7343 in_f: usize,
7344 out_f: usize,
7345 row_bytes: usize,
7346 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7347 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7348 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7350 const BN: u32 = 256;
7351 let cfg = LaunchConfig {
7352 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7353 block_dim: (32, 4, 1),
7354 shared_mem_bytes: 0,
7355 };
7356 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7357 let __s_b = self.gpu.stream();
7358 let mut b = __s_b.launch_builder(&f);
7359 b.arg(bytes)
7360 .arg(aq4)
7361 .arg(ad4)
7362 .arg(&mut y)
7363 .arg(&inf)
7364 .arg(&outf)
7365 .arg(&mi)
7366 .arg(&rb);
7367 unsafe {
7368 b.launch(cfg)?;
7369 }
7370 Ok(y)
7371 }
7372
7373 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7375 &self,
7376 bytes: &CudaSlice<u8>,
7377 x: &CudaSlice<f32>,
7378 m: usize,
7379 in_f: usize,
7380 out_f: usize,
7381 row_bytes: usize,
7382 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7383 assert!(
7384 in_f % 64 == 0,
7385 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7386 );
7387 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7388 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7389 }
7390
7391 pub fn qmatvec_q8_0_fast(
7393 &self,
7394 w: &CudaSlice<u8>,
7395 x: &CudaSlice<f32>,
7396 m: usize,
7397 in_f: usize,
7398 out_f: usize,
7399 row_bytes: usize,
7400 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7401 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7402 let f = self.func("qmatvec_q8_0_dp4a");
7403 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7405 grid_dim: (out_f as u32, m as u32, 1),
7406 block_dim: (128, 1, 1),
7407 shared_mem_bytes: 0,
7408 };
7409 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7410 let __s_b = self.gpu.stream();
7411 let mut b = __s_b.launch_builder(&f);
7412 b.arg(w)
7413 .arg(&aq)
7414 .arg(&ad)
7415 .arg(&mut y)
7416 .arg(&inf)
7417 .arg(&outf)
7418 .arg(&mi)
7419 .arg(&rb);
7420 unsafe {
7421 b.launch(cfg)?;
7422 }
7423 Ok(y)
7424 }
7425
7426 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7429 &self,
7430 w: &CudaSlice<u8>,
7431 x: &CudaSlice<f32>,
7432 m: usize,
7433 in_f: usize,
7434 out_f: usize,
7435 row_bytes: usize,
7436 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7437 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7438 let f = self.func("qmatvec_q4_K_dp4a");
7439 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7441 grid_dim: (out_f as u32, m as u32, 1),
7442 block_dim: (128, 1, 1),
7443 shared_mem_bytes: 0,
7444 };
7445 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7446 let __s_b = self.gpu.stream();
7447 let mut b = __s_b.launch_builder(&f);
7448 b.arg(w)
7449 .arg(&aq)
7450 .arg(&ad)
7451 .arg(&mut y)
7452 .arg(&inf)
7453 .arg(&outf)
7454 .arg(&mi)
7455 .arg(&rb);
7456 unsafe {
7457 b.launch(cfg)?;
7458 }
7459 Ok(y)
7460 }
7461
7462 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7465 &self,
7466 w: &CudaSlice<u8>,
7467 x: &CudaSlice<f32>,
7468 m: usize,
7469 in_f: usize,
7470 out_f: usize,
7471 row_bytes: usize,
7472 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7473 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7474 let f = self.func("qmatvec_q6_K_dp4a");
7475 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7477 grid_dim: (out_f as u32, m as u32, 1),
7478 block_dim: (128, 1, 1),
7479 shared_mem_bytes: 0,
7480 };
7481 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7482 let __s_b = self.gpu.stream();
7483 let mut b = __s_b.launch_builder(&f);
7484 b.arg(w)
7485 .arg(&aq)
7486 .arg(&ad)
7487 .arg(&mut y)
7488 .arg(&inf)
7489 .arg(&outf)
7490 .arg(&mi)
7491 .arg(&rb);
7492 unsafe {
7493 b.launch(cfg)?;
7494 }
7495 Ok(y)
7496 }
7497
7498 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7501 &self,
7502 w: &CudaSlice<u8>,
7503 x: &CudaSlice<f32>,
7504 m: usize,
7505 in_f: usize,
7506 out_f: usize,
7507 row_bytes: usize,
7508 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7509 self.qmatvec_dp4a_named(
7510 "qmatvec_q5_K_dp4a",
7511 &w.slice(0..w.len()),
7512 x,
7513 m,
7514 in_f,
7515 out_f,
7516 row_bytes,
7517 )
7518 }
7519 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7522 &self,
7523 w: &CudaSlice<u8>,
7524 x: &CudaSlice<f32>,
7525 m: usize,
7526 in_f: usize,
7527 out_f: usize,
7528 row_bytes: usize,
7529 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7530 self.qmatvec_dp4a_named(
7531 "qmatvec_q3_K_dp4a",
7532 &w.slice(0..w.len()),
7533 x,
7534 m,
7535 in_f,
7536 out_f,
7537 row_bytes,
7538 )
7539 }
7540 pub fn qmatvec_nvfp4_fast_rp(
7542 &self,
7543 w: &CudaSlice<u8>,
7544 x: &CudaSlice<f32>,
7545 m: usize,
7546 in_f: usize,
7547 out_f: usize,
7548 row_bytes: usize,
7549 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7550 assert!(
7551 in_f % 64 == 0,
7552 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7553 );
7554 self.qmatvec_dp4a_named(
7555 "qmatvec_nvfp4_dp4a_rp",
7556 &w.slice(0..w.len()),
7557 x,
7558 m,
7559 in_f,
7560 out_f,
7561 row_bytes,
7562 )
7563 }
7564 pub fn qmatvec_nvfp4_fast(
7566 &self,
7567 w: &cudarc::driver::CudaView<'_, u8>,
7568 x: &CudaSlice<f32>,
7569 m: usize,
7570 in_f: usize,
7571 out_f: usize,
7572 row_bytes: usize,
7573 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7574 assert!(
7577 in_f % 64 == 0,
7578 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7579 );
7580 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7581 }
7582 pub fn qmatvec_nvfp4_fast_v2(
7585 &self,
7586 w: &cudarc::driver::CudaView<'_, u8>,
7587 x: &CudaSlice<f32>,
7588 m: usize,
7589 in_f: usize,
7590 out_f: usize,
7591 row_bytes: usize,
7592 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7593 assert!(
7594 in_f % 64 == 0,
7595 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7596 );
7597 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7598 }
7599 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7602 &self,
7603 w: &CudaSlice<u8>,
7604 x: &CudaSlice<f32>,
7605 m: usize,
7606 in_f: usize,
7607 out_f: usize,
7608 row_bytes: usize,
7609 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7610 self.qmatvec_dp4a_named(
7611 "qmatvec_iq4_XS_dp4a",
7612 &w.slice(0..w.len()),
7613 x,
7614 m,
7615 in_f,
7616 out_f,
7617 row_bytes,
7618 )
7619 }
7620
7621 fn qmatvec_dp4a_named(
7623 &self,
7624 name: &str,
7625 w: &cudarc::driver::CudaView<'_, u8>,
7626 x: &CudaSlice<f32>,
7627 m: usize,
7628 in_f: usize,
7629 out_f: usize,
7630 row_bytes: usize,
7631 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7632 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7633 let f = self.func(name);
7634 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7636 grid_dim: (out_f as u32, m as u32, 1),
7637 block_dim: (128, 1, 1),
7638 shared_mem_bytes: 0,
7639 };
7640 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7641 let __s_b = self.gpu.stream();
7642 let mut b = __s_b.launch_builder(&f);
7643 b.arg(w)
7644 .arg(&aq)
7645 .arg(&ad)
7646 .arg(&mut y)
7647 .arg(&inf)
7648 .arg(&outf)
7649 .arg(&mi)
7650 .arg(&rb);
7651 unsafe {
7652 b.launch(cfg)?;
7653 }
7654 Ok(y)
7655 }
7656
7657 #[allow(clippy::too_many_arguments)]
7663 pub fn qmatvec_nvfp4_fast_prequant_into(
7664 &self,
7665 w: &CudaSlice<u8>,
7666 aq: &CudaSlice<i8>,
7667 ad: &CudaSlice<f32>,
7668 y: &mut CudaSlice<f32>,
7669 m: usize,
7670 in_f: usize,
7671 out_f: usize,
7672 row_bytes: usize,
7673 ) -> Result<(), Box<dyn std::error::Error>> {
7674 assert!(
7675 in_f % 64 == 0,
7676 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7677 );
7678 if y.len() < m * out_f {
7679 return Err(format!(
7680 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7681 y.len()
7682 )
7683 .into());
7684 }
7685 let f = self.func("qmatvec_nvfp4_dp4a");
7686 let cfg = LaunchConfig {
7687 grid_dim: (out_f as u32, m as u32, 1),
7688 block_dim: (128, 1, 1),
7689 shared_mem_bytes: 0,
7690 };
7691 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7692 let __s_b = self.gpu.stream();
7693 let mut b = __s_b.launch_builder(&f);
7694 b.arg(w)
7695 .arg(aq)
7696 .arg(ad)
7697 .arg(y)
7698 .arg(&inf)
7699 .arg(&outf)
7700 .arg(&mi)
7701 .arg(&rb);
7702 unsafe {
7703 b.launch(cfg)?;
7704 }
7705 Ok(())
7706 }
7707
7708 #[allow(clippy::too_many_arguments)]
7711 pub fn matvec_f32_qkv_into(
7712 &self,
7713 wq: &CudaSlice<f32>,
7714 wk: &CudaSlice<f32>,
7715 wv: &CudaSlice<f32>,
7716 wg: &CudaSlice<f32>,
7717 x: &CudaSlice<f32>,
7718 yq: &mut CudaSlice<f32>,
7719 yk: &mut CudaSlice<f32>,
7720 yv: &mut CudaSlice<f32>,
7721 yg: &mut CudaSlice<f32>,
7722 in_f: usize,
7723 out_q: usize,
7724 out_kv: usize,
7725 out_g: usize,
7726 ) -> Result<(), Box<dyn std::error::Error>> {
7727 if in_f % 4 != 0
7728 || wq.len() != out_q * in_f
7729 || wk.len() != out_kv * in_f
7730 || wv.len() != out_kv * in_f
7731 || wg.len() < out_g * in_f
7732 || x.len() < in_f
7733 || yq.len() < out_q
7734 || yk.len() < out_kv
7735 || yv.len() < out_kv
7736 || (out_g > 0 && yg.len() < out_g)
7737 {
7738 return Err(format!(
7739 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7740 wq={} wk={} wv={} wg={}",
7741 wq.len(),
7742 wk.len(),
7743 wv.len(),
7744 wg.len()
7745 )
7746 .into());
7747 }
7748 let f = self.func("matvec_f32_qkv");
7749 let cfg = LaunchConfig {
7750 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7751 block_dim: (128, 1, 1),
7752 shared_mem_bytes: 0,
7753 };
7754 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7755 let __s_b = self.gpu.stream();
7756 let mut b = __s_b.launch_builder(&f);
7757 b.arg(wq)
7758 .arg(wk)
7759 .arg(wv)
7760 .arg(wg)
7761 .arg(x)
7762 .arg(yq)
7763 .arg(yk)
7764 .arg(yv)
7765 .arg(yg)
7766 .arg(&inf)
7767 .arg(&oq)
7768 .arg(&okv)
7769 .arg(&og);
7770 unsafe {
7771 b.launch(cfg)?;
7772 }
7773 Ok(())
7774 }
7775
7776 #[allow(clippy::too_many_arguments)]
7782 #[allow(clippy::too_many_arguments)]
7785 pub fn qmatvec_nvfp4_sel_gu_into(
7786 &self,
7787 gate_bank: &CudaSlice<u8>,
7788 up_bank: &CudaSlice<u8>,
7789 sel: &CudaSlice<i32>,
7790 aq: &CudaSlice<i8>,
7791 ad: &CudaSlice<f32>,
7792 yg: &mut CudaSlice<f32>,
7793 yu: &mut CudaSlice<f32>,
7794 n_sel: usize,
7795 in_f: usize,
7796 out_f: usize,
7797 row_bytes: usize,
7798 expert_stride: usize,
7799 ) -> Result<(), Box<dyn std::error::Error>> {
7800 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7801 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7802 return Err("NVFP4 gu sel geometry".into());
7803 }
7804 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7807 let rpw = *RPW.get_or_init(|| {
7808 std::env::var("MEMRA_SEL_GU_RPW")
7809 .ok()
7810 .and_then(|v| v.parse().ok())
7811 .filter(|r| *r == 2 || *r == 4)
7812 .unwrap_or(1)
7813 });
7814 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
7815 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7818 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
7819 let f = self.func(match (wpr, rpw) {
7820 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
7821 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
7822 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
7823 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
7824 });
7825 let cfg = LaunchConfig {
7826 grid_dim: if wpr {
7827 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
7828 } else if rpw == 1 {
7829 ((2 * out_f) as u32, n_sel as u32, 1)
7830 } else {
7831 ((out_f / rpw) as u32, n_sel as u32, 1)
7832 },
7833 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
7834 shared_mem_bytes: 0,
7835 };
7836 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7837 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7838 let (ars, adrs) = (0i64, 0i64);
7839 let __s_b = self.gpu.stream();
7840 let mut b = __s_b.launch_builder(&f);
7841 b.arg(gate_bank)
7842 .arg(up_bank)
7843 .arg(sel)
7844 .arg(aq)
7845 .arg(ad)
7846 .arg(yg)
7847 .arg(yu)
7848 .arg(&inf)
7849 .arg(&outf)
7850 .arg(&ns)
7851 .arg(&rb)
7852 .arg(&es)
7853 .arg(&ars)
7854 .arg(&adrs);
7855 unsafe {
7856 b.launch(cfg)?;
7857 }
7858 Ok(())
7859 }
7860
7861 #[allow(clippy::too_many_arguments)]
7869 pub fn qmatvec_nvfp4_sel_down8_into(
7870 &self,
7871 bank: &CudaSlice<u8>,
7872 sel: &CudaSlice<i32>,
7873 aq: &CudaSlice<i8>,
7874 ad: &CudaSlice<f32>,
7875 route_w: &CudaSlice<f32>,
7876 md: &CudaSlice<f32>,
7877 dst: &mut CudaSlice<f32>,
7878 n_sel: usize,
7879 in_f: usize,
7880 out_f: usize,
7881 row_bytes: usize,
7882 expert_stride: usize,
7883 act_row_stride: usize,
7884 ad_row_stride: usize,
7885 ) -> Result<(), Box<dyn std::error::Error>> {
7886 if in_f % 64 != 0
7887 || n_sel == 0
7888 || n_sel > 8
7889 || (in_f >> 5) > 32
7890 || dst.len() < out_f
7891 || sel.len() < n_sel
7892 || route_w.len() < n_sel
7893 {
7894 return Err(format!(
7895 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
7896 dst.len()
7897 )
7898 .into());
7899 }
7900 if !crate::tp::nvfp4_bank_v2_on() {
7901 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
7902 }
7903 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
7904 let cfg = LaunchConfig {
7905 grid_dim: (out_f as u32, 1, 1),
7906 block_dim: (32, n_sel as u32, 1),
7907 shared_mem_bytes: 0,
7908 };
7909 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7910 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7911 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7912 let __s_b = self.gpu.stream();
7913 let mut b = __s_b.launch_builder(&f);
7914 b.arg(bank)
7915 .arg(sel)
7916 .arg(aq)
7917 .arg(ad)
7918 .arg(route_w)
7919 .arg(md)
7920 .arg(dst)
7921 .arg(&inf)
7922 .arg(&outf)
7923 .arg(&ns)
7924 .arg(&rb)
7925 .arg(&es)
7926 .arg(&ars)
7927 .arg(&adrs);
7928 unsafe {
7929 b.launch(cfg)?;
7930 }
7931 Ok(())
7932 }
7933
7934 #[allow(clippy::too_many_arguments)]
7938 pub fn qmatvec_nvfp4_sel_down8_rows_into(
7939 &self,
7940 bank: &CudaSlice<u8>,
7941 sel: &CudaSlice<i32>,
7942 aq: &CudaSlice<i8>,
7943 ad: &CudaSlice<f32>,
7944 route_w: &CudaSlice<f32>,
7945 md: &CudaSlice<f32>,
7946 dst: &mut CudaSlice<f32>,
7947 t: usize,
7948 n_sel_col: usize,
7949 in_f: usize,
7950 out_f: usize,
7951 row_bytes: usize,
7952 expert_stride: usize,
7953 act_row_stride: usize,
7954 ad_row_stride: usize,
7955 ) -> Result<(), Box<dyn std::error::Error>> {
7956 let n_sel = t * n_sel_col;
7957 if in_f % 64 != 0
7958 || n_sel_col == 0
7959 || n_sel_col > 8
7960 || t == 0
7961 || t > 64
7962 || (in_f >> 5) > 32
7963 || dst.len() < t * out_f
7964 || sel.len() < n_sel
7965 || route_w.len() < n_sel
7966 {
7967 return Err(format!(
7968 "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
7969 dst.len()
7970 )
7971 .into());
7972 }
7973 if !crate::tp::nvfp4_bank_v2_on() {
7974 return Err(
7975 "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
7976 );
7977 }
7978 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
7979 let cfg = LaunchConfig {
7980 grid_dim: (out_f as u32, t as u32, 1),
7981 block_dim: (32, n_sel_col as u32, 1),
7982 shared_mem_bytes: 0,
7983 };
7984 let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
7985 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7986 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7987 let __s_b = self.gpu.stream();
7988 let mut b = __s_b.launch_builder(&f);
7989 b.arg(bank)
7990 .arg(sel)
7991 .arg(aq)
7992 .arg(ad)
7993 .arg(route_w)
7994 .arg(md)
7995 .arg(dst)
7996 .arg(&inf)
7997 .arg(&outf)
7998 .arg(&nsc)
7999 .arg(&rb)
8000 .arg(&es)
8001 .arg(&ars)
8002 .arg(&adrs);
8003 unsafe {
8004 b.launch(cfg)?;
8005 }
8006 Ok(())
8007 }
8008
8009 #[allow(clippy::too_many_arguments)]
8012 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8013 &self,
8014 gate_bank: &CudaSlice<u8>,
8015 up_bank: &CudaSlice<u8>,
8016 sel: &CudaSlice<i32>,
8017 aq: &CudaSlice<i8>,
8018 ad: &CudaSlice<f32>,
8019 yg: &mut CudaSlice<f32>,
8020 yu: &mut CudaSlice<f32>,
8021 n_sel: usize,
8022 in_f: usize,
8023 out_f: usize,
8024 row_bytes: usize,
8025 expert_stride: usize,
8026 owner: usize,
8027 ) -> Result<(), Box<dyn std::error::Error>> {
8028 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8029 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8030 return Err("NVFP4 gu ep geometry".into());
8031 }
8032 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8033 let cfg = LaunchConfig {
8034 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8035 block_dim: (128, 1, 1),
8036 shared_mem_bytes: 0,
8037 };
8038 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8039 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8040 let (ars, adrs) = (0i64, 0i64);
8041 let __s_b = self.gpu.stream();
8042 let mut b = __s_b.launch_builder(&f);
8043 b.arg(gate_bank)
8044 .arg(up_bank)
8045 .arg(sel)
8046 .arg(aq)
8047 .arg(ad)
8048 .arg(yg)
8049 .arg(yu)
8050 .arg(&inf)
8051 .arg(&outf)
8052 .arg(&ns)
8053 .arg(&rb)
8054 .arg(&es)
8055 .arg(&ars)
8056 .arg(&adrs)
8057 .arg(&own);
8058 unsafe {
8059 b.launch(cfg)?;
8060 }
8061 Ok(())
8062 }
8063
8064 #[allow(clippy::too_many_arguments)]
8066 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8067 &self,
8068 gate: &CudaSlice<f32>,
8069 up: &CudaSlice<f32>,
8070 gmac: &CudaSlice<f32>,
8071 umac: &CudaSlice<f32>,
8072 sel: &CudaSlice<i32>,
8073 limit: Option<f32>,
8074 out_q: &mut CudaSlice<i8>,
8075 out_d: &mut CudaSlice<f32>,
8076 n_per: usize,
8077 n_sel: usize,
8078 owner: usize,
8079 ) -> Result<(), Box<dyn std::error::Error>> {
8080 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8081 return Err("NVFP4 silu ep geometry".into());
8082 }
8083 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8084 let warps = n_sel * n_per / 32;
8085 let cfg = LaunchConfig {
8086 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8087 block_dim: (128, 1, 1),
8088 shared_mem_bytes: 0,
8089 };
8090 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8091 let (lim, has) = match limit {
8092 Some(l) => (l, 1i32),
8093 None => (0.0f32, 0i32),
8094 };
8095 let __s_b = self.gpu.stream();
8096 let mut b = __s_b.launch_builder(&f);
8097 b.arg(gate)
8098 .arg(up)
8099 .arg(gmac)
8100 .arg(umac)
8101 .arg(sel)
8102 .arg(&lim)
8103 .arg(&has)
8104 .arg(out_q)
8105 .arg(out_d)
8106 .arg(&np)
8107 .arg(&ns)
8108 .arg(&own);
8109 unsafe {
8110 b.launch(cfg)?;
8111 }
8112 Ok(())
8113 }
8114
8115 #[allow(clippy::too_many_arguments)]
8117 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8118 &self,
8119 bank: &CudaSlice<u8>,
8120 sel: &CudaSlice<i32>,
8121 aq: &CudaSlice<i8>,
8122 ad: &CudaSlice<f32>,
8123 route_w: &CudaSlice<f32>,
8124 md: &CudaSlice<f32>,
8125 dst: &mut CudaSlice<f32>,
8126 n_sel: usize,
8127 in_f: usize,
8128 out_f: usize,
8129 row_bytes: usize,
8130 expert_stride: usize,
8131 act_row_stride: usize,
8132 ad_row_stride: usize,
8133 owner: usize,
8134 ) -> Result<(), Box<dyn std::error::Error>> {
8135 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8136 return Err("NVFP4 down8 ep geometry".into());
8137 }
8138 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8139 let cfg = LaunchConfig {
8140 grid_dim: (out_f as u32, 1, 1),
8141 block_dim: (32, n_sel as u32, 1),
8142 shared_mem_bytes: 0,
8143 };
8144 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8145 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8146 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8147 let __s_b = self.gpu.stream();
8148 let mut b = __s_b.launch_builder(&f);
8149 b.arg(bank)
8150 .arg(sel)
8151 .arg(aq)
8152 .arg(ad)
8153 .arg(route_w)
8154 .arg(md)
8155 .arg(dst)
8156 .arg(&inf)
8157 .arg(&outf)
8158 .arg(&ns)
8159 .arg(&rb)
8160 .arg(&es)
8161 .arg(&ars)
8162 .arg(&adrs)
8163 .arg(&own);
8164 unsafe {
8165 b.launch(cfg)?;
8166 }
8167 Ok(())
8168 }
8169
8170 pub fn qmatvec_nvfp4_sel_into(
8171 &self,
8172 bank: &CudaSlice<u8>,
8173 sel: &CudaSlice<i32>,
8174 aq: &CudaSlice<i8>,
8175 ad: &CudaSlice<f32>,
8176 y: &mut CudaSlice<f32>,
8177 n_sel: usize,
8178 in_f: usize,
8179 out_f: usize,
8180 row_bytes: usize,
8181 expert_stride: usize,
8182 act_row_stride: usize,
8183 ad_row_stride: usize,
8184 ) -> Result<(), Box<dyn std::error::Error>> {
8185 assert!(
8186 in_f % 64 == 0,
8187 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8188 );
8189 if y.len() < n_sel * out_f || sel.len() < n_sel {
8190 return Err(format!(
8191 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8192 y.len(),
8193 sel.len()
8194 )
8195 .into());
8196 }
8197 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8204 let mode = *MR.get_or_init(|| {
8205 if crate::tp::nvfp4_bank_v2_on() {
8206 3
8207 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8208 2
8209 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8210 1
8211 } else {
8212 0
8213 }
8214 });
8215 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8216 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8220 let v2s = mode == 3
8221 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8222 && row_bytes % 16 == 0
8223 && in_f <= 4096;
8224 let f = match (mode, v2s) {
8225 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8226 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8227 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8228 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8229 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8230 };
8231 let nsb = in_f >> 5;
8236 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8237 32
8238 } else if mode == 1 {
8239 512
8240 } else {
8241 128
8242 };
8243 let cfg = LaunchConfig {
8244 grid_dim: (
8245 if v2s {
8246 (out_f as u32).div_ceil(8)
8247 } else {
8248 match mode {
8249 2 => (out_f as u32).div_ceil(16),
8250 1 => (out_f as u32).div_ceil(4),
8251 _ => out_f as u32,
8252 }
8253 },
8254 n_sel as u32,
8255 1,
8256 ),
8257 block_dim: (fit_block, 1, 1),
8258 shared_mem_bytes: 0,
8259 };
8260 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8261 let (rb, es, ars, adrs) = (
8262 row_bytes as i64,
8263 expert_stride as i64,
8264 act_row_stride as i64,
8265 ad_row_stride as i64,
8266 );
8267 let __s_b = self.gpu.stream();
8268 let mut b = __s_b.launch_builder(&f);
8269 b.arg(bank)
8270 .arg(sel)
8271 .arg(aq)
8272 .arg(ad)
8273 .arg(y)
8274 .arg(&inf)
8275 .arg(&outf)
8276 .arg(&ns)
8277 .arg(&rb)
8278 .arg(&es)
8279 .arg(&ars)
8280 .arg(&adrs);
8281 unsafe {
8282 b.launch(cfg)?;
8283 }
8284 Ok(())
8285 }
8286
8287 #[allow(clippy::too_many_arguments)]
8292 pub fn silu_mul_scaled_q8_1_sel_into(
8293 &self,
8294 gate: &CudaSlice<f32>,
8295 up: &CudaSlice<f32>,
8296 gmac: &CudaSlice<f32>,
8297 umac: &CudaSlice<f32>,
8298 sel: &CudaSlice<i32>,
8299 limit: Option<f32>,
8300 out_q: &mut CudaSlice<i8>,
8301 out_d: &mut CudaSlice<f32>,
8302 n_per: usize,
8303 n_sel: usize,
8304 ) -> Result<(), Box<dyn std::error::Error>> {
8305 let n = n_per * n_sel;
8306 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8307 return Err(format!(
8308 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8309 out_q.len(),
8310 out_d.len()
8311 )
8312 .into());
8313 }
8314 if let Some(limit) = limit {
8315 if limit <= 1e-6 {
8316 return Err(format!(
8317 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8318 )
8319 .into());
8320 }
8321 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8322 let cfg = LaunchConfig::for_num_elems(n as u32);
8323 let (np, ns) = (n_per as i32, n_sel as i32);
8324 let __s_b = self.gpu.stream();
8325 let mut b = __s_b.launch_builder(&f);
8326 b.arg(gate)
8327 .arg(up)
8328 .arg(gmac)
8329 .arg(umac)
8330 .arg(sel)
8331 .arg(&limit)
8332 .arg(out_q)
8333 .arg(out_d)
8334 .arg(&np)
8335 .arg(&ns);
8336 unsafe {
8337 b.launch(cfg)?;
8338 }
8339 return Ok(());
8340 }
8341 let f = self.func("silu_mul_scaled_q8_1_sel");
8342 let cfg = LaunchConfig::for_num_elems(n as u32);
8343 let (np, ns) = (n_per as i32, n_sel as i32);
8344 let __s_b = self.gpu.stream();
8345 let mut b = __s_b.launch_builder(&f);
8346 b.arg(gate)
8347 .arg(up)
8348 .arg(gmac)
8349 .arg(umac)
8350 .arg(sel)
8351 .arg(out_q)
8352 .arg(out_d)
8353 .arg(&np)
8354 .arg(&ns);
8355 unsafe {
8356 b.launch(cfg)?;
8357 }
8358 Ok(())
8359 }
8360
8361 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8362 Ok(self.gpu.stream().clone_htod(v)?)
8363 }
8364 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8365 Ok(self.gpu.stream().clone_htod(v)?)
8366 }
8367 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8369 Ok(self.gpu.stream().clone_htod(v)?)
8370 }
8371 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8372 Ok(self.gpu.stream().clone_htod(v)?)
8373 }
8374 pub fn dtoh_view(
8376 &self,
8377 d: &cudarc::driver::CudaView<f32>,
8378 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8379 let v = self.gpu.stream().clone_dtoh(d)?;
8380 self.gpu.stream().synchronize()?;
8381 Ok(v)
8382 }
8383 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8384 let v = self.gpu.stream().clone_dtoh(d)?;
8385 self.gpu.stream().synchronize()?;
8386 Ok(v)
8387 }
8388 pub fn dtoh_pair(
8392 &self,
8393 a: &CudaSlice<f32>,
8394 b: &CudaSlice<f32>,
8395 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8396 let av = self.gpu.stream().clone_dtoh(a)?;
8397 let bv = self.gpu.stream().clone_dtoh(b)?;
8398 self.gpu.stream().synchronize()?;
8399 Ok((av, bv))
8400 }
8401 pub fn dtoh_pair_views(
8404 &self,
8405 a: &cudarc::driver::CudaView<f32>,
8406 b: &cudarc::driver::CudaView<f32>,
8407 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8408 let av = self.gpu.stream().clone_dtoh(a)?;
8409 let bv = self.gpu.stream().clone_dtoh(b)?;
8410 self.gpu.stream().synchronize()?;
8411 Ok((av, bv))
8412 }
8413 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8415 let v = self.gpu.stream().clone_dtoh(d)?;
8416 self.gpu.stream().synchronize()?;
8417 Ok(v)
8418 }
8419 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8421 let v = self.gpu.stream().clone_dtoh(d)?;
8422 self.gpu.stream().synchronize()?;
8423 Ok(v)
8424 }
8425 pub fn dtoh_u8_view(
8426 &self,
8427 d: &cudarc::driver::CudaView<u8>,
8428 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8429 let v = self.gpu.stream().clone_dtoh(d)?;
8430 self.gpu.stream().synchronize()?;
8431 Ok(v)
8432 }
8433 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8434 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8435 self.keep_if_capturing(&s);
8436 Ok(s)
8437 }
8438
8439 pub fn prob_of_token_device(
8448 &self,
8449 logits: &CudaSlice<f32>,
8450 tok: &CudaSlice<u32>,
8451 n_vocab: usize,
8452 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8453 let nb = ARGMAX_NB;
8454 let mut part = self.alloc_uninit::<f32>(nb)?;
8455 let mut p = self.alloc_uninit::<f32>(1)?;
8456 let f1 = self.func("prob_of_token_partial_f32");
8457 let cfg1 = LaunchConfig {
8458 grid_dim: (nb as u32, 1, 1),
8459 block_dim: (256, 1, 1),
8460 shared_mem_bytes: 0,
8461 };
8462 let nv = n_vocab as i32;
8463 let __s_b1 = self.gpu.stream();
8464 let mut b1 = __s_b1.launch_builder(&f1);
8465 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8466 unsafe {
8467 b1.launch(cfg1)?;
8468 }
8469 let f2 = self.func("prob_of_token_final_f32");
8470 let cfg2 = LaunchConfig {
8471 grid_dim: (1, 1, 1),
8472 block_dim: (256, 1, 1),
8473 shared_mem_bytes: 0,
8474 };
8475 let nbi = nb as i32;
8476 let __s_b2 = self.gpu.stream();
8477 let mut b2 = __s_b2.launch_builder(&f2);
8478 b2.arg(&part).arg(&mut p).arg(&nbi);
8479 unsafe {
8480 b2.launch(cfg2)?;
8481 }
8482 Ok(p)
8483 }
8484
8485 pub fn prob_of_token_device_col(
8492 &self,
8493 logits: &CudaSlice<f32>,
8494 tok_all: &CudaSlice<u32>,
8495 tok_idx: usize,
8496 p_out: &mut CudaSlice<f32>,
8497 p_idx: usize,
8498 n_vocab: usize,
8499 ) -> Result<(), Box<dyn std::error::Error>> {
8500 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8501 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8502 let nb = ARGMAX_NB;
8503 let mut part = self.alloc_uninit::<f32>(nb)?;
8504 let f1 = self.func("prob_of_token_partial_f32");
8505 let cfg1 = LaunchConfig {
8506 grid_dim: (nb as u32, 1, 1),
8507 block_dim: (256, 1, 1),
8508 shared_mem_bytes: 0,
8509 };
8510 let nv = n_vocab as i32;
8511 let __s_b1 = self.gpu.stream();
8512 let mut b1 = __s_b1.launch_builder(&f1);
8513 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8514 unsafe {
8515 b1.launch(cfg1)?;
8516 }
8517 let f2 = self.func("prob_of_token_final_f32");
8518 let cfg2 = LaunchConfig {
8519 grid_dim: (1, 1, 1),
8520 block_dim: (256, 1, 1),
8521 shared_mem_bytes: 0,
8522 };
8523 let nbi = nb as i32;
8524 let __s_b2 = self.gpu.stream();
8525 let mut b2 = __s_b2.launch_builder(&f2);
8526 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8527 unsafe {
8528 b2.launch(cfg2)?;
8529 }
8530 Ok(())
8531 }
8532
8533 pub fn prob_of_token_device_into(
8534 &self,
8535 logits: &CudaSlice<f32>,
8536 tok: &CudaSlice<u32>,
8537 p_out: &mut CudaSlice<f32>,
8538 n_vocab: usize,
8539 ) -> Result<(), Box<dyn std::error::Error>> {
8540 let nb = ARGMAX_NB;
8541 let mut part = self.alloc_uninit::<f32>(nb)?;
8542 let f1 = self.func("prob_of_token_partial_f32");
8543 let cfg1 = LaunchConfig {
8544 grid_dim: (nb as u32, 1, 1),
8545 block_dim: (256, 1, 1),
8546 shared_mem_bytes: 0,
8547 };
8548 let nv = n_vocab as i32;
8549 let __s_b1 = self.gpu.stream();
8550 let mut b1 = __s_b1.launch_builder(&f1);
8551 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8552 unsafe {
8553 b1.launch(cfg1)?;
8554 }
8555 let f2 = self.func("prob_of_token_final_f32");
8556 let cfg2 = LaunchConfig {
8557 grid_dim: (1, 1, 1),
8558 block_dim: (256, 1, 1),
8559 shared_mem_bytes: 0,
8560 };
8561 let nbi = nb as i32;
8562 let __s_b2 = self.gpu.stream();
8563 let mut b2 = __s_b2.launch_builder(&f2);
8564 b2.arg(&part).arg(p_out).arg(&nbi);
8565 unsafe {
8566 b2.launch(cfg2)?;
8567 }
8568 Ok(())
8569 }
8570
8571 pub fn u32_hist_append(
8574 &self,
8575 tok: &CudaSlice<u32>,
8576 hist: &mut CudaSlice<u32>,
8577 idx: &mut CudaSlice<i32>,
8578 ) -> Result<(), Box<dyn std::error::Error>> {
8579 let f = self.func("u32_hist_append");
8580 let cfg = LaunchConfig {
8581 grid_dim: (1, 1, 1),
8582 block_dim: (32, 1, 1),
8583 shared_mem_bytes: 0,
8584 };
8585 let __s_b = self.gpu.stream();
8586 let mut b = __s_b.launch_builder(&f);
8587 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8588 unsafe {
8589 b.launch(cfg)?;
8590 }
8591 Ok(())
8592 }
8593
8594 pub fn argmax_token_device(
8595 &self,
8596 logits: &CudaSlice<f32>,
8597 n_vocab: usize,
8598 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8599 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8600 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8601 Ok(tok)
8602 }
8603 pub fn argmax_token_device_into(
8610 &self,
8611 logits: &CudaSlice<f32>,
8612 tok: &mut CudaSlice<u32>,
8613 n_vocab: usize,
8614 ) -> Result<(), Box<dyn std::error::Error>> {
8615 let nb = ARGMAX_NB;
8616 let f1 = self.func("argmax_partial_f32");
8617 let f2 = self.func("argmax_final_f32");
8618 let mut guard = self.argmax_partials.lock().unwrap();
8619 if guard.is_none() {
8620 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8623 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8624 *guard = Some((pv, pi));
8625 }
8626 let (part_v, part_i) = guard.as_mut().unwrap();
8627 let nv = n_vocab as i32;
8628 let nbi = nb as i32;
8629 let cfg1 = LaunchConfig {
8631 grid_dim: (nb as u32, 1, 1),
8632 block_dim: (256, 1, 1),
8633 shared_mem_bytes: 0,
8634 };
8635 let __s_b1 = self.gpu.stream();
8636 let mut b1 = __s_b1.launch_builder(&f1);
8637 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8638 unsafe {
8639 b1.launch(cfg1)?;
8640 }
8641 let cfg2 = LaunchConfig {
8643 grid_dim: (1, 1, 1),
8644 block_dim: (256, 1, 1),
8645 shared_mem_bytes: 0,
8646 };
8647 let __s_b2 = self.gpu.stream();
8648 let mut b2 = __s_b2.launch_builder(&f2);
8649 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8650 unsafe {
8651 b2.launch(cfg2)?;
8652 }
8653 Ok(())
8654 }
8655 pub fn argmax_token_device_col(
8661 &self,
8662 logits: &CudaSlice<f32>,
8663 col: usize,
8664 n_vocab: usize,
8665 toks: &mut CudaSlice<u32>,
8666 out_idx: usize,
8667 ) -> Result<(), Box<dyn std::error::Error>> {
8668 let nb = ARGMAX_NB;
8669 let f1 = self.func("argmax_partial_f32");
8670 let f2 = self.func("argmax_final_f32");
8671 let mut guard = self.argmax_partials.lock().unwrap();
8672 if guard.is_none() {
8673 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8674 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8675 *guard = Some((pv, pi));
8676 }
8677 let (part_v, part_i) = guard.as_mut().unwrap();
8678 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8679 let nv = n_vocab as i32;
8680 let nbi = nb as i32;
8681 let cfg1 = LaunchConfig {
8682 grid_dim: (nb as u32, 1, 1),
8683 block_dim: (256, 1, 1),
8684 shared_mem_bytes: 0,
8685 };
8686 let __s_b1 = self.gpu.stream();
8687 let mut b1 = __s_b1.launch_builder(&f1);
8688 b1.arg(&col_view)
8689 .arg(&mut *part_v)
8690 .arg(&mut *part_i)
8691 .arg(&nv);
8692 unsafe {
8693 b1.launch(cfg1)?;
8694 }
8695 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8696 let cfg2 = LaunchConfig {
8697 grid_dim: (1, 1, 1),
8698 block_dim: (256, 1, 1),
8699 shared_mem_bytes: 0,
8700 };
8701 let __s_b2 = self.gpu.stream();
8702 let mut b2 = __s_b2.launch_builder(&f2);
8703 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8704 unsafe {
8705 b2.launch(cfg2)?;
8706 }
8707 Ok(())
8708 }
8709 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8711 Ok(self.gpu.stream().clone_htod(v)?)
8712 }
8713 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8714 let v = self.gpu.stream().clone_dtoh(d)?;
8715 self.gpu.stream().synchronize()?;
8716 Ok(v)
8717 }
8718
8719 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8720 let v = self.gpu.stream().clone_dtoh(d)?;
8721 self.gpu.stream().synchronize()?;
8722 Ok(v)
8723 }
8724 pub fn htod_u32_into(
8728 &self,
8729 dst: &mut CudaSlice<u32>,
8730 src: &[u32],
8731 ) -> Result<(), Box<dyn std::error::Error>> {
8732 let mut view = dst.slice_mut(0..src.len());
8733 self.gpu.stream().memcpy_htod(src, &mut view)?;
8734 Ok(())
8735 }
8736
8737 pub fn htod_i32_into(
8740 &self,
8741 dst: &mut CudaSlice<i32>,
8742 src: &[i32],
8743 ) -> Result<(), Box<dyn std::error::Error>> {
8744 let mut view = dst.slice_mut(0..src.len());
8745 self.gpu.stream().memcpy_htod(src, &mut view)?;
8746 Ok(())
8747 }
8748
8749 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8750 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8751 self.keep_if_capturing(&s);
8752 Ok(s)
8753 }
8754 pub fn embed_gather_device_into(
8757 &self,
8758 embd: &CudaSlice<u8>,
8759 token_d: &CudaSlice<u32>,
8760 x_out: &mut CudaSlice<f32>,
8761 n_embd: usize,
8762 qtype: i32,
8763 row_bytes: usize,
8764 ) -> Result<(), Box<dyn std::error::Error>> {
8765 let f = self.func("embed_gather_u32");
8766 let cfg = LaunchConfig {
8767 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8768 block_dim: (256, 1, 1),
8769 shared_mem_bytes: 0,
8770 };
8771 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8772 let __s_b = self.gpu.stream();
8773 let mut b = __s_b.launch_builder(&f);
8774 b.arg(embd)
8775 .arg(token_d)
8776 .arg(x_out)
8777 .arg(&ne)
8778 .arg(&qt)
8779 .arg(&rb);
8780 unsafe {
8781 b.launch(cfg)?;
8782 }
8783 Ok(())
8784 }
8785 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8787 let v = self.gpu.stream().clone_dtoh(d)?;
8788 self.gpu.stream().synchronize()?;
8789 Ok(v[0])
8790 }
8791 pub fn i32_set_k(
8798 &self,
8799 dst: &mut CudaSlice<i32>,
8800 v: i32,
8801 ) -> Result<(), Box<dyn std::error::Error>> {
8802 let f = self.func("i32_set_k");
8803 let cfg = LaunchConfig {
8804 grid_dim: (1, 1, 1),
8805 block_dim: (1, 1, 1),
8806 shared_mem_bytes: 0,
8807 };
8808 let idx = 0i32;
8809 let __s_b = self.gpu.stream();
8810 let mut b = __s_b.launch_builder(&f);
8811 b.arg(dst).arg(&v).arg(&idx);
8812 unsafe {
8813 b.launch(cfg)?;
8814 }
8815 Ok(())
8816 }
8817
8818 pub fn set_i32_one(
8819 &self,
8820 d: &mut CudaSlice<i32>,
8821 v: i32,
8822 ) -> Result<(), Box<dyn std::error::Error>> {
8823 self.gpu.stream().memcpy_htod(&[v], d)?;
8824 Ok(())
8825 }
8826 pub fn set_u32_one(
8829 &self,
8830 d: &mut CudaSlice<u32>,
8831 v: u32,
8832 ) -> Result<(), Box<dyn std::error::Error>> {
8833 self.gpu.stream().memcpy_htod(&[v], d)?;
8834 Ok(())
8835 }
8836 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8838 let v = self.gpu.stream().clone_dtoh(d)?;
8839 self.gpu.stream().synchronize()?;
8840 Ok(v[0])
8841 }
8842 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8844 Ok(self.gpu.stream().clone_htod(bytes)?)
8845 }
8846 pub fn embed_gather_device(
8850 &self,
8851 embd: &CudaSlice<u8>,
8852 token_d: &CudaSlice<u32>,
8853 n_embd: usize,
8854 qtype: i32,
8855 row_bytes: usize,
8856 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8857 let f = self.func("embed_gather_u32");
8858 let mut x = self.alloc_uninit::<f32>(n_embd)?;
8859 let cfg = LaunchConfig {
8860 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8861 block_dim: (256, 1, 1),
8862 shared_mem_bytes: 0,
8863 };
8864 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8865 let __s_b = self.gpu.stream();
8866 let mut b = __s_b.launch_builder(&f);
8867 b.arg(embd)
8868 .arg(token_d)
8869 .arg(&mut x)
8870 .arg(&ne)
8871 .arg(&qt)
8872 .arg(&rb);
8873 unsafe {
8874 b.launch(cfg)?;
8875 }
8876 Ok(x)
8877 }
8878
8879 pub fn embed_gather_device_t(
8883 &self,
8884 embd: &CudaSlice<u8>,
8885 tokens: &[u32],
8886 n_embd: usize,
8887 qtype: i32,
8888 row_bytes: usize,
8889 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8890 let t = tokens.len();
8891 let tok_d = self.gpu.stream().clone_htod(tokens)?;
8892 let f = self.func("embed_gather_u32_t");
8893 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8894 let cfg = LaunchConfig {
8895 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8896 block_dim: (256, 1, 1),
8897 shared_mem_bytes: 0,
8898 };
8899 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8900 let __s_b = self.gpu.stream();
8901 let mut b = __s_b.launch_builder(&f);
8902 b.arg(embd)
8903 .arg(&tok_d)
8904 .arg(&mut x)
8905 .arg(&ne)
8906 .arg(&qt)
8907 .arg(&rb)
8908 .arg(&ti);
8909 unsafe {
8910 b.launch(cfg)?;
8911 }
8912 Ok(x)
8913 }
8914
8915 pub fn embed_gather_device_tv(
8920 &self,
8921 embd: &CudaSlice<u8>,
8922 tok_v: &cudarc::driver::CudaView<u32>,
8923 t: usize,
8924 n_embd: usize,
8925 qtype: i32,
8926 row_bytes: usize,
8927 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8928 let f = self.func("embed_gather_u32_t");
8929 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8930 let cfg = LaunchConfig {
8931 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8932 block_dim: (256, 1, 1),
8933 shared_mem_bytes: 0,
8934 };
8935 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8936 let __s_b = self.gpu.stream();
8937 let mut b = __s_b.launch_builder(&f);
8938 b.arg(embd)
8939 .arg(tok_v)
8940 .arg(&mut x)
8941 .arg(&ne)
8942 .arg(&qt)
8943 .arg(&rb)
8944 .arg(&ti);
8945 unsafe {
8946 b.launch(cfg)?;
8947 }
8948 Ok(x)
8949 }
8950
8951 pub fn embed_gather_device_td(
8952 &self,
8953 embd: &CudaSlice<u8>,
8954 tok_d: &CudaSlice<u32>,
8955 t: usize,
8956 n_embd: usize,
8957 qtype: i32,
8958 row_bytes: usize,
8959 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8960 let f = self.func("embed_gather_u32_t");
8961 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8962 let cfg = LaunchConfig {
8963 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8964 block_dim: (256, 1, 1),
8965 shared_mem_bytes: 0,
8966 };
8967 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8968 let __s_b = self.gpu.stream();
8969 let mut b = __s_b.launch_builder(&f);
8970 b.arg(embd)
8971 .arg(tok_d)
8972 .arg(&mut x)
8973 .arg(&ne)
8974 .arg(&qt)
8975 .arg(&rb)
8976 .arg(&ti);
8977 unsafe {
8978 b.launch(cfg)?;
8979 }
8980 Ok(x)
8981 }
8982
8983 #[inline]
8989 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8991 if self
8992 .capture_keep_on
8993 .load(std::sync::atomic::Ordering::Relaxed)
8994 {
8995 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8996 }
8997 }
8998
8999 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9000 &self,
9001 n: usize,
9002 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9003 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9004 {
9008 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9009 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9010 use cudarc::driver::DevicePtrMut;
9012 let n_bytes = s.len() * std::mem::size_of::<T>();
9013 let stream = self.gpu.stream();
9014 let (p_, _g) = s.device_ptr_mut(&stream);
9015 unsafe {
9016 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9017 .result()?;
9018 }
9019 }
9020 }
9021 self.keep_if_capturing(&s);
9022 Ok(s)
9023 }
9024
9025 pub fn uninit_q8_pair(
9030 &self,
9031 n: usize,
9032 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9033 Ok((
9034 self.alloc_uninit::<i8>(n)?,
9035 self.alloc_uninit::<f32>(n / 32)?,
9036 ))
9037 }
9038
9039 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9040 self.alloc_uninit::<f32>(n)
9041 }
9042
9043 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9045 self.alloc_uninit::<i8>(n)
9046 }
9047
9048 #[allow(clippy::too_many_arguments)]
9052 pub fn rms_norm3(
9053 &self,
9054 x: &CudaSlice<f32>,
9055 w0: &CudaSlice<f32>,
9056 w1: &CudaSlice<f32>,
9057 w2: &CudaSlice<f32>,
9058 d0: &mut CudaSlice<f32>,
9059 d1: &mut CudaSlice<f32>,
9060 d2: &mut CudaSlice<f32>,
9061 ncols: usize,
9062 nrows: usize,
9063 eps: f32,
9064 ) -> Result<(), Box<dyn std::error::Error>> {
9065 let f = self.func("rms_norm3_f32");
9066 let cfg = LaunchConfig {
9067 grid_dim: (nrows as u32, 1, 1),
9068 block_dim: (rms_block(), 1, 1),
9069 shared_mem_bytes: 0,
9070 };
9071 let (nc, e) = (ncols as i32, eps);
9072 let __s_b = self.gpu.stream();
9073 let mut b = __s_b.launch_builder(&f);
9074 b.arg(x)
9075 .arg(w0)
9076 .arg(w1)
9077 .arg(w2)
9078 .arg(d0)
9079 .arg(d1)
9080 .arg(d2)
9081 .arg(&nc)
9082 .arg(&e);
9083 unsafe {
9084 b.launch(cfg)?;
9085 }
9086 Ok(())
9087 }
9088
9089 #[allow(clippy::too_many_arguments)]
9091 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9094 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9095 *WARP_ON.get_or_init(|| {
9096 std::env::var("MEMRA_QKVNORM_W")
9097 .map(|v| v != "0")
9098 .unwrap_or(true)
9099 }) && ncols % 4 == 0
9100 && rows >= 64
9101 }
9102
9103 #[allow(clippy::too_many_arguments)]
9106 pub fn rms_norm_qkv_w4b(
9107 &self,
9108 q: &CudaSlice<f32>,
9109 k: &CudaSlice<f32>,
9110 v: &CudaSlice<f32>,
9111 wq: &CudaSlice<f32>,
9112 wk: &CudaSlice<f32>,
9113 wv: &CudaSlice<f32>,
9114 dq: &mut CudaSlice<f32>,
9115 dk: &mut CudaSlice<f32>,
9116 dv: &mut CudaSlice<f32>,
9117 dvb: &mut CudaSlice<u8>,
9118 ncols: usize,
9119 rq: usize,
9120 rk: usize,
9121 eps: f32,
9122 vf16: bool,
9123 ) -> Result<(), Box<dyn std::error::Error>> {
9124 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9125 let f = self.func("rms_norm_qkv_w4b_f32");
9126 let rows = (rq + 2 * rk) as u32;
9127 let cfg = LaunchConfig {
9128 grid_dim: (rows.div_ceil(8), 1, 1),
9129 block_dim: (256, 1, 1),
9130 shared_mem_bytes: 0,
9131 };
9132 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9133 let vf = vf16 as i32;
9134 let __s_b = self.gpu.stream();
9135 let mut b = __s_b.launch_builder(&f);
9136 b.arg(q)
9137 .arg(k)
9138 .arg(v)
9139 .arg(wq)
9140 .arg(wk)
9141 .arg(wv)
9142 .arg(dq)
9143 .arg(dk)
9144 .arg(dv)
9145 .arg(&mut *dvb)
9146 .arg(&nc)
9147 .arg(&rqi)
9148 .arg(&rki)
9149 .arg(&rvi)
9150 .arg(&e)
9151 .arg(&vf);
9152 unsafe {
9153 b.launch(cfg)?;
9154 }
9155 Ok(())
9156 }
9157
9158 pub fn rms_norm_qkv(
9159 &self,
9160 q: &CudaSlice<f32>,
9161 k: &CudaSlice<f32>,
9162 v: &CudaSlice<f32>,
9163 wq: &CudaSlice<f32>,
9164 wk: &CudaSlice<f32>,
9165 wv: &CudaSlice<f32>,
9166 dq: &mut CudaSlice<f32>,
9167 dk: &mut CudaSlice<f32>,
9168 dv: &mut CudaSlice<f32>,
9169 ncols: usize,
9170 rq: usize,
9171 rk: usize,
9172 eps: f32,
9173 ) -> Result<(), Box<dyn std::error::Error>> {
9174 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9178 let warp_on = *WARP_ON.get_or_init(|| {
9179 std::env::var("MEMRA_QKVNORM_W")
9180 .map(|v| v != "0")
9181 .unwrap_or(true)
9182 });
9183 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9186 let f = self.func("rms_norm_qkv_w4_f32");
9187 let rows = (rq + 2 * rk) as u32;
9188 let cfg = LaunchConfig {
9189 grid_dim: (rows.div_ceil(8), 1, 1),
9190 block_dim: (256, 1, 1),
9191 shared_mem_bytes: 0,
9192 };
9193 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9194 let __s_b = self.gpu.stream();
9195 let mut b = __s_b.launch_builder(&f);
9196 b.arg(q)
9197 .arg(k)
9198 .arg(v)
9199 .arg(wq)
9200 .arg(wk)
9201 .arg(wv)
9202 .arg(dq)
9203 .arg(dk)
9204 .arg(dv)
9205 .arg(&nc)
9206 .arg(&rqi)
9207 .arg(&rki)
9208 .arg(&rvi)
9209 .arg(&e);
9210 unsafe {
9211 b.launch(cfg)?;
9212 }
9213 return Ok(());
9214 }
9215 let f = self.func("rms_norm_qkv_f32");
9216 let grid = (rq + 2 * rk) as u32;
9217 let cfg = LaunchConfig {
9218 grid_dim: (grid, 1, 1),
9219 block_dim: (rms_block(), 1, 1),
9220 shared_mem_bytes: 0,
9221 };
9222 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9223 let __s_b = self.gpu.stream();
9224 let mut b = __s_b.launch_builder(&f);
9225 b.arg(q)
9226 .arg(k)
9227 .arg(v)
9228 .arg(wq)
9229 .arg(wk)
9230 .arg(wv)
9231 .arg(dq)
9232 .arg(dk)
9233 .arg(dv)
9234 .arg(&nc)
9235 .arg(&rqi)
9236 .arg(&rki)
9237 .arg(&e);
9238 unsafe {
9239 b.launch(cfg)?;
9240 }
9241 Ok(())
9242 }
9243
9244 #[allow(clippy::too_many_arguments)]
9246 pub fn rms_norm2x(
9247 &self,
9248 a: &CudaSlice<f32>,
9249 bb: &CudaSlice<f32>,
9250 wa: &CudaSlice<f32>,
9251 wb: &CudaSlice<f32>,
9252 da: &mut CudaSlice<f32>,
9253 db: &mut CudaSlice<f32>,
9254 ncols: usize,
9255 nrows: usize,
9256 eps: f32,
9257 ) -> Result<(), Box<dyn std::error::Error>> {
9258 let f = self.func("rms_norm2x_f32");
9259 let cfg = LaunchConfig {
9260 grid_dim: (2 * nrows as u32, 1, 1),
9261 block_dim: (rms_block(), 1, 1),
9262 shared_mem_bytes: 0,
9263 };
9264 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9265 let __s_b = self.gpu.stream();
9266 let mut b = __s_b.launch_builder(&f);
9267 b.arg(a)
9268 .arg(bb)
9269 .arg(wa)
9270 .arg(wb)
9271 .arg(da)
9272 .arg(db)
9273 .arg(&nc)
9274 .arg(&nr)
9275 .arg(&e);
9276 unsafe {
9277 b.launch(cfg)?;
9278 }
9279 Ok(())
9280 }
9281
9282 pub fn softcap(
9284 &self,
9285 y: &mut CudaSlice<f32>,
9286 cap: f32,
9287 n: usize,
9288 ) -> Result<(), Box<dyn std::error::Error>> {
9289 let f = self.func("softcap_f32");
9290 let cfg = LaunchConfig::for_num_elems(n as u32);
9291 let ni = n as i32;
9292 let __s_b = self.gpu.stream();
9293 let mut b = __s_b.launch_builder(&f);
9294 b.arg(y).arg(&cap).arg(&ni);
9295 unsafe {
9296 b.launch(cfg)?;
9297 }
9298 Ok(())
9299 }
9300
9301 pub fn mask_ids_rows(
9304 &self,
9305 y: &mut CudaSlice<f32>,
9306 ids: &CudaSlice<i32>,
9307 n_ids: usize,
9308 n_vocab: usize,
9309 t: usize,
9310 ) -> Result<(), Box<dyn std::error::Error>> {
9311 let f = self.func("mask_ids_rows_f32");
9312 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9313 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9314 let __s_b = self.gpu.stream();
9315 let mut b = __s_b.launch_builder(&f);
9316 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9317 unsafe {
9318 b.launch(cfg)?;
9319 }
9320 Ok(())
9321 }
9322
9323 #[allow(clippy::too_many_arguments)]
9325 pub fn add_scale_rms_norm(
9326 &self,
9327 a: &CudaSlice<f32>,
9328 b_in: &CudaSlice<f32>,
9329 c: f32,
9330 w: &CudaSlice<f32>,
9331 res: &mut CudaSlice<f32>,
9332 dst: &mut CudaSlice<f32>,
9333 ncols: usize,
9334 nrows: usize,
9335 eps: f32,
9336 ) -> Result<(), Box<dyn std::error::Error>> {
9337 let f = self.func("add_scale_rms_norm_f32");
9338 let cfg = LaunchConfig {
9339 grid_dim: (nrows as u32, 1, 1),
9340 block_dim: (rms_block(), 1, 1),
9341 shared_mem_bytes: 0,
9342 };
9343 let (nc, e2) = (ncols as i32, eps);
9344 let __s_b = self.gpu.stream();
9345 let mut b = __s_b.launch_builder(&f);
9346 b.arg(a)
9347 .arg(b_in)
9348 .arg(&c)
9349 .arg(w)
9350 .arg(res)
9351 .arg(dst)
9352 .arg(&nc)
9353 .arg(&e2);
9354 unsafe {
9355 b.launch(cfg)?;
9356 }
9357 Ok(())
9358 }
9359
9360 #[allow(clippy::too_many_arguments)]
9363 pub fn add_scale_rms_norm_q8_1(
9364 &self,
9365 a: &CudaSlice<f32>,
9366 b_in: &CudaSlice<f32>,
9367 c: f32,
9368 w: &CudaSlice<f32>,
9369 res: &mut CudaSlice<f32>,
9370 ncols: usize,
9371 nrows: usize,
9372 eps: f32,
9373 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9374 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9375 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9376 let (nc, e2) = (ncols as i32, eps);
9377 if Self::pdl_on() && Self::pdl_wb_on() {
9378 {
9379 use cudarc::driver::{DevicePtr, DevicePtrMut};
9380 let s = &self.gpu.stream();
9381 let (pa, _g0) = a.device_ptr(s);
9382 let (pb, _g1) = b_in.device_ptr(s);
9383 let (pw, _g2) = w.device_ptr(s);
9384 let (pr, _g3) = res.device_ptr_mut(s);
9385 let (pq, _g4) = out_q.device_ptr_mut(s);
9386 let (pd, _g5) = out_d.device_ptr_mut(s);
9387 let mut ps = [
9388 &pa as *const _ as *mut std::ffi::c_void,
9389 &pb as *const _ as *mut _,
9390 &c as *const _ as *mut _,
9391 &pw as *const _ as *mut _,
9392 &pr as *const _ as *mut _,
9393 &pq as *const _ as *mut _,
9394 &pd as *const _ as *mut _,
9395 &nc as *const _ as *mut _,
9396 &e2 as *const _ as *mut _,
9397 ];
9398 unsafe {
9399 self.launch_pdl(
9400 "add_scale_rms_norm_q8_1",
9401 (nrows as u32, 1, 1),
9402 (rms_block(), 1, 1),
9403 &mut ps,
9404 )?;
9405 }
9406 }
9407 return Ok((out_q, out_d));
9408 }
9409 let f = self.func("add_scale_rms_norm_q8_1");
9410 let cfg = LaunchConfig {
9411 grid_dim: (nrows as u32, 1, 1),
9412 block_dim: (rms_block(), 1, 1),
9413 shared_mem_bytes: 0,
9414 };
9415 let __s_b = self.gpu.stream();
9416 let mut b = __s_b.launch_builder(&f);
9417 b.arg(a)
9418 .arg(b_in)
9419 .arg(&c)
9420 .arg(w)
9421 .arg(res)
9422 .arg(&mut out_q)
9423 .arg(&mut out_d)
9424 .arg(&nc)
9425 .arg(&e2);
9426 unsafe {
9427 b.launch(cfg)?;
9428 }
9429 Ok((out_q, out_d))
9430 }
9431
9432 #[allow(clippy::too_many_arguments)]
9434 pub fn add_scale_rms_norm_q8_1_into(
9435 &self,
9436 a: &CudaSlice<f32>,
9437 b_in: &CudaSlice<f32>,
9438 c: f32,
9439 w: &CudaSlice<f32>,
9440 res: &mut CudaSlice<f32>,
9441 ncols: usize,
9442 nrows: usize,
9443 eps: f32,
9444 out_q: &mut CudaSlice<i8>,
9445 out_d: &mut CudaSlice<f32>,
9446 ) -> Result<(), Box<dyn std::error::Error>> {
9447 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9448 let (nc, e2) = (ncols as i32, eps);
9449 if Self::pdl_on() && Self::pdl_wb_on() {
9450 use cudarc::driver::{DevicePtr, DevicePtrMut};
9451 let s = &self.gpu.stream();
9452 let (pa, _g0) = a.device_ptr(s);
9453 let (pb, _g1) = b_in.device_ptr(s);
9454 let (pw, _g2) = w.device_ptr(s);
9455 let (pr, _g3) = res.device_ptr_mut(s);
9456 let (pq, _g4) = out_q.device_ptr_mut(s);
9457 let (pd, _g5) = out_d.device_ptr_mut(s);
9458 let mut ps = [
9459 &pa as *const _ as *mut std::ffi::c_void,
9460 &pb as *const _ as *mut _,
9461 &c as *const _ as *mut _,
9462 &pw as *const _ as *mut _,
9463 &pr as *const _ as *mut _,
9464 &pq as *const _ as *mut _,
9465 &pd as *const _ as *mut _,
9466 &nc as *const _ as *mut _,
9467 &e2 as *const _ as *mut _,
9468 ];
9469 unsafe {
9470 self.launch_pdl(
9471 "add_scale_rms_norm_q8_1",
9472 (nrows as u32, 1, 1),
9473 (rms_block(), 1, 1),
9474 &mut ps,
9475 )?;
9476 }
9477 return Ok(());
9478 }
9479 let f = self.func("add_scale_rms_norm_q8_1");
9480 let cfg = LaunchConfig {
9481 grid_dim: (nrows as u32, 1, 1),
9482 block_dim: (rms_block(), 1, 1),
9483 shared_mem_bytes: 0,
9484 };
9485 let __s_b = self.gpu.stream();
9486 let mut b = __s_b.launch_builder(&f);
9487 b.arg(a)
9488 .arg(b_in)
9489 .arg(&c)
9490 .arg(w)
9491 .arg(res)
9492 .arg(&mut *out_q)
9493 .arg(&mut *out_d)
9494 .arg(&nc)
9495 .arg(&e2);
9496 unsafe {
9497 b.launch(cfg)?;
9498 }
9499 Ok(())
9500 }
9501
9502 #[allow(clippy::too_many_arguments)]
9505 pub fn rms_pre_add_scale_rms_norm_q8_1(
9506 &self,
9507 a: &CudaSlice<f32>,
9508 wa: &CudaSlice<f32>,
9509 b_in: &CudaSlice<f32>,
9510 c: f32,
9511 w: &CudaSlice<f32>,
9512 res: &mut CudaSlice<f32>,
9513 ncols: usize,
9514 nrows: usize,
9515 eps: f32,
9516 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9517 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9518 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9519 let (nc, e2) = (ncols as i32, eps);
9520 if Self::pdl_on() {
9521 {
9522 use cudarc::driver::{DevicePtr, DevicePtrMut};
9523 let s = &self.gpu.stream();
9524 let (pa, _g0) = a.device_ptr(s);
9525 let (pwa, _g1) = wa.device_ptr(s);
9526 let (pb, _g2) = b_in.device_ptr(s);
9527 let (pw, _g3) = w.device_ptr(s);
9528 let (pr, _g4) = res.device_ptr_mut(s);
9529 let (pq, _g5) = out_q.device_ptr_mut(s);
9530 let (pd, _g6) = out_d.device_ptr_mut(s);
9531 let mut ps = [
9532 &pa as *const _ as *mut std::ffi::c_void,
9533 &pwa as *const _ as *mut _,
9534 &pb as *const _ as *mut _,
9535 &c as *const _ as *mut _,
9536 &pw as *const _ as *mut _,
9537 &pr as *const _ as *mut _,
9538 &pq as *const _ as *mut _,
9539 &pd as *const _ as *mut _,
9540 &nc as *const _ as *mut _,
9541 &e2 as *const _ as *mut _,
9542 ];
9543 unsafe {
9544 self.launch_pdl(
9545 "rms_pre_add_scale_rms_norm_q8_1",
9546 (nrows as u32, 1, 1),
9547 (rms_block(), 1, 1),
9548 &mut ps,
9549 )?;
9550 }
9551 }
9552 return Ok((out_q, out_d));
9553 }
9554 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9555 let cfg = LaunchConfig {
9556 grid_dim: (nrows as u32, 1, 1),
9557 block_dim: (rms_block(), 1, 1),
9558 shared_mem_bytes: 0,
9559 };
9560 let __s_b = self.gpu.stream();
9561 let mut b = __s_b.launch_builder(&f);
9562 b.arg(a)
9563 .arg(wa)
9564 .arg(b_in)
9565 .arg(&c)
9566 .arg(w)
9567 .arg(res)
9568 .arg(&mut out_q)
9569 .arg(&mut out_d)
9570 .arg(&nc)
9571 .arg(&e2);
9572 unsafe {
9573 b.launch(cfg)?;
9574 }
9575 Ok((out_q, out_d))
9576 }
9577
9578 pub fn gelu_tanh_mul_q8_1(
9581 &self,
9582 gate: &CudaSlice<f32>,
9583 up: &cudarc::driver::CudaView<f32>,
9584 act: &mut CudaSlice<f32>,
9585 ncols: usize,
9586 nrows: usize,
9587 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9588 debug_assert!(ncols % 128 == 0);
9589 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9590 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9591 let nc = ncols as i32;
9592 if Self::pdl_on() {
9593 {
9594 use cudarc::driver::{DevicePtr, DevicePtrMut};
9595 let s = &self.gpu.stream();
9596 let (pg, _g0) = gate.device_ptr(s);
9597 let (pu, _g1) = up.device_ptr(s);
9598 let (pact, _g2) = act.device_ptr_mut(s);
9599 let (pq, _g3) = out_q.device_ptr_mut(s);
9600 let (pd, _g4) = out_d.device_ptr_mut(s);
9601 let mut ps = [
9602 &pg as *const _ as *mut std::ffi::c_void,
9603 &pu as *const _ as *mut _,
9604 &pact as *const _ as *mut _,
9605 &pq as *const _ as *mut _,
9606 &pd as *const _ as *mut _,
9607 &nc as *const _ as *mut _,
9608 ];
9609 unsafe {
9610 self.launch_pdl(
9611 "gelu_tanh_mul_q8_1",
9612 (nrows as u32, 1, 1),
9613 (rms_block(), 1, 1),
9614 &mut ps,
9615 )?;
9616 }
9617 }
9618 return Ok((out_q, out_d));
9619 }
9620 let f = self.func("gelu_tanh_mul_q8_1");
9621 let cfg = LaunchConfig {
9622 grid_dim: (nrows as u32, 1, 1),
9623 block_dim: (rms_block(), 1, 1),
9624 shared_mem_bytes: 0,
9625 };
9626 let __s_b = self.gpu.stream();
9627 let mut b = __s_b.launch_builder(&f);
9628 b.arg(gate)
9629 .arg(up)
9630 .arg(act)
9631 .arg(&mut out_q)
9632 .arg(&mut out_d)
9633 .arg(&nc);
9634 unsafe {
9635 b.launch(cfg)?;
9636 }
9637 Ok((out_q, out_d))
9638 }
9639
9640 #[allow(clippy::too_many_arguments)]
9642 pub fn gelu_tanh_mul_q8_1_into(
9643 &self,
9644 gate: &CudaSlice<f32>,
9645 up: &cudarc::driver::CudaView<f32>,
9646 act: &mut CudaSlice<f32>,
9647 ncols: usize,
9648 nrows: usize,
9649 out_q: &mut CudaSlice<i8>,
9650 out_d: &mut CudaSlice<f32>,
9651 ) -> Result<(), Box<dyn std::error::Error>> {
9652 debug_assert!(ncols % 128 == 0);
9653 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9654 let nc = ncols as i32;
9655 if Self::pdl_on() {
9656 use cudarc::driver::{DevicePtr, DevicePtrMut};
9657 let s = &self.gpu.stream();
9658 let (pg, _g0) = gate.device_ptr(s);
9659 let (pu, _g1) = up.device_ptr(s);
9660 let (pact, _g2) = act.device_ptr_mut(s);
9661 let (pq, _g3) = out_q.device_ptr_mut(s);
9662 let (pd, _g4) = out_d.device_ptr_mut(s);
9663 let mut ps = [
9664 &pg as *const _ as *mut std::ffi::c_void,
9665 &pu as *const _ as *mut _,
9666 &pact as *const _ as *mut _,
9667 &pq as *const _ as *mut _,
9668 &pd as *const _ as *mut _,
9669 &nc as *const _ as *mut _,
9670 ];
9671 unsafe {
9672 self.launch_pdl(
9673 "gelu_tanh_mul_q8_1",
9674 (nrows as u32, 1, 1),
9675 (rms_block(), 1, 1),
9676 &mut ps,
9677 )?;
9678 }
9679 return Ok(());
9680 }
9681 let f = self.func("gelu_tanh_mul_q8_1");
9682 let cfg = LaunchConfig {
9683 grid_dim: (nrows as u32, 1, 1),
9684 block_dim: (rms_block(), 1, 1),
9685 shared_mem_bytes: 0,
9686 };
9687 let __s_b = self.gpu.stream();
9688 let mut b = __s_b.launch_builder(&f);
9689 b.arg(gate)
9690 .arg(up)
9691 .arg(&mut *act)
9692 .arg(&mut *out_q)
9693 .arg(&mut *out_d)
9694 .arg(&nc);
9695 unsafe {
9696 b.launch(cfg)?;
9697 }
9698 Ok(())
9699 }
9700
9701 #[allow(clippy::too_many_arguments)]
9703 pub fn add_rms_norm3_q8z(
9704 &self,
9705 a: &CudaSlice<f32>,
9706 b_in: &CudaSlice<f32>,
9707 w0: &CudaSlice<f32>,
9708 w1: &CudaSlice<f32>,
9709 w2: &CudaSlice<f32>,
9710 res: &mut CudaSlice<f32>,
9711 out1: &mut CudaSlice<f32>,
9712 ncols: usize,
9713 nrows: usize,
9714 eps: f32,
9715 ) -> Result<
9716 (
9717 (CudaSlice<i8>, CudaSlice<f32>),
9718 (CudaSlice<i8>, CudaSlice<f32>),
9719 ),
9720 Box<dyn std::error::Error>,
9721 > {
9722 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9723 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9724 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9725 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9726 let f = self.func("add_rms_norm3_q8z_f32");
9727 let cfg = LaunchConfig {
9728 grid_dim: (nrows as u32, 1, 1),
9729 block_dim: (rms_block(), 1, 1),
9730 shared_mem_bytes: 0,
9731 };
9732 let (nc, e2) = (ncols as i32, eps);
9733 let __s_b = self.gpu.stream();
9734 let mut b = __s_b.launch_builder(&f);
9735 b.arg(a)
9736 .arg(b_in)
9737 .arg(w0)
9738 .arg(w1)
9739 .arg(w2)
9740 .arg(res)
9741 .arg(&mut q0)
9742 .arg(&mut d0)
9743 .arg(out1)
9744 .arg(&mut q2)
9745 .arg(&mut d2)
9746 .arg(&nc)
9747 .arg(&e2);
9748 unsafe {
9749 b.launch(cfg)?;
9750 }
9751 Ok(((q0, d0), (q2, d2)))
9752 }
9753
9754 #[allow(clippy::too_many_arguments)]
9756 pub fn add_rms_norm3(
9757 &self,
9758 a: &CudaSlice<f32>,
9759 b_in: &CudaSlice<f32>,
9760 w0: &CudaSlice<f32>,
9761 w1: &CudaSlice<f32>,
9762 w2: &CudaSlice<f32>,
9763 res: &mut CudaSlice<f32>,
9764 d0: &mut CudaSlice<f32>,
9765 d1: &mut CudaSlice<f32>,
9766 d2: &mut CudaSlice<f32>,
9767 ncols: usize,
9768 nrows: usize,
9769 eps: f32,
9770 ) -> Result<(), Box<dyn std::error::Error>> {
9771 let f = self.func("add_rms_norm3_f32");
9772 let cfg = LaunchConfig {
9773 grid_dim: (nrows as u32, 1, 1),
9774 block_dim: (rms_block(), 1, 1),
9775 shared_mem_bytes: 0,
9776 };
9777 let (nc, e2) = (ncols as i32, eps);
9778 let __s_b = self.gpu.stream();
9779 let mut b = __s_b.launch_builder(&f);
9780 b.arg(a)
9781 .arg(b_in)
9782 .arg(w0)
9783 .arg(w1)
9784 .arg(w2)
9785 .arg(res)
9786 .arg(d0)
9787 .arg(d1)
9788 .arg(d2)
9789 .arg(&nc)
9790 .arg(&e2);
9791 unsafe {
9792 b.launch(cfg)?;
9793 }
9794 Ok(())
9795 }
9796
9797 pub fn add_scale(
9799 &self,
9800 a: &CudaSlice<f32>,
9801 b_in: &CudaSlice<f32>,
9802 c: f32,
9803 dst: &mut CudaSlice<f32>,
9804 n: usize,
9805 ) -> Result<(), Box<dyn std::error::Error>> {
9806 let f = self.func("add_scale_f32");
9807 let cfg = LaunchConfig::for_num_elems(n as u32);
9808 let ni = n as i32;
9809 let __s_b = self.gpu.stream();
9810 let mut b = __s_b.launch_builder(&f);
9811 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9812 unsafe {
9813 b.launch(cfg)?;
9814 }
9815 Ok(())
9816 }
9817
9818 pub fn layer_norm_bias(
9820 &self,
9821 x: &CudaSlice<f32>,
9822 w: &CudaSlice<f32>,
9823 b: &CudaSlice<f32>,
9824 dst: &mut CudaSlice<f32>,
9825 ncols: usize,
9826 nrows: usize,
9827 eps: f32,
9828 ) -> Result<(), Box<dyn std::error::Error>> {
9829 let f = self.func("layer_norm_bias_f32");
9830 let (nc, e) = (ncols as i32, eps);
9831 let cfg = LaunchConfig {
9832 grid_dim: (nrows as u32, 1, 1),
9833 block_dim: (256, 1, 1),
9834 shared_mem_bytes: 0,
9835 };
9836 let __s_b = self.gpu.stream();
9837 let mut lb = __s_b.launch_builder(&f);
9838 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9839 unsafe {
9840 lb.launch(cfg)?;
9841 }
9842 Ok(())
9843 }
9844
9845 pub fn gelu_tanh(
9847 &self,
9848 x: &CudaSlice<f32>,
9849 dst: &mut CudaSlice<f32>,
9850 n: usize,
9851 ) -> Result<(), Box<dyn std::error::Error>> {
9852 let f = self.func("gelu_tanh_f32");
9853 let ni = n as i64;
9854 let cfg = LaunchConfig {
9855 grid_dim: (n.div_ceil(256) as u32, 1, 1),
9856 block_dim: (256, 1, 1),
9857 shared_mem_bytes: 0,
9858 };
9859 let __s_b = self.gpu.stream();
9860 let mut lb = __s_b.launch_builder(&f);
9861 lb.arg(x).arg(&mut *dst).arg(&ni);
9862 unsafe {
9863 lb.launch(cfg)?;
9864 }
9865 Ok(())
9866 }
9867
9868 pub fn row_softmax(
9870 &self,
9871 x: &mut CudaSlice<f32>,
9872 ncols: usize,
9873 nrows: usize,
9874 ) -> Result<(), Box<dyn std::error::Error>> {
9875 let f = self.func("row_softmax_f32");
9876 let nc = ncols as i32;
9877 let cfg = LaunchConfig {
9878 grid_dim: (nrows as u32, 1, 1),
9879 block_dim: (256, 1, 1),
9880 shared_mem_bytes: 0,
9881 };
9882 let __s_b = self.gpu.stream();
9883 let mut lb = __s_b.launch_builder(&f);
9884 lb.arg(&mut *x).arg(&nc);
9885 unsafe {
9886 lb.launch(cfg)?;
9887 }
9888 Ok(())
9889 }
9890
9891 pub fn rms_norm(
9892 &self,
9893 x: &CudaSlice<f32>,
9894 w: &CudaSlice<f32>,
9895 dst: &mut CudaSlice<f32>,
9896 ncols: usize,
9897 nrows: usize,
9898 eps: f32,
9899 ) -> Result<(), Box<dyn std::error::Error>> {
9900 let (nc, e) = (ncols as i32, eps);
9901 let kname = if Self::norm_ilp_on() {
9902 "rms_norm_f32_v2"
9903 } else {
9904 "rms_norm_f32"
9905 };
9906 if Self::pdl_on() && Self::pdl_wb_on() {
9907 use cudarc::driver::{DevicePtr, DevicePtrMut};
9908 let s = &self.gpu.stream();
9909 let (px, _g0) = x.device_ptr(s);
9910 let (pw, _g1) = w.device_ptr(s);
9911 let (pd, _g2) = dst.device_ptr_mut(s);
9912 let mut ps = [
9913 &px as *const _ as *mut std::ffi::c_void,
9914 &pw as *const _ as *mut _,
9915 &pd as *const _ as *mut _,
9916 &nc as *const _ as *mut _,
9917 &e as *const _ as *mut _,
9918 ];
9919 unsafe {
9920 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9921 }
9922 return Ok(());
9923 }
9924 let f = self.func(kname);
9925 let cfg = LaunchConfig {
9926 grid_dim: (nrows as u32, 1, 1),
9927 block_dim: (rms_block(), 1, 1),
9928 shared_mem_bytes: 0,
9929 };
9930 let __s_b = self.gpu.stream();
9931 let mut b = __s_b.launch_builder(&f);
9932 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9933 unsafe {
9934 b.launch(cfg)?;
9935 }
9936 Ok(())
9937 }
9938
9939 pub fn rms_norm_decode(
9947 &self,
9948 x: &CudaSlice<f32>,
9949 w: &CudaSlice<f32>,
9950 dst: &mut CudaSlice<f32>,
9951 ncols: usize,
9952 nrows: usize,
9953 eps: f32,
9954 ) -> Result<(), Box<dyn std::error::Error>> {
9955 let f = self.func(if Self::norm_ilp_on() {
9956 "rms_norm_f32_v2"
9957 } else {
9958 "rms_norm_f32"
9959 });
9960 let cfg = LaunchConfig {
9961 grid_dim: (nrows as u32, 1, 1),
9962 block_dim: (1024, 1, 1),
9963 shared_mem_bytes: 0,
9964 };
9965 let (nc, e) = (ncols as i32, eps);
9966 let __s_b = self.gpu.stream();
9967 let mut b = __s_b.launch_builder(&f);
9968 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9969 unsafe {
9970 b.launch(cfg)?;
9971 }
9972 Ok(())
9973 }
9974
9975 pub fn rms_norm_q8_1(
9979 &self,
9980 x: &CudaSlice<f32>,
9981 w: &CudaSlice<f32>,
9982 ncols: usize,
9983 nrows: usize,
9984 eps: f32,
9985 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9986 let nblk = ncols / 32;
9987 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9988 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9989 let (nc, e) = (ncols as i32, eps);
9990 if Self::pdl_on() {
9991 {
9992 use cudarc::driver::{DevicePtr, DevicePtrMut};
9993 let s = &self.gpu.stream();
9994 let (px, _g0) = x.device_ptr(s);
9995 let (pw, _g1) = w.device_ptr(s);
9996 let (pq, _g2) = q.device_ptr_mut(s);
9997 let (pd, _g3) = d.device_ptr_mut(s);
9998 let mut ps = [
9999 &px as *const _ as *mut std::ffi::c_void,
10000 &pw as *const _ as *mut _,
10001 &pq as *const _ as *mut _,
10002 &pd as *const _ as *mut _,
10003 &nc as *const _ as *mut _,
10004 &e as *const _ as *mut _,
10005 ];
10006 unsafe {
10007 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10008 }
10009 }
10010 return Ok((q, d));
10011 }
10012 let f = self.func("rms_norm_q8_1");
10013 let cfg = LaunchConfig {
10016 grid_dim: (nrows as u32, 1, 1),
10017 block_dim: (1024, 1, 1),
10018 shared_mem_bytes: 0,
10019 };
10020 let __s_b = self.gpu.stream();
10021 let mut b = __s_b.launch_builder(&f);
10022 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10023 unsafe {
10024 b.launch(cfg)?;
10025 }
10026 Ok((q, d))
10027 }
10028
10029 pub fn rms_norm_q8_1_into(
10032 &self,
10033 x: &CudaSlice<f32>,
10034 w: &CudaSlice<f32>,
10035 ncols: usize,
10036 nrows: usize,
10037 eps: f32,
10038 q: &mut CudaSlice<i8>,
10039 d: &mut CudaSlice<f32>,
10040 ) -> Result<(), Box<dyn std::error::Error>> {
10041 let nblk = ncols / 32;
10042 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10043 let (nc, e) = (ncols as i32, eps);
10044 if Self::pdl_on() {
10045 use cudarc::driver::{DevicePtr, DevicePtrMut};
10046 let s = &self.gpu.stream();
10047 let (px, _g0) = x.device_ptr(s);
10048 let (pw, _g1) = w.device_ptr(s);
10049 let (pq, _g2) = q.device_ptr_mut(s);
10050 let (pd, _g3) = d.device_ptr_mut(s);
10051 let mut ps = [
10052 &px as *const _ as *mut std::ffi::c_void,
10053 &pw as *const _ as *mut _,
10054 &pq as *const _ as *mut _,
10055 &pd as *const _ as *mut _,
10056 &nc as *const _ as *mut _,
10057 &e as *const _ as *mut _,
10058 ];
10059 unsafe {
10060 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10061 }
10062 return Ok(());
10063 }
10064 let f = self.func("rms_norm_q8_1");
10065 let cfg = LaunchConfig {
10066 grid_dim: (nrows as u32, 1, 1),
10067 block_dim: (1024, 1, 1),
10068 shared_mem_bytes: 0,
10069 };
10070 let __s_b = self.gpu.stream();
10071 let mut b = __s_b.launch_builder(&f);
10072 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10073 unsafe {
10074 b.launch(cfg)?;
10075 }
10076 Ok(())
10077 }
10078
10079 pub fn quantize_q8_1_into(
10081 &self,
10082 x: &CudaSlice<f32>,
10083 m: usize,
10084 in_f: usize,
10085 q: &mut CudaSlice<i8>,
10086 d: &mut CudaSlice<f32>,
10087 ) -> Result<(), Box<dyn std::error::Error>> {
10088 let nblk = in_f / 32;
10089 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10090 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10091 let (inf, mi) = (in_f as i32, m as i32);
10092 if Self::pdl_on() && Self::pdl_wb_on() {
10093 use cudarc::driver::{DevicePtr, DevicePtrMut};
10094 let s = &self.gpu.stream();
10095 let (px, _g0) = x.device_ptr(s);
10096 let (pq, _g1) = q.device_ptr_mut(s);
10097 let (pd, _g2) = d.device_ptr_mut(s);
10098 let mut ps = [
10099 &px as *const _ as *mut std::ffi::c_void,
10100 &pq as *const _ as *mut _,
10101 &pd as *const _ as *mut _,
10102 &inf as *const _ as *mut _,
10103 &mi as *const _ as *mut _,
10104 ];
10105 unsafe {
10106 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10107 }
10108 return Ok(());
10109 }
10110 let f = self.func("quantize_q8_1");
10111 let __s_b = self.gpu.stream();
10112 let mut b = __s_b.launch_builder(&f);
10113 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10114 unsafe {
10115 b.launch(cfg)?;
10116 }
10117 Ok(())
10118 }
10119
10120 pub fn add_rms_norm_q8_1(
10124 &self,
10125 a: &CudaSlice<f32>,
10126 b_in: &CudaSlice<f32>,
10127 w: &CudaSlice<f32>,
10128 res: &mut CudaSlice<f32>,
10129 ncols: usize,
10130 nrows: usize,
10131 eps: f32,
10132 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10133 let nblk = ncols / 32;
10134 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10135 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10136 let f = self.func("add_rms_norm_q8_1");
10137 let cfg = LaunchConfig {
10139 grid_dim: (nrows as u32, 1, 1),
10140 block_dim: (1024, 1, 1),
10141 shared_mem_bytes: 0,
10142 };
10143 let (nc, e) = (ncols as i32, eps);
10144 let __s_bld = self.gpu.stream();
10145 let mut bld = __s_bld.launch_builder(&f);
10146 bld.arg(a)
10147 .arg(b_in)
10148 .arg(w)
10149 .arg(res)
10150 .arg(&mut q)
10151 .arg(&mut d)
10152 .arg(&nc)
10153 .arg(&e);
10154 unsafe {
10155 bld.launch(cfg)?;
10156 }
10157 Ok((q, d))
10158 }
10159
10160 #[allow(clippy::too_many_arguments)]
10166 pub fn join_add_rms_norm_raw(
10167 &self,
10168 a0_raw: u64,
10169 a1_raw: u64,
10170 x: &CudaSlice<f32>,
10171 w: &CudaSlice<f32>,
10172 res: &mut CudaSlice<f32>,
10173 dst: &mut CudaSlice<f32>,
10174 ncols: usize,
10175 eps: f32,
10176 ) -> Result<(), Box<dyn std::error::Error>> {
10177 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10178 return Err("join_add_rms_norm geometry".into());
10179 }
10180 let f = self.func("join_add_rms_norm_f32");
10181 let cfg = LaunchConfig {
10182 grid_dim: (1, 1, 1),
10183 block_dim: (rms_block(), 1, 1),
10184 shared_mem_bytes: 0,
10185 };
10186 let (nc, e) = (ncols as i32, eps);
10187 let __s_b = self.gpu.stream();
10188 let mut b = __s_b.launch_builder(&f);
10189 b.arg(&a0_raw)
10190 .arg(&a1_raw)
10191 .arg(x)
10192 .arg(w)
10193 .arg(&mut *res)
10194 .arg(&mut *dst)
10195 .arg(&nc)
10196 .arg(&e);
10197 unsafe {
10198 b.launch(cfg)?;
10199 }
10200 Ok(())
10201 }
10202
10203 pub fn add_rms_norm(
10204 &self,
10205 a: &CudaSlice<f32>,
10206 b: &CudaSlice<f32>,
10207 w: &CudaSlice<f32>,
10208 res: &mut CudaSlice<f32>,
10209 dst: &mut CudaSlice<f32>,
10210 ncols: usize,
10211 nrows: usize,
10212 eps: f32,
10213 ) -> Result<(), Box<dyn std::error::Error>> {
10214 let (nc, e) = (ncols as i32, eps);
10215 let kname = if Self::norm_ilp_on() {
10216 "add_rms_norm_f32_v2"
10217 } else {
10218 "add_rms_norm_f32"
10219 };
10220 if Self::pdl_on() && Self::pdl_wb_on() {
10221 use cudarc::driver::{DevicePtr, DevicePtrMut};
10222 let s = &self.gpu.stream();
10223 let (pa, _g0) = a.device_ptr(s);
10224 let (pb, _g1) = b.device_ptr(s);
10225 let (pw, _g2) = w.device_ptr(s);
10226 let (pr, _g3) = res.device_ptr_mut(s);
10227 let (pd, _g4) = dst.device_ptr_mut(s);
10228 let mut ps = [
10229 &pa as *const _ as *mut std::ffi::c_void,
10230 &pb as *const _ as *mut _,
10231 &pw as *const _ as *mut _,
10232 &pr as *const _ as *mut _,
10233 &pd as *const _ as *mut _,
10234 &nc as *const _ as *mut _,
10235 &e as *const _ as *mut _,
10236 ];
10237 unsafe {
10238 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10239 }
10240 return Ok(());
10241 }
10242 let f = self.func(kname);
10243 let cfg = LaunchConfig {
10244 grid_dim: (nrows as u32, 1, 1),
10245 block_dim: (rms_block(), 1, 1),
10246 shared_mem_bytes: 0,
10247 };
10248 let __s_b2 = self.gpu.stream();
10249 let mut b2 = __s_b2.launch_builder(&f);
10250 b2.arg(a)
10251 .arg(b)
10252 .arg(w)
10253 .arg(&mut *res)
10254 .arg(&mut *dst)
10255 .arg(&nc)
10256 .arg(&e);
10257 unsafe {
10258 b2.launch(cfg)?;
10259 }
10260 Ok(())
10261 }
10262
10263 #[allow(clippy::too_many_arguments)]
10266 pub fn rms_pre_add_rms_norm(
10267 &self,
10268 a: &CudaSlice<f32>,
10269 wa: &CudaSlice<f32>,
10270 b: &CudaSlice<f32>,
10271 w: &CudaSlice<f32>,
10272 res: &mut CudaSlice<f32>,
10273 dst: &mut CudaSlice<f32>,
10274 ncols: usize,
10275 nrows: usize,
10276 eps: f32,
10277 ) -> Result<(), Box<dyn std::error::Error>> {
10278 let f = self.func("rms_pre_add_rms_norm_f32");
10279 let cfg = LaunchConfig {
10280 grid_dim: (nrows as u32, 1, 1),
10281 block_dim: (rms_block(), 1, 1),
10282 shared_mem_bytes: 0,
10283 };
10284 let (nc, e) = (ncols as i32, eps);
10285 let __s_b2 = self.gpu.stream();
10286 let mut b2 = __s_b2.launch_builder(&f);
10287 b2.arg(a)
10288 .arg(wa)
10289 .arg(b)
10290 .arg(w)
10291 .arg(&mut *res)
10292 .arg(&mut *dst)
10293 .arg(&nc)
10294 .arg(&e);
10295 unsafe {
10296 b2.launch(cfg)?;
10297 }
10298 Ok(())
10299 }
10300
10301 #[allow(clippy::too_many_arguments)]
10303 pub fn rms_pre_add_rms_norm_q8z(
10304 &self,
10305 a: &CudaSlice<f32>,
10306 wa: &CudaSlice<f32>,
10307 b: &CudaSlice<f32>,
10308 w: &CudaSlice<f32>,
10309 res: &mut CudaSlice<f32>,
10310 dst: &mut CudaSlice<f32>,
10311 ncols: usize,
10312 nrows: usize,
10313 eps: f32,
10314 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10315 debug_assert!(ncols % 128 == 0);
10316 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10317 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10318 let (nc, e) = (ncols as i32, eps);
10319 if Self::pdl_on() {
10320 {
10321 use cudarc::driver::{DevicePtr, DevicePtrMut};
10322 let s = &self.gpu.stream();
10323 let (pa, _g0) = a.device_ptr(s);
10324 let (pwa, _g1) = wa.device_ptr(s);
10325 let (pb, _g2) = b.device_ptr(s);
10326 let (pw, _g3) = w.device_ptr(s);
10327 let (pr, _g4) = res.device_ptr_mut(s);
10328 let (pdst, _g5) = dst.device_ptr_mut(s);
10329 let (pq, _g6) = out_q.device_ptr_mut(s);
10330 let (pd, _g7) = out_d.device_ptr_mut(s);
10331 let mut ps = [
10332 &pa as *const _ as *mut std::ffi::c_void,
10333 &pwa as *const _ as *mut _,
10334 &pb as *const _ as *mut _,
10335 &pw as *const _ as *mut _,
10336 &pr as *const _ as *mut _,
10337 &pdst as *const _ as *mut _,
10338 &pq as *const _ as *mut _,
10339 &pd as *const _ as *mut _,
10340 &nc as *const _ as *mut _,
10341 &e as *const _ as *mut _,
10342 ];
10343 unsafe {
10344 self.launch_pdl(
10345 "rms_pre_add_rms_norm_q8z_f32",
10346 (nrows as u32, 1, 1),
10347 (rms_block(), 1, 1),
10348 &mut ps,
10349 )?;
10350 }
10351 }
10352 return Ok((out_q, out_d));
10353 }
10354 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10355 let cfg = LaunchConfig {
10356 grid_dim: (nrows as u32, 1, 1),
10357 block_dim: (rms_block(), 1, 1),
10358 shared_mem_bytes: 0,
10359 };
10360 let __s_b2 = self.gpu.stream();
10361 let mut b2 = __s_b2.launch_builder(&f);
10362 b2.arg(a)
10363 .arg(wa)
10364 .arg(b)
10365 .arg(w)
10366 .arg(&mut *res)
10367 .arg(&mut *dst)
10368 .arg(&mut out_q)
10369 .arg(&mut out_d)
10370 .arg(&nc)
10371 .arg(&e);
10372 unsafe {
10373 b2.launch(cfg)?;
10374 }
10375 Ok((out_q, out_d))
10376 }
10377
10378 #[allow(clippy::too_many_arguments)]
10382 pub fn rms_pre_add_rms_norm_q8z_into(
10383 &self,
10384 a: &CudaSlice<f32>,
10385 wa: &CudaSlice<f32>,
10386 b: &CudaSlice<f32>,
10387 w: &CudaSlice<f32>,
10388 res: &mut CudaSlice<f32>,
10389 dst: &mut CudaSlice<f32>,
10390 ncols: usize,
10391 nrows: usize,
10392 eps: f32,
10393 out_q: &mut CudaSlice<i8>,
10394 out_d: &mut CudaSlice<f32>,
10395 ) -> Result<(), Box<dyn std::error::Error>> {
10396 debug_assert!(ncols % 128 == 0);
10397 let (nc, e) = (ncols as i32, eps);
10398 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10399 let cfg = LaunchConfig {
10400 grid_dim: (nrows as u32, 1, 1),
10401 block_dim: (rms_block(), 1, 1),
10402 shared_mem_bytes: 0,
10403 };
10404 let __s_b = self.gpu.stream();
10405 let mut b2 = __s_b.launch_builder(&f);
10406 b2.arg(a)
10407 .arg(wa)
10408 .arg(b)
10409 .arg(w)
10410 .arg(&mut *res)
10411 .arg(&mut *dst)
10412 .arg(&mut *out_q)
10413 .arg(&mut *out_d)
10414 .arg(&nc)
10415 .arg(&e);
10416 unsafe {
10417 b2.launch(cfg)?;
10418 }
10419 Ok(())
10420 }
10421
10422 #[allow(clippy::too_many_arguments)]
10425 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10426 &self,
10427 a: &CudaSlice<f32>,
10428 wa: &CudaSlice<f32>,
10429 b_in: &CudaSlice<f32>,
10430 c: f32,
10431 w: &CudaSlice<f32>,
10432 res: &mut CudaSlice<f32>,
10433 ncols: usize,
10434 nrows: usize,
10435 eps: f32,
10436 out_q: &mut CudaSlice<i8>,
10437 out_d: &mut CudaSlice<f32>,
10438 ) -> Result<(), Box<dyn std::error::Error>> {
10439 debug_assert!(ncols % 128 == 0);
10440 let (nc, e2) = (ncols as i32, eps);
10441 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10442 let cfg = LaunchConfig {
10443 grid_dim: (nrows as u32, 1, 1),
10444 block_dim: (rms_block(), 1, 1),
10445 shared_mem_bytes: 0,
10446 };
10447 let __s_b = self.gpu.stream();
10448 let mut b2 = __s_b.launch_builder(&f);
10449 b2.arg(a)
10450 .arg(wa)
10451 .arg(b_in)
10452 .arg(&c)
10453 .arg(w)
10454 .arg(&mut *res)
10455 .arg(&mut *out_q)
10456 .arg(&mut *out_d)
10457 .arg(&nc)
10458 .arg(&e2);
10459 unsafe {
10460 b2.launch(cfg)?;
10461 }
10462 Ok(())
10463 }
10464
10465 pub fn g4_pnfold_on() -> bool {
10473 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10474 *ON.get_or_init(|| {
10475 std::env::var("MEMRA_G4_PNFOLD")
10476 .map(|v| v != "0")
10477 .unwrap_or(true)
10478 })
10479 }
10480
10481 pub fn build_q4_out_concat3(
10485 &self,
10486 w0: &crate::model::GpuTensor,
10487 w1: &crate::model::GpuTensor,
10488 w2: &crate::model::GpuTensor,
10489 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10490 use crate::model::GpuTensor;
10491 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10492 match w {
10493 GpuTensor::Quant {
10494 qtype,
10495 row_bytes,
10496 rp,
10497 ..
10498 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10499 _ => None,
10500 }
10501 };
10502 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10503 else {
10504 return Ok(None);
10505 };
10506 if rb0 != rb1
10507 || rb0 != rb2
10508 || w0.in_features() != w1.in_features()
10509 || w0.in_features() != w2.in_features()
10510 {
10511 return Ok(None);
10512 }
10513 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10514 match w {
10515 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10516 _ => unreachable!(),
10517 }
10518 }
10519 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10520 let total = rb0 * (o0 + o1 + o2);
10521 let mut cat = self.alloc_u8(total)?;
10522 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10523 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10524 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10525 Ok(Some(GpuTensor::Quant {
10526 bytes: cat,
10527 qtype: QT_Q4_0,
10528 row_bytes: rb0,
10529 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10530 scale: 1.0,
10531 rp: false,
10532 #[cfg(memra_cutlass)]
10533 cutlass: None,
10534 fp8: None,
10535 blk: None,
10536 rp4: None,
10537 f16: None,
10538 }))
10539 }
10540
10541 fn full_width_rope_only(
10559 kernel: &str,
10560 n_rot: usize,
10561 head_dim: usize,
10562 ) -> Result<(), Box<dyn std::error::Error>> {
10563 if n_rot == head_dim {
10564 return Ok(());
10565 }
10566 Err(format!(
10567 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10568 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10569 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10570 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10571 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10572 )
10573 .into())
10574 }
10575
10576 #[allow(clippy::too_many_arguments)]
10580 pub fn rms_norm_qkv_rope_cat(
10581 &self,
10582 qkv: &CudaSlice<f32>,
10583 wq: &CudaSlice<f32>,
10584 wk: &CudaSlice<f32>,
10585 wv: &CudaSlice<f32>,
10586 q: &mut CudaSlice<f32>,
10587 k: &mut CudaSlice<f32>,
10588 v: &mut CudaSlice<f32>,
10589 head_dim: usize,
10590 n_rot: usize,
10591 rq: usize,
10592 rk: usize,
10593 pos: &CudaSlice<i32>,
10594 nh_q: usize,
10595 nh_k: usize,
10596 base: f32,
10597 freq_scale: f32,
10598 ff: Option<&CudaSlice<f32>>,
10599 eps: f32,
10600 ) -> Result<(), Box<dyn std::error::Error>> {
10601 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10602 let rows = rq + rk + rk;
10603 let theta_scale = base.powf(-2.0 / head_dim as f32);
10604 let (nc, rqi, rki, nhq, nhk) = (
10605 head_dim as i32,
10606 rq as i32,
10607 rk as i32,
10608 nh_q as i32,
10609 nh_k as i32,
10610 );
10611 if Self::pdl_on() {
10612 use cudarc::driver::{DevicePtr, DevicePtrMut};
10613 let s = &self.gpu.stream();
10614 let (pqkv, _g0) = qkv.device_ptr(s);
10615 let (pwq, _g1) = wq.device_ptr(s);
10616 let (pwk, _g2) = wk.device_ptr(s);
10617 let (pwv, _g3) = wv.device_ptr(s);
10618 let (pq, _g4) = q.device_ptr_mut(s);
10619 let (pk, _g5) = k.device_ptr_mut(s);
10620 let (pv, _g6) = v.device_ptr_mut(s);
10621 let (ppos, _g7) = pos.device_ptr(s);
10622 let (pff, _g8) = match ff {
10623 Some(t) => {
10624 let (p, g) = t.device_ptr(s);
10625 (p, Some(g))
10626 }
10627 None => (0, None),
10628 };
10629 let mut ps = [
10630 &pqkv as *const _ as *mut std::ffi::c_void,
10631 &pwq as *const _ as *mut _,
10632 &pwk as *const _ as *mut _,
10633 &pwv as *const _ as *mut _,
10634 &pq as *const _ as *mut _,
10635 &pk as *const _ as *mut _,
10636 &pv as *const _ as *mut _,
10637 &nc as *const _ as *mut _,
10638 &rqi as *const _ as *mut _,
10639 &rki as *const _ as *mut _,
10640 &ppos as *const _ as *mut _,
10641 &nhq as *const _ as *mut _,
10642 &nhk as *const _ as *mut _,
10643 &theta_scale as *const _ as *mut _,
10644 &freq_scale as *const _ as *mut _,
10645 &pff as *const _ as *mut _,
10646 &eps as *const _ as *mut _,
10647 ];
10648 unsafe {
10649 self.launch_pdl(
10650 "rms_norm_qkv_rope_cat_f32",
10651 (rows as u32, 1, 1),
10652 (rms_block(), 1, 1),
10653 &mut ps,
10654 )?;
10655 }
10656 return Ok(());
10657 }
10658 let f = self.func("rms_norm_qkv_rope_cat_f32");
10659 let cfg = LaunchConfig {
10660 grid_dim: (rows as u32, 1, 1),
10661 block_dim: (rms_block(), 1, 1),
10662 shared_mem_bytes: 0,
10663 };
10664 let __s_b = self.gpu.stream();
10665 let mut b = __s_b.launch_builder(&f);
10666 match ff {
10667 Some(t) => {
10668 b.arg(qkv)
10669 .arg(wq)
10670 .arg(wk)
10671 .arg(wv)
10672 .arg(&mut *q)
10673 .arg(&mut *k)
10674 .arg(&mut *v)
10675 .arg(&nc)
10676 .arg(&rqi)
10677 .arg(&rki)
10678 .arg(pos)
10679 .arg(&nhq)
10680 .arg(&nhk)
10681 .arg(&theta_scale)
10682 .arg(&freq_scale)
10683 .arg(t)
10684 .arg(&eps);
10685 unsafe {
10686 b.launch(cfg)?;
10687 }
10688 }
10689 None => {
10690 let null: u64 = 0;
10691 b.arg(qkv)
10692 .arg(wq)
10693 .arg(wk)
10694 .arg(wv)
10695 .arg(&mut *q)
10696 .arg(&mut *k)
10697 .arg(&mut *v)
10698 .arg(&nc)
10699 .arg(&rqi)
10700 .arg(&rki)
10701 .arg(pos)
10702 .arg(&nhq)
10703 .arg(&nhk)
10704 .arg(&theta_scale)
10705 .arg(&freq_scale)
10706 .arg(&null)
10707 .arg(&eps);
10708 unsafe {
10709 b.launch(cfg)?;
10710 }
10711 }
10712 }
10713 Ok(())
10714 }
10715
10716 #[allow(clippy::too_many_arguments)]
10720 pub fn rms_norm_qkv_rope(
10721 &self,
10722 q0: &CudaSlice<f32>,
10723 k0: &CudaSlice<f32>,
10724 v0: &CudaSlice<f32>,
10725 wq: &CudaSlice<f32>,
10726 wk: &CudaSlice<f32>,
10727 wv: &CudaSlice<f32>,
10728 q: &mut CudaSlice<f32>,
10729 k: &mut CudaSlice<f32>,
10730 v: &mut CudaSlice<f32>,
10731 head_dim: usize,
10732 n_rot: usize,
10733 rq: usize,
10734 rk: usize,
10735 pos: &CudaSlice<i32>,
10736 nh_q: usize,
10737 nh_k: usize,
10738 base: f32,
10739 freq_scale: f32,
10740 ff: Option<&CudaSlice<f32>>,
10741 eps: f32,
10742 ) -> Result<(), Box<dyn std::error::Error>> {
10743 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10744 let f = self.func("rms_norm_qkv_rope_f32");
10745 let rows = rq + rk + rk; let cfg = LaunchConfig {
10747 grid_dim: (rows as u32, 1, 1),
10748 block_dim: (rms_block(), 1, 1),
10749 shared_mem_bytes: 0,
10750 };
10751 let theta_scale = base.powf(-2.0 / head_dim as f32);
10752 let (nc, rqi, rki, nhq, nhk) = (
10753 head_dim as i32,
10754 rq as i32,
10755 rk as i32,
10756 nh_q as i32,
10757 nh_k as i32,
10758 );
10759 let __s_b = self.gpu.stream();
10760 let mut b = __s_b.launch_builder(&f);
10761 match ff {
10762 Some(t) => {
10763 b.arg(q0)
10764 .arg(k0)
10765 .arg(v0)
10766 .arg(wq)
10767 .arg(wk)
10768 .arg(wv)
10769 .arg(&mut *q)
10770 .arg(&mut *k)
10771 .arg(&mut *v)
10772 .arg(&nc)
10773 .arg(&rqi)
10774 .arg(&rki)
10775 .arg(pos)
10776 .arg(&nhq)
10777 .arg(&nhk)
10778 .arg(&theta_scale)
10779 .arg(&freq_scale)
10780 .arg(t)
10781 .arg(&eps);
10782 unsafe {
10783 b.launch(cfg)?;
10784 }
10785 }
10786 None => {
10787 let null: u64 = 0;
10788 b.arg(q0)
10789 .arg(k0)
10790 .arg(v0)
10791 .arg(wq)
10792 .arg(wk)
10793 .arg(wv)
10794 .arg(&mut *q)
10795 .arg(&mut *k)
10796 .arg(&mut *v)
10797 .arg(&nc)
10798 .arg(&rqi)
10799 .arg(&rki)
10800 .arg(pos)
10801 .arg(&nhq)
10802 .arg(&nhk)
10803 .arg(&theta_scale)
10804 .arg(&freq_scale)
10805 .arg(&null)
10806 .arg(&eps);
10807 unsafe {
10808 b.launch(cfg)?;
10809 }
10810 }
10811 }
10812 Ok(())
10813 }
10814
10815 #[allow(clippy::too_many_arguments)]
10821 pub fn rms_norm_qkv_rope_append_dc(
10822 &self,
10823 q0: &CudaSlice<f32>,
10824 k0: &CudaSlice<f32>,
10825 v0: &CudaSlice<f32>,
10826 wq: &CudaSlice<f32>,
10827 wk: &CudaSlice<f32>,
10828 wv: &CudaSlice<f32>,
10829 q: &mut CudaSlice<f32>,
10830 k: &mut CudaSlice<f32>,
10831 v: &mut CudaSlice<f32>,
10832 head_dim: usize,
10833 n_rot: usize,
10834 rq: usize,
10835 rk: usize,
10836 pos: &CudaSlice<i32>,
10837 nh_q: usize,
10838 nh_k: usize,
10839 base: f32,
10840 freq_scale: f32,
10841 ff: Option<&CudaSlice<f32>>,
10842 eps: f32,
10843 kc: &mut CudaSlice<u8>,
10844 vc: &mut CudaSlice<u8>,
10845 t_dev: &CudaSlice<i32>,
10846 k_tok_bytes: usize,
10847 v_tok_bytes: usize,
10848 g: bool,
10849 ) -> Result<(), Box<dyn std::error::Error>> {
10850 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10851 let rows = rq + rk + rk;
10852 let theta_scale = base.powf(-2.0 / head_dim as f32);
10853 let (nc, rqi, rki, nhq, nhk) = (
10854 head_dim as i32,
10855 rq as i32,
10856 rk as i32,
10857 nh_q as i32,
10858 nh_k as i32,
10859 );
10860 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10861 if Self::pdl_on() && Self::pdl_wb_on() {
10862 use cudarc::driver::{DevicePtr, DevicePtrMut};
10863 let s = &self.gpu.stream();
10864 let (p0, _a0) = q0.device_ptr(s);
10865 let (p1, _a1) = k0.device_ptr(s);
10866 let (p2, _a2) = v0.device_ptr(s);
10867 let (pwq, _a3) = wq.device_ptr(s);
10868 let (pwk, _a4) = wk.device_ptr(s);
10869 let (pwv, _a5) = wv.device_ptr(s);
10870 let (pq, _a6) = q.device_ptr_mut(s);
10871 let (pk, _a7) = k.device_ptr_mut(s);
10872 let (pv, _a8) = v.device_ptr_mut(s);
10873 let (pp, _a9) = pos.device_ptr(s);
10874 let pff: u64 = match ff {
10875 Some(t) => {
10876 let (p, _gg) = t.device_ptr(s);
10877 p as u64
10878 }
10879 None => 0,
10880 };
10881 let (pkc, _a10) = kc.device_ptr_mut(s);
10882 let (pvc, _a11) = vc.device_ptr_mut(s);
10883 let (pt, _a12) = t_dev.device_ptr(s);
10884 let mut ps = [
10885 &p0 as *const _ as *mut std::ffi::c_void,
10886 &p1 as *const _ as *mut _,
10887 &p2 as *const _ as *mut _,
10888 &pwq as *const _ as *mut _,
10889 &pwk as *const _ as *mut _,
10890 &pwv as *const _ as *mut _,
10891 &pq as *const _ as *mut _,
10892 &pk as *const _ as *mut _,
10893 &pv as *const _ as *mut _,
10894 &nc as *const _ as *mut _,
10895 &rqi as *const _ as *mut _,
10896 &rki as *const _ as *mut _,
10897 &pp as *const _ as *mut _,
10898 &nhq as *const _ as *mut _,
10899 &nhk as *const _ as *mut _,
10900 &theta_scale as *const _ as *mut _,
10901 &freq_scale as *const _ as *mut _,
10902 &pff as *const _ as *mut _,
10903 &eps as *const _ as *mut _,
10904 &pkc as *const _ as *mut _,
10905 &pvc as *const _ as *mut _,
10906 &pt as *const _ as *mut _,
10907 &ktb as *const _ as *mut _,
10908 &vtb as *const _ as *mut _,
10909 ];
10910 unsafe {
10911 self.launch_pdl_flash(
10912 g,
10913 "rms_norm_qkv_rope_append_dc_f32",
10914 (rows as u32, 1, 1),
10915 (rms_block(), 1, 1),
10916 0,
10917 &mut ps,
10918 )?;
10919 }
10920 return Ok(());
10921 }
10922 let f = if g {
10923 self.func_g("rms_norm_qkv_rope_append_dc_f32")
10924 } else {
10925 self.func("rms_norm_qkv_rope_append_dc_f32")
10926 };
10927 let cfg = LaunchConfig {
10928 grid_dim: (rows as u32, 1, 1),
10929 block_dim: (rms_block(), 1, 1),
10930 shared_mem_bytes: 0,
10931 };
10932 let __s_b = self.gpu.stream();
10933 let mut b = __s_b.launch_builder(&f);
10934 match ff {
10935 Some(t) => {
10936 b.arg(q0)
10937 .arg(k0)
10938 .arg(v0)
10939 .arg(wq)
10940 .arg(wk)
10941 .arg(wv)
10942 .arg(&mut *q)
10943 .arg(&mut *k)
10944 .arg(&mut *v)
10945 .arg(&nc)
10946 .arg(&rqi)
10947 .arg(&rki)
10948 .arg(pos)
10949 .arg(&nhq)
10950 .arg(&nhk)
10951 .arg(&theta_scale)
10952 .arg(&freq_scale)
10953 .arg(t)
10954 .arg(&eps)
10955 .arg(&mut *kc)
10956 .arg(&mut *vc)
10957 .arg(t_dev)
10958 .arg(&ktb)
10959 .arg(&vtb);
10960 unsafe {
10961 b.launch(cfg)?;
10962 }
10963 }
10964 None => {
10965 let null: u64 = 0;
10966 b.arg(q0)
10967 .arg(k0)
10968 .arg(v0)
10969 .arg(wq)
10970 .arg(wk)
10971 .arg(wv)
10972 .arg(&mut *q)
10973 .arg(&mut *k)
10974 .arg(&mut *v)
10975 .arg(&nc)
10976 .arg(&rqi)
10977 .arg(&rki)
10978 .arg(pos)
10979 .arg(&nhq)
10980 .arg(&nhk)
10981 .arg(&theta_scale)
10982 .arg(&freq_scale)
10983 .arg(&null)
10984 .arg(&eps)
10985 .arg(&mut *kc)
10986 .arg(&mut *vc)
10987 .arg(t_dev)
10988 .arg(&ktb)
10989 .arg(&vtb);
10990 unsafe {
10991 b.launch(cfg)?;
10992 }
10993 }
10994 }
10995 Ok(())
10996 }
10997
10998 #[allow(clippy::too_many_arguments)]
11006 pub fn rms_norm_qkv_rope_append(
11007 &self,
11008 q0: &CudaSlice<f32>,
11009 k0: &CudaSlice<f32>,
11010 v0: &CudaSlice<f32>,
11011 wq: &CudaSlice<f32>,
11012 wk: &CudaSlice<f32>,
11013 wv: &CudaSlice<f32>,
11014 q: &mut CudaSlice<f32>,
11015 k: &mut CudaSlice<f32>,
11016 v: &mut CudaSlice<f32>,
11017 head_dim: usize,
11018 n_rot: usize,
11019 rq: usize,
11020 rk: usize,
11021 pos: &CudaSlice<i32>,
11022 nh_q: usize,
11023 nh_k: usize,
11024 base: f32,
11025 freq_scale: f32,
11026 ff: Option<&CudaSlice<f32>>,
11027 eps: f32,
11028 kc: &mut CudaSlice<u8>,
11029 vc: &mut CudaSlice<u8>,
11030 t: usize,
11031 k_tok_bytes: usize,
11032 v_tok_bytes: usize,
11033 g: bool,
11034 ) -> Result<(), Box<dyn std::error::Error>> {
11035 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11036 let rows = rq + rk + rk;
11037 let theta_scale = base.powf(-2.0 / head_dim as f32);
11038 let (nc, rqi, rki, nhq, nhk) = (
11039 head_dim as i32,
11040 rq as i32,
11041 rk as i32,
11042 nh_q as i32,
11043 nh_k as i32,
11044 );
11045 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11046 let ti = t as i32;
11047 if Self::pdl_on() && Self::pdl_wb_on() {
11048 use cudarc::driver::{DevicePtr, DevicePtrMut};
11049 let s = &self.gpu.stream();
11050 let (p0, _a0) = q0.device_ptr(s);
11051 let (p1, _a1) = k0.device_ptr(s);
11052 let (p2, _a2) = v0.device_ptr(s);
11053 let (pwq, _a3) = wq.device_ptr(s);
11054 let (pwk, _a4) = wk.device_ptr(s);
11055 let (pwv, _a5) = wv.device_ptr(s);
11056 let (pq, _a6) = q.device_ptr_mut(s);
11057 let (pk, _a7) = k.device_ptr_mut(s);
11058 let (pv, _a8) = v.device_ptr_mut(s);
11059 let (pp, _a9) = pos.device_ptr(s);
11060 let pff: u64 = match ff {
11061 Some(t) => {
11062 let (p, _gg) = t.device_ptr(s);
11063 p as u64
11064 }
11065 None => 0,
11066 };
11067 let (pkc, _a10) = kc.device_ptr_mut(s);
11068 let (pvc, _a11) = vc.device_ptr_mut(s);
11069 let mut ps = [
11070 &p0 as *const _ as *mut std::ffi::c_void,
11071 &p1 as *const _ as *mut _,
11072 &p2 as *const _ as *mut _,
11073 &pwq as *const _ as *mut _,
11074 &pwk as *const _ as *mut _,
11075 &pwv as *const _ as *mut _,
11076 &pq as *const _ as *mut _,
11077 &pk as *const _ as *mut _,
11078 &pv as *const _ as *mut _,
11079 &nc as *const _ as *mut _,
11080 &rqi as *const _ as *mut _,
11081 &rki as *const _ as *mut _,
11082 &pp as *const _ as *mut _,
11083 &nhq as *const _ as *mut _,
11084 &nhk as *const _ as *mut _,
11085 &theta_scale as *const _ as *mut _,
11086 &freq_scale as *const _ as *mut _,
11087 &pff as *const _ as *mut _,
11088 &eps as *const _ as *mut _,
11089 &pkc as *const _ as *mut _,
11090 &pvc as *const _ as *mut _,
11091 &ti as *const _ as *mut _,
11092 &ktb as *const _ as *mut _,
11093 &vtb as *const _ as *mut _,
11094 ];
11095 unsafe {
11096 self.launch_pdl_flash(
11097 g,
11098 "rms_norm_qkv_rope_append_f32",
11099 (rows as u32, 1, 1),
11100 (rms_block(), 1, 1),
11101 0,
11102 &mut ps,
11103 )?;
11104 }
11105 return Ok(());
11106 }
11107 let f = if g {
11108 self.func_g("rms_norm_qkv_rope_append_f32")
11109 } else {
11110 self.func("rms_norm_qkv_rope_append_f32")
11111 };
11112 let cfg = LaunchConfig {
11113 grid_dim: (rows as u32, 1, 1),
11114 block_dim: (rms_block(), 1, 1),
11115 shared_mem_bytes: 0,
11116 };
11117 let __s_b = self.gpu.stream();
11118 let mut b = __s_b.launch_builder(&f);
11119 let null: u64 = 0;
11120 b.arg(q0)
11121 .arg(k0)
11122 .arg(v0)
11123 .arg(wq)
11124 .arg(wk)
11125 .arg(wv)
11126 .arg(&mut *q)
11127 .arg(&mut *k)
11128 .arg(&mut *v)
11129 .arg(&nc)
11130 .arg(&rqi)
11131 .arg(&rki)
11132 .arg(pos)
11133 .arg(&nhq)
11134 .arg(&nhk)
11135 .arg(&theta_scale)
11136 .arg(&freq_scale);
11137 match ff {
11138 Some(t) => {
11139 b.arg(t);
11140 }
11141 None => {
11142 b.arg(&null);
11143 }
11144 }
11145 b.arg(&eps)
11146 .arg(&mut *kc)
11147 .arg(&mut *vc)
11148 .arg(&ti)
11149 .arg(&ktb)
11150 .arg(&vtb);
11151 unsafe {
11152 b.launch(cfg)?;
11153 }
11154 Ok(())
11155 }
11156
11157 pub fn add_q8_1(
11158 &self,
11159 a: &CudaSlice<f32>,
11160 b: &CudaSlice<f32>,
11161 res: &mut CudaSlice<f32>,
11162 ncols: usize,
11163 nrows: usize,
11164 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11165 debug_assert!(ncols % 128 == 0);
11166 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11167 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11168 let f = self.func("add_q8_1_f32");
11169 let cfg = LaunchConfig {
11170 grid_dim: (nrows as u32, 1, 1),
11171 block_dim: (rms_block(), 1, 1),
11172 shared_mem_bytes: 0,
11173 };
11174 let nc = ncols as i32;
11175 let __s_b2 = self.gpu.stream();
11176 let mut b2 = __s_b2.launch_builder(&f);
11177 b2.arg(a)
11178 .arg(b)
11179 .arg(&mut *res)
11180 .arg(&mut out_q)
11181 .arg(&mut out_d)
11182 .arg(&nc);
11183 unsafe {
11184 b2.launch(cfg)?;
11185 }
11186 Ok((out_q, out_d))
11187 }
11188
11189 pub fn rms_pre_add_q8_1(
11193 &self,
11194 a: &CudaSlice<f32>,
11195 wa: &CudaSlice<f32>,
11196 b: &CudaSlice<f32>,
11197 res: &mut CudaSlice<f32>,
11198 ncols: usize,
11199 nrows: usize,
11200 eps: f32,
11201 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11202 debug_assert!(ncols % 128 == 0);
11203 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11204 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11205 let f = self.func("rms_pre_add_q8_1_f32");
11206 let cfg = LaunchConfig {
11207 grid_dim: (nrows as u32, 1, 1),
11208 block_dim: (rms_block(), 1, 1),
11209 shared_mem_bytes: 0,
11210 };
11211 let (nc, ep) = (ncols as i32, eps);
11212 let __s_b2 = self.gpu.stream();
11213 let mut b2 = __s_b2.launch_builder(&f);
11214 b2.arg(a)
11215 .arg(wa)
11216 .arg(b)
11217 .arg(&mut *res)
11218 .arg(&mut out_q)
11219 .arg(&mut out_d)
11220 .arg(&nc)
11221 .arg(&ep);
11222 unsafe {
11223 b2.launch(cfg)?;
11224 }
11225 Ok((out_q, out_d))
11226 }
11227
11228 pub fn l2_v2_on(ncols: usize) -> bool {
11232 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11233 }
11234
11235 pub fn l2_norm_pp(
11236 &self,
11237 x: &CudaSlice<f32>,
11238 dst: &mut CudaSlice<f32>,
11239 dst16: Option<&mut CudaSlice<u8>>,
11240 ncols: usize,
11241 nrows: usize,
11242 eps: f32,
11243 ) -> Result<(), Box<dyn std::error::Error>> {
11244 if Self::l2_v2_on(ncols) {
11245 let f = self.func("l2_norm_pp_v2_f32");
11246 let rows_per_block = 8u32; let cfg = LaunchConfig {
11248 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11249 block_dim: (256, 1, 1),
11250 shared_mem_bytes: 0,
11251 };
11252 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11253 let d16: u64 = match dst16 {
11255 Some(d) => self.addr_u8(d),
11256 None => 0,
11257 };
11258 let __s_b = self.gpu.stream();
11259 let mut b = __s_b.launch_builder(&f);
11260 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11261 unsafe {
11262 b.launch(cfg)?;
11263 }
11264 return Ok(());
11265 }
11266 self.l2_norm(x, dst, ncols, nrows, eps)
11267 }
11268
11269 pub fn l2_norm(
11270 &self,
11271 x: &CudaSlice<f32>,
11272 dst: &mut CudaSlice<f32>,
11273 ncols: usize,
11274 nrows: usize,
11275 eps: f32,
11276 ) -> Result<(), Box<dyn std::error::Error>> {
11277 let f = self.func("l2_norm_f32");
11278 let cfg = LaunchConfig {
11279 grid_dim: (nrows as u32, 1, 1),
11280 block_dim: (256, 1, 1),
11281 shared_mem_bytes: 0,
11282 };
11283 let (nc, e) = (ncols as i32, eps);
11284 let __s_b = self.gpu.stream();
11285 let mut b = __s_b.launch_builder(&f);
11286 b.arg(x).arg(dst).arg(&nc).arg(&e);
11287 unsafe {
11288 b.launch(cfg)?;
11289 }
11290 Ok(())
11291 }
11292
11293 pub fn l2_norm_decode(
11299 &self,
11300 x: &CudaSlice<f32>,
11301 dst: &mut CudaSlice<f32>,
11302 ncols: usize,
11303 nrows: usize,
11304 eps: f32,
11305 ) -> Result<(), Box<dyn std::error::Error>> {
11306 let f = self.func("l2_norm_f32");
11307 let cfg = LaunchConfig {
11308 grid_dim: (nrows as u32, 1, 1),
11309 block_dim: (32, 1, 1),
11310 shared_mem_bytes: 0,
11311 };
11312 let (nc, e) = (ncols as i32, eps);
11313 let __s_b = self.gpu.stream();
11314 let mut b = __s_b.launch_builder(&f);
11315 b.arg(x).arg(dst).arg(&nc).arg(&e);
11316 unsafe {
11317 b.launch(cfg)?;
11318 }
11319 Ok(())
11320 }
11321
11322 pub fn rope_neox(
11324 &self,
11325 x: &mut CudaSlice<f32>,
11326 pos: &CudaSlice<i32>,
11327 head_dim: usize,
11328 n_dims: usize,
11329 n_heads: usize,
11330 n_tokens: usize,
11331 freq_base: f32,
11332 freq_scale: f32,
11333 ) -> Result<(), Box<dyn std::error::Error>> {
11334 let f = self.func("rope_neox_f32");
11335 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11336 let grid = (n_heads * n_tokens) as u32;
11337 let cfg = LaunchConfig {
11338 grid_dim: (grid, 1, 1),
11339 block_dim: ((head_dim / 2) as u32, 1, 1),
11340 shared_mem_bytes: 0,
11341 };
11342 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11343 let __s_b = self.gpu.stream();
11344 let mut b = __s_b.launch_builder(&f);
11345 b.arg(x)
11346 .arg(pos)
11347 .arg(&hd)
11348 .arg(&nd)
11349 .arg(&nh)
11350 .arg(&theta_scale)
11351 .arg(&freq_scale);
11352 unsafe {
11353 b.launch(cfg)?;
11354 }
11355 Ok(())
11356 }
11357
11358 pub fn rope_neox_ff(
11360 &self,
11361 x: &mut CudaSlice<f32>,
11362 pos: &CudaSlice<i32>,
11363 head_dim: usize,
11364 n_dims: usize,
11365 n_heads: usize,
11366 n_tokens: usize,
11367 freq_base: f32,
11368 freq_scale: f32,
11369 ff: &CudaSlice<f32>,
11370 ) -> Result<(), Box<dyn std::error::Error>> {
11371 let f = self.func("rope_neox_ff_f32");
11372 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11373 let grid = (n_heads * n_tokens) as u32;
11374 let cfg = LaunchConfig {
11375 grid_dim: (grid, 1, 1),
11376 block_dim: ((head_dim / 2) as u32, 1, 1),
11377 shared_mem_bytes: 0,
11378 };
11379 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11380 let __s_b = self.gpu.stream();
11381 let mut b = __s_b.launch_builder(&f);
11382 b.arg(x)
11383 .arg(pos)
11384 .arg(&hd)
11385 .arg(&nd)
11386 .arg(&nh)
11387 .arg(&theta_scale)
11388 .arg(&freq_scale)
11389 .arg(ff);
11390 unsafe {
11391 b.launch(cfg)?;
11392 }
11393 Ok(())
11394 }
11395
11396 #[allow(clippy::too_many_arguments)]
11398 pub fn rope_neox2(
11399 &self,
11400 q: &mut CudaSlice<f32>,
11401 k: &mut CudaSlice<f32>,
11402 pos: &CudaSlice<i32>,
11403 head_dim: usize,
11404 n_dims: usize,
11405 nh_q: usize,
11406 nh_k: usize,
11407 n_tokens: usize,
11408 freq_base: f32,
11409 freq_scale: f32,
11410 ff: Option<&CudaSlice<f32>>,
11411 ) -> Result<(), Box<dyn std::error::Error>> {
11412 let f = self.func("rope_neox2_f32");
11413 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11414 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11415 let cfg = LaunchConfig {
11416 grid_dim: (grid, 1, 1),
11417 block_dim: ((head_dim / 2) as u32, 1, 1),
11418 shared_mem_bytes: 0,
11419 };
11420 let (hd, nd, nq, nk, nt) = (
11421 head_dim as i32,
11422 n_dims as i32,
11423 nh_q as i32,
11424 nh_k as i32,
11425 n_tokens as i32,
11426 );
11427 let __s_b = self.gpu.stream();
11428 let mut b = __s_b.launch_builder(&f);
11429 b.arg(q)
11430 .arg(k)
11431 .arg(pos)
11432 .arg(&hd)
11433 .arg(&nd)
11434 .arg(&nq)
11435 .arg(&nk)
11436 .arg(&nt)
11437 .arg(&theta_scale)
11438 .arg(&freq_scale);
11439 match ff {
11440 Some(ffv) => {
11441 b.arg(ffv);
11442 unsafe {
11443 b.launch(cfg)?;
11444 }
11445 }
11446 None => {
11447 let null: u64 = 0;
11448 b.arg(&null);
11449 unsafe {
11450 b.launch(cfg)?;
11451 }
11452 }
11453 }
11454 Ok(())
11455 }
11456
11457 pub fn gelu_tanh_mul(
11459 &self,
11460 gate: &CudaSlice<f32>,
11461 up: &CudaSlice<f32>,
11462 dst: &mut CudaSlice<f32>,
11463 n: usize,
11464 ) -> Result<(), Box<dyn std::error::Error>> {
11465 let f = self.func("gelu_tanh_mul_f32");
11466 let cfg = LaunchConfig::for_num_elems(n as u32);
11467 let ni = n as i32;
11468 let __s_b = self.gpu.stream();
11469 let mut b = __s_b.launch_builder(&f);
11470 b.arg(gate).arg(up).arg(dst).arg(&ni);
11471 unsafe {
11472 b.launch(cfg)?;
11473 }
11474 Ok(())
11475 }
11476
11477 pub fn silu_mul(
11478 &self,
11479 gate: &CudaSlice<f32>,
11480 up: &CudaSlice<f32>,
11481 dst: &mut CudaSlice<f32>,
11482 n: usize,
11483 ) -> Result<(), Box<dyn std::error::Error>> {
11484 let f = self.func("silu_mul_f32");
11485 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11487 let ni = n as i32;
11488 let __s_b = self.gpu.stream();
11489 let mut b = __s_b.launch_builder(&f);
11490 b.arg(gate).arg(up).arg(dst).arg(&ni);
11491 unsafe {
11492 b.launch(cfg)?;
11493 }
11494 Ok(())
11495 }
11496
11497 pub fn silu_mul_host_expf(
11499 &self,
11500 gate: &CudaSlice<f32>,
11501 up: &CudaSlice<f32>,
11502 dst: &mut CudaSlice<f32>,
11503 n: usize,
11504 ) -> Result<(), Box<dyn std::error::Error>> {
11505 let f = self.func("silu_mul_host_expf_f32");
11506 let cfg = LaunchConfig::for_num_elems(n as u32);
11507 let ni = n as i32;
11508 let __s_b = self.gpu.stream();
11509 let mut b = __s_b.launch_builder(&f);
11510 b.arg(gate).arg(up).arg(dst).arg(&ni);
11511 unsafe {
11512 b.launch(cfg)?;
11513 }
11514 Ok(())
11515 }
11516
11517 pub fn silu_clamped_mul_host_expf(
11519 &self,
11520 gate: &CudaSlice<f32>,
11521 up: &CudaSlice<f32>,
11522 limit: f32,
11523 dst: &mut CudaSlice<f32>,
11524 n: usize,
11525 ) -> Result<(), Box<dyn std::error::Error>> {
11526 if !limit.is_finite() || limit <= 0.0 {
11527 return Err(
11528 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11529 );
11530 }
11531 let f = self.func("silu_clamped_mul_host_expf_f32");
11532 let cfg = LaunchConfig::for_num_elems(n as u32);
11533 let ni = n as i32;
11534 let __s_b = self.gpu.stream();
11535 let mut b = __s_b.launch_builder(&f);
11536 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11537 unsafe {
11538 b.launch(cfg)?;
11539 }
11540 Ok(())
11541 }
11542
11543 pub fn silu_mul_f16out(
11546 &self,
11547 gate: &CudaSlice<f32>,
11548 up: &CudaSlice<f32>,
11549 dst: &mut CudaSlice<f32>,
11550 dst16: &mut CudaSlice<u8>,
11551 n: usize,
11552 ) -> Result<(), Box<dyn std::error::Error>> {
11553 let f = self.func("silu_mul_f16out_f32");
11554 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11555 let ni = n as i32;
11556 let __s_b = self.gpu.stream();
11557 let mut b = __s_b.launch_builder(&f);
11558 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11559 unsafe {
11560 b.launch(cfg)?;
11561 }
11562 Ok(())
11563 }
11564
11565 pub fn silu_mul_scaled(
11572 &self,
11573 gate: &CudaSlice<f32>,
11574 up: &CudaSlice<f32>,
11575 gs: f32,
11576 us: f32,
11577 dst: &mut CudaSlice<f32>,
11578 n: usize,
11579 ) -> Result<(), Box<dyn std::error::Error>> {
11580 let f = self.func("silu_mul_scaled_f32");
11581 let cfg = LaunchConfig::for_num_elems(n as u32);
11582 let ni = n as i32;
11583 let (gsf, usf) = (gs, us);
11584 let __s_b = self.gpu.stream();
11585 let mut b = __s_b.launch_builder(&f);
11586 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11587 unsafe {
11588 b.launch(cfg)?;
11589 }
11590 Ok(())
11591 }
11592
11593 #[allow(clippy::too_many_arguments)]
11597 pub fn swigluoai_mul_scaled(
11598 &self,
11599 gate: &CudaSlice<f32>,
11600 up: &CudaSlice<f32>,
11601 gs: f32,
11602 us: f32,
11603 alpha: f32,
11604 limit: f32,
11605 dst: &mut CudaSlice<f32>,
11606 n: usize,
11607 ) -> Result<(), Box<dyn std::error::Error>> {
11608 let f = self.func("swigluoai_mul_scaled_f32");
11609 let cfg = LaunchConfig::for_num_elems(n as u32);
11610 let ni = n as i32;
11611 let __s_b = self.gpu.stream();
11612 let mut b = __s_b.launch_builder(&f);
11613 b.arg(gate)
11614 .arg(up)
11615 .arg(&gs)
11616 .arg(&us)
11617 .arg(&alpha)
11618 .arg(&limit)
11619 .arg(dst)
11620 .arg(&ni);
11621 unsafe {
11622 b.launch(cfg)?;
11623 }
11624 Ok(())
11625 }
11626
11627 pub fn silu_mul_scaled_q8_1(
11635 &self,
11636 gate: &CudaSlice<f32>,
11637 up: &CudaSlice<f32>,
11638 gs: f32,
11639 us: f32,
11640 n: usize,
11641 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11642 let f = self.func("silu_mul_scaled_q8_1");
11643 let nblk = n / 32;
11644 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);
11648 let (gsf, usf, ni) = (gs, us, n as i32);
11649 let __s_b = self.gpu.stream();
11650 let mut b = __s_b.launch_builder(&f);
11651 b.arg(gate)
11652 .arg(up)
11653 .arg(&gsf)
11654 .arg(&usf)
11655 .arg(&mut aq)
11656 .arg(&mut ad)
11657 .arg(&ni);
11658 unsafe {
11659 b.launch(cfg)?;
11660 }
11661 Ok((aq, ad))
11662 }
11663
11664 pub fn add(
11665 &self,
11666 a: &CudaSlice<f32>,
11667 b_in: &CudaSlice<f32>,
11668 dst: &mut CudaSlice<f32>,
11669 n: usize,
11670 ) -> Result<(), Box<dyn std::error::Error>> {
11671 let f = self.func("add_f32");
11672 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11674 let ni = n as i32;
11675 let __s_bld = self.gpu.stream();
11676 let mut bld = __s_bld.launch_builder(&f);
11677 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11678 unsafe {
11679 bld.launch(cfg)?;
11680 }
11681 Ok(())
11682 }
11683
11684 pub fn mul(
11685 &self,
11686 a: &CudaSlice<f32>,
11687 b_in: &CudaSlice<f32>,
11688 dst: &mut CudaSlice<f32>,
11689 n: usize,
11690 ) -> Result<(), Box<dyn std::error::Error>> {
11691 let f = self.func("mul_f32");
11692 let cfg = LaunchConfig::for_num_elems(n as u32);
11693 let ni = n as i32;
11694 let __s_bld = self.gpu.stream();
11695 let mut bld = __s_bld.launch_builder(&f);
11696 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11697 unsafe {
11698 bld.launch(cfg)?;
11699 }
11700 Ok(())
11701 }
11702
11703 pub fn matmul(
11706 &self,
11707 w: &crate::model::GpuTensor,
11708 x: &CudaSlice<f32>,
11709 m: usize,
11710 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11711 use crate::model::GpuTensor;
11712 let in_f = w.in_features();
11713 let out_f = w.out_features();
11714 #[allow(non_snake_case)]
11722 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11725 usize::MAX
11726 } else {
11727 16usize
11728 };
11729
11730 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11755 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11756 return Ok(y);
11757 }
11758 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11765 return Ok(y);
11766 }
11767 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11770 return Ok(y);
11771 }
11772 }
11773 if let GpuTensor::Quant { qtype, .. } = w {
11788 if *qtype == QT_F8_E4M3_BLK {
11789 if m >= GEMM_M_THRESHOLD {
11790 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11791 return Ok(y);
11792 }
11793 }
11794 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11795 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11796 return Ok(y);
11797 }
11798 }
11799 }
11800 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11801 return self.qmatvec_mmq(w, x, m);
11802 }
11803 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11804 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11805 return self.qmatvec_gemm(w, &aq, &ad, m);
11806 }
11807 if m >= GEMM_M_THRESHOLD {
11810 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11811 return Ok(y);
11812 }
11813 }
11814 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11818 if m == 1 && fast {
11823 if let GpuTensor::Quant {
11824 bytes,
11825 qtype,
11826 row_bytes,
11827 rp,
11828 rp4,
11829 scale,
11830 ..
11831 } = w
11832 {
11833 if self.mmvq_supports(*qtype) {
11834 let (bytes, rp) = match rp4 {
11838 Some(m4) => (m4, true),
11839 None => (bytes, *rp),
11840 };
11841 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11842 return self.qmatvec_mmvq(
11843 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11844 );
11845 }
11846 }
11847 }
11848 if (2..=16).contains(&m)
11864 && fast
11865 && std::env::var("MEMRA_NO_BATCHED").is_err()
11866 && (m <= 4 || Self::b8_enabled())
11867 {
11868 let m_ok = m <= 8
11878 || matches!(w, GpuTensor::Quant { qtype, .. }
11879 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11880 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11881 if m_ok {
11882 if let GpuTensor::Quant {
11883 bytes,
11884 qtype,
11885 row_bytes,
11886 rp,
11887 rp4,
11888 ..
11889 } = w
11890 {
11891 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11892 let (bytes, rp) = match rp4 {
11893 Some(m4) => (m4, true),
11894 None => (bytes, *rp),
11895 };
11896 let mcols = Self::batched_mcols(m);
11897 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11898 let mut y = self.qmatvec_mmvq_batched(
11899 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11900 )?;
11901 if let GpuTensor::Quant { scale, .. } = w {
11902 if *scale != 1.0 {
11903 self.scale_inplace(&mut y, *scale, m * out_f)?;
11904 }
11905 }
11906 return Ok(y);
11907 }
11908 }
11909 }
11910 }
11911 if fast {
11917 if let GpuTensor::Quant {
11918 bytes,
11919 qtype,
11920 row_bytes,
11921 scale,
11922 ..
11923 } = w
11924 {
11925 if *qtype == QT_F8_E4M3 {
11926 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11927 return self.qmatvec_mmvq(
11928 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11929 );
11930 }
11931 }
11932 }
11933 let mut y = match w {
11934 GpuTensor::Quant {
11935 bytes,
11936 qtype,
11937 row_bytes,
11938 ..
11939 } if fast && *qtype == QT_Q8_0 => {
11940 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11941 }
11942 GpuTensor::Quant {
11943 bytes,
11944 qtype,
11945 row_bytes,
11946 ..
11947 } if fast && *qtype == QT_Q4_K => {
11948 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11949 }
11950 GpuTensor::Quant {
11951 bytes,
11952 qtype,
11953 row_bytes,
11954 ..
11955 } if fast && *qtype == QT_Q6_K => {
11956 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11957 }
11958 GpuTensor::Quant {
11959 bytes,
11960 qtype,
11961 row_bytes,
11962 ..
11963 } if fast && *qtype == QT_Q5_K => {
11964 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11965 }
11966 GpuTensor::Quant {
11967 bytes,
11968 qtype,
11969 row_bytes,
11970 ..
11971 } if fast && *qtype == QT_Q3_K => {
11972 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11973 }
11974 GpuTensor::Quant {
11975 bytes,
11976 qtype,
11977 row_bytes,
11978 rp,
11979 ..
11980 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11981 if *rp {
11982 "qmatvec_nvfp4_dp4a_rp"
11983 } else {
11984 "qmatvec_nvfp4_dp4a"
11985 },
11986 &bytes.slice(0..bytes.len()),
11987 x,
11988 m,
11989 in_f,
11990 out_f,
11991 *row_bytes,
11992 )?,
11993 GpuTensor::Quant {
11997 bytes,
11998 qtype,
11999 row_bytes,
12000 ..
12001 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12002 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12003 }
12004 GpuTensor::Quant {
12009 bytes,
12010 qtype,
12011 row_bytes,
12012 rp,
12013 ..
12014 } =>
12015 {
12018 self.qmatvec(
12019 bytes,
12020 x,
12021 m,
12022 in_f,
12023 out_f,
12024 if *rp && *qtype == QT_NVFP4 {
12025 QT_NVFP4_RP
12026 } else {
12027 *qtype
12028 },
12029 *row_bytes,
12030 )?
12031 }
12032 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12033 GpuTensor::FloatBf16 { data, .. } => {
12036 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12043 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12044 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12045 y
12046 } else {
12047 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12048 }
12049 }
12050 };
12051 if let GpuTensor::Quant { scale, .. } = w {
12053 if *scale != 1.0 {
12054 self.scale_inplace(&mut y, *scale, m * out_f)?;
12055 }
12056 }
12057 Ok(y)
12058 }
12059
12060 pub fn stage_a_raw_needed() -> bool {
12070 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12071 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12072 }
12073
12074 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12077 use crate::model::GpuTensor;
12078 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12079 return false;
12080 }
12081 match w {
12082 GpuTensor::Quant { qtype, .. } => {
12089 matches!(
12090 *qtype,
12091 QT_Q8_0
12092 | QT_Q4_K
12093 | QT_Q6_K
12094 | QT_Q5_K
12095 | QT_Q3_K
12096 | QT_NVFP4
12097 | QT_F8_E4M3
12098 | QT_F8_E4M3_BLK
12099 | QT_Q4_0
12100 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12101 }
12102 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12103 }
12104 }
12105
12106 pub fn matmul_pre(
12111 &self,
12112 w: &crate::model::GpuTensor,
12113 aq: &CudaSlice<i8>,
12114 ad: &CudaSlice<f32>,
12115 x_fallback: &CudaSlice<f32>,
12116 m: usize,
12117 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12118 use crate::model::GpuTensor;
12119 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12125 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12128 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12129 return Ok(y);
12130 }
12131 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12134 return Ok(y);
12135 }
12136 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12138 return Ok(y);
12139 }
12140 }
12141 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12147 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12148 return Ok(y);
12149 }
12150 }
12151 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12152 return Ok(y);
12153 }
12154 if m >= 16
12159 && w.out_features() >= 128
12160 && self.mmq_supports(w)
12161 && !self.verify_exact_on()
12162 && x_raw_ok
12163 {
12164 return self.qmatvec_mmq(w, x_fallback, m);
12165 }
12166 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12169 if let Some(y) =
12170 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12171 {
12172 return Ok(y);
12173 }
12174 }
12175 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12178 return self.qmatvec_gemm(w, aq, ad, m);
12179 }
12180 if !self.uses_q8_1_fast(w) {
12199 if !x_raw_ok {
12200 return Err(format!(
12201 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12202 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12203 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12204 activation (see Engine::rms_norm_decode, which is bit-identical to \
12205 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12206 x_fallback.len(),
12207 m,
12208 w.in_features(),
12209 m * w.in_features()
12210 )
12211 .into());
12212 }
12213 return self.matmul(w, x_fallback, m);
12214 }
12215 let in_f = w.in_features();
12216 let out_f = w.out_features();
12217 let (bytes, qtype, row_bytes, scale, rp) = match w {
12218 GpuTensor::Quant {
12219 bytes,
12220 qtype,
12221 row_bytes,
12222 scale,
12223 rp,
12224 ..
12225 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12226 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12227 };
12228 let (mbytes, mrp) = match w {
12231 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12232 _ => (bytes, rp),
12233 };
12234 if m == 1 && self.mmvq_supports(qtype) {
12238 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12239 }
12240 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12253 && std::env::var("MEMRA_NO_BATCHED").is_err()
12254 && (m <= 4 || Self::b8_enabled())
12255 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12259 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12260 {
12261 let mcols = Self::batched_mcols(m);
12262 return self.qmatvec_mmvq_batched(
12263 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12264 );
12265 }
12266 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12272 let (b2, r2) = if qtype == QT_Q4_0 {
12273 (mbytes, mrp)
12274 } else {
12275 (bytes, rp)
12276 };
12277 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12278 }
12279 let name = match qtype {
12280 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12281 QT_Q4_K => "qmatvec_q4_K_dp4a",
12282 QT_Q6_K => "qmatvec_q6_K_dp4a",
12283 QT_Q5_K => "qmatvec_q5_K_dp4a",
12284 QT_Q3_K => "qmatvec_q3_K_dp4a",
12285 QT_NVFP4 => {
12286 if rp {
12287 "qmatvec_nvfp4_dp4a_rp"
12288 } else {
12289 "qmatvec_nvfp4_dp4a"
12290 }
12291 }
12292 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12293 _ => unreachable!(),
12294 };
12295 let f = self.func(name);
12296 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12298 grid_dim: (out_f as u32, m as u32, 1),
12299 block_dim: (128, 1, 1),
12300 shared_mem_bytes: 0,
12301 };
12302 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12303 let __s_b = self.gpu.stream();
12304 let mut b = __s_b.launch_builder(&f);
12305 b.arg(bytes)
12306 .arg(aq)
12307 .arg(ad)
12308 .arg(&mut y)
12309 .arg(&inf)
12310 .arg(&outf)
12311 .arg(&mi)
12312 .arg(&rb);
12313 unsafe {
12314 b.launch(cfg)?;
12315 }
12316 if scale != 1.0 {
12317 self.scale_inplace(&mut y, scale, m * out_f)?;
12318 }
12319 Ok(y)
12320 }
12321
12322 pub fn matmul_decode_exact(
12330 &self,
12331 w: &crate::model::GpuTensor,
12332 x: &CudaSlice<f32>,
12333 m: usize,
12334 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12335 use crate::model::GpuTensor;
12336 if let GpuTensor::Float { data, .. } = w {
12344 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12345 }
12346 if let GpuTensor::FloatBf16 { data, .. } = w {
12349 let (in_f, out_f) = (w.in_features(), w.out_features());
12350 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12353 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12354 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12355 return Ok(y);
12356 }
12357 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12358 }
12359 if !self.uses_q8_1_fast(w) {
12360 return self.matmul(w, x, m);
12361 }
12362 let in_f = w.in_features();
12363 let out_f = w.out_features();
12364 let (bytes, qtype, row_bytes, scale, rp) = match w {
12365 GpuTensor::Quant {
12366 bytes,
12367 qtype,
12368 row_bytes,
12369 scale,
12370 rp,
12371 ..
12372 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12373 _ => return self.matmul(w, x, m),
12374 };
12375 let (bytes, rp) = match w {
12378 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12379 _ => (bytes, rp),
12380 };
12381 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12382 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12386 return Ok(y);
12387 }
12388 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12397 && std::env::var("MEMRA_NO_BATCHED").is_err()
12398 && (m <= 4 || Self::b8_enabled())
12399 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12402 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12403 {
12404 let mcols = Self::batched_mcols(m);
12405 return self.qmatvec_mmvq_batched(
12406 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12407 );
12408 }
12409 if self.mmvq_supports(qtype) {
12410 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12413 }
12414 self.matmul_pre(w, &aq, &ad, x, m)
12417 }
12418
12419 pub fn matmul_decode_exact_pre(
12429 &self,
12430 w: &crate::model::GpuTensor,
12431 aq: &CudaSlice<i8>,
12432 ad: &CudaSlice<f32>,
12433 m: usize,
12434 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12435 use crate::model::GpuTensor;
12436 debug_assert!(
12437 self.uses_q8_1_fast(w),
12438 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12439 );
12440 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12442 return Ok(y);
12443 }
12444 let in_f = w.in_features();
12445 let out_f = w.out_features();
12446 let (bytes, qtype, row_bytes, scale, rp) = match w {
12447 GpuTensor::Quant {
12448 bytes,
12449 qtype,
12450 row_bytes,
12451 scale,
12452 rp,
12453 ..
12454 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12455 _ => {
12456 return Err(
12457 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12458 );
12459 }
12460 };
12461 let (bytes, rp) = match w {
12463 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12464 _ => (bytes, rp),
12465 };
12466 if (2..=16).contains(&m)
12468 && self.batched_supports(qtype)
12469 && self.mmvq_supports(qtype)
12470 && std::env::var("MEMRA_NO_BATCHED").is_err()
12471 && (m <= 4 || Self::b8_enabled())
12472 && (m <= 8
12473 || qtype == QT_Q4_0
12474 || qtype == QT_Q6_K
12475 || qtype == QT_F8_E4M3
12476 || qtype == QT_NVFP4
12477 || qtype == QT_Q4_K
12478 || qtype == QT_Q5_K
12479 || qtype == QT_Q8_0)
12480 {
12481 let mcols = Self::batched_mcols(m);
12482 return self.qmatvec_mmvq_batched(
12483 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12484 );
12485 }
12486 if self.mmvq_supports(qtype) {
12487 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12488 }
12489 let x0 = self.zeros(0)?;
12492 self.matmul_pre(w, aq, ad, &x0, m)
12493 }
12494
12495 pub fn matmul_decode_exact_dual_pre(
12504 &self,
12505 w0: &crate::model::GpuTensor,
12506 w1: &crate::model::GpuTensor,
12507 aq: &CudaSlice<i8>,
12508 ad: &CudaSlice<f32>,
12509 m: usize,
12510 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12511 {
12512 use crate::model::GpuTensor;
12513 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12514 let on = *ON.get_or_init(|| {
12515 std::env::var("MEMRA_SPEC_DUAL_T")
12516 .map(|v| v != "0")
12517 .unwrap_or(true)
12518 });
12519 if !on
12520 || !(2..=7).contains(&m)
12521 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12522 || !self.uses_q8_1_fast(w0)
12523 || !self.uses_q8_1_fast(w1)
12524 {
12525 return Ok(None);
12526 }
12527 if !self.mmvq_supports(QT_NVFP4) {
12532 return Ok(None);
12533 }
12534 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12535 if w1.in_features() != in_f || w1.out_features() != out_f {
12536 return Ok(None);
12537 }
12538 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12539 (
12540 GpuTensor::Quant {
12541 bytes: b0,
12542 qtype: q0,
12543 row_bytes: rb0,
12544 scale: s0,
12545 rp: rp0,
12546 rp4: None,
12547 ..
12548 },
12549 GpuTensor::Quant {
12550 bytes: b1,
12551 qtype: q1,
12552 row_bytes: rb1,
12553 scale: s1,
12554 rp: rp1,
12555 rp4: None,
12556 ..
12557 },
12558 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12559 (b0, b1, *rb0, *s0, *s1, *rp0)
12560 }
12561 _ => return Ok(None),
12562 };
12563 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12566 {
12567 return Ok(None);
12568 }
12569 let (y0, y1) =
12570 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12571 Ok(Some(((y0, s0), (y1, s1))))
12572 }
12573
12574 pub fn matmul_decode_exact_group4_pre(
12586 &self,
12587 ws: [&crate::model::GpuTensor; 4],
12588 aq: &CudaSlice<i8>,
12589 ad: &CudaSlice<f32>,
12590 m: usize,
12591 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12592 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12593 let on = *ON.get_or_init(|| {
12594 std::env::var("MEMRA_TK_GDN_GROUP")
12595 .map(|v| v != "0")
12596 .unwrap_or(true)
12597 });
12598 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12599 }
12600
12601 pub fn matmul_decode_exact_group3_pre(
12606 &self,
12607 ws: [&crate::model::GpuTensor; 3],
12608 aq: &CudaSlice<i8>,
12609 ad: &CudaSlice<f32>,
12610 m: usize,
12611 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12612 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12613 let on = *ON.get_or_init(|| {
12614 std::env::var("MEMRA_TK_FA_GROUP")
12615 .map(|v| v != "0")
12616 .unwrap_or(true)
12617 });
12618 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12619 }
12620
12621 fn matmul_decode_exact_group_pre(
12625 &self,
12626 ws: &[&crate::model::GpuTensor],
12627 aq: &CudaSlice<i8>,
12628 ad: &CudaSlice<f32>,
12629 m: usize,
12630 on: bool,
12631 tag: &'static str,
12632 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12633 use crate::model::GpuTensor;
12634 if !on
12635 || !(2..=16).contains(&m)
12636 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12637 || (m > 4 && !Self::b8_enabled())
12638 || !self.mmvq_supports(QT_NVFP4)
12639 || !self.batched_supports(QT_NVFP4)
12640 {
12641 return Ok(None);
12642 }
12643 let in_f = ws[0].in_features();
12644 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12645 for w in ws {
12646 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12647 return Ok(None);
12648 }
12649 match w {
12650 GpuTensor::Quant {
12651 bytes,
12652 qtype,
12653 scale,
12654 rp: true,
12655 rp4: None,
12656 ..
12657 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12658 parts.push((bytes, w.out_features(), *scale));
12659 }
12660 _ => return Ok(None),
12661 }
12662 }
12663 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12665 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12666 let mcols = if (5..=7).contains(&m) && b567 {
12667 m
12668 } else {
12669 Self::batched_mcols(m)
12670 };
12671 let kname: &'static str = match mcols {
12672 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12673 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12674 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12675 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12676 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12677 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12678 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12679 _ => return Ok(None),
12680 };
12681 if std::env::var("MEMRA_DEBUG").is_ok() {
12684 use std::sync::Mutex;
12685 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12686 let mut seen = SEEN.lock().unwrap();
12687 if !seen.contains(&tag) {
12688 seen.push(tag);
12689 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12690 }
12691 }
12692 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12695 let three = parts.len() == 3;
12696 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12697 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12698 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12699 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12702 let cfg = LaunchConfig {
12703 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12704 block_dim: (32, ROWS_PER_BLOCK, 1),
12705 shared_mem_bytes: 0,
12706 };
12707 let (inf, mi) = (in_f as i32, m as i32);
12708 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12709 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12710 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12711 let s3 = if three { 1.0f32 } else { parts[3].2 };
12712 let w3 = if three { parts[0].0 } else { parts[3].0 };
12713 let f = self.func(kname);
12714 let __s_b = self.gpu.stream();
12715 let mut b = __s_b.launch_builder(&f);
12716 b.arg(parts[0].0)
12717 .arg(parts[1].0)
12718 .arg(parts[2].0)
12719 .arg(w3)
12720 .arg(aq)
12721 .arg(ad)
12722 .arg(&mut y0)
12723 .arg(&mut y1)
12724 .arg(&mut y2)
12725 .arg(&mut y3)
12726 .arg(&inf)
12727 .arg(&n0)
12728 .arg(&n1)
12729 .arg(&n2)
12730 .arg(&n3)
12731 .arg(&mi)
12732 .arg(&s0)
12733 .arg(&s1)
12734 .arg(&s2)
12735 .arg(&s3);
12736 unsafe {
12737 b.launch(cfg)?;
12738 }
12739 Ok(Some(if three {
12740 vec![y0, y1, y2]
12741 } else {
12742 vec![y0, y1, y2, y3]
12743 }))
12744 }
12745
12746 pub fn matmul_decode_exact_dual(
12762 &self,
12763 w0: &crate::model::GpuTensor,
12764 w1: &crate::model::GpuTensor,
12765 x: &CudaSlice<f32>,
12766 m: usize,
12767 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12768 use crate::model::GpuTensor;
12769 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12770 let on = *ON.get_or_init(|| {
12771 std::env::var("MEMRA_SPEC_DUAL_T")
12772 .map(|v| v != "0")
12773 .unwrap_or(true)
12774 });
12775 if !on
12776 || !(2..=4).contains(&m)
12777 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12778 || !self.uses_q8_1_fast(w0)
12779 || !self.uses_q8_1_fast(w1)
12780 {
12781 return Ok(None);
12782 }
12783 if !self.mmvq_supports(QT_NVFP4) {
12788 return Ok(None);
12789 }
12790 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12791 if w1.in_features() != in_f || w1.out_features() != out_f {
12792 return Ok(None);
12793 }
12794 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12795 (
12796 GpuTensor::Quant {
12797 bytes: b0,
12798 qtype: q0,
12799 row_bytes: rb0,
12800 scale: s0,
12801 rp: rp0,
12802 rp4: None,
12803 ..
12804 },
12805 GpuTensor::Quant {
12806 bytes: b1,
12807 qtype: q1,
12808 row_bytes: rb1,
12809 scale: s1,
12810 rp: rp1,
12811 rp4: None,
12812 ..
12813 },
12814 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12815 (b0, b1, *rb0, *s0, *s1, *rp0)
12816 }
12817 _ => return Ok(None),
12818 };
12819 if std::env::var("MEMRA_DEBUG").is_ok() {
12822 static ONCE: std::sync::Once = std::sync::Once::new();
12823 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12824 }
12825 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12826 let (y0, y1) =
12827 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12828 let mut y0 = y0;
12829 let mut y1 = y1;
12830 if s0 != 1.0 {
12831 self.scale_inplace(&mut y0, s0, m * out_f)?;
12832 }
12833 if s1 != 1.0 {
12834 self.scale_inplace(&mut y1, s1, m * out_f)?;
12835 }
12836 Ok(Some((y0, y1)))
12837 }
12838
12839 #[allow(clippy::too_many_arguments)]
12844 pub fn qmatvec_batched_dual_raw(
12845 &self,
12846 b0: &CudaSlice<u8>,
12847 b1: &CudaSlice<u8>,
12848 aq: &CudaSlice<i8>,
12849 ad: &CudaSlice<f32>,
12850 m: usize,
12851 in_f: usize,
12852 out_f: usize,
12853 row_bytes: usize,
12854 rp: bool,
12855 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12856 const ROWS_PER_BLOCK: u32 = 4;
12857 let mcols = Self::batched_mcols(m);
12858 let tiny_rp1 = rp
12861 && mcols == 4
12862 && out_f <= 128
12863 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12864 let (name, rows_per_block) = if tiny_rp1 {
12865 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12866 } else {
12867 match (mcols, rp, m) {
12868 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12869 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12870 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12871 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12872 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12873 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12874 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12875 _ => {
12876 return Err(
12877 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12878 );
12879 }
12880 }
12881 };
12882 let f = self.func(name);
12883 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12884 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12885 let cfg = LaunchConfig {
12886 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12887 block_dim: (32, ROWS_PER_BLOCK, 1),
12888 shared_mem_bytes: 0,
12889 };
12890 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12891 let __s_b = self.gpu.stream();
12892 let mut b = __s_b.launch_builder(&f);
12893 b.arg(b0)
12894 .arg(b1)
12895 .arg(aq)
12896 .arg(ad)
12897 .arg(&mut y0)
12898 .arg(&mut y1)
12899 .arg(&inf)
12900 .arg(&outf)
12901 .arg(&mi)
12902 .arg(&rb);
12903 unsafe {
12904 b.launch(cfg)?;
12905 }
12906 Ok((y0, y1))
12907 }
12908
12909 pub fn matmul_pre_dual_noscale(
12921 &self,
12922 w0: &crate::model::GpuTensor,
12923 w1: &crate::model::GpuTensor,
12924 aq: &CudaSlice<i8>,
12925 ad: &CudaSlice<f32>,
12926 m: usize,
12927 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12928 {
12929 use crate::model::GpuTensor;
12930 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12931 return Ok(None);
12932 }
12933 if !self.mmvq_supports(QT_NVFP4) {
12943 return Ok(None);
12944 }
12945 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12946 if w1.in_features() != in_f || w1.out_features() != out_f {
12947 return Ok(None);
12948 }
12949 let no_mirror =
12962 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12963 if self.q8_ffn_fuse2_on()
12964 && no_mirror(w0)
12965 && no_mirror(w1)
12966 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12967 {
12968 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12969 return Ok(Some(((y0, 1.0), (y1, 1.0))));
12970 }
12971 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12981 let (y0, y1) =
12982 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12983 return Ok(Some(((y0, p0.3), (y1, p1.3))));
12984 }
12985 let (b0, q0, rb0, s0, rp0) = match w0 {
12986 GpuTensor::Quant {
12987 bytes,
12988 qtype,
12989 row_bytes,
12990 scale,
12991 rp,
12992 ..
12993 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12994 _ => return Ok(None),
12995 };
12996 let (b1, q1, rb1, s1, rp1) = match w1 {
12997 GpuTensor::Quant {
12998 bytes,
12999 qtype,
13000 row_bytes,
13001 scale,
13002 rp,
13003 ..
13004 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13005 _ => return Ok(None),
13006 };
13007 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13008 return Ok(None);
13009 }
13010 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13012 let rows_per_block = ROWS_PER_BLOCK * RPW;
13013 let f = self.func(if rp0 {
13014 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13015 } else {
13016 "qmatvec_nvfp4_mmvq_dual_mr2"
13017 });
13018 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13019 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13020 let cfg = LaunchConfig {
13021 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13022 block_dim: (32, ROWS_PER_BLOCK, 1),
13023 shared_mem_bytes: 0,
13024 };
13025 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13026 let one = 1.0f32;
13029 let __s_b = self.gpu.stream();
13030 let mut b = __s_b.launch_builder(&f);
13031 b.arg(b0)
13032 .arg(b1)
13033 .arg(aq)
13034 .arg(ad)
13035 .arg(&mut y0)
13036 .arg(&mut y1)
13037 .arg(&inf)
13038 .arg(&outf)
13039 .arg(&mi)
13040 .arg(&rb)
13041 .arg(&one)
13042 .arg(&one);
13043 unsafe {
13044 b.launch(cfg)?;
13045 }
13046 Ok(Some(((y0, s0), (y1, s1))))
13047 }
13048
13049 #[allow(clippy::too_many_arguments)]
13057 pub fn matmul_nvfp4_fused3(
13058 &self,
13059 w0: &crate::model::GpuTensor,
13060 w1: &crate::model::GpuTensor,
13061 w2: &crate::model::GpuTensor,
13062 aq: &CudaSlice<i8>,
13063 ad: &CudaSlice<f32>,
13064 m: usize,
13065 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13066 {
13067 use crate::model::GpuTensor;
13068 if !self.mmvq_supports(QT_NVFP4)
13075 || !self.uses_q8_1_fast(w0)
13076 || !self.uses_q8_1_fast(w1)
13077 || !self.uses_q8_1_fast(w2)
13078 {
13079 return Ok(None);
13080 }
13081 if (9..=16).contains(&m) {
13084 return Ok(
13085 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13086 Some(mut ys) => {
13087 let y2 = ys.pop().unwrap();
13088 let y1 = ys.pop().unwrap();
13089 let y0 = ys.pop().unwrap();
13090 Some((y0, y1, y2))
13091 }
13092 None => None,
13093 },
13094 );
13095 }
13096 if !(1..=8).contains(&m) {
13097 return Ok(None);
13098 }
13099 if m > 1 {
13100 let in_f = w0.in_features();
13101 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13102 || !self.batched_supports(QT_NVFP4)
13103 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13104 || (m > 4 && !Self::b8_enabled())
13105 || in_f % 512 != 0
13106 || in_f / 64 > 272
13107 {
13108 return Ok(None);
13109 }
13110 }
13111 let unpack = |w: &crate::model::GpuTensor| match w {
13112 GpuTensor::Quant {
13113 bytes,
13114 qtype,
13115 scale,
13116 rp,
13117 ..
13118 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13119 _ => None,
13120 };
13121 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13122 return Ok(None);
13123 };
13124 let in_f = w0.in_features();
13125 if w1.in_features() != in_f || w2.in_features() != in_f {
13126 return Ok(None);
13127 }
13128 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13129 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13131 let rows_pb = ROWS_PER_BLOCK * RPW;
13132 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13133 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13134 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13135 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13136 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13137 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13140 if m > 1 {
13141 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13143 return Ok(None);
13144 }
13145 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13146 let cfg = LaunchConfig {
13147 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13148 block_dim: (32, ROWS_PER_BLOCK, 1),
13149 shared_mem_bytes: 0,
13150 };
13151 let __s_b = self.gpu.stream();
13152 let mut b = __s_b.launch_builder(&f);
13153 b.arg(b0)
13154 .arg(b1)
13155 .arg(b2)
13156 .arg(aq)
13157 .arg(ad)
13158 .arg(&mut y0)
13159 .arg(&mut y1)
13160 .arg(&mut y2)
13161 .arg(&inf)
13162 .arg(&oi0)
13163 .arg(&oi1)
13164 .arg(&oi2)
13165 .arg(&mi);
13166 unsafe {
13167 b.launch(cfg)?;
13168 }
13169 return Ok(Some((y0, y1, y2)));
13170 }
13171 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13172 let cfg = LaunchConfig {
13173 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13174 block_dim: (32, ROWS_PER_BLOCK, 1),
13175 shared_mem_bytes: 0,
13176 };
13177 let __s_b = self.gpu.stream();
13178 let mut b = __s_b.launch_builder(&f);
13179 b.arg(b0)
13180 .arg(b1)
13181 .arg(b2)
13182 .arg(aq)
13183 .arg(ad)
13184 .arg(&mut y0)
13185 .arg(&mut y1)
13186 .arg(&mut y2)
13187 .arg(&inf)
13188 .arg(&oi0)
13189 .arg(&oi1)
13190 .arg(&oi2)
13191 .arg(&mi)
13192 .arg(&p0.1)
13193 .arg(&p1.1)
13194 .arg(&p2.1);
13195 unsafe {
13196 b.launch(cfg)?;
13197 }
13198 Ok(Some((y0, y1, y2)))
13199 }
13200
13201 pub fn matmul_nvfp4_fused2(
13210 &self,
13211 w0: &crate::model::GpuTensor,
13212 w1: &crate::model::GpuTensor,
13213 aq: &CudaSlice<i8>,
13214 ad: &CudaSlice<f32>,
13215 m: usize,
13216 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13217 use crate::model::GpuTensor;
13218 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13219 let off =
13220 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13221 if off
13224 || m != 1
13225 || !self.mmvq_supports(QT_NVFP4)
13226 || !self.uses_q8_1_fast(w0)
13227 || !self.uses_q8_1_fast(w1)
13228 {
13229 return Ok(None);
13230 }
13231 let unpack = |w: &crate::model::GpuTensor| match w {
13232 GpuTensor::Quant {
13233 bytes,
13234 qtype,
13235 scale,
13236 rp,
13237 ..
13238 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13239 _ => None,
13240 };
13241 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13242 return Ok(None);
13243 };
13244 let in_f = w0.in_features();
13245 if w1.in_features() != in_f {
13246 return Ok(None);
13247 }
13248 let (o0, o1) = (w0.out_features(), w1.out_features());
13249 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13251 let rows_pb = ROWS_PER_BLOCK * RPW;
13252 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13253 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13254 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13255 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13256 let cfg = LaunchConfig {
13257 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13258 block_dim: (32, ROWS_PER_BLOCK, 1),
13259 shared_mem_bytes: 0,
13260 };
13261 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13262 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13265 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13268 {
13269 use cudarc::driver::{DevicePtr, DevicePtrMut};
13270 let s = &self.gpu.stream();
13271 let (pw0, _g0) = b0.device_ptr(s);
13272 let (pw1, _g1) = b1.device_ptr(s);
13273 let (paq, _g2) = aq.device_ptr(s);
13274 let (pad, _g3) = ad.device_ptr(s);
13275 let (py0, _g4) = y0.device_ptr_mut(s);
13276 let (py1, _g5) = y1.device_ptr_mut(s);
13277 let (s0, s1) = (p0.1, p1.1);
13278 let mut ps = [
13279 &pw0 as *const _ as *mut std::ffi::c_void,
13280 &pw1 as *const _ as *mut _,
13281 &paq as *const _ as *mut _,
13282 &pad as *const _ as *mut _,
13283 &py0 as *const _ as *mut _,
13284 &py1 as *const _ as *mut _,
13285 &inf as *const _ as *mut _,
13286 &oi0 as *const _ as *mut _,
13287 &oi1 as *const _ as *mut _,
13288 &mi as *const _ as *mut _,
13289 &s0 as *const _ as *mut _,
13290 &s1 as *const _ as *mut _,
13291 ];
13292 unsafe {
13293 self.launch_pdl(
13294 "qmatvec_nvfp4_mmvq_fused2_rp",
13295 cfg.grid_dim,
13296 cfg.block_dim,
13297 &mut ps,
13298 )?;
13299 }
13300 }
13301 return Ok(Some((y0, y1)));
13302 }
13303 let __s_b = self.gpu.stream();
13304 let mut b = __s_b.launch_builder(&f);
13305 b.arg(b0)
13306 .arg(b1)
13307 .arg(aq)
13308 .arg(ad)
13309 .arg(&mut y0)
13310 .arg(&mut y1)
13311 .arg(&inf)
13312 .arg(&oi0)
13313 .arg(&oi1)
13314 .arg(&mi)
13315 .arg(&p0.1)
13316 .arg(&p1.1);
13317 unsafe {
13318 b.launch(cfg)?;
13319 }
13320 Ok(Some((y0, y1)))
13321 }
13322
13323 pub fn matmul_nvfp4_fused2_into(
13328 &self,
13329 w0: &crate::model::GpuTensor,
13330 w1: &crate::model::GpuTensor,
13331 aq: &CudaSlice<i8>,
13332 ad: &CudaSlice<f32>,
13333 y0: &mut CudaSlice<f32>,
13334 y1: &mut CudaSlice<f32>,
13335 ) -> Result<bool, Box<dyn std::error::Error>> {
13336 use crate::model::GpuTensor;
13337 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13338 let off =
13339 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13340 if off
13341 || !self.mmvq_supports(QT_NVFP4)
13342 || !self.uses_q8_1_fast(w0)
13343 || !self.uses_q8_1_fast(w1)
13344 {
13345 return Ok(false);
13346 }
13347 let unpack = |w: &crate::model::GpuTensor| match w {
13348 GpuTensor::Quant {
13349 bytes,
13350 qtype,
13351 scale,
13352 rp,
13353 ..
13354 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13355 _ => None,
13356 };
13357 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13358 return Ok(false);
13359 };
13360 let in_f = w0.in_features();
13361 if w1.in_features() != in_f {
13362 return Ok(false);
13363 }
13364 let (o0, o1) = (w0.out_features(), w1.out_features());
13365 if y0.len() < o0 || y1.len() < o1 {
13366 return Ok(false);
13367 }
13368 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13370 let rows_pb = ROWS_PER_BLOCK * RPW;
13371 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13372 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13373 let cfg = LaunchConfig {
13374 grid_dim: (nb(o0) + nb(o1), 1, 1),
13375 block_dim: (32, ROWS_PER_BLOCK, 1),
13376 shared_mem_bytes: 0,
13377 };
13378 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13379 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13382 let __s_b = self.gpu.stream();
13383 let mut b = __s_b.launch_builder(&f);
13384 b.arg(b0)
13385 .arg(b1)
13386 .arg(aq)
13387 .arg(ad)
13388 .arg(&mut *y0)
13389 .arg(&mut *y1)
13390 .arg(&inf)
13391 .arg(&oi0)
13392 .arg(&oi1)
13393 .arg(&mi)
13394 .arg(&p0.1)
13395 .arg(&p1.1);
13396 unsafe {
13397 b.launch(cfg)?;
13398 }
13399 Ok(true)
13400 }
13401
13402 #[allow(clippy::type_complexity)]
13407 pub fn matmul_nvfp4_fused4(
13408 &self,
13409 w0: &crate::model::GpuTensor,
13410 w1: &crate::model::GpuTensor,
13411 w2: &crate::model::GpuTensor,
13412 w3: &crate::model::GpuTensor,
13413 aq: &CudaSlice<i8>,
13414 ad: &CudaSlice<f32>,
13415 m: usize,
13416 ) -> Result<
13417 Option<(
13418 CudaSlice<f32>,
13419 CudaSlice<f32>,
13420 CudaSlice<f32>,
13421 CudaSlice<f32>,
13422 )>,
13423 Box<dyn std::error::Error>,
13424 > {
13425 use crate::model::GpuTensor;
13426 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13433 || !self.mmvq_supports(QT_NVFP4)
13434 || !self.uses_q8_1_fast(w0)
13435 || !self.uses_q8_1_fast(w1)
13436 || !self.uses_q8_1_fast(w2)
13437 || !self.uses_q8_1_fast(w3)
13438 {
13439 return Ok(None);
13440 }
13441 if (9..=16).contains(&m) {
13446 return Ok(
13447 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13448 Some(mut ys) => {
13449 let y3 = ys.pop().unwrap();
13450 let y2 = ys.pop().unwrap();
13451 let y1 = ys.pop().unwrap();
13452 let y0 = ys.pop().unwrap();
13453 Some((y0, y1, y2, y3))
13454 }
13455 None => None,
13456 },
13457 );
13458 }
13459 if !(1..=8).contains(&m) {
13460 return Ok(None);
13461 }
13462 if m > 1 {
13463 let in_f = w0.in_features();
13466 if !self.batched_supports(QT_NVFP4)
13467 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13468 || (m > 4 && !Self::b8_enabled())
13469 || in_f % 512 != 0
13470 || in_f / 64 > 272
13471 {
13472 return Ok(None);
13473 }
13474 }
13475 let unpack = |w: &crate::model::GpuTensor| match w {
13476 GpuTensor::Quant {
13477 bytes,
13478 qtype,
13479 scale,
13480 rp,
13481 ..
13482 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13483 _ => None,
13484 };
13485 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13486 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13487 else {
13488 return Ok(None);
13489 };
13490 let in_f = w0.in_features();
13491 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13492 return Ok(None);
13493 }
13494 let (o0, o1, o2, o3) = (
13495 w0.out_features(),
13496 w1.out_features(),
13497 w2.out_features(),
13498 w3.out_features(),
13499 );
13500 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13502 let rows_pb = ROWS_PER_BLOCK * RPW;
13503 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13504 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13505 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13506 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13507 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13508 let (inf, oi0, oi1, oi2, oi3, mi) = (
13509 in_f as i32,
13510 o0 as i32,
13511 o1 as i32,
13512 o2 as i32,
13513 o3 as i32,
13514 m as i32,
13515 );
13516 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13519 if m > 1 {
13520 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13523 return Ok(None);
13524 }
13525 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13526 let cfg = LaunchConfig {
13527 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13528 block_dim: (32, ROWS_PER_BLOCK, 1),
13529 shared_mem_bytes: 0,
13530 };
13531 let __s_b = self.gpu.stream();
13532 let mut b = __s_b.launch_builder(&f);
13533 b.arg(b0)
13534 .arg(b1)
13535 .arg(b2)
13536 .arg(b3)
13537 .arg(aq)
13538 .arg(ad)
13539 .arg(&mut y0)
13540 .arg(&mut y1)
13541 .arg(&mut y2)
13542 .arg(&mut y3)
13543 .arg(&inf)
13544 .arg(&oi0)
13545 .arg(&oi1)
13546 .arg(&oi2)
13547 .arg(&oi3)
13548 .arg(&mi);
13549 unsafe {
13550 b.launch(cfg)?;
13551 }
13552 return Ok(Some((y0, y1, y2, y3)));
13553 }
13554 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13555 let cfg = LaunchConfig {
13556 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13557 block_dim: (32, ROWS_PER_BLOCK, 1),
13558 shared_mem_bytes: 0,
13559 };
13560 let __s_b = self.gpu.stream();
13561 let mut b = __s_b.launch_builder(&f);
13562 b.arg(b0)
13563 .arg(b1)
13564 .arg(b2)
13565 .arg(b3)
13566 .arg(aq)
13567 .arg(ad)
13568 .arg(&mut y0)
13569 .arg(&mut y1)
13570 .arg(&mut y2)
13571 .arg(&mut y3)
13572 .arg(&inf)
13573 .arg(&oi0)
13574 .arg(&oi1)
13575 .arg(&oi2)
13576 .arg(&oi3)
13577 .arg(&mi)
13578 .arg(&p0.1)
13579 .arg(&p1.1)
13580 .arg(&p2.1)
13581 .arg(&p3.1);
13582 unsafe {
13583 b.launch(cfg)?;
13584 }
13585 Ok(Some((y0, y1, y2, y3)))
13586 }
13587
13588 pub fn matmul_q8_fused2(
13596 &self,
13597 w0: &crate::model::GpuTensor,
13598 w1: &crate::model::GpuTensor,
13599 aq: &CudaSlice<i8>,
13600 ad: &CudaSlice<f32>,
13601 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13602 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13608 return Ok(Some(self.e4m3_fused2_core(
13609 p0.0,
13610 p1.0,
13611 aq,
13612 ad,
13613 w0.in_features(),
13614 p0.1,
13615 p1.1,
13616 p0.2,
13617 p0.3,
13618 p1.3,
13619 )?));
13620 }
13621 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13622 return Ok(None);
13623 };
13624 Ok(Some(self.q8_fused2_core(
13625 p0.0,
13626 p1.0,
13627 aq,
13628 ad,
13629 w0.in_features(),
13630 p0.1,
13631 p1.1,
13632 p0.2,
13633 )?))
13634 }
13635
13636 #[allow(clippy::too_many_arguments)]
13637 fn q8_fused2_core(
13638 &self,
13639 b0: &CudaSlice<u8>,
13640 b1: &CudaSlice<u8>,
13641 aq: &CudaSlice<i8>,
13642 ad: &CudaSlice<f32>,
13643 in_f: usize,
13644 out0: usize,
13645 out1: usize,
13646 row_bytes: usize,
13647 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13648 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13650 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13651 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13652 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13653 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13654 let cfg = LaunchConfig {
13655 grid_dim: (nb0 + nb1, 1, 1),
13656 block_dim: (32, ROWS_PER_BLOCK, 1),
13657 shared_mem_bytes: 0,
13658 };
13659 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13660 let __s_b = self.gpu.stream();
13661 let mut b = __s_b.launch_builder(&f);
13662 b.arg(b0)
13663 .arg(b1)
13664 .arg(aq)
13665 .arg(ad)
13666 .arg(&mut y0)
13667 .arg(&mut y1)
13668 .arg(&inf)
13669 .arg(&o0)
13670 .arg(&o1)
13671 .arg(&rbl);
13672 unsafe {
13673 b.launch(cfg)?;
13674 }
13675 Ok((y0, y1))
13676 }
13677
13678 pub fn matmul_q8_fused2_x(
13684 &self,
13685 w0: &crate::model::GpuTensor,
13686 w1: &crate::model::GpuTensor,
13687 x: &CudaSlice<f32>,
13688 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13689 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13690 return Ok(None);
13691 }
13692 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13693 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13694 return Ok(Some(self.e4m3_fused2_core(
13695 p0.0,
13696 p1.0,
13697 &aq,
13698 &ad,
13699 w0.in_features(),
13700 p0.1,
13701 p1.1,
13702 p0.2,
13703 p0.3,
13704 p1.3,
13705 )?));
13706 }
13707 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13708 return Ok(None);
13709 };
13710 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13711 Ok(Some(self.q8_fused2_core(
13712 p0.0,
13713 p1.0,
13714 &aq,
13715 &ad,
13716 w0.in_features(),
13717 p0.1,
13718 p1.1,
13719 p0.2,
13720 )?))
13721 }
13722
13723 #[allow(clippy::too_many_arguments)]
13726 pub fn qmatvec_q8_fused2_raw(
13727 &self,
13728 b0: &CudaSlice<u8>,
13729 b1: &CudaSlice<u8>,
13730 x: &CudaSlice<f32>,
13731 in_f: usize,
13732 out0: usize,
13733 out1: usize,
13734 row_bytes: usize,
13735 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13736 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13737 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13738 }
13739
13740 pub fn matmul_q4_fused3(
13746 &self,
13747 w0: &crate::model::GpuTensor,
13748 w1: &crate::model::GpuTensor,
13749 w2: &crate::model::GpuTensor,
13750 aq: &CudaSlice<i8>,
13751 ad: &CudaSlice<f32>,
13752 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13753 {
13754 use crate::model::GpuTensor;
13755 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13756 match w {
13757 GpuTensor::Quant {
13758 qtype, row_bytes, ..
13759 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13760 _ => None,
13761 }
13762 };
13763 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13764 return Ok(None);
13765 };
13766 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13767 return Ok(None);
13768 }
13769 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13773 match w {
13774 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13775 Some(m) => (m, true),
13776 None => (bytes, *rp),
13777 },
13778 _ => unreachable!(),
13779 }
13780 }
13781 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13782 if rp0 != rp1 || rp1 != rp2 {
13783 return Ok(None);
13784 }
13785 let rp = rp0;
13786 let rpb: u32 = 4;
13787 let mr1 = rp && Self::q40_mr1_on();
13791 let nb = |o: usize| {
13792 if mr1 {
13793 (o as u32).div_ceil(rpb)
13794 } else {
13795 (o as u32).div_ceil(2).div_ceil(rpb)
13796 }
13797 };
13798 let grid = nb(o0) + nb(o1) + nb(o2);
13799 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13800 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13801 let mut y2 = self.alloc_uninit::<f32>(o2)?;
13802 let f = self.func(if mr1 {
13803 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13804 } else if rp {
13805 "qmatvec_q4_0_mmvq_fused3_rp"
13806 } else {
13807 "qmatvec_q4_0_mmvq_fused3"
13808 });
13809 let cfg = LaunchConfig {
13810 grid_dim: (grid, 1, 1),
13811 block_dim: (32, rpb, 1),
13812 shared_mem_bytes: 0,
13813 };
13814 let inf = w0.in_features() as i32;
13815 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13816 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13817 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13820 {
13821 use cudarc::driver::{DevicePtr, DevicePtrMut};
13822 let s = &self.gpu.stream();
13823 let (p0, _g0) = b0.device_ptr(s);
13824 let (p1, _g1) = b1.device_ptr(s);
13825 let (p2, _g2) = b2.device_ptr(s);
13826 let (paq, _g3) = aq.device_ptr(s);
13827 let (pad, _g4) = ad.device_ptr(s);
13828 let (py0, _g5) = y0.device_ptr_mut(s);
13829 let (py1, _g6) = y1.device_ptr_mut(s);
13830 let (py2, _g7) = y2.device_ptr_mut(s);
13831 let mut ps = [
13832 &p0 as *const _ as *mut std::ffi::c_void,
13833 &p1 as *const _ as *mut _,
13834 &p2 as *const _ as *mut _,
13835 &paq as *const _ as *mut _,
13836 &pad as *const _ as *mut _,
13837 &py0 as *const _ as *mut _,
13838 &py1 as *const _ as *mut _,
13839 &py2 as *const _ as *mut _,
13840 &inf as *const _ as *mut _,
13841 &oo0 as *const _ as *mut _,
13842 &oo1 as *const _ as *mut _,
13843 &oo2 as *const _ as *mut _,
13844 &r0 as *const _ as *mut _,
13845 &r1 as *const _ as *mut _,
13846 &r2 as *const _ as *mut _,
13847 ];
13848 unsafe {
13849 self.launch_pdl(
13850 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13851 (grid, 1, 1),
13852 (32, rpb, 1),
13853 &mut ps,
13854 )?;
13855 }
13856 }
13857 return Ok(Some((y0, y1, y2)));
13858 }
13859 let __s_b = self.gpu.stream();
13860 let mut b = __s_b.launch_builder(&f);
13861 b.arg(b0)
13862 .arg(b1)
13863 .arg(b2)
13864 .arg(aq)
13865 .arg(ad)
13866 .arg(&mut y0)
13867 .arg(&mut y1)
13868 .arg(&mut y2)
13869 .arg(&inf)
13870 .arg(&oo0)
13871 .arg(&oo1)
13872 .arg(&oo2)
13873 .arg(&r0)
13874 .arg(&r1)
13875 .arg(&r2);
13876 unsafe {
13877 b.launch(cfg)?;
13878 }
13879 Ok(Some((y0, y1, y2)))
13880 }
13881
13882 #[allow(clippy::too_many_arguments)]
13885 pub fn matmul_q4_fused3_into(
13886 &self,
13887 w0: &crate::model::GpuTensor,
13888 w1: &crate::model::GpuTensor,
13889 w2: &crate::model::GpuTensor,
13890 aq: &CudaSlice<i8>,
13891 ad: &CudaSlice<f32>,
13892 y0: &mut CudaSlice<f32>,
13893 y1: &mut CudaSlice<f32>,
13894 y2: &mut CudaSlice<f32>,
13895 ) -> Result<bool, Box<dyn std::error::Error>> {
13896 use crate::model::GpuTensor;
13897 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13898 match w {
13899 GpuTensor::Quant {
13900 qtype, row_bytes, ..
13901 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13902 _ => None,
13903 }
13904 };
13905 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13906 return Ok(false);
13907 };
13908 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13909 return Ok(false);
13910 }
13911 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13912 match w {
13913 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13914 Some(m) => (m, true),
13915 None => (bytes, *rp),
13916 },
13917 _ => unreachable!(),
13918 }
13919 }
13920 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13921 if rp0 != rp1 || rp1 != rp2 {
13922 return Ok(false);
13923 }
13924 let rp = rp0;
13925 let rpb: u32 = 4;
13926 let mr1 = rp && Self::q40_mr1_on();
13927 let nb = |o: usize| {
13928 if mr1 {
13929 (o as u32).div_ceil(rpb)
13930 } else {
13931 (o as u32).div_ceil(2).div_ceil(rpb)
13932 }
13933 };
13934 let grid = nb(o0) + nb(o1) + nb(o2);
13935 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13936 let f = self.func(if mr1 {
13937 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13938 } else if rp {
13939 "qmatvec_q4_0_mmvq_fused3_rp"
13940 } else {
13941 "qmatvec_q4_0_mmvq_fused3"
13942 });
13943 let cfg = LaunchConfig {
13944 grid_dim: (grid, 1, 1),
13945 block_dim: (32, rpb, 1),
13946 shared_mem_bytes: 0,
13947 };
13948 let inf = w0.in_features() as i32;
13949 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13950 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13951 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13953 use cudarc::driver::{DevicePtr, DevicePtrMut};
13954 let s = &self.gpu.stream();
13955 let (p0, _g0) = b0.device_ptr(s);
13956 let (p1, _g1) = b1.device_ptr(s);
13957 let (p2, _g2) = b2.device_ptr(s);
13958 let (paq, _g3) = aq.device_ptr(s);
13959 let (pad, _g4) = ad.device_ptr(s);
13960 let (py0, _g5) = y0.device_ptr_mut(s);
13961 let (py1, _g6) = y1.device_ptr_mut(s);
13962 let (py2, _g7) = y2.device_ptr_mut(s);
13963 let mut ps = [
13964 &p0 as *const _ as *mut std::ffi::c_void,
13965 &p1 as *const _ as *mut _,
13966 &p2 as *const _ as *mut _,
13967 &paq as *const _ as *mut _,
13968 &pad as *const _ as *mut _,
13969 &py0 as *const _ as *mut _,
13970 &py1 as *const _ as *mut _,
13971 &py2 as *const _ as *mut _,
13972 &inf as *const _ as *mut _,
13973 &oo0 as *const _ as *mut _,
13974 &oo1 as *const _ as *mut _,
13975 &oo2 as *const _ as *mut _,
13976 &r0 as *const _ as *mut _,
13977 &r1 as *const _ as *mut _,
13978 &r2 as *const _ as *mut _,
13979 ];
13980 unsafe {
13981 self.launch_pdl(
13982 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13983 (grid, 1, 1),
13984 (32, rpb, 1),
13985 &mut ps,
13986 )?;
13987 }
13988 return Ok(true);
13989 }
13990 let __s_b = self.gpu.stream();
13991 let mut b = __s_b.launch_builder(&f);
13992 b.arg(b0)
13993 .arg(b1)
13994 .arg(b2)
13995 .arg(aq)
13996 .arg(ad)
13997 .arg(&mut *y0)
13998 .arg(&mut *y1)
13999 .arg(&mut *y2)
14000 .arg(&inf)
14001 .arg(&oo0)
14002 .arg(&oo1)
14003 .arg(&oo2)
14004 .arg(&r0)
14005 .arg(&r1)
14006 .arg(&r2);
14007 unsafe {
14008 b.launch(cfg)?;
14009 }
14010 Ok(true)
14011 }
14012
14013 pub fn matmul_q4_fused2(
14015 &self,
14016 w0: &crate::model::GpuTensor,
14017 w1: &crate::model::GpuTensor,
14018 aq: &CudaSlice<i8>,
14019 ad: &CudaSlice<f32>,
14020 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14021 use crate::model::GpuTensor;
14022 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14023 match w {
14024 GpuTensor::Quant {
14025 qtype, row_bytes, ..
14026 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14027 _ => None,
14028 }
14029 };
14030 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14031 return Ok(None);
14032 };
14033 if w0.in_features() != w1.in_features() {
14034 return Ok(None);
14035 }
14036 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14038 match w {
14039 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14040 Some(m) => (m, true),
14041 None => (bytes, *rp),
14042 },
14043 _ => unreachable!(),
14044 }
14045 }
14046 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14047 if rp0 != rp1 {
14048 return Ok(None);
14049 }
14050 let rp = rp0;
14051 let rpb: u32 = 4;
14052 let mr1 = rp && Self::q40_mr1_on();
14054 let nb = |o: usize| {
14055 if mr1 {
14056 (o as u32).div_ceil(rpb)
14057 } else {
14058 (o as u32).div_ceil(2).div_ceil(rpb)
14059 }
14060 };
14061 let grid = nb(o0) + nb(o1);
14062 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14063 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14064 let f = self.func(if mr1 {
14065 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14066 } else if rp {
14067 "qmatvec_q4_0_mmvq_fused2_rp"
14068 } else {
14069 "qmatvec_q4_0_mmvq_fused2"
14070 });
14071 let cfg = LaunchConfig {
14072 grid_dim: (grid, 1, 1),
14073 block_dim: (32, rpb, 1),
14074 shared_mem_bytes: 0,
14075 };
14076 let inf = w0.in_features() as i32;
14077 let (oo0, oo1) = (o0 as i32, o1 as i32);
14078 let (r0, r1) = (rb0 as i64, rb1 as i64);
14079 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14081 {
14082 use cudarc::driver::{DevicePtr, DevicePtrMut};
14083 let s = &self.gpu.stream();
14084 let (p0, _g0) = b0.device_ptr(s);
14085 let (p1, _g1) = b1.device_ptr(s);
14086 let (paq, _g2) = aq.device_ptr(s);
14087 let (pad, _g3) = ad.device_ptr(s);
14088 let (py0, _g4) = y0.device_ptr_mut(s);
14089 let (py1, _g5) = y1.device_ptr_mut(s);
14090 let mut ps = [
14091 &p0 as *const _ as *mut std::ffi::c_void,
14092 &p1 as *const _ as *mut _,
14093 &paq as *const _ as *mut _,
14094 &pad as *const _ as *mut _,
14095 &py0 as *const _ as *mut _,
14096 &py1 as *const _ as *mut _,
14097 &inf as *const _ as *mut _,
14098 &oo0 as *const _ as *mut _,
14099 &oo1 as *const _ as *mut _,
14100 &r0 as *const _ as *mut _,
14101 &r1 as *const _ as *mut _,
14102 ];
14103 unsafe {
14104 self.launch_pdl(
14105 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14106 (grid, 1, 1),
14107 (32, rpb, 1),
14108 &mut ps,
14109 )?;
14110 }
14111 }
14112 return Ok(Some((y0, y1)));
14113 }
14114 let __s_b = self.gpu.stream();
14115 let mut b = __s_b.launch_builder(&f);
14116 b.arg(b0)
14117 .arg(b1)
14118 .arg(aq)
14119 .arg(ad)
14120 .arg(&mut y0)
14121 .arg(&mut y1)
14122 .arg(&inf)
14123 .arg(&oo0)
14124 .arg(&oo1)
14125 .arg(&r0)
14126 .arg(&r1);
14127 unsafe {
14128 b.launch(cfg)?;
14129 }
14130 Ok(Some((y0, y1)))
14131 }
14132
14133 pub fn matmul_q4_fused2_into(
14135 &self,
14136 w0: &crate::model::GpuTensor,
14137 w1: &crate::model::GpuTensor,
14138 aq: &CudaSlice<i8>,
14139 ad: &CudaSlice<f32>,
14140 y0: &mut CudaSlice<f32>,
14141 y1: &mut CudaSlice<f32>,
14142 ) -> Result<bool, Box<dyn std::error::Error>> {
14143 use crate::model::GpuTensor;
14144 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14145 match w {
14146 GpuTensor::Quant {
14147 qtype, row_bytes, ..
14148 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14149 _ => None,
14150 }
14151 };
14152 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14153 return Ok(false);
14154 };
14155 if w0.in_features() != w1.in_features() {
14156 return Ok(false);
14157 }
14158 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14159 match w {
14160 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14161 Some(m) => (m, true),
14162 None => (bytes, *rp),
14163 },
14164 _ => unreachable!(),
14165 }
14166 }
14167 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14168 if rp0 != rp1 {
14169 return Ok(false);
14170 }
14171 let rp = rp0;
14172 let rpb: u32 = 4;
14173 let mr1 = rp && Self::q40_mr1_on();
14174 let nb = |o: usize| {
14175 if mr1 {
14176 (o as u32).div_ceil(rpb)
14177 } else {
14178 (o as u32).div_ceil(2).div_ceil(rpb)
14179 }
14180 };
14181 let grid = nb(o0) + nb(o1);
14182 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14183 let f = self.func(if mr1 {
14184 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14185 } else if rp {
14186 "qmatvec_q4_0_mmvq_fused2_rp"
14187 } else {
14188 "qmatvec_q4_0_mmvq_fused2"
14189 });
14190 let cfg = LaunchConfig {
14191 grid_dim: (grid, 1, 1),
14192 block_dim: (32, rpb, 1),
14193 shared_mem_bytes: 0,
14194 };
14195 let inf = w0.in_features() as i32;
14196 let (oo0, oo1) = (o0 as i32, o1 as i32);
14197 let (r0, r1) = (rb0 as i64, rb1 as i64);
14198 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14200 use cudarc::driver::{DevicePtr, DevicePtrMut};
14201 let s = &self.gpu.stream();
14202 let (p0, _g0) = b0.device_ptr(s);
14203 let (p1, _g1) = b1.device_ptr(s);
14204 let (paq, _g2) = aq.device_ptr(s);
14205 let (pad, _g3) = ad.device_ptr(s);
14206 let (py0, _g4) = y0.device_ptr_mut(s);
14207 let (py1, _g5) = y1.device_ptr_mut(s);
14208 let mut ps = [
14209 &p0 as *const _ as *mut std::ffi::c_void,
14210 &p1 as *const _ as *mut _,
14211 &paq as *const _ as *mut _,
14212 &pad as *const _ as *mut _,
14213 &py0 as *const _ as *mut _,
14214 &py1 as *const _ as *mut _,
14215 &inf as *const _ as *mut _,
14216 &oo0 as *const _ as *mut _,
14217 &oo1 as *const _ as *mut _,
14218 &r0 as *const _ as *mut _,
14219 &r1 as *const _ as *mut _,
14220 ];
14221 unsafe {
14222 self.launch_pdl(
14223 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14224 (grid, 1, 1),
14225 (32, rpb, 1),
14226 &mut ps,
14227 )?;
14228 }
14229 return Ok(true);
14230 }
14231 let __s_b = self.gpu.stream();
14232 let mut b = __s_b.launch_builder(&f);
14233 b.arg(b0)
14234 .arg(b1)
14235 .arg(aq)
14236 .arg(ad)
14237 .arg(&mut *y0)
14238 .arg(&mut *y1)
14239 .arg(&inf)
14240 .arg(&oo0)
14241 .arg(&oo1)
14242 .arg(&r0)
14243 .arg(&r1);
14244 unsafe {
14245 b.launch(cfg)?;
14246 }
14247 Ok(true)
14248 }
14249
14250 pub fn matmul_q4_fused2_batched(
14255 &self,
14256 w0: &crate::model::GpuTensor,
14257 w1: &crate::model::GpuTensor,
14258 aq: &CudaSlice<i8>,
14259 ad: &CudaSlice<f32>,
14260 m: usize,
14261 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14262 use crate::model::GpuTensor;
14263 if m < 2 || m > 8 {
14264 return Ok(None);
14265 }
14266 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14267 match w {
14268 GpuTensor::Quant {
14269 qtype, row_bytes, ..
14270 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14271 _ => None,
14272 }
14273 };
14274 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14275 return Ok(None);
14276 };
14277 if w0.in_features() != w1.in_features() {
14278 return Ok(None);
14279 }
14280 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14281 match w {
14282 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14283 Some(mr) => (mr, true),
14284 None => (bytes, *rp),
14285 },
14286 _ => unreachable!(),
14287 }
14288 }
14289 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14290 if !rp0 || !rp1 {
14291 return Ok(None);
14292 }
14293 let mcols = Self::batched_mcols(m);
14294 let rpb: u32 = 4;
14295 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14296 let grid = nb(o0) + nb(o1);
14297 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14298 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14299 let f = self.func(match mcols {
14300 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14301 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14302 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14303 });
14304 let cfg = LaunchConfig {
14305 grid_dim: (grid, 1, 1),
14306 block_dim: (32, rpb, 1),
14307 shared_mem_bytes: 0,
14308 };
14309 let inf = w0.in_features() as i32;
14310 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14311 let rb = rb0 as i64;
14312 let __s_b = self.gpu.stream();
14313 let mut b = __s_b.launch_builder(&f);
14314 b.arg(b0)
14315 .arg(b1)
14316 .arg(aq)
14317 .arg(ad)
14318 .arg(&mut y0)
14319 .arg(&mut y1)
14320 .arg(&inf)
14321 .arg(&oo0)
14322 .arg(&oo1)
14323 .arg(&mi)
14324 .arg(&rb);
14325 unsafe {
14326 b.launch(cfg)?;
14327 }
14328 Ok(Some((y0, y1)))
14329 }
14330
14331 #[allow(clippy::too_many_arguments)]
14334 pub fn matmul_q4_fused3_batched(
14335 &self,
14336 w0: &crate::model::GpuTensor,
14337 w1: &crate::model::GpuTensor,
14338 w2: &crate::model::GpuTensor,
14339 aq: &CudaSlice<i8>,
14340 ad: &CudaSlice<f32>,
14341 m: usize,
14342 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14343 {
14344 use crate::model::GpuTensor;
14345 if m < 2 || m > 8 {
14346 return Ok(None);
14347 }
14348 let q4 = |w: &GpuTensor| -> Option<usize> {
14349 match w {
14350 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14351 _ => None,
14352 }
14353 };
14354 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14355 return Ok(None);
14356 };
14357 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14358 return Ok(None);
14359 }
14360 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14361 match w {
14362 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14363 Some(mr) => (mr, true),
14364 None => (bytes, *rp),
14365 },
14366 _ => unreachable!(),
14367 }
14368 }
14369 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14370 if !rp0 || !rp1 || !rp2 {
14371 return Ok(None);
14372 }
14373 let mcols = Self::batched_mcols(m);
14374 let rpb: u32 = 4;
14375 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14376 let grid = nb(o0) + nb(o1) + nb(o2);
14377 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14378 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14379 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14380 let f = self.func(match mcols {
14381 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14382 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14383 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14384 });
14385 let cfg = LaunchConfig {
14386 grid_dim: (grid, 1, 1),
14387 block_dim: (32, rpb, 1),
14388 shared_mem_bytes: 0,
14389 };
14390 let inf = w0.in_features() as i32;
14391 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14392 let rb = 0i64;
14393 let __s_b = self.gpu.stream();
14394 let mut b = __s_b.launch_builder(&f);
14395 b.arg(b0)
14396 .arg(b1)
14397 .arg(b2)
14398 .arg(aq)
14399 .arg(ad)
14400 .arg(&mut y0)
14401 .arg(&mut y1)
14402 .arg(&mut y2)
14403 .arg(&inf)
14404 .arg(&oo0)
14405 .arg(&oo1)
14406 .arg(&oo2)
14407 .arg(&mi)
14408 .arg(&rb);
14409 unsafe {
14410 b.launch(cfg)?;
14411 }
14412 Ok(Some((y0, y1, y2)))
14413 }
14414
14415 pub fn matmul_q8_fused3(
14416 &self,
14417 w0: &crate::model::GpuTensor,
14418 w1: &crate::model::GpuTensor,
14419 w2: &crate::model::GpuTensor,
14420 aq: &CudaSlice<i8>,
14421 ad: &CudaSlice<f32>,
14422 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14423 {
14424 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14427 return Ok(Some(self.e4m3_fused3_core(
14428 p0.0,
14429 p1.0,
14430 p2.0,
14431 aq,
14432 ad,
14433 w0.in_features(),
14434 p0.1,
14435 p1.1,
14436 p2.1,
14437 p0.2,
14438 p0.3,
14439 p1.3,
14440 p2.3,
14441 )?));
14442 }
14443 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14444 return Ok(None);
14445 };
14446 Ok(Some(self.q8_fused3_core(
14447 p0.0,
14448 p1.0,
14449 p2.0,
14450 aq,
14451 ad,
14452 w0.in_features(),
14453 p0.1,
14454 p1.1,
14455 p2.1,
14456 p0.2,
14457 )?))
14458 }
14459
14460 #[allow(clippy::too_many_arguments)]
14461 fn q8_fused3_core(
14462 &self,
14463 b0: &CudaSlice<u8>,
14464 b1: &CudaSlice<u8>,
14465 b2: &CudaSlice<u8>,
14466 aq: &CudaSlice<i8>,
14467 ad: &CudaSlice<f32>,
14468 in_f: usize,
14469 out0: usize,
14470 out1: usize,
14471 out2: usize,
14472 row_bytes: usize,
14473 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14474 const ROWS_PER_BLOCK: u32 = 4;
14475 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14476 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14477 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14478 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14479 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14480 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14481 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14482 let cfg = LaunchConfig {
14483 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14484 block_dim: (32, ROWS_PER_BLOCK, 1),
14485 shared_mem_bytes: 0,
14486 };
14487 let (inf, o0, o1, o2, rbl) = (
14488 in_f as i32,
14489 out0 as i32,
14490 out1 as i32,
14491 out2 as i32,
14492 row_bytes as i64,
14493 );
14494 let __s_b = self.gpu.stream();
14495 let mut b = __s_b.launch_builder(&f);
14496 b.arg(b0)
14497 .arg(b1)
14498 .arg(b2)
14499 .arg(aq)
14500 .arg(ad)
14501 .arg(&mut y0)
14502 .arg(&mut y1)
14503 .arg(&mut y2)
14504 .arg(&inf)
14505 .arg(&o0)
14506 .arg(&o1)
14507 .arg(&o2)
14508 .arg(&rbl);
14509 unsafe {
14510 b.launch(cfg)?;
14511 }
14512 Ok((y0, y1, y2))
14513 }
14514
14515 #[allow(clippy::too_many_arguments)]
14517 pub fn qmatvec_q8_fused3_raw(
14518 &self,
14519 b0: &CudaSlice<u8>,
14520 b1: &CudaSlice<u8>,
14521 b2: &CudaSlice<u8>,
14522 x: &CudaSlice<f32>,
14523 in_f: usize,
14524 out0: usize,
14525 out1: usize,
14526 out2: usize,
14527 row_bytes: usize,
14528 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14529 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14530 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14531 }
14532
14533 pub fn matmul_q8_fused2_t(
14544 &self,
14545 w0: &crate::model::GpuTensor,
14546 w1: &crate::model::GpuTensor,
14547 aq: &CudaSlice<i8>,
14548 ad: &CudaSlice<f32>,
14549 m: usize,
14550 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14551 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14555 return Ok(None);
14556 }
14557 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14560 if m > 4 && !Self::b8_enabled() {
14561 return Ok(None);
14562 }
14563 return Ok(Some(self.e4m3_fused2_t_core(
14564 p0.0,
14565 p1.0,
14566 aq,
14567 ad,
14568 m,
14569 w0.in_features(),
14570 p0.1,
14571 p1.1,
14572 p0.2,
14573 p0.3,
14574 p1.3,
14575 )?));
14576 }
14577 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14578 return Ok(None);
14579 };
14580 Ok(Some(self.q8_fused2_t_core(
14581 p0.0,
14582 p1.0,
14583 aq,
14584 ad,
14585 m,
14586 w0.in_features(),
14587 p0.1,
14588 p1.1,
14589 p0.2,
14590 )?))
14591 }
14592
14593 #[allow(clippy::too_many_arguments)]
14594 fn q8_fused2_t_core(
14595 &self,
14596 b0: &CudaSlice<u8>,
14597 b1: &CudaSlice<u8>,
14598 aq: &CudaSlice<i8>,
14599 ad: &CudaSlice<f32>,
14600 m: usize,
14601 in_f: usize,
14602 out0: usize,
14603 out1: usize,
14604 row_bytes: usize,
14605 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14606 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14608 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14609 let f = self.func(match Self::batched_mcols(m) {
14610 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14611 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14612 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14614 });
14615 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14616 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14617 let cfg = LaunchConfig {
14618 grid_dim: (nb0 + nb1, 1, 1),
14619 block_dim: (32, ROWS_PER_BLOCK, 1),
14620 shared_mem_bytes: 0,
14621 };
14622 let (inf, o0, o1, mi, rbl) = (
14623 in_f as i32,
14624 out0 as i32,
14625 out1 as i32,
14626 m as i32,
14627 row_bytes as i64,
14628 );
14629 let __s_b = self.gpu.stream();
14630 let mut b = __s_b.launch_builder(&f);
14631 b.arg(b0)
14632 .arg(b1)
14633 .arg(aq)
14634 .arg(ad)
14635 .arg(&mut y0)
14636 .arg(&mut y1)
14637 .arg(&inf)
14638 .arg(&o0)
14639 .arg(&o1)
14640 .arg(&mi)
14641 .arg(&rbl);
14642 unsafe {
14643 b.launch(cfg)?;
14644 }
14645 Ok((y0, y1))
14646 }
14647
14648 #[allow(clippy::too_many_arguments)]
14651 pub fn qmatvec_q8_fused2_t_raw(
14652 &self,
14653 b0: &CudaSlice<u8>,
14654 b1: &CudaSlice<u8>,
14655 x: &CudaSlice<f32>,
14656 m: usize,
14657 in_f: usize,
14658 out0: usize,
14659 out1: usize,
14660 row_bytes: usize,
14661 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14662 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14663 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14664 }
14665
14666 #[allow(clippy::too_many_arguments)]
14669 pub fn matmul_q8_fused3_t(
14670 &self,
14671 w0: &crate::model::GpuTensor,
14672 w1: &crate::model::GpuTensor,
14673 w2: &crate::model::GpuTensor,
14674 aq: &CudaSlice<i8>,
14675 ad: &CudaSlice<f32>,
14676 m: usize,
14677 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14678 {
14679 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14680 return Ok(None);
14681 }
14682 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14683 return Ok(Some(self.e4m3_fused3_t_core(
14684 p0.0,
14685 p1.0,
14686 p2.0,
14687 aq,
14688 ad,
14689 m,
14690 w0.in_features(),
14691 p0.1,
14692 p1.1,
14693 p2.1,
14694 p0.2,
14695 p0.3,
14696 p1.3,
14697 p2.3,
14698 )?));
14699 }
14700 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14701 return Ok(None);
14702 };
14703 Ok(Some(self.q8_fused3_t_core(
14704 p0.0,
14705 p1.0,
14706 p2.0,
14707 aq,
14708 ad,
14709 m,
14710 w0.in_features(),
14711 p0.1,
14712 p1.1,
14713 p2.1,
14714 p0.2,
14715 )?))
14716 }
14717
14718 #[allow(clippy::too_many_arguments)]
14719 fn q8_fused3_t_core(
14720 &self,
14721 b0: &CudaSlice<u8>,
14722 b1: &CudaSlice<u8>,
14723 b2: &CudaSlice<u8>,
14724 aq: &CudaSlice<i8>,
14725 ad: &CudaSlice<f32>,
14726 m: usize,
14727 in_f: usize,
14728 out0: usize,
14729 out1: usize,
14730 out2: usize,
14731 row_bytes: usize,
14732 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14733 const ROWS_PER_BLOCK: u32 = 4;
14734 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14735 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14736 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14737 let f = self.func(if Self::batched_mcols(m) == 2 {
14738 "qmatvec_q8_0_mmvq_fused3_b2"
14739 } else {
14740 "qmatvec_q8_0_mmvq_fused3_b4"
14741 });
14742 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14743 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14744 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14745 let cfg = LaunchConfig {
14746 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14747 block_dim: (32, ROWS_PER_BLOCK, 1),
14748 shared_mem_bytes: 0,
14749 };
14750 let (inf, o0, o1, o2, mi, rbl) = (
14751 in_f as i32,
14752 out0 as i32,
14753 out1 as i32,
14754 out2 as i32,
14755 m as i32,
14756 row_bytes as i64,
14757 );
14758 let __s_b = self.gpu.stream();
14759 let mut b = __s_b.launch_builder(&f);
14760 b.arg(b0)
14761 .arg(b1)
14762 .arg(b2)
14763 .arg(aq)
14764 .arg(ad)
14765 .arg(&mut y0)
14766 .arg(&mut y1)
14767 .arg(&mut y2)
14768 .arg(&inf)
14769 .arg(&o0)
14770 .arg(&o1)
14771 .arg(&o2)
14772 .arg(&mi)
14773 .arg(&rbl);
14774 unsafe {
14775 b.launch(cfg)?;
14776 }
14777 Ok((y0, y1, y2))
14778 }
14779
14780 #[allow(clippy::too_many_arguments)]
14782 pub fn qmatvec_q8_fused3_t_raw(
14783 &self,
14784 b0: &CudaSlice<u8>,
14785 b1: &CudaSlice<u8>,
14786 b2: &CudaSlice<u8>,
14787 x: &CudaSlice<f32>,
14788 m: usize,
14789 in_f: usize,
14790 out0: usize,
14791 out1: usize,
14792 out2: usize,
14793 row_bytes: usize,
14794 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14795 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14796 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14797 }
14798
14799 pub fn q8_ffn_fuse2_on(&self) -> bool {
14803 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14804 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14805 }
14806
14807 #[allow(clippy::type_complexity)]
14813 fn q8_fused_params<'w, const N: usize>(
14814 &self,
14815 ws: &[&'w crate::model::GpuTensor; N],
14816 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14817 use crate::model::GpuTensor;
14818 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14819 return None;
14820 }
14821 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14822 return None;
14823 }
14824 let in_f = ws[0].in_features();
14825 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14826 for (i, w) in ws.iter().enumerate() {
14827 match w {
14828 GpuTensor::Quant {
14829 bytes,
14830 qtype,
14831 row_bytes,
14832 scale,
14833 ..
14834 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14835 out[i] = Some((bytes, w.out_features(), *row_bytes))
14836 }
14837 _ => return None,
14838 }
14839 }
14840 Some(out.map(|o| o.unwrap()))
14841 }
14842
14843 pub fn e4m3_dual_on(&self) -> bool {
14846 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14847 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14848 }
14849
14850 #[allow(clippy::type_complexity)]
14862 fn e4m3_fused_params<'w, const N: usize>(
14863 &self,
14864 ws: &[&'w crate::model::GpuTensor; N],
14865 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14866 use crate::model::GpuTensor;
14867 if !self.e4m3_dual_on() {
14868 return None;
14869 }
14870 let in_f = ws[0].in_features();
14871 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14872 for (i, w) in ws.iter().enumerate() {
14873 match w {
14874 GpuTensor::Quant {
14875 bytes,
14876 qtype,
14877 row_bytes,
14878 scale,
14879 rp,
14880 rp4,
14881 ..
14882 } if *qtype == QT_F8_E4M3
14883 && w.in_features() == in_f
14884 && *row_bytes == in_f
14885 && !*rp
14886 && rp4.is_none() =>
14887 {
14888 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14889 }
14890 _ => return None,
14891 }
14892 }
14893 Some(out.map(|o| o.unwrap()))
14894 }
14895
14896 #[allow(clippy::too_many_arguments)]
14900 fn e4m3_fused2_core(
14901 &self,
14902 b0: &CudaSlice<u8>,
14903 b1: &CudaSlice<u8>,
14904 aq: &CudaSlice<i8>,
14905 ad: &CudaSlice<f32>,
14906 in_f: usize,
14907 out0: usize,
14908 out1: usize,
14909 row_bytes: usize,
14910 ws0: f32,
14911 ws1: f32,
14912 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14913 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14915 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14916 let f = self.func("qmatvec_e4m3_mmvq_fused2");
14917 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14918 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14919 let cfg = LaunchConfig {
14920 grid_dim: (nb0 + nb1, 1, 1),
14921 block_dim: (32, ROWS_PER_BLOCK, 1),
14922 shared_mem_bytes: 0,
14923 };
14924 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14925 let __s_b = self.gpu.stream();
14926 let mut b = __s_b.launch_builder(&f);
14927 b.arg(b0)
14928 .arg(b1)
14929 .arg(aq)
14930 .arg(ad)
14931 .arg(&mut y0)
14932 .arg(&mut y1)
14933 .arg(&inf)
14934 .arg(&o0)
14935 .arg(&o1)
14936 .arg(&rbl)
14937 .arg(&ws0)
14938 .arg(&ws1);
14939 unsafe {
14940 b.launch(cfg)?;
14941 }
14942 Ok((y0, y1))
14943 }
14944
14945 #[allow(clippy::too_many_arguments)]
14947 fn e4m3_fused3_core(
14948 &self,
14949 b0: &CudaSlice<u8>,
14950 b1: &CudaSlice<u8>,
14951 b2: &CudaSlice<u8>,
14952 aq: &CudaSlice<i8>,
14953 ad: &CudaSlice<f32>,
14954 in_f: usize,
14955 out0: usize,
14956 out1: usize,
14957 out2: usize,
14958 row_bytes: usize,
14959 ws0: f32,
14960 ws1: f32,
14961 ws2: f32,
14962 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14963 const ROWS_PER_BLOCK: u32 = 4;
14964 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14965 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14966 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14967 let f = self.func("qmatvec_e4m3_mmvq_fused3");
14968 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14969 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14970 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14971 let cfg = LaunchConfig {
14972 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14973 block_dim: (32, ROWS_PER_BLOCK, 1),
14974 shared_mem_bytes: 0,
14975 };
14976 let (inf, o0, o1, o2, rbl) = (
14977 in_f as i32,
14978 out0 as i32,
14979 out1 as i32,
14980 out2 as i32,
14981 row_bytes as i64,
14982 );
14983 let __s_b = self.gpu.stream();
14984 let mut b = __s_b.launch_builder(&f);
14985 b.arg(b0)
14986 .arg(b1)
14987 .arg(b2)
14988 .arg(aq)
14989 .arg(ad)
14990 .arg(&mut y0)
14991 .arg(&mut y1)
14992 .arg(&mut y2)
14993 .arg(&inf)
14994 .arg(&o0)
14995 .arg(&o1)
14996 .arg(&o2)
14997 .arg(&rbl)
14998 .arg(&ws0)
14999 .arg(&ws1)
15000 .arg(&ws2);
15001 unsafe {
15002 b.launch(cfg)?;
15003 }
15004 Ok((y0, y1, y2))
15005 }
15006
15007 #[allow(clippy::too_many_arguments)]
15011 fn e4m3_fused2_t_core(
15012 &self,
15013 b0: &CudaSlice<u8>,
15014 b1: &CudaSlice<u8>,
15015 aq: &CudaSlice<i8>,
15016 ad: &CudaSlice<f32>,
15017 m: usize,
15018 in_f: usize,
15019 out0: usize,
15020 out1: usize,
15021 row_bytes: usize,
15022 ws0: f32,
15023 ws1: f32,
15024 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15025 const ROWS_PER_BLOCK: u32 = 4;
15026 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15027 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15028 let f = self.func(match Self::batched_mcols(m) {
15029 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15030 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15031 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15032 });
15033 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15034 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15035 let cfg = LaunchConfig {
15036 grid_dim: (nb0 + nb1, 1, 1),
15037 block_dim: (32, ROWS_PER_BLOCK, 1),
15038 shared_mem_bytes: 0,
15039 };
15040 let (inf, o0, o1, mi, rbl) = (
15041 in_f as i32,
15042 out0 as i32,
15043 out1 as i32,
15044 m as i32,
15045 row_bytes as i64,
15046 );
15047 let __s_b = self.gpu.stream();
15048 let mut b = __s_b.launch_builder(&f);
15049 b.arg(b0)
15050 .arg(b1)
15051 .arg(aq)
15052 .arg(ad)
15053 .arg(&mut y0)
15054 .arg(&mut y1)
15055 .arg(&inf)
15056 .arg(&o0)
15057 .arg(&o1)
15058 .arg(&mi)
15059 .arg(&rbl);
15060 unsafe {
15061 b.launch(cfg)?;
15062 }
15063 if ws0 != 1.0 {
15064 self.scale_inplace(&mut y0, ws0, m * out0)?;
15065 }
15066 if ws1 != 1.0 {
15067 self.scale_inplace(&mut y1, ws1, m * out1)?;
15068 }
15069 Ok((y0, y1))
15070 }
15071
15072 #[allow(clippy::too_many_arguments)]
15074 fn e4m3_fused3_t_core(
15075 &self,
15076 b0: &CudaSlice<u8>,
15077 b1: &CudaSlice<u8>,
15078 b2: &CudaSlice<u8>,
15079 aq: &CudaSlice<i8>,
15080 ad: &CudaSlice<f32>,
15081 m: usize,
15082 in_f: usize,
15083 out0: usize,
15084 out1: usize,
15085 out2: usize,
15086 row_bytes: usize,
15087 ws0: f32,
15088 ws1: f32,
15089 ws2: f32,
15090 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15091 const ROWS_PER_BLOCK: u32 = 4;
15092 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15093 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15094 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15095 let f = self.func(if Self::batched_mcols(m) == 2 {
15096 "qmatvec_e4m3_mmvq_fused3_b2"
15097 } else {
15098 "qmatvec_e4m3_mmvq_fused3_b4"
15099 });
15100 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15101 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15102 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15103 let cfg = LaunchConfig {
15104 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15105 block_dim: (32, ROWS_PER_BLOCK, 1),
15106 shared_mem_bytes: 0,
15107 };
15108 let (inf, o0, o1, o2, mi, rbl) = (
15109 in_f as i32,
15110 out0 as i32,
15111 out1 as i32,
15112 out2 as i32,
15113 m as i32,
15114 row_bytes as i64,
15115 );
15116 let __s_b = self.gpu.stream();
15117 let mut b = __s_b.launch_builder(&f);
15118 b.arg(b0)
15119 .arg(b1)
15120 .arg(b2)
15121 .arg(aq)
15122 .arg(ad)
15123 .arg(&mut y0)
15124 .arg(&mut y1)
15125 .arg(&mut y2)
15126 .arg(&inf)
15127 .arg(&o0)
15128 .arg(&o1)
15129 .arg(&o2)
15130 .arg(&mi)
15131 .arg(&rbl);
15132 unsafe {
15133 b.launch(cfg)?;
15134 }
15135 if ws0 != 1.0 {
15136 self.scale_inplace(&mut y0, ws0, m * out0)?;
15137 }
15138 if ws1 != 1.0 {
15139 self.scale_inplace(&mut y1, ws1, m * out1)?;
15140 }
15141 if ws2 != 1.0 {
15142 self.scale_inplace(&mut y2, ws2, m * out2)?;
15143 }
15144 Ok((y0, y1, y2))
15145 }
15146
15147 pub fn qmatvec_e4m3_blk_mmvq(
15157 &self,
15158 bytes: &CudaSlice<u8>,
15159 aq: &CudaSlice<i8>,
15160 ad: &CudaSlice<f32>,
15161 scales: &CudaSlice<f32>,
15162 m: usize,
15163 in_f: usize,
15164 out_f: usize,
15165 row_bytes: usize,
15166 scale_cols: usize,
15167 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15168 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15170 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15171 )?;
15172 Ok(y)
15173 }
15174
15175 #[allow(clippy::too_many_arguments)]
15177 pub fn qmatvec_e4m3_blk_mmvq_into(
15178 &self,
15179 bytes: &CudaSlice<u8>,
15180 aq: &CudaSlice<i8>,
15181 ad: &CudaSlice<f32>,
15182 scales: &CudaSlice<f32>,
15183 m: usize,
15184 in_f: usize,
15185 out_f: usize,
15186 row_bytes: usize,
15187 scale_cols: usize,
15188 y: &mut CudaSlice<f32>,
15189 ) -> Result<(), Box<dyn std::error::Error>> {
15190 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15192 let cfg = LaunchConfig {
15193 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15194 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15197 let (inf, outf, mi, rb, sc) = (
15198 in_f as i32,
15199 out_f as i32,
15200 m as i32,
15201 row_bytes as i64,
15202 scale_cols as i32,
15203 );
15204 let __s_b = self.gpu.stream();
15205 let mut b = __s_b.launch_builder(&f);
15206 b.arg(bytes)
15207 .arg(aq)
15208 .arg(ad)
15209 .arg(scales)
15210 .arg(&mut *y)
15211 .arg(&inf)
15212 .arg(&outf)
15213 .arg(&mi)
15214 .arg(&rb)
15215 .arg(&sc);
15216 unsafe {
15217 b.launch(cfg)?;
15218 }
15219 Ok(())
15220 }
15221
15222 #[allow(clippy::too_many_arguments)]
15228 pub fn qmatvec_e4m3_blk_mmvq_batched(
15229 &self,
15230 bytes: &CudaSlice<u8>,
15231 aq: &CudaSlice<i8>,
15232 ad: &CudaSlice<f32>,
15233 scales: &CudaSlice<f32>,
15234 m: usize,
15235 in_f: usize,
15236 out_f: usize,
15237 row_bytes: usize,
15238 scale_cols: usize,
15239 mcols: usize,
15240 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15241 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15243 let name = match mcols {
15244 2 => "qmatvec_e4m3_blk_mmvq_b2",
15245 4 => "qmatvec_e4m3_blk_mmvq_b4",
15246 8 => "qmatvec_e4m3_blk_mmvq_b8",
15247 16 => "qmatvec_e4m3_blk_mmvq_b16",
15248 _ => {
15249 return Err(
15250 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15251 );
15252 }
15253 };
15254 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15255 let f = self.func(name);
15256 let cfg = LaunchConfig {
15257 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15258 block_dim: (32, ROWS_PER_BLOCK, 1),
15259 shared_mem_bytes: 0,
15260 };
15261 let (inf, outf, mi, rb, sc) = (
15262 in_f as i32,
15263 out_f as i32,
15264 m as i32,
15265 row_bytes as i64,
15266 scale_cols as i32,
15267 );
15268 let __s_b = self.gpu.stream();
15269 let mut b = __s_b.launch_builder(&f);
15270 b.arg(bytes)
15271 .arg(aq)
15272 .arg(ad)
15273 .arg(scales)
15274 .arg(&mut y)
15275 .arg(&inf)
15276 .arg(&outf)
15277 .arg(&mi)
15278 .arg(&rb)
15279 .arg(&sc);
15280 unsafe {
15281 b.launch(cfg)?;
15282 }
15283 Ok(y)
15284 }
15285
15286 #[allow(clippy::too_many_arguments)]
15289 pub fn qmatvec_e4m3_blk_batched_raw(
15290 &self,
15291 bytes: &CudaSlice<u8>,
15292 x: &CudaSlice<f32>,
15293 scales: &CudaSlice<f32>,
15294 m: usize,
15295 in_f: usize,
15296 out_f: usize,
15297 row_bytes: usize,
15298 scale_cols: usize,
15299 mcols: usize,
15300 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15301 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15302 self.qmatvec_e4m3_blk_mmvq_batched(
15303 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15304 )
15305 }
15306
15307 #[allow(clippy::too_many_arguments)]
15310 pub fn qmatvec_e4m3_blk_mmvq_raw(
15311 &self,
15312 bytes: &CudaSlice<u8>,
15313 x: &CudaSlice<f32>,
15314 scales: &CudaSlice<f32>,
15315 m: usize,
15316 in_f: usize,
15317 out_f: usize,
15318 row_bytes: usize,
15319 scale_cols: usize,
15320 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15321 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15322 self.qmatvec_e4m3_blk_mmvq(
15323 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15324 )
15325 }
15326
15327 #[allow(clippy::too_many_arguments)]
15330 pub fn qmatvec_e4m3_fused2_raw(
15331 &self,
15332 b0: &CudaSlice<u8>,
15333 b1: &CudaSlice<u8>,
15334 x: &CudaSlice<f32>,
15335 in_f: usize,
15336 out0: usize,
15337 out1: usize,
15338 row_bytes: usize,
15339 ws0: f32,
15340 ws1: f32,
15341 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15342 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15343 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15344 }
15345
15346 #[allow(clippy::too_many_arguments)]
15347 pub fn qmatvec_e4m3_fused3_raw(
15348 &self,
15349 b0: &CudaSlice<u8>,
15350 b1: &CudaSlice<u8>,
15351 b2: &CudaSlice<u8>,
15352 x: &CudaSlice<f32>,
15353 in_f: usize,
15354 out0: usize,
15355 out1: usize,
15356 out2: usize,
15357 row_bytes: usize,
15358 ws0: f32,
15359 ws1: f32,
15360 ws2: f32,
15361 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15362 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15363 self.e4m3_fused3_core(
15364 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15365 )
15366 }
15367
15368 #[allow(clippy::too_many_arguments)]
15369 pub fn qmatvec_e4m3_fused2_t_raw(
15370 &self,
15371 b0: &CudaSlice<u8>,
15372 b1: &CudaSlice<u8>,
15373 x: &CudaSlice<f32>,
15374 m: usize,
15375 in_f: usize,
15376 out0: usize,
15377 out1: usize,
15378 row_bytes: usize,
15379 ws0: f32,
15380 ws1: f32,
15381 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15382 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15383 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15384 }
15385
15386 #[allow(clippy::too_many_arguments)]
15387 pub fn qmatvec_e4m3_fused3_t_raw(
15388 &self,
15389 b0: &CudaSlice<u8>,
15390 b1: &CudaSlice<u8>,
15391 b2: &CudaSlice<u8>,
15392 x: &CudaSlice<f32>,
15393 m: usize,
15394 in_f: usize,
15395 out0: usize,
15396 out1: usize,
15397 out2: usize,
15398 row_bytes: usize,
15399 ws0: f32,
15400 ws1: f32,
15401 ws2: f32,
15402 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15403 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15404 self.e4m3_fused3_t_core(
15405 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15406 )
15407 }
15408
15409 fn try_e4m3_blk_pre(
15420 &self,
15421 w: &crate::model::GpuTensor,
15422 aq: &CudaSlice<i8>,
15423 ad: &CudaSlice<f32>,
15424 m: usize,
15425 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15426 use crate::model::GpuTensor;
15427 if let GpuTensor::Quant {
15428 bytes,
15429 qtype,
15430 row_bytes,
15431 blk: Some(g),
15432 ..
15433 } = w
15434 {
15435 if *qtype == QT_F8_E4M3_BLK {
15436 if (2..=16).contains(&m)
15442 && std::env::var("MEMRA_NO_BATCHED").is_err()
15443 && (m <= 4 || Self::b8_enabled())
15444 {
15445 let mcols = Self::batched_mcols(m);
15446 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15447 bytes,
15448 aq,
15449 ad,
15450 &g.scales,
15451 m,
15452 w.in_features(),
15453 w.out_features(),
15454 *row_bytes,
15455 g.cols,
15456 mcols,
15457 )?));
15458 }
15459 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15460 bytes,
15461 aq,
15462 ad,
15463 &g.scales,
15464 m,
15465 w.in_features(),
15466 w.out_features(),
15467 *row_bytes,
15468 g.cols,
15469 )?));
15470 }
15471 }
15472 Ok(None)
15473 }
15474
15475 fn try_e4m3_blk_prefill(
15522 &self,
15523 w: &crate::model::GpuTensor,
15524 x: &CudaSlice<f32>,
15525 m: usize,
15526 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15527 use crate::model::GpuTensor;
15528 let GpuTensor::Quant {
15529 bytes,
15530 qtype,
15531 blk: Some(g),
15532 ..
15533 } = w
15534 else {
15535 return Ok(None);
15536 };
15537 if *qtype != QT_F8_E4M3_BLK {
15538 return Ok(None);
15539 }
15540 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15545 return Ok(Some(y));
15546 }
15547 let (in_f, out_f) = (w.in_features(), w.out_features());
15548 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15549 let tmp = GpuTensor::Quant {
15550 bytes: slab,
15551 qtype: QT_Q8_0,
15552 row_bytes: in_f / 32 * 34,
15553 ne: vec![in_f as u64, out_f as u64],
15554 scale: 1.0,
15555 rp: false,
15556 #[cfg(memra_cutlass)]
15557 cutlass: None,
15558 fp8: None,
15559 blk: None,
15560 f16: None,
15561 rp4: None,
15562 };
15563 Ok(Some(self.matmul(&tmp, x, m)?))
15565 }
15566
15567 pub fn matmul_pre_noscale(
15568 &self,
15569 w: &crate::model::GpuTensor,
15570 aq: &CudaSlice<i8>,
15571 ad: &CudaSlice<f32>,
15572 m: usize,
15573 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15574 use crate::model::GpuTensor;
15575 if m == 1 {
15579 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15580 return Ok(Some((y, 1.0)));
15581 }
15582 }
15583 if m != 1 || !self.uses_q8_1_fast(w) {
15585 return Ok(None);
15586 }
15587 let in_f = w.in_features();
15588 let out_f = w.out_features();
15589 let (bytes, qtype, row_bytes, scale, rp) = match w {
15590 GpuTensor::Quant {
15591 bytes,
15592 qtype,
15593 row_bytes,
15594 scale,
15595 rp,
15596 ..
15597 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15598 _ => return Ok(None),
15599 };
15600 if self.mmvq_supports(qtype) {
15602 let (mbytes, mrp) = match w {
15604 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15605 _ => (bytes, rp),
15606 };
15607 let y = self.qmatvec_mmvq(
15608 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15609 )?;
15610 return Ok(Some((y, scale)));
15611 }
15612 let name = match qtype {
15614 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15615 QT_Q4_K => "qmatvec_q4_K_dp4a",
15616 QT_Q6_K => "qmatvec_q6_K_dp4a",
15617 QT_Q5_K => "qmatvec_q5_K_dp4a",
15618 QT_Q3_K => "qmatvec_q3_K_dp4a",
15619 QT_NVFP4 => {
15620 if rp {
15621 "qmatvec_nvfp4_dp4a_rp"
15622 } else {
15623 "qmatvec_nvfp4_dp4a"
15624 }
15625 }
15626 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15627 _ => return Ok(None),
15628 };
15629 let f = self.func(name);
15630 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15631 let cfg = LaunchConfig {
15632 grid_dim: (out_f as u32, m as u32, 1),
15633 block_dim: (128, 1, 1),
15634 shared_mem_bytes: 0,
15635 };
15636 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15637 let __s_b = self.gpu.stream();
15638 let mut b = __s_b.launch_builder(&f);
15639 b.arg(bytes)
15640 .arg(aq)
15641 .arg(ad)
15642 .arg(&mut y)
15643 .arg(&inf)
15644 .arg(&outf)
15645 .arg(&mi)
15646 .arg(&rb);
15647 unsafe {
15648 b.launch(cfg)?;
15649 }
15650 Ok(Some((y, scale)))
15651 }
15652
15653 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15656 if qtype == QT_F8_E4M3 {
15661 return true;
15662 }
15663 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15664 return false;
15665 }
15666 matches!(
15667 qtype,
15668 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15669 )
15670 }
15671
15672 pub fn qmatvec_mmvq(
15677 &self,
15678 bytes: &CudaSlice<u8>,
15679 aq: &CudaSlice<i8>,
15680 ad: &CudaSlice<f32>,
15681 m: usize,
15682 in_f: usize,
15683 out_f: usize,
15684 qtype: i32,
15685 row_bytes: usize,
15686 scale: f32,
15687 rp: bool,
15688 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15689 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15691 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15692 )?;
15693 Ok(y)
15694 }
15695
15696 #[allow(clippy::too_many_arguments)]
15698 pub fn qmatvec_mmvq_into(
15699 &self,
15700 bytes: &CudaSlice<u8>,
15701 aq: &CudaSlice<i8>,
15702 ad: &CudaSlice<f32>,
15703 m: usize,
15704 in_f: usize,
15705 out_f: usize,
15706 qtype: i32,
15707 row_bytes: usize,
15708 scale: f32,
15709 rp: bool,
15710 y: &mut CudaSlice<f32>,
15711 ) -> Result<(), Box<dyn std::error::Error>> {
15712 debug_assert!(y.len() >= m * out_f);
15713 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15719 && rp
15720 && m == 1
15721 && out_f >= 64
15722 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15723 && {
15724 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15725 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15726 }
15727 {
15728 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15729 let cfg = LaunchConfig {
15730 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15731 block_dim: (32, 2, 1),
15732 shared_mem_bytes: 0,
15733 };
15734 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15735 let __s_b = self.gpu.stream();
15736 let mut b = __s_b.launch_builder(&f);
15737 b.arg(bytes)
15738 .arg(aq)
15739 .arg(ad)
15740 .arg(&mut *y)
15741 .arg(&inf)
15742 .arg(&outf)
15743 .arg(&mi)
15744 .arg(&rb);
15745 unsafe {
15746 b.launch(cfg)?;
15747 }
15748 if scale != 1.0 {
15749 self.scale_inplace(y, scale, out_f)?;
15750 }
15751 return Ok(());
15752 }
15753 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15762 2
15763 } else {
15764 1
15765 };
15766 if m == 1 && qtype == QT_Q4_0 {
15771 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15772 mr = *Q40MR.get_or_init(|| {
15775 std::env::var("MEMRA_Q40_MR")
15776 .ok()
15777 .and_then(|v| v.parse().ok())
15778 .unwrap_or(1)
15779 });
15780 }
15781 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15792 let q5_force = q5_mode.as_deref() == Some("2");
15793 let q5_il = qtype == QT_Q5_K
15796 && m == 1
15797 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15798 if q5_il && !q5_force && out_f > 65536 {
15799 mr = 1;
15800 }
15801 if qtype == QT_Q4_0 && rp && mr != 1 {
15804 mr = 2;
15805 }
15806 if qtype == QT_Q8_0 && rp {
15810 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15811 mr = *Q80MR.get_or_init(|| {
15812 std::env::var("MEMRA_Q80_MR")
15813 .ok()
15814 .and_then(|v| v.parse().ok())
15815 .unwrap_or(1)
15816 });
15817 }
15818 let name = match (qtype, mr, rp) {
15819 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15820 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15821 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15822 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15823 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15824 (QT_Q5_K, 2, _) => {
15825 if q5_il {
15826 "qmatvec_q5_K_mmvq_mr2_il"
15827 } else {
15828 "qmatvec_q5_K_mmvq_mr2"
15829 }
15830 }
15831 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15832 (QT_Q8_0, _, true)
15837 if in_f % 1024 == 0 && {
15838 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15839 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15840 } =>
15841 {
15842 "qmatvec_q8_0_mmvq_rpca"
15843 }
15844 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15845 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15846 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15850 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15851 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15852 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15853 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15854 (QT_Q5_K, _, _) => {
15855 if q5_il {
15856 "qmatvec_q5_K_mmvq_il"
15857 } else {
15858 "qmatvec_q5_K_mmvq"
15859 }
15860 }
15861 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15862 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15863 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15864 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15865 };
15866 let f = self.func(name);
15867 let rows_per_block = ROWS_PER_BLOCK * mr;
15869 let cfg = LaunchConfig {
15870 grid_dim: (
15871 (out_f as u32 + rows_per_block - 1) / rows_per_block,
15872 m as u32,
15873 1,
15874 ),
15875 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15878 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15879 let __s_b = self.gpu.stream();
15880 let mut b = __s_b.launch_builder(&f);
15881 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15886 if Self::pdl_on()
15889 && Self::pdl_mmvq_on()
15890 && Self::pdl_nvfp4q8_on()
15891 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15892 {
15893 use cudarc::driver::{DevicePtr, DevicePtrMut};
15894 let s = &self.gpu.stream();
15895 let (pw, _g0) = bytes.device_ptr(s);
15896 let (paq, _g1) = aq.device_ptr(s);
15897 let (pad, _g2) = ad.device_ptr(s);
15898 let (py, _g3) = y.device_ptr_mut(s);
15899 let mut ps = [
15900 &pw as *const _ as *mut std::ffi::c_void,
15901 &paq as *const _ as *mut _,
15902 &pad as *const _ as *mut _,
15903 &py as *const _ as *mut _,
15904 &inf as *const _ as *mut _,
15905 &outf as *const _ as *mut _,
15906 &mi as *const _ as *mut _,
15907 &rb as *const _ as *mut _,
15908 &scale as *const _ as *mut _,
15909 ];
15910 unsafe {
15911 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15912 }
15913 return Ok(());
15914 }
15915 b.arg(bytes)
15916 .arg(aq)
15917 .arg(ad)
15918 .arg(&mut *y)
15919 .arg(&inf)
15920 .arg(&outf)
15921 .arg(&mi)
15922 .arg(&rb)
15923 .arg(&scale);
15924 unsafe {
15925 b.launch(cfg)?;
15926 }
15927 } else if Self::pdl_on()
15928 && Self::pdl_mmvq_on()
15929 && (matches!(
15930 name,
15931 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15932 ) || (Self::pdl_nvfp4q8_on()
15933 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15934 {
15935 {
15939 use cudarc::driver::{DevicePtr, DevicePtrMut};
15940 let s = &self.gpu.stream();
15941 let (pw, _g0) = bytes.device_ptr(s);
15942 let (paq, _g1) = aq.device_ptr(s);
15943 let (pad, _g2) = ad.device_ptr(s);
15944 let (py, _g3) = y.device_ptr_mut(s);
15945 let mut ps = [
15946 &pw as *const _ as *mut std::ffi::c_void,
15947 &paq as *const _ as *mut _,
15948 &pad as *const _ as *mut _,
15949 &py as *const _ as *mut _,
15950 &inf as *const _ as *mut _,
15951 &outf as *const _ as *mut _,
15952 &mi as *const _ as *mut _,
15953 &rb as *const _ as *mut _,
15954 ];
15955 unsafe {
15956 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15957 }
15958 }
15959 if scale != 1.0 {
15960 self.scale_inplace(y, scale, m * out_f)?;
15961 }
15962 } else {
15963 b.arg(bytes)
15964 .arg(aq)
15965 .arg(ad)
15966 .arg(&mut *y)
15967 .arg(&inf)
15968 .arg(&outf)
15969 .arg(&mi)
15970 .arg(&rb);
15971 unsafe {
15972 b.launch(cfg)?;
15973 }
15974 if scale != 1.0 {
15975 self.scale_inplace(y, scale, m * out_f)?;
15976 }
15977 }
15978 Ok(())
15979 }
15980
15981 pub fn qmatvec_mmvq_raw(
15985 &self,
15986 bytes: &CudaSlice<u8>,
15987 x: &CudaSlice<f32>,
15988 m: usize,
15989 in_f: usize,
15990 out_f: usize,
15991 qtype: i32,
15992 row_bytes: usize,
15993 rp: bool,
15994 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15995 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15996 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15997 }
15998
15999 pub fn batched_supports(&self, qtype: i32) -> bool {
16003 matches!(
16004 qtype,
16005 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16006 )
16007 }
16008
16009 pub fn iq_fast_enabled() -> bool {
16017 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16018 *ON.get_or_init(|| {
16019 std::env::var("MEMRA_IQ_FAST")
16020 .map(|v| v != "0")
16021 .unwrap_or(true)
16022 })
16023 }
16024
16025 pub fn b8_enabled() -> bool {
16028 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16029 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16030 }
16031
16032 pub fn batched_mcols(m: usize) -> usize {
16034 if m == 2 {
16035 2
16036 } else if m <= 4 {
16037 4
16038 } else if m <= 8 {
16039 8
16040 } else {
16041 16
16042 }
16043 }
16044
16045 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16050 Some(match (qtype, mcols) {
16051 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16052 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16053 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16054 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16060 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16061 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16062 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16063 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16066 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16067 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16068 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16069 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16072 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16073 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16074 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16075 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16076 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16077 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16078 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16079 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16083 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16084 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16085 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16086 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16090 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16091 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16092 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16093 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16094 _ => return None,
16095 })
16096 }
16097
16098 pub fn sm_count(&self) -> i32 {
16133 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16134 *SMS.get_or_init(|| {
16135 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16136 self.gpu
16137 .ctx
16138 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16139 .unwrap_or(82)
16140 })
16141 }
16142
16143 pub fn batched_variant(
16144 &self,
16145 _m: usize,
16146 in_f: usize,
16147 out_f: usize,
16148 qtype: i32,
16149 row_bytes: usize,
16150 mcols: usize,
16151 rp: bool,
16152 ) -> &'static str {
16153 if qtype == QT_Q8_0 {
16158 return if rp { "rp" } else { "base" };
16159 }
16160 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16161 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16162 Ok("base") => "base",
16163 Ok("pf") => "pf",
16164 Ok("r2") => "r2",
16165 Ok("r2w8") => "r2w8",
16166 Ok("pfr2") => "pfr2",
16167 Ok("ca") => "ca",
16168 Ok("car2") => "car2",
16169 Ok("rp") => "rp",
16172 Ok("rpr2") => "rpr2",
16173 Ok("rpr2w8") => "rpr2w8",
16174 Ok("rpca") => "rpca",
16177 Ok("rpcar2") => "rpcar2",
16178 Ok("rpsc") => "rpsc",
16185 Ok("rpms") => "rpms",
16186 Ok("rpmsc") => "rpmsc",
16187 Ok("rpks") => "rpks",
16188 Ok("rpksc") => "rpksc",
16189 _ => "auto",
16190 });
16191 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16195 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16200 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16201 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16202 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16203 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16204 let sms = *SMS.get_or_init(|| {
16205 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16206 self.gpu
16207 .ctx
16208 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16209 .unwrap_or(82)
16210 });
16211 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16231 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16234 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16235 Ok("base") => "base",
16236 Ok("r2") => "r2",
16237 Ok("r2w8") => "r2w8",
16238 _ => "auto",
16239 });
16240 let variant: &'static str = if qtype == QT_Q4_0 {
16241 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16245 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16246 Ok("base") => "base",
16252 Ok("r2") => "r2",
16253 Ok("ms") => "ms",
16254 Ok("sm") => "sm",
16255 Ok("la") => "la",
16256 _ => "auto",
16257 });
16258 let v = if q40 != "auto" {
16259 q40
16260 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16261 "r2"
16262 } else {
16263 "base"
16264 };
16265 if rp {
16270 match v {
16271 "ms" => "r2ms_rp",
16272 "sm" => "r2sm_rp",
16273 "la" => "r2la_rp",
16274 "r2" => "r2_rp",
16275 _ => "rp",
16276 }
16277 } else if matches!(v, "ms" | "sm" | "la") {
16278 "r2"
16279 } else {
16280 v
16281 }
16282 } else if qtype != QT_NVFP4 && !kq_r2 {
16283 "base"
16284 } else if kq_r2 && rp {
16285 "rp"
16289 } else if kq_r2 {
16290 if kq_bv != "auto" {
16293 if kq_bv == "r2w8" && mcols != 4 {
16294 "r2"
16295 } else {
16296 kq_bv
16297 }
16298 } else if bv != "auto" {
16299 match bv {
16300 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16301 "r2w8" | "rpr2w8" => {
16302 if mcols != 4 {
16303 "r2"
16304 } else {
16305 "r2w8"
16306 }
16307 }
16308 _ => "base", }
16310 } else {
16311 let blocks = (out_f + 7) / 8;
16312 let waves = blocks as f64 / (7 * sms as usize) as f64;
16313 let filled = blocks >= 4 * sms as usize;
16314 let use_r2 = if qtype == QT_Q4_K {
16315 filled
16316 } else {
16317 waves >= 2.0
16318 };
16319 if use_r2 { "r2" } else { "base" }
16320 }
16321 } else if bv != "auto" {
16322 let v = if bv == "r2w8" && mcols == 2 {
16327 "r2"
16328 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16329 "pf"
16330 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16331 "r2"
16332 } else if bv == "pfr2" && mcols == 8 {
16333 "r2"
16334 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16335 "rpr2"
16336 }
16337 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16339 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16340 } else if bv == "rpcar2" && mcols == 2 {
16341 "rpca"
16342 }
16343 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16346 "rpr2"
16347 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16348 "rpr2"
16349 } else {
16350 bv
16351 };
16352 if rp {
16353 match v {
16354 "base" | "pf" | "ca" | "rp" => "rp",
16355 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16356 "r2w8" | "rpr2w8" => {
16357 if mcols == 2 {
16358 "rpr2"
16359 } else {
16360 "rpr2w8"
16361 }
16362 }
16363 other => other, }
16365 } else {
16366 v
16367 }
16368 } else if mcols == 8 {
16369 if rp {
16380 if sc_ok { "rpsc" } else { "rpr2w8" }
16381 } else {
16382 "r2w8"
16383 }
16384 } else if mcols >= 4 {
16385 let blocks = (out_f + 7) / 8;
16389 let r7 = 7 * sms as usize;
16390 let r8 = 8 * sms as usize;
16391 let waves = blocks as f64 / r7 as f64;
16392 let filled = blocks >= 4 * sms as usize;
16393 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16397 if rp { "rpr2w8" } else { "r2w8" }
16401 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16402 if rp { "rpr2" } else { "r2" }
16405 } else {
16406 if rp { "rp" } else { "pf" }
16410 }
16411 } else if in_f >= 6144 {
16412 if rp { "rpr2" } else { "r2" }
16416 } else if rp {
16417 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16422 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16423 "rpsc"
16424 } else {
16425 "rp"
16426 }
16427 } else {
16428 "base"
16429 };
16430 variant
16431 }
16432
16433 pub fn qmatvec_mmvq_batched(
16434 &self,
16435 bytes: &CudaSlice<u8>,
16436 aq: &CudaSlice<i8>,
16437 ad: &CudaSlice<f32>,
16438 m: usize,
16439 in_f: usize,
16440 out_f: usize,
16441 qtype: i32,
16442 row_bytes: usize,
16443 mcols: usize,
16444 scale: f32,
16445 rp: bool,
16446 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16447 const ROWS_PER_BLOCK: u32 = 4;
16448 let forced: Option<&'static str> = {
16453 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16454 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16455 .as_deref()
16456 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16457 };
16458 let variant = match forced {
16459 Some(v) if !rp || v.contains("rp") => v,
16460 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16461 };
16462 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16463 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16464 })?;
16465 let variant = if mcols == 16 {
16469 if rp { "rp" } else { "base" }
16470 } else {
16471 variant
16472 };
16473 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16480 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16481 if b567
16482 && qtype == QT_NVFP4
16483 && rp
16484 && mcols == 8
16485 && (5..=7).contains(&m)
16486 && matches!(variant, "rpsc" | "rpr2w8")
16487 {
16488 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16489 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16491 let cfg = LaunchConfig {
16492 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16493 block_dim: (32, ROWS_PER_BLOCK, 1),
16494 shared_mem_bytes: 0,
16495 };
16496 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16497 let __s_b = self.gpu.stream();
16498 let mut b = __s_b.launch_builder(&f);
16499 b.arg(bytes)
16500 .arg(aq)
16501 .arg(ad)
16502 .arg(&mut y)
16503 .arg(&inf)
16504 .arg(&outf)
16505 .arg(&mi)
16506 .arg(&rb);
16507 unsafe {
16508 b.launch(cfg)?;
16509 }
16510 if scale != 1.0 {
16511 self.scale_inplace(&mut y, scale, m * out_f)?;
16512 }
16513 return Ok(y);
16514 }
16515 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16516 "base" => (base_name.into(), ROWS_PER_BLOCK),
16517 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16518 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16519 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16520 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16524 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16525 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16526 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16527 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16528 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16529 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16530 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16532 debug_assert!(
16533 !rp || name.contains("_rp"),
16534 "rp weight dispatched to a GGUF-layout kernel"
16535 );
16536 let f = self.func(&name);
16537 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16538 let smem = if name.contains("_r2sm_rp") {
16540 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16541 } else {
16542 0
16543 };
16544 let cfg = LaunchConfig {
16545 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16546 block_dim: (32, ROWS_PER_BLOCK, 1),
16547 shared_mem_bytes: smem,
16548 };
16549 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16550 let __s_b = self.gpu.stream();
16551 let mut b = __s_b.launch_builder(&f);
16552 b.arg(bytes)
16553 .arg(aq)
16554 .arg(ad)
16555 .arg(&mut y)
16556 .arg(&inf)
16557 .arg(&outf)
16558 .arg(&mi)
16559 .arg(&rb);
16560 unsafe {
16561 b.launch(cfg)?;
16562 }
16563 if scale != 1.0 {
16564 self.scale_inplace(&mut y, scale, m * out_f)?;
16565 }
16566 Ok(y)
16567 }
16568
16569 pub fn qmatvec_batched_raw(
16573 &self,
16574 bytes: &CudaSlice<u8>,
16575 x: &CudaSlice<f32>,
16576 m: usize,
16577 in_f: usize,
16578 out_f: usize,
16579 qtype: i32,
16580 row_bytes: usize,
16581 mcols: usize,
16582 rp: bool,
16583 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16584 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16585 self.qmatvec_mmvq_batched(
16586 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16587 )
16588 }
16589
16590 pub fn qmatvec_nvfp4_batched_raw(
16592 &self,
16593 bytes: &CudaSlice<u8>,
16594 x: &CudaSlice<f32>,
16595 m: usize,
16596 in_f: usize,
16597 out_f: usize,
16598 row_bytes: usize,
16599 mcols: usize,
16600 rp: bool,
16601 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16602 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16603 }
16604
16605 fn try_fp4_gemm(
16609 &self,
16610 w: &crate::model::GpuTensor,
16611 x: &CudaSlice<f32>,
16612 m: usize,
16613 in_f: usize,
16614 out_f: usize,
16615 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16616 use crate::model::GpuTensor;
16617 if cfg!(memra_portable_cuda) {
16618 return Ok(None);
16619 }
16620 if std::env::var("MEMRA_FP4").is_ok() {
16623 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16624 }
16625 if std::env::var("MEMRA_FP4").is_err() {
16626 return Ok(None);
16627 }
16628 #[cfg(memra_cutlass)]
16637 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16638 if let GpuTensor::Quant {
16639 bytes,
16640 qtype,
16641 scale,
16642 row_bytes,
16643 cutlass,
16644 ..
16645 } = w
16646 {
16647 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16648 if let Some(cw) = cutlass {
16649 let y = self.cutlass_fp4_gemm(
16651 &cw.b_packed,
16652 &cw.sfb_swizzled,
16653 x,
16654 *scale,
16655 m,
16656 out_f,
16657 in_f,
16658 )?;
16659 return Ok(Some(y));
16660 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16661 let (b_packed, sfb_sw) =
16666 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16667 let y =
16668 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16669 return Ok(Some(y));
16670 }
16671 }
16672 }
16673 }
16674 if let GpuTensor::Quant {
16675 bytes,
16676 qtype,
16677 row_bytes,
16678 scale,
16679 rp,
16680 ..
16681 } = w
16682 {
16683 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16686 let y =
16687 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16688 return Ok(Some(y));
16689 }
16690 }
16691 Ok(None)
16692 }
16693
16694 pub fn rms_norm_f16out(
16698 &self,
16699 x: &CudaSlice<f32>,
16700 w: &CudaSlice<f32>,
16701 dst: &mut CudaSlice<f32>,
16702 dst16: &mut CudaSlice<u8>,
16703 ncols: usize,
16704 nrows: usize,
16705 eps: f32,
16706 ) -> Result<(), Box<dyn std::error::Error>> {
16707 let f = self.func("rms_norm_f16out_f32");
16708 let cfg = LaunchConfig {
16709 grid_dim: (nrows as u32, 1, 1),
16710 block_dim: (rms_block(), 1, 1),
16711 shared_mem_bytes: 0,
16712 };
16713 let (nc, e) = (ncols as i32, eps);
16714 let __s_b = self.gpu.stream();
16715 let mut b = __s_b.launch_builder(&f);
16716 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16717 unsafe {
16718 b.launch(cfg)?;
16719 }
16720 Ok(())
16721 }
16722
16723 #[allow(clippy::too_many_arguments)]
16726 pub fn add_rms_norm_f16out(
16727 &self,
16728 a: &CudaSlice<f32>,
16729 b: &CudaSlice<f32>,
16730 w: &CudaSlice<f32>,
16731 res: &mut CudaSlice<f32>,
16732 dst: &mut CudaSlice<f32>,
16733 dst16: &mut CudaSlice<u8>,
16734 ncols: usize,
16735 nrows: usize,
16736 eps: f32,
16737 ) -> Result<(), Box<dyn std::error::Error>> {
16738 let f = self.func("add_rms_norm_f16out_f32");
16739 let cfg = LaunchConfig {
16740 grid_dim: (nrows as u32, 1, 1),
16741 block_dim: (rms_block(), 1, 1),
16742 shared_mem_bytes: 0,
16743 };
16744 let (nc, e) = (ncols as i32, eps);
16745 let __s_lb = self.gpu.stream();
16746 let mut lb = __s_lb.launch_builder(&f);
16747 lb.arg(a)
16748 .arg(b)
16749 .arg(w)
16750 .arg(res)
16751 .arg(dst)
16752 .arg(dst16)
16753 .arg(&nc)
16754 .arg(&e);
16755 unsafe {
16756 lb.launch(cfg)?;
16757 }
16758 Ok(())
16759 }
16760
16761 pub fn matmul_group_xh(
16764 &self,
16765 ws: &[&crate::model::GpuTensor],
16766 x: &CudaSlice<f32>,
16767 xh: &CudaSlice<u8>,
16768 m: usize,
16769 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16770 let mut out = Vec::with_capacity(ws.len());
16771 let in_f = ws[0].in_features();
16772 for w in ws {
16773 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16774 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16775 out.push(y);
16776 continue;
16777 }
16778 }
16779 out.push(self.matmul(w, x, m)?);
16780 }
16781 Ok(out)
16782 }
16783
16784 pub fn gdn_pad_mask(
16787 &self,
16788 beta: &mut CudaSlice<f32>,
16789 g_log: &mut CudaSlice<f32>,
16790 len_d: &CudaSlice<i32>,
16791 h: usize,
16792 t: usize,
16793 ) -> Result<(), Box<dyn std::error::Error>> {
16794 let f = self.func("gdn_pad_mask_f32");
16795 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16796 let (hi, ti) = (h as i32, t as i32);
16797 let __s_b = self.gpu.stream();
16798 let mut b = __s_b.launch_builder(&f);
16799 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16800 unsafe {
16801 b.launch(cfg)?;
16802 }
16803 Ok(())
16804 }
16805
16806 pub fn row_gather_dev(
16809 &self,
16810 src: &CudaSlice<f32>,
16811 dst: &mut CudaSlice<f32>,
16812 len_d: &CudaSlice<i32>,
16813 ncols: usize,
16814 ) -> Result<(), Box<dyn std::error::Error>> {
16815 let f = self.func("row_gather_dev_f32");
16816 let cfg = LaunchConfig::for_num_elems(ncols as u32);
16817 let nc = ncols as i32;
16818 let __s_b = self.gpu.stream();
16819 let mut b = __s_b.launch_builder(&f);
16820 b.arg(src).arg(dst).arg(len_d).arg(&nc);
16821 unsafe {
16822 b.launch(cfg)?;
16823 }
16824 Ok(())
16825 }
16826
16827 pub fn matmul_group(
16834 &self,
16835 ws: &[&crate::model::GpuTensor],
16836 x: &CudaSlice<f32>,
16837 m: usize,
16838 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16839 use crate::model::GpuTensor;
16840 let mut out = Vec::with_capacity(ws.len());
16841 let any_mirror = ws
16842 .iter()
16843 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16844 if m >= 16 && any_mirror && !self.verify_exact_on() {
16845 let in_f = ws[0].in_features();
16846 let xh = self.f16_act(x, m * in_f, in_f)?;
16847 for w in ws {
16848 if w.in_features() == in_f {
16849 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16850 out.push(y);
16851 continue;
16852 }
16853 }
16854 out.push(self.matmul(w, x, m)?);
16855 }
16856 return Ok(out);
16857 }
16858 for w in ws {
16859 out.push(self.matmul(w, x, m)?);
16860 }
16861 Ok(out)
16862 }
16863
16864 pub fn matmul_group_multi(
16871 &self,
16872 ws: &[&crate::model::GpuTensor],
16873 xs: &[&CudaSlice<f32>],
16874 ms: &[usize],
16875 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16876 assert_eq!(xs.len(), ms.len());
16877 let in_f = ws[0].in_features();
16878 let total: usize = ms.iter().sum();
16879 let mut xcat = self.uninit(total * in_f)?;
16880 let mut off = 0usize;
16881 for (x, &m) in xs.iter().zip(ms) {
16882 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16883 off += m;
16884 }
16885 let ys = self.matmul_group(ws, &xcat, total)?;
16886 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16887 for (w, y) in ws.iter().zip(ys) {
16888 let out_f = w.out_features();
16889 let mut off = 0usize;
16890 for (s, &m) in ms.iter().enumerate() {
16891 let mut ys_s = self.uninit(m * out_f)?;
16892 let src = y.slice(off * out_f..(off + m) * out_f);
16893 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16894 out[s].push(ys_s);
16895 off += m;
16896 }
16897 }
16898 Ok(out)
16899 }
16900
16901 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16911 use crate::model::GpuTensor;
16912 if !legacy_quant_gemm_allowed(
16913 cfg!(memra_portable_cuda),
16914 cfg!(memra_hopper_mma),
16915 std::env::var_os("MEMRA_NO_GEMM").is_some(),
16916 ) {
16917 return false;
16918 }
16919 match w {
16920 GpuTensor::Quant { qtype, .. } => {
16921 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16922 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16923 }
16924 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16925 }
16926 }
16927
16928 pub fn qmatvec_gemm(
16935 &self,
16936 w: &crate::model::GpuTensor,
16937 aq: &CudaSlice<i8>,
16938 ad: &CudaSlice<f32>,
16939 m: usize,
16940 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16941 use crate::model::GpuTensor;
16942 let in_f = w.in_features();
16943 let out_f = w.out_features();
16944 let (bytes, qtype, row_bytes, scale, rp) = match w {
16945 GpuTensor::Quant {
16946 bytes,
16947 qtype,
16948 row_bytes,
16949 scale,
16950 rp,
16951 ..
16952 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16953 _ => unreachable!("gemm_supports guaranteed Quant"),
16954 };
16955 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16961 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16962 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16963 if scale != 1.0 {
16964 self.scale_inplace(&mut y, scale, m * out_f)?;
16965 }
16966 return Ok(y);
16967 }
16968 }
16969 let name = match qtype {
16970 QT_Q8_0 => "qmatvec_gemm_q8_0",
16971 QT_Q4_K => "qmatvec_gemm_q4_K",
16972 QT_Q4_0 => {
16973 if rp {
16974 "qmatvec_gemm_q4_0_rp"
16975 } else {
16976 "qmatvec_gemm_q4_0"
16977 }
16978 }
16979 QT_Q5_K => "qmatvec_gemm_q5_K",
16980 QT_Q6_K => "qmatvec_gemm_q6_K",
16981 QT_NVFP4 => {
16982 if rp {
16983 "qmatvec_gemm_nvfp4_rp"
16984 } else {
16985 "qmatvec_gemm_nvfp4"
16986 }
16987 }
16988 _ => unreachable!(),
16989 };
16990 let f = self.func(name);
16991 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);
16996 let k1_tile = if is_k1 {
16998 k1_launch_override().unwrap_or((128, 128, 8))
16999 } else {
17000 (128, 128, 8)
17001 };
17002 let (bm, bn): (u32, u32) = if is_k1 {
17003 (k1_tile.0, k1_tile.1)
17004 } else {
17005 (64, 256)
17006 };
17007 let warps: u32 = if is_k1 {
17008 k1_tile.2
17009 } else {
17010 match qtype {
17011 QT_NVFP4 => 8,
17012 _ => 4,
17013 }
17014 };
17015 let cfg = LaunchConfig {
17016 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17017 block_dim: (32, warps, 1),
17018 shared_mem_bytes: 0,
17019 };
17020 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17021 let __s_b = self.gpu.stream();
17022 let mut b = __s_b.launch_builder(&f);
17023 b.arg(bytes)
17024 .arg(aq)
17025 .arg(ad)
17026 .arg(&mut y)
17027 .arg(&inf)
17028 .arg(&outf)
17029 .arg(&mi)
17030 .arg(&rb);
17031 unsafe {
17032 b.launch(cfg)?;
17033 }
17034 if scale != 1.0 {
17035 self.scale_inplace(&mut y, scale, m * out_f)?;
17036 }
17037 Ok(y)
17038 }
17039
17040 pub fn qmatvec_gemm_raw(
17045 &self,
17046 bytes: &CudaSlice<u8>,
17047 x: &CudaSlice<f32>,
17048 m: usize,
17049 in_f: usize,
17050 out_f: usize,
17051 qtype: i32,
17052 row_bytes: usize,
17053 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17054 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17055 let name = match qtype {
17056 QT_Q8_0 => "qmatvec_gemm_q8_0",
17057 QT_Q4_K => "qmatvec_gemm_q4_K",
17058 QT_Q4_0 => "qmatvec_gemm_q4_0",
17059 QT_Q5_K => "qmatvec_gemm_q5_K",
17060 QT_Q6_K => "qmatvec_gemm_q6_K",
17061 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17062 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17063 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17064 };
17065 let f = self.func(name);
17066 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);
17070 let k1_tile = if is_k1 {
17072 k1_launch_override().unwrap_or((128, 128, 8))
17073 } else {
17074 (128, 128, 8)
17075 };
17076 let (bm, bn): (u32, u32) = if is_k1 {
17077 (k1_tile.0, k1_tile.1)
17078 } else {
17079 (64, 256)
17080 };
17081 let warps: u32 = if is_k1 {
17082 k1_tile.2
17083 } else {
17084 match qtype {
17085 QT_NVFP4 | QT_NVFP4_RP => 8,
17086 _ => 4,
17087 }
17088 };
17089 let cfg = LaunchConfig {
17090 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17091 block_dim: (32, warps, 1),
17092 shared_mem_bytes: 0,
17093 };
17094 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17095 let __s_b = self.gpu.stream();
17096 let mut b = __s_b.launch_builder(&f);
17097 b.arg(bytes)
17098 .arg(&aq)
17099 .arg(&ad)
17100 .arg(&mut y)
17101 .arg(&inf)
17102 .arg(&outf)
17103 .arg(&mi)
17104 .arg(&rb);
17105 unsafe {
17106 b.launch(cfg)?;
17107 }
17108 Ok(y)
17109 }
17110
17111 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17118 &self,
17119 rp4: &CudaSlice<u8>,
17120 aq: &CudaSlice<i8>,
17121 ad: &CudaSlice<f32>,
17122 m: usize,
17123 in_f: usize,
17124 out_f: usize,
17125 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17126 assert!(
17127 out_f % 64 == 0 && in_f % 32 == 0,
17128 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17129 );
17130 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17131 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17133 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17134 block_dim: (128, 1, 1),
17135 shared_mem_bytes: 0,
17136 };
17137 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17138 let __s_b = self.gpu.stream();
17139 let mut b = __s_b.launch_builder(&f);
17140 b.arg(rp4)
17141 .arg(aq)
17142 .arg(ad)
17143 .arg(&mut y)
17144 .arg(&inf)
17145 .arg(&outf)
17146 .arg(&mi);
17147 unsafe {
17148 b.launch(cfg)?;
17149 }
17150 Ok(y)
17151 }
17152
17153 pub fn scale_inplace(
17155 &self,
17156 y: &mut CudaSlice<f32>,
17157 s: f32,
17158 n: usize,
17159 ) -> Result<(), Box<dyn std::error::Error>> {
17160 let f = self.func("scale_f32");
17161 let cfg = LaunchConfig::for_num_elems(n as u32);
17162 let (sf, ni) = (s, n as i32);
17163 let __s_b = self.gpu.stream();
17164 let mut b = __s_b.launch_builder(&f);
17165 b.arg(y).arg(&sf).arg(&ni);
17166 unsafe {
17167 b.launch(cfg)?;
17168 }
17169 Ok(())
17170 }
17171
17172 pub fn bf16_to_f32(
17177 &self,
17178 data: &cudarc::driver::CudaView<'_, u8>,
17179 n: usize,
17180 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17181 let mut out = self.alloc_uninit::<f32>(n)?;
17182 let f = self.func("bf16_to_f32");
17183 let cfg = LaunchConfig::for_num_elems(n as u32);
17184 let ni = n as i32;
17185 let __s_b = self.gpu.stream();
17186 let mut b = __s_b.launch_builder(&f);
17187 b.arg(data).arg(&mut out).arg(&ni);
17188 unsafe {
17189 b.launch(cfg)?;
17190 }
17191 Ok(out)
17192 }
17193
17194 fn linear_bf16_chunked(
17201 &self,
17202 x: &CudaSlice<f32>,
17203 data: &CudaSlice<u8>,
17204 m: usize,
17205 in_f: usize,
17206 out_f: usize,
17207 exact: bool,
17208 canonical_chunk_rows: Option<usize>,
17209 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17210 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17214 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17215 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17216 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17217 let started = timing.then(std::time::Instant::now);
17218 let result =
17219 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17220 if let Some(started) = started {
17221 use std::sync::atomic::Ordering;
17222 self.stream().synchronize()?;
17223 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17224 + started.elapsed().as_nanos() as u64;
17225 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17226 + (in_f * out_f * 2) as u64;
17227 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17228 if calls % 1024 == 0 {
17229 eprintln!(
17230 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17231 weight_gb={:.2}",
17232 ns as f64 / 1.0e6,
17233 ns as f64 / calls as f64 / 1.0e3,
17234 wb as f64 / 1.0e9,
17235 );
17236 }
17237 }
17238 result
17239 }
17240
17241 pub(crate) fn bf16_mmv_on() -> bool {
17246 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17247 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17248 }
17249
17250 fn matvec_bf16(
17253 &self,
17254 data: &CudaSlice<u8>,
17255 x: &CudaSlice<f32>,
17256 in_f: usize,
17257 out_f: usize,
17258 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17259 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17260 return Err(format!(
17261 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17262 data.len(),
17263 x.len()
17264 )
17265 .into());
17266 }
17267 let mut y = self.alloc_uninit::<f32>(out_f)?;
17268 let f = self.func("matvec_bf16_f32acc");
17269 let cfg = LaunchConfig {
17270 grid_dim: (out_f as u32, 1, 1),
17271 block_dim: (mmv_block(), 1, 1),
17272 shared_mem_bytes: 0,
17273 };
17274 let ini = in_f as i32;
17275 let __s_bld = self.gpu.stream();
17276 let mut bld = __s_bld.launch_builder(&f);
17277 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17278 unsafe {
17279 bld.launch(cfg)?;
17280 }
17281 Ok(y)
17282 }
17283
17284 #[allow(clippy::too_many_arguments)]
17288 #[allow(clippy::too_many_arguments)]
17293 #[allow(clippy::too_many_arguments)]
17298 pub fn qk_norm_rope_append_inc_dcw_rows(
17299 &self,
17300 q_raw_t: &CudaSlice<f32>,
17301 k_raw_t: &CudaSlice<f32>,
17302 v_raw_t: &CudaSlice<f32>,
17303 qw: &CudaSlice<f32>,
17304 kw: &CudaSlice<f32>,
17305 q_out_t: &mut CudaSlice<f32>,
17306 k_out_t: &mut CudaSlice<f32>,
17307 tab: &CudaSlice<u64>,
17308 pos_t: &CudaSlice<i32>,
17309 same_session: bool,
17310 t: usize,
17311 kv_dim_k: usize,
17312 kv_dim_v: usize,
17313 k_tok_bytes: usize,
17314 v_tok_bytes: usize,
17315 head_dim: usize,
17316 n_dims: usize,
17317 nh_q: usize,
17318 nh_k: usize,
17319 eps: f32,
17320 freq_base: f32,
17321 freq_scale: f32,
17322 ff: Option<&CudaSlice<f32>>,
17323 ) -> Result<(), Box<dyn std::error::Error>> {
17324 if head_dim != 128
17325 || kv_dim_v != kv_dim_k
17326 || kv_dim_k != nh_k * head_dim
17327 || t == 0
17328 || t > 32
17329 || tab.len() < t * 6
17330 || pos_t.len() < t
17331 || q_raw_t.len() < t * nh_q * head_dim
17332 || k_raw_t.len() < t * nh_k * head_dim
17333 || v_raw_t.len() < t * kv_dim_v
17334 || q_out_t.len() < t * nh_q * head_dim
17335 || k_out_t.len() < t * nh_k * head_dim
17336 {
17337 return Err(format!(
17338 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17339 nh_q={nh_q} nh_k={nh_k}"
17340 )
17341 .into());
17342 }
17343 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17344 let same_t: i32 = if same_session { t as i32 } else { 0 };
17345 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17346 let cfg = LaunchConfig {
17347 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17348 block_dim: (128, 1, 1),
17349 shared_mem_bytes: 0,
17350 };
17351 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17352 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17353 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17354 let null: u64 = 0;
17355 let __s_b = self.gpu.stream();
17356 let mut b = __s_b.launch_builder(&f);
17357 b.arg(q_raw_t)
17358 .arg(k_raw_t)
17359 .arg(v_raw_t)
17360 .arg(qw)
17361 .arg(kw)
17362 .arg(q_out_t)
17363 .arg(k_out_t)
17364 .arg(tab)
17365 .arg(pos_t)
17366 .arg(&same_t)
17367 .arg(&kvk)
17368 .arg(&kvv)
17369 .arg(&ktb)
17370 .arg(&vtb)
17371 .arg(&hd)
17372 .arg(&nd)
17373 .arg(&nq)
17374 .arg(&nk)
17375 .arg(&eps)
17376 .arg(&theta_scale)
17377 .arg(&freq_scale);
17378 match ff {
17379 Some(freqs) => {
17380 b.arg(freqs);
17381 }
17382 None => {
17383 b.arg(&null);
17384 }
17385 }
17386 unsafe {
17387 b.launch(cfg)?;
17388 }
17389 Ok(())
17390 }
17391
17392 pub fn qk_norm_rope_append_inc_dcw(
17393 &self,
17394 q_raw: &CudaSlice<f32>,
17395 k_raw: &CudaSlice<f32>,
17396 v_raw: &CudaSlice<f32>,
17397 qw: &CudaSlice<f32>,
17398 kw: &CudaSlice<f32>,
17399 q_out: &mut CudaSlice<f32>,
17400 k_out: &mut CudaSlice<f32>,
17401 pos: &CudaSlice<i32>,
17402 k_plane: &mut CudaSlice<u8>,
17403 v_plane: &mut CudaSlice<u8>,
17404 len_dev: &CudaSlice<i32>,
17407 base_dev: Option<&CudaSlice<i32>>,
17408 done_ctr: &mut CudaSlice<u32>,
17409 kv_dim_k: usize,
17410 kv_dim_v: usize,
17411 k_tok_bytes: usize,
17412 v_tok_bytes: usize,
17413 head_dim: usize,
17414 n_dims: usize,
17415 nh_q: usize,
17416 nh_k: usize,
17417 eps: f32,
17418 freq_base: f32,
17419 freq_scale: f32,
17420 ff: Option<&CudaSlice<f32>>,
17421 ) -> Result<(), Box<dyn std::error::Error>> {
17422 if head_dim != 128
17423 || kv_dim_v != kv_dim_k
17424 || kv_dim_k != nh_k * head_dim
17425 || q_raw.len() < nh_q * head_dim
17426 || k_raw.len() < nh_k * head_dim
17427 || v_raw.len() < kv_dim_v
17428 || q_out.len() < nh_q * head_dim
17429 || k_out.len() < nh_k * head_dim
17430 || pos.is_empty()
17431 || done_ctr.is_empty()
17432 {
17433 return Err(format!(
17434 "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}"
17435 )
17436 .into());
17437 }
17438 let f = self.func("qk_norm_rope_append_inc_dcw");
17439 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17440 let cfg = LaunchConfig {
17441 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17442 block_dim: (128, 1, 1),
17443 shared_mem_bytes: 0,
17444 };
17445 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17446 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17447 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17448 let null: u64 = 0;
17449 let __s_b = self.gpu.stream();
17450 let mut b = __s_b.launch_builder(&f);
17451 b.arg(q_raw)
17452 .arg(k_raw)
17453 .arg(v_raw)
17454 .arg(qw)
17455 .arg(kw)
17456 .arg(q_out)
17457 .arg(k_out)
17458 .arg(pos)
17459 .arg(&mut *k_plane)
17460 .arg(&mut *v_plane)
17461 .arg(len_dev);
17462 match base_dev {
17463 Some(base) => {
17464 b.arg(base);
17465 }
17466 None => {
17467 b.arg(&null);
17468 }
17469 }
17470 b.arg(&mut *done_ctr)
17471 .arg(&kvk)
17472 .arg(&kvv)
17473 .arg(&ktb)
17474 .arg(&vtb)
17475 .arg(&hd)
17476 .arg(&nd)
17477 .arg(&nq)
17478 .arg(&eps)
17479 .arg(&theta_scale)
17480 .arg(&freq_scale);
17481 match ff {
17482 Some(freqs) => {
17483 b.arg(freqs);
17484 }
17485 None => {
17486 b.arg(&null);
17487 }
17488 }
17489 unsafe {
17490 b.launch(cfg)?;
17491 }
17492 Ok(())
17493 }
17494
17495 pub fn qk_norm_rope_into(
17496 &self,
17497 q_raw: &CudaSlice<f32>,
17498 k_raw: &CudaSlice<f32>,
17499 qw: &CudaSlice<f32>,
17500 kw: &CudaSlice<f32>,
17501 q_out: &mut CudaSlice<f32>,
17502 k_out: &mut CudaSlice<f32>,
17503 pos: &CudaSlice<i32>,
17504 head_dim: usize,
17505 n_dims: usize,
17506 nh_q: usize,
17507 nh_k: usize,
17508 eps: f32,
17509 freq_base: f32,
17510 freq_scale: f32,
17511 ff: Option<&CudaSlice<f32>>,
17512 ) -> Result<(), Box<dyn std::error::Error>> {
17513 if head_dim > 512
17514 || q_raw.len() < nh_q * head_dim
17515 || k_raw.len() < nh_k * head_dim
17516 || q_out.len() < nh_q * head_dim
17517 || k_out.len() < nh_k * head_dim
17518 || qw.len() < head_dim
17519 || kw.len() < head_dim
17520 || pos.is_empty()
17521 {
17522 return Err(format!(
17523 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17524 )
17525 .into());
17526 }
17527 let f = self.func("qk_norm_rope_f32");
17528 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17529 let cfg = LaunchConfig {
17530 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17531 block_dim: (128, 1, 1),
17532 shared_mem_bytes: 0,
17533 };
17534 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17535 let __s_b = self.gpu.stream();
17536 let mut b = __s_b.launch_builder(&f);
17537 b.arg(q_raw)
17538 .arg(k_raw)
17539 .arg(qw)
17540 .arg(kw)
17541 .arg(q_out)
17542 .arg(k_out)
17543 .arg(pos)
17544 .arg(&hd)
17545 .arg(&nd)
17546 .arg(&nq)
17547 .arg(&eps)
17548 .arg(&theta_scale)
17549 .arg(&freq_scale);
17550 match ff {
17551 Some(ffv) => {
17552 b.arg(ffv);
17553 unsafe {
17554 b.launch(cfg)?;
17555 }
17556 }
17557 None => {
17558 let null: u64 = 0;
17559 b.arg(&null);
17560 unsafe {
17561 b.launch(cfg)?;
17562 }
17563 }
17564 }
17565 Ok(())
17566 }
17567
17568 #[allow(clippy::too_many_arguments)]
17571 pub fn matvec_f32_b4_into(
17572 &self,
17573 w: [&CudaSlice<f32>; 4],
17574 x: &CudaSlice<f32>,
17575 y: &mut CudaSlice<f32>,
17576 block_cols: usize,
17577 out_f: usize,
17578 ) -> Result<(), Box<dyn std::error::Error>> {
17579 if block_cols % 4 != 0
17580 || x.len() < 4 * block_cols
17581 || y.len() < out_f
17582 || w.iter().any(|w| w.len() != out_f * block_cols)
17583 {
17584 return Err(format!(
17585 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17586 x.len()
17587 )
17588 .into());
17589 }
17590 let f = self.func("matvec_f32_b4");
17591 let cfg = LaunchConfig {
17592 grid_dim: (out_f as u32, 1, 1),
17593 block_dim: (128, 1, 1),
17594 shared_mem_bytes: 0,
17595 };
17596 let (bc, of) = (block_cols as i32, out_f as i32);
17597 let __s_b = self.gpu.stream();
17598 let mut b = __s_b.launch_builder(&f);
17599 b.arg(w[0])
17600 .arg(w[1])
17601 .arg(w[2])
17602 .arg(w[3])
17603 .arg(x)
17604 .arg(y)
17605 .arg(&bc)
17606 .arg(&of);
17607 unsafe {
17608 b.launch(cfg)?;
17609 }
17610 Ok(())
17611 }
17612
17613 pub fn axpy_rows_seq_into(
17616 &self,
17617 x: &CudaSlice<f32>,
17618 w: &CudaSlice<f32>,
17619 y: &mut CudaSlice<f32>,
17620 width: usize,
17621 n_rows: usize,
17622 ) -> Result<(), Box<dyn std::error::Error>> {
17623 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17624 return Err(format!(
17625 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17626 x.len(),
17627 w.len(),
17628 y.len()
17629 )
17630 .into());
17631 }
17632 let f = self.func("axpy_rows_seq_f32");
17633 let cfg = LaunchConfig::for_num_elems(width as u32);
17634 let (wi, nr) = (width as i32, n_rows as i32);
17635 let __s_b = self.gpu.stream();
17636 let mut b = __s_b.launch_builder(&f);
17637 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17638 unsafe {
17639 b.launch(cfg)?;
17640 }
17641 Ok(())
17642 }
17643
17644 #[allow(clippy::too_many_arguments)]
17648 pub fn axpy_rows_seq_md_off_into(
17649 &self,
17650 x: &CudaSlice<f32>,
17651 w_route: &CudaSlice<f32>,
17652 md: &CudaSlice<f32>,
17653 sel: &CudaSlice<i32>,
17654 y: &mut CudaSlice<f32>,
17655 width: usize,
17656 n_rows: usize,
17657 row0: usize,
17658 ) -> Result<(), Box<dyn std::error::Error>> {
17659 if x.len() < (row0 + n_rows) * width
17660 || w_route.len() < row0 + n_rows
17661 || sel.len() < row0 + n_rows
17662 || y.len() < width
17663 {
17664 return Err(format!(
17665 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17666 rows={n_rows} row0={row0}",
17667 x.len(),
17668 w_route.len(),
17669 sel.len(),
17670 y.len()
17671 )
17672 .into());
17673 }
17674 let f = self.func("axpy_rows_seq_md_off_f32");
17675 let cfg = LaunchConfig::for_num_elems(width as u32);
17676 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17677 let __s_b = self.gpu.stream();
17678 let mut b = __s_b.launch_builder(&f);
17679 b.arg(x)
17680 .arg(w_route)
17681 .arg(md)
17682 .arg(sel)
17683 .arg(y)
17684 .arg(&wi)
17685 .arg(&nr)
17686 .arg(&r0);
17687 unsafe {
17688 b.launch(cfg)?;
17689 }
17690 Ok(())
17691 }
17692
17693 #[allow(clippy::too_many_arguments)]
17698 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17699 &self,
17700 gate_bank: &CudaSlice<u8>,
17701 up_bank: &CudaSlice<u8>,
17702 sel: &CudaSlice<i32>,
17703 aq: &CudaSlice<i8>,
17704 ad: &CudaSlice<f32>,
17705 yg: &mut CudaSlice<f32>,
17706 yu: &mut CudaSlice<f32>,
17707 n_sel: usize,
17708 n_sel_col: usize,
17709 in_f: usize,
17710 out_f: usize,
17711 row_bytes: usize,
17712 expert_stride: usize,
17713 act_row_stride: usize,
17714 ad_row_stride: usize,
17715 ) -> Result<(), Box<dyn std::error::Error>> {
17716 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17717 if yg.len() < n_sel * out_f
17718 || yu.len() < n_sel * out_f
17719 || sel.len() < n_sel
17720 || n_sel_col == 0
17721 || n_sel % n_sel_col != 0
17722 {
17723 return Err("NVFP4 gu tcol geometry".into());
17724 }
17725 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17726 let cfg = LaunchConfig {
17727 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17728 block_dim: (128, 1, 1),
17729 shared_mem_bytes: 0,
17730 };
17731 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17732 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17733 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17734 let __s_b = self.gpu.stream();
17735 let mut b = __s_b.launch_builder(&f);
17736 b.arg(gate_bank)
17737 .arg(up_bank)
17738 .arg(sel)
17739 .arg(aq)
17740 .arg(ad)
17741 .arg(yg)
17742 .arg(yu)
17743 .arg(&inf)
17744 .arg(&outf)
17745 .arg(&ns)
17746 .arg(&rb)
17747 .arg(&es)
17748 .arg(&ars)
17749 .arg(&adrs)
17750 .arg(&nsc);
17751 unsafe {
17752 b.launch(cfg)?;
17753 }
17754 Ok(())
17755 }
17756
17757 #[allow(clippy::too_many_arguments)]
17760 pub fn axpy_rows_seq_md_into(
17761 &self,
17762 x: &CudaSlice<f32>,
17763 w_route: &CudaSlice<f32>,
17764 md: &CudaSlice<f32>,
17765 sel: &CudaSlice<i32>,
17766 y: &mut CudaSlice<f32>,
17767 width: usize,
17768 n_rows: usize,
17769 ) -> Result<(), Box<dyn std::error::Error>> {
17770 if x.len() < n_rows * width
17771 || w_route.len() < n_rows
17772 || sel.len() < n_rows
17773 || y.len() < width
17774 {
17775 return Err(format!(
17776 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17777 x.len(),
17778 w_route.len(),
17779 sel.len(),
17780 y.len()
17781 )
17782 .into());
17783 }
17784 let f = self.func("axpy_rows_seq_md_f32");
17785 let cfg = LaunchConfig::for_num_elems(width as u32);
17786 let (wi, nr) = (width as i32, n_rows as i32);
17787 let __s_b = self.gpu.stream();
17788 let mut b = __s_b.launch_builder(&f);
17789 b.arg(x)
17790 .arg(w_route)
17791 .arg(md)
17792 .arg(sel)
17793 .arg(y)
17794 .arg(&wi)
17795 .arg(&nr);
17796 unsafe {
17797 b.launch(cfg)?;
17798 }
17799 Ok(())
17800 }
17801
17802 #[allow(clippy::too_many_arguments)]
17804 #[allow(clippy::too_many_arguments)]
17808 pub fn matvec_bf16_qkvg_tcol_into(
17809 &self,
17810 wq: &CudaSlice<u8>,
17811 wk: &CudaSlice<u8>,
17812 wv: &CudaSlice<u8>,
17813 wg: &CudaSlice<u8>,
17814 x_t: &CudaSlice<f32>,
17815 yq: &mut CudaSlice<f32>,
17816 yk: &mut CudaSlice<f32>,
17817 yv: &mut CudaSlice<f32>,
17818 yg: &mut CudaSlice<f32>,
17819 in_f: usize,
17820 out_q: usize,
17821 out_kv: usize,
17822 out_g: usize,
17823 t: usize,
17824 ) -> Result<(), Box<dyn std::error::Error>> {
17825 if t == 0
17826 || t > 8
17827 || in_f % 8 != 0
17828 || x_t.len() < t * in_f
17829 || yq.len() < t * out_q
17830 || yk.len() < t * out_kv
17831 || yv.len() < t * out_kv
17832 || (out_g > 0 && yg.len() < t * out_g)
17833 {
17834 return Err("matvec_bf16_qkvg_tcol geometry".into());
17835 }
17836 let grid = out_q + 2 * out_kv + out_g;
17837 let cfg = LaunchConfig {
17838 grid_dim: (grid as u32, 1, 1),
17839 block_dim: (mmv_block(), 1, 1),
17840 shared_mem_bytes: 0,
17841 };
17842 let (ini, oq, okv, og, ti) = (
17843 in_f as i32,
17844 out_q as i32,
17845 out_kv as i32,
17846 out_g as i32,
17847 t as i32,
17848 );
17849 let __s_b = self.gpu.stream();
17850 if let Some(name) = match t {
17852 2 => Some("matvec_bf16_qkvg_tcol_t2"),
17853 4 => Some("matvec_bf16_qkvg_tcol_t4"),
17854 8 => Some("matvec_bf16_qkvg_tcol_t8"),
17855 _ => None,
17856 } {
17857 let f = self.func(name);
17858 let mut b = __s_b.launch_builder(&f);
17859 b.arg(wq)
17860 .arg(wk)
17861 .arg(wv)
17862 .arg(wg)
17863 .arg(x_t)
17864 .arg(yq)
17865 .arg(yk)
17866 .arg(yv)
17867 .arg(yg)
17868 .arg(&ini)
17869 .arg(&oq)
17870 .arg(&okv)
17871 .arg(&og);
17872 unsafe {
17873 b.launch(cfg)?;
17874 }
17875 return Ok(());
17876 }
17877 let f = self.func("matvec_bf16_qkvg_tcol");
17878 let mut b = __s_b.launch_builder(&f);
17879 b.arg(wq)
17880 .arg(wk)
17881 .arg(wv)
17882 .arg(wg)
17883 .arg(x_t)
17884 .arg(yq)
17885 .arg(yk)
17886 .arg(yv)
17887 .arg(yg)
17888 .arg(&ini)
17889 .arg(&oq)
17890 .arg(&okv)
17891 .arg(&og)
17892 .arg(&ti);
17893 unsafe {
17894 b.launch(cfg)?;
17895 }
17896 Ok(())
17897 }
17898
17899 pub fn matvec_bf16_qkvg_into(
17900 &self,
17901 wq: &CudaSlice<u8>,
17902 wk: &CudaSlice<u8>,
17903 wv: &CudaSlice<u8>,
17904 wg: &CudaSlice<u8>,
17905 x: &CudaSlice<f32>,
17906 yq: &mut CudaSlice<f32>,
17907 yk: &mut CudaSlice<f32>,
17908 yv: &mut CudaSlice<f32>,
17909 yg: &mut CudaSlice<f32>,
17910 in_f: usize,
17911 out_q: usize,
17912 out_kv: usize,
17913 out_g: usize,
17914 ) -> Result<(), Box<dyn std::error::Error>> {
17915 if in_f % 8 != 0
17916 || wq.len() != out_q * in_f * 2
17917 || wk.len() != out_kv * in_f * 2
17918 || wv.len() != out_kv * in_f * 2
17919 || wg.len() < out_g * in_f * 2
17920 || x.len() < in_f
17921 || yq.len() < out_q
17922 || yk.len() < out_kv
17923 || yv.len() < out_kv
17924 || (out_g > 0 && yg.len() < out_g)
17925 {
17926 return Err(format!(
17927 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17928 )
17929 .into());
17930 }
17931 let f = self.func("matvec_bf16_qkvg");
17932 let cfg = LaunchConfig {
17933 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17934 block_dim: (mmv_block(), 1, 1),
17935 shared_mem_bytes: 0,
17936 };
17937 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17938 let __s_b = self.gpu.stream();
17939 let mut b = __s_b.launch_builder(&f);
17940 b.arg(wq)
17941 .arg(wk)
17942 .arg(wv)
17943 .arg(wg)
17944 .arg(x)
17945 .arg(yq)
17946 .arg(yk)
17947 .arg(yv)
17948 .arg(yg)
17949 .arg(&inf)
17950 .arg(&oq)
17951 .arg(&okv)
17952 .arg(&og);
17953 unsafe {
17954 b.launch(cfg)?;
17955 }
17956 Ok(())
17957 }
17958
17959 pub fn matvec_bf16_b4_into(
17961 &self,
17962 w: [&CudaSlice<u8>; 4],
17963 x: &CudaSlice<f32>,
17964 y: &mut CudaSlice<f32>,
17965 block_cols: usize,
17966 out_f: usize,
17967 ) -> Result<(), Box<dyn std::error::Error>> {
17968 if block_cols % 8 != 0
17969 || x.len() < 4 * block_cols
17970 || y.len() < out_f
17971 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17972 {
17973 return Err(format!(
17974 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17975 x.len()
17976 )
17977 .into());
17978 }
17979 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17982 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17983 let f = self.func(if x2 {
17984 "matvec_bf16_b4_x2"
17985 } else {
17986 "matvec_bf16_b4"
17987 });
17988 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17989 let cfg = LaunchConfig {
17990 grid_dim: (grid as u32, 1, 1),
17991 block_dim: (mmv_block(), 1, 1),
17992 shared_mem_bytes: 0,
17993 };
17994 let (bc, of) = (block_cols as i32, out_f as i32);
17995 let __s_b = self.gpu.stream();
17996 let mut b = __s_b.launch_builder(&f);
17997 b.arg(w[0])
17998 .arg(w[1])
17999 .arg(w[2])
18000 .arg(w[3])
18001 .arg(x)
18002 .arg(y)
18003 .arg(&bc)
18004 .arg(&of);
18005 unsafe {
18006 b.launch(cfg)?;
18007 }
18008 Ok(())
18009 }
18010
18011 pub fn matvec_bf16_b4_tcol_into(
18017 &self,
18018 w: [&CudaSlice<u8>; 4],
18019 x_t: &CudaSlice<f32>,
18020 y_t: &mut CudaSlice<f32>,
18021 block_cols: usize,
18022 out_f: usize,
18023 t: usize,
18024 ) -> Result<(), Box<dyn std::error::Error>> {
18025 if block_cols % 8 != 0
18026 || t == 0
18027 || t > 8
18028 || x_t.len() < t * 4 * block_cols
18029 || y_t.len() < t * out_f
18030 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18031 {
18032 return Err(format!(
18033 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18034 x_t.len()
18035 )
18036 .into());
18037 }
18038 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18039 return Err(
18040 "b4 tcol verify is qualified against the plain b4 kernel only \
18041 (MEMRA_B4_X2=1 is a different t=1 program)"
18042 .into(),
18043 );
18044 }
18045 let cfg = LaunchConfig {
18049 grid_dim: (out_f as u32, 1, 1),
18050 block_dim: (mmv_block(), 1, 1),
18051 shared_mem_bytes: 0,
18052 };
18053 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18054 let __s_b = self.gpu.stream();
18055 if let Some(name) = match t {
18056 2 => Some("matvec_bf16_b4_tcol_t2"),
18057 4 => Some("matvec_bf16_b4_tcol_t4"),
18058 8 => Some("matvec_bf16_b4_tcol_t8"),
18059 _ => None,
18060 } {
18061 let f = self.func(name);
18062 let mut b = __s_b.launch_builder(&f);
18063 b.arg(w[0])
18064 .arg(w[1])
18065 .arg(w[2])
18066 .arg(w[3])
18067 .arg(x_t)
18068 .arg(y_t)
18069 .arg(&bc)
18070 .arg(&of);
18071 unsafe {
18072 b.launch(cfg)?;
18073 }
18074 return Ok(());
18075 }
18076 let f = self.func("matvec_bf16_b4_tcol");
18077 let mut b = __s_b.launch_builder(&f);
18078 b.arg(w[0])
18079 .arg(w[1])
18080 .arg(w[2])
18081 .arg(w[3])
18082 .arg(x_t)
18083 .arg(y_t)
18084 .arg(&bc)
18085 .arg(&of)
18086 .arg(&ti);
18087 unsafe {
18088 b.launch(cfg)?;
18089 }
18090 Ok(())
18091 }
18092
18093 pub fn matvec_bf16_into(
18095 &self,
18096 data: &CudaSlice<u8>,
18097 x: &CudaSlice<f32>,
18098 y: &mut CudaSlice<f32>,
18099 in_f: usize,
18100 out_f: usize,
18101 ) -> Result<(), Box<dyn std::error::Error>> {
18102 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18103 return Err(format!(
18104 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18105 data.len(),
18106 x.len(),
18107 y.len()
18108 )
18109 .into());
18110 }
18111 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18115 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18116 && in_f <= 2048;
18117 if x4 {
18118 let f = self.func("matvec_bf16_f32acc_x4");
18119 let cfg = LaunchConfig {
18120 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18121 block_dim: (mmv_block(), 1, 1),
18122 shared_mem_bytes: 0,
18123 };
18124 let (ini, outi) = (in_f as i32, out_f as i32);
18125 let __s_b = self.gpu.stream();
18126 let mut b = __s_b.launch_builder(&f);
18127 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18128 unsafe {
18129 b.launch(cfg)?;
18130 }
18131 return Ok(());
18132 }
18133 let f = self.func("matvec_bf16_f32acc");
18134 let cfg = LaunchConfig {
18135 grid_dim: (out_f as u32, 1, 1),
18136 block_dim: (mmv_block(), 1, 1),
18137 shared_mem_bytes: 0,
18138 };
18139 let ini = in_f as i32;
18140 let __s_b = self.gpu.stream();
18141 let mut b = __s_b.launch_builder(&f);
18142 b.arg(data).arg(x).arg(y).arg(&ini);
18143 unsafe {
18144 b.launch(cfg)?;
18145 }
18146 Ok(())
18147 }
18148
18149 pub fn matvec_bf16_view_into(
18152 &self,
18153 data: &cudarc::driver::CudaView<'_, u8>,
18154 x: &CudaSlice<f32>,
18155 y: &mut CudaSlice<f32>,
18156 in_f: usize,
18157 out_f: usize,
18158 ) -> Result<(), Box<dyn std::error::Error>> {
18159 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18160 return Err(format!(
18161 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18162 data.len(),
18163 x.len(),
18164 y.len()
18165 )
18166 .into());
18167 }
18168 let f = self.func("matvec_bf16_f32acc");
18169 let cfg = LaunchConfig {
18170 grid_dim: (out_f as u32, 1, 1),
18171 block_dim: (mmv_block(), 1, 1),
18172 shared_mem_bytes: 0,
18173 };
18174 let ini = in_f as i32;
18175 let __s_b = self.gpu.stream();
18176 let mut b = __s_b.launch_builder(&f);
18177 b.arg(data).arg(x).arg(y).arg(&ini);
18178 unsafe {
18179 b.launch(cfg)?;
18180 }
18181 Ok(())
18182 }
18183
18184 pub fn matvec_bf16_raw_out(
18187 &self,
18188 w: &CudaSlice<u8>,
18189 x: &CudaSlice<f32>,
18190 y_raw: u64,
18191 in_f: usize,
18192 out_f: usize,
18193 ) -> Result<(), Box<dyn std::error::Error>> {
18194 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18195 return Err("matvec_bf16_raw_out geometry".into());
18196 }
18197 let f = self.func("matvec_bf16_f32acc");
18198 let cfg = LaunchConfig {
18199 grid_dim: (out_f as u32, 1, 1),
18200 block_dim: (mmv_block(), 1, 1),
18201 shared_mem_bytes: 0,
18202 };
18203 let ini = in_f as i32;
18204 let __s_b = self.gpu.stream();
18205 let mut b = __s_b.launch_builder(&f);
18206 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18207 unsafe {
18208 b.launch(cfg)?;
18209 }
18210 Ok(())
18211 }
18212
18213 pub fn add3_raw(
18217 &self,
18218 a: &CudaSlice<f32>,
18219 b: &CudaSlice<f32>,
18220 sh_raw: u64,
18221 scale_raw: u64,
18222 dst: &mut CudaSlice<f32>,
18223 n: usize,
18224 ) -> Result<(), Box<dyn std::error::Error>> {
18225 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18226 return Err("add3_raw geometry".into());
18227 }
18228 let f = self.func("add3_f32");
18229 let cfg = LaunchConfig {
18230 grid_dim: ((n as u32).div_ceil(256), 1, 1),
18231 block_dim: (256, 1, 1),
18232 shared_mem_bytes: 0,
18233 };
18234 let ni = n as i32;
18235 let __s_b = self.gpu.stream();
18236 let mut bld = __s_b.launch_builder(&f);
18237 bld.arg(a)
18238 .arg(b)
18239 .arg(&sh_raw)
18240 .arg(&scale_raw)
18241 .arg(dst)
18242 .arg(&ni);
18243 unsafe {
18244 bld.launch(cfg)?;
18245 }
18246 Ok(())
18247 }
18248
18249 pub fn matvec_bf16_down_addscale_into(
18252 &self,
18253 w: &CudaSlice<u8>,
18254 x: &CudaSlice<f32>,
18255 scale: &CudaSlice<f32>,
18256 dst: &mut CudaSlice<f32>,
18257 in_f: usize,
18258 out_f: usize,
18259 ) -> Result<(), Box<dyn std::error::Error>> {
18260 if w.len() != in_f * out_f * 2
18261 || x.len() < in_f
18262 || in_f % 8 != 0
18263 || dst.len() < out_f
18264 || scale.is_empty()
18265 {
18266 return Err("matvec_bf16_down_addscale geometry".into());
18267 }
18268 let f = self.func("matvec_bf16_down_addscale");
18269 let cfg = LaunchConfig {
18270 grid_dim: (out_f as u32, 1, 1),
18271 block_dim: (mmv_block(), 1, 1),
18272 shared_mem_bytes: 0,
18273 };
18274 let ini = in_f as i32;
18275 let __s_b = self.gpu.stream();
18276 let mut b = __s_b.launch_builder(&f);
18277 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18278 unsafe {
18279 b.launch(cfg)?;
18280 }
18281 Ok(())
18282 }
18283
18284 #[allow(clippy::too_many_arguments)]
18288 pub fn matvec_bf16_dual_silu_rows_into(
18289 &self,
18290 wg: &CudaSlice<u8>,
18291 wu: &CudaSlice<u8>,
18292 x: &CudaSlice<f32>,
18293 act: &mut CudaSlice<f32>,
18294 in_f: usize,
18295 out_f: usize,
18296 limit: Option<f32>,
18297 t: usize,
18298 ) -> Result<(), Box<dyn std::error::Error>> {
18299 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18300 return Err("matvec_bf16_dual_silu_rows geometry".into());
18301 }
18302 let f = self.func("matvec_bf16_dual_silu_rows");
18303 let cfg = LaunchConfig {
18304 grid_dim: (out_f as u32, t as u32, 1),
18305 block_dim: (mmv_block(), 1, 1),
18306 shared_mem_bytes: 0,
18307 };
18308 let (ini, outi) = (in_f as i32, out_f as i32);
18309 let lim = limit.unwrap_or(0.0);
18310 let __s_b = self.gpu.stream();
18311 let mut b = __s_b.launch_builder(&f);
18312 b.arg(wg)
18313 .arg(wu)
18314 .arg(x)
18315 .arg(&mut *act)
18316 .arg(&ini)
18317 .arg(&outi)
18318 .arg(&lim);
18319 unsafe {
18320 b.launch(cfg)?;
18321 }
18322 Ok(())
18323 }
18324
18325 pub fn matvec_bf16_rows_into(
18327 &self,
18328 w: &CudaSlice<u8>,
18329 x: &CudaSlice<f32>,
18330 y: &mut CudaSlice<f32>,
18331 in_f: usize,
18332 out_f: usize,
18333 t: usize,
18334 ) -> Result<(), Box<dyn std::error::Error>> {
18335 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18336 return Err("matvec_bf16_rows geometry".into());
18337 }
18338 let f = self.func("matvec_bf16_f32acc_x4_rows");
18339 let cfg = LaunchConfig {
18340 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18341 block_dim: (mmv_block(), 1, 1),
18342 shared_mem_bytes: 0,
18343 };
18344 let (ini, outi) = (in_f as i32, out_f as i32);
18345 let __s_b = self.gpu.stream();
18346 let mut b = __s_b.launch_builder(&f);
18347 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18348 unsafe {
18349 b.launch(cfg)?;
18350 }
18351 Ok(())
18352 }
18353
18354 pub fn matvec_bf16_dual_silu_into(
18355 &self,
18356 wg: &CudaSlice<u8>,
18357 wu: &CudaSlice<u8>,
18358 x: &CudaSlice<f32>,
18359 act: &mut CudaSlice<f32>,
18360 in_f: usize,
18361 out_f: usize,
18362 limit: Option<f32>,
18363 ) -> Result<(), Box<dyn std::error::Error>> {
18364 if wg.len() != in_f * out_f * 2
18365 || wu.len() != in_f * out_f * 2
18366 || x.len() < in_f
18367 || in_f % 8 != 0
18368 || act.len() < out_f
18369 {
18370 return Err("matvec_bf16_dual_silu geometry".into());
18371 }
18372 let f = self.func("matvec_bf16_dual_silu");
18373 let cfg = LaunchConfig {
18374 grid_dim: (out_f as u32, 1, 1),
18375 block_dim: (mmv_block(), 1, 1),
18376 shared_mem_bytes: 0,
18377 };
18378 let (ini, outi) = (in_f as i32, out_f as i32);
18379 let lim = limit.unwrap_or(0.0);
18380 let __s_b = self.gpu.stream();
18381 let mut b = __s_b.launch_builder(&f);
18382 b.arg(wg)
18383 .arg(wu)
18384 .arg(x)
18385 .arg(act)
18386 .arg(&ini)
18387 .arg(&outi)
18388 .arg(&lim);
18389 unsafe {
18390 b.launch(cfg)?;
18391 }
18392 Ok(())
18393 }
18394
18395 #[allow(clippy::too_many_arguments)]
18398 pub fn matvec_bf16_dual_view_into(
18399 &self,
18400 wg: &cudarc::driver::CudaView<'_, u8>,
18401 wu: &cudarc::driver::CudaView<'_, u8>,
18402 x: &CudaSlice<f32>,
18403 yg: &mut CudaSlice<f32>,
18404 yu: &mut CudaSlice<f32>,
18405 in_f: usize,
18406 out_f: usize,
18407 ) -> Result<(), Box<dyn std::error::Error>> {
18408 if wg.len() != in_f * out_f * 2
18409 || wu.len() != in_f * out_f * 2
18410 || x.len() < in_f
18411 || in_f % 8 != 0
18412 || yg.len() < out_f
18413 || yu.len() < out_f
18414 {
18415 return Err(format!(
18416 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18417 wg.len(),
18418 wu.len(),
18419 x.len()
18420 )
18421 .into());
18422 }
18423 let f = self.func("matvec_bf16_dual");
18424 let cfg = LaunchConfig {
18425 grid_dim: ((2 * out_f) as u32, 1, 1),
18426 block_dim: (mmv_block(), 1, 1),
18427 shared_mem_bytes: 0,
18428 };
18429 let (ini, outi) = (in_f as i32, out_f as i32);
18430 let __s_b = self.gpu.stream();
18431 let mut b = __s_b.launch_builder(&f);
18432 b.arg(wg)
18433 .arg(wu)
18434 .arg(x)
18435 .arg(yg)
18436 .arg(yu)
18437 .arg(&ini)
18438 .arg(&outi);
18439 unsafe {
18440 b.launch(cfg)?;
18441 }
18442 Ok(())
18443 }
18444
18445 #[allow(clippy::too_many_arguments)]
18447 pub fn matvec_bf16_dual_into(
18448 &self,
18449 wg: &CudaSlice<u8>,
18450 wu: &CudaSlice<u8>,
18451 x: &CudaSlice<f32>,
18452 yg: &mut CudaSlice<f32>,
18453 yu: &mut CudaSlice<f32>,
18454 in_f: usize,
18455 out_f: usize,
18456 ) -> Result<(), Box<dyn std::error::Error>> {
18457 if wg.len() != in_f * out_f * 2
18458 || wu.len() != in_f * out_f * 2
18459 || x.len() < in_f
18460 || in_f % 8 != 0
18461 || yg.len() < out_f
18462 || yu.len() < out_f
18463 {
18464 return Err(format!(
18465 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18466 wg.len(),
18467 wu.len(),
18468 x.len()
18469 )
18470 .into());
18471 }
18472 let f = self.func("matvec_bf16_dual");
18473 let cfg = LaunchConfig {
18474 grid_dim: ((2 * out_f) as u32, 1, 1),
18475 block_dim: (mmv_block(), 1, 1),
18476 shared_mem_bytes: 0,
18477 };
18478 let (ini, outi) = (in_f as i32, out_f as i32);
18479 let __s_b = self.gpu.stream();
18480 let mut b = __s_b.launch_builder(&f);
18481 b.arg(wg)
18482 .arg(wu)
18483 .arg(x)
18484 .arg(yg)
18485 .arg(yu)
18486 .arg(&ini)
18487 .arg(&outi);
18488 unsafe {
18489 b.launch(cfg)?;
18490 }
18491 Ok(())
18492 }
18493
18494 pub(crate) fn matvec_bf16_dual(
18497 &self,
18498 wg: &CudaSlice<u8>,
18499 wu: &CudaSlice<u8>,
18500 x: &CudaSlice<f32>,
18501 in_f: usize,
18502 out_f: usize,
18503 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18504 if wg.len() != in_f * out_f * 2
18505 || wu.len() != in_f * out_f * 2
18506 || x.len() < in_f
18507 || in_f % 8 != 0
18508 {
18509 return Err(format!(
18510 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18511 wg.len(),
18512 wu.len(),
18513 x.len()
18514 )
18515 .into());
18516 }
18517 let mut yg = self.alloc_uninit::<f32>(out_f)?;
18518 let mut yu = self.alloc_uninit::<f32>(out_f)?;
18519 let f = self.func("matvec_bf16_dual");
18520 let cfg = LaunchConfig {
18521 grid_dim: ((2 * out_f) as u32, 1, 1),
18522 block_dim: (mmv_block(), 1, 1),
18523 shared_mem_bytes: 0,
18524 };
18525 let (ini, outi) = (in_f as i32, out_f as i32);
18526 let __s_b = self.gpu.stream();
18527 let mut b = __s_b.launch_builder(&f);
18528 b.arg(wg)
18529 .arg(wu)
18530 .arg(x)
18531 .arg(&mut yg)
18532 .arg(&mut yu)
18533 .arg(&ini)
18534 .arg(&outi);
18535 unsafe {
18536 b.launch(cfg)?;
18537 }
18538 Ok((yg, yu))
18539 }
18540
18541 #[allow(clippy::too_many_arguments)]
18542 fn linear_bf16_chunked_inner(
18543 &self,
18544 x: &CudaSlice<f32>,
18545 data: &CudaSlice<u8>,
18546 m: usize,
18547 in_f: usize,
18548 out_f: usize,
18549 exact: bool,
18550 canonical_chunk_rows: Option<usize>,
18551 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18552 const CHUNK_BYTES: usize = 256 << 20;
18553 if m == 1
18556 && !exact
18557 && canonical_chunk_rows.is_none()
18558 && in_f % 8 == 0
18559 && Self::bf16_mmv_on()
18560 {
18561 return self.matvec_bf16(data, x, in_f, out_f);
18562 }
18563 let row_bytes = in_f
18564 .checked_mul(std::mem::size_of::<f32>())
18565 .ok_or("BF16 chunk row byte count overflow")?;
18566 if row_bytes == 0 || out_f == 0 {
18567 return Err("BF16 chunk dimensions must be nonzero".into());
18568 }
18569 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18570 let chunk_rows = match canonical_chunk_rows {
18571 Some(rows) if rows == 0 => {
18572 return Err("canonical BF16 chunk rows must be nonzero".into());
18573 }
18574 Some(rows) if rows > max_chunk_rows => {
18575 return Err(format!(
18576 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18577 )
18578 .into());
18579 }
18580 Some(rows) if out_f % rows != 0 => {
18581 return Err(format!(
18582 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18583 )
18584 .into());
18585 }
18586 Some(rows) => rows,
18587 None => max_chunk_rows,
18588 };
18589 if chunk_rows >= out_f {
18590 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18591 return if exact {
18592 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
18593 } else {
18594 self.linear(x, &wf32, m, in_f, out_f)
18595 };
18596 }
18597 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18598 let mut r0 = 0usize;
18599 while r0 < out_f {
18600 let rows = chunk_rows.min(out_f - r0);
18601 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
18602 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
18603 let yc = if exact {
18604 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
18605 } else {
18606 self.linear(x, &wf32, m, in_f, rows)?
18607 };
18608 for mi in 0..m {
18610 let src = yc.slice(mi * rows..(mi + 1) * rows);
18611 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
18612 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
18613 }
18614 r0 += rows;
18615 }
18616 Ok(y)
18617 }
18618
18619 pub fn linear_bf16_resident(
18623 &self,
18624 x: &CudaSlice<f32>,
18625 data: &CudaSlice<u8>,
18626 m: usize,
18627 in_f: usize,
18628 out_f: usize,
18629 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18630 if data.len() != in_f * out_f * 2 {
18631 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18632 }
18633 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
18634 }
18635
18636 pub fn linear_bf16_resident_canonical_rows(
18642 &self,
18643 x: &CudaSlice<f32>,
18644 data: &CudaSlice<u8>,
18645 m: usize,
18646 in_f: usize,
18647 out_f: usize,
18648 canonical_chunk_rows: usize,
18649 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18650 if data.len() != in_f * out_f * 2 {
18651 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18652 }
18653 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18654 }
18655
18656 pub fn linear_f32_resident_canonical_rows(
18661 &self,
18662 x: &CudaSlice<f32>,
18663 data: &CudaSlice<f32>,
18664 m: usize,
18665 in_f: usize,
18666 out_f: usize,
18667 canonical_chunk_rows: usize,
18668 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18669 self.linear_f32_resident_canonical_rows_inner(
18670 x,
18671 data,
18672 m,
18673 in_f,
18674 out_f,
18675 canonical_chunk_rows,
18676 false,
18677 )
18678 }
18679
18680 pub fn linear_f32_resident_canonical_rows_strided(
18686 &self,
18687 x: &CudaSlice<f32>,
18688 data: &CudaSlice<f32>,
18689 m: usize,
18690 in_f: usize,
18691 out_f: usize,
18692 canonical_chunk_rows: usize,
18693 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18694 self.linear_f32_resident_canonical_rows_inner(
18695 x,
18696 data,
18697 m,
18698 in_f,
18699 out_f,
18700 canonical_chunk_rows,
18701 true,
18702 )
18703 }
18704
18705 fn linear_f32_resident_canonical_rows_inner(
18706 &self,
18707 x: &CudaSlice<f32>,
18708 data: &CudaSlice<f32>,
18709 m: usize,
18710 in_f: usize,
18711 out_f: usize,
18712 canonical_chunk_rows: usize,
18713 strided_output: bool,
18714 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18715 if data.len() != in_f * out_f {
18716 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18717 }
18718 if canonical_chunk_rows == 0
18719 || canonical_chunk_rows > out_f
18720 || out_f % canonical_chunk_rows != 0
18721 {
18722 return Err(format!(
18723 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18724 )
18725 .into());
18726 }
18727 if canonical_chunk_rows == out_f {
18728 return self.linear(x, data, m, in_f, out_f);
18729 }
18730
18731 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18732 let input = x.slice(0..x.len());
18733 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18734 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18735 if m == 1 {
18736 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18737 self.linear_device_into(
18738 &input,
18739 &weights,
18740 &mut destination,
18741 1,
18742 in_f,
18743 canonical_chunk_rows,
18744 )?;
18745 continue;
18746 }
18747 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18748 if strided_output {
18749 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18750 } else {
18751 for token in 0..m {
18752 let source = chunk
18753 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18754 let mut destination =
18755 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18756 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18757 }
18758 }
18759 }
18760 Ok(y)
18761 }
18762
18763 pub fn linear_f32_resident_canonical_rows_t1_into(
18768 &self,
18769 x: &CudaSlice<f32>,
18770 data: &CudaSlice<f32>,
18771 y: &mut CudaSlice<f32>,
18772 in_f: usize,
18773 out_f: usize,
18774 canonical_chunk_rows: usize,
18775 ) -> Result<(), Box<dyn std::error::Error>> {
18776 if data.len() != in_f * out_f {
18777 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18778 }
18779 if y.len() != out_f || x.len() != in_f {
18780 return Err(format!(
18781 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18782 x.len(),
18783 y.len()
18784 )
18785 .into());
18786 }
18787 if canonical_chunk_rows == 0
18788 || canonical_chunk_rows > out_f
18789 || out_f % canonical_chunk_rows != 0
18790 {
18791 return Err(format!(
18792 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18793 )
18794 .into());
18795 }
18796 let input = x.slice(0..x.len());
18797 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18798 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18799 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18800 self.linear_device_into(
18801 &input,
18802 &weights,
18803 &mut destination,
18804 1,
18805 in_f,
18806 canonical_chunk_rows,
18807 )?;
18808 }
18809 Ok(())
18810 }
18811
18812 pub fn linear_t1_into(
18815 &self,
18816 x: &cudarc::driver::CudaView<'_, f32>,
18817 w: &cudarc::driver::CudaView<'_, f32>,
18818 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18819 in_f: usize,
18820 out_f: usize,
18821 ) -> Result<(), Box<dyn std::error::Error>> {
18822 self.linear_device_into(x, w, y, 1, in_f, out_f)
18823 }
18824
18825 pub fn linear_decode_exact(
18832 &self,
18833 x: &CudaSlice<f32>,
18834 w: &CudaSlice<f32>,
18835 m_tokens: usize,
18836 in_f: usize,
18837 out_f: usize,
18838 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18839 if m_tokens == 1 {
18840 return self.linear(x, w, 1, in_f, out_f);
18841 }
18842 let xv = self.view(x, m_tokens * in_f);
18843 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18844 for t in 0..m_tokens {
18845 let row = xv.slice(t * in_f..(t + 1) * in_f);
18846 let mut xr = self.alloc_uninit::<f32>(in_f)?;
18847 self.copy_view_into(&mut xr, 0, &row, in_f)?;
18848 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18849 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18850 }
18851 Ok(y)
18852 }
18853
18854 pub fn linear(
18855 &self,
18856 x: &CudaSlice<f32>,
18857 w: &CudaSlice<f32>,
18858 m_tokens: usize,
18859 in_f: usize,
18860 out_f: usize,
18861 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18862 self.linear_device(x, w, m_tokens, in_f, out_f)
18863 }
18864
18865 fn linear_device<I>(
18866 &self,
18867 x: &I,
18868 w: &I,
18869 m_tokens: usize,
18870 in_f: usize,
18871 out_f: usize,
18872 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18873 where
18874 I: cudarc::driver::DevicePtr<f32>,
18875 {
18876 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)?;
18878 Ok(c)
18879 }
18880
18881 fn linear_device_into<I, O>(
18882 &self,
18883 x: &I,
18884 w: &I,
18885 c: &mut O,
18886 m_tokens: usize,
18887 in_f: usize,
18888 out_f: usize,
18889 ) -> Result<(), Box<dyn std::error::Error>>
18890 where
18891 I: cudarc::driver::DevicePtr<f32>,
18892 O: cudarc::driver::DevicePtrMut<f32>,
18893 {
18894 use cudarc::cublaslt::{Matmul, MatmulConfig};
18895 let cfg = MatmulConfig {
18896 transa: true,
18897 transb: false,
18898 transc: false,
18899 m: out_f as u64,
18900 n: m_tokens as u64,
18901 k: in_f as u64,
18902 alpha: 1.0,
18903 lda: in_f as i64,
18904 ldb: in_f as i64,
18905 beta: 0.0,
18906 ldc: out_f as i64,
18907 stride_a: None,
18908 stride_b: None,
18909 stride_c: None,
18910 stride_bias: None,
18911 batch_size: None,
18912 };
18913 let blas = self.gpu.blas();
18914 unsafe {
18915 blas.matmul(cfg, w, x, c, None, None)?;
18916 }
18917 Ok(())
18918 }
18919
18920 pub fn sdpa_naive(
18929 &self,
18930 q: &CudaSlice<f32>,
18931 k: &CudaSlice<f32>,
18932 v: &CudaSlice<f32>,
18933 o: &mut CudaSlice<f32>,
18934 head_dim: usize,
18935 n_head: usize,
18936 n_head_kv: usize,
18937 t: usize,
18938 t_kv: usize,
18939 scale: f32,
18940 causal: bool,
18941 ) -> Result<(), Box<dyn std::error::Error>> {
18942 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
18943 return self.sdpa_naive_gmem(
18944 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18945 );
18946 }
18947 let f = self.func("sdpa_naive_f32");
18948 let cfg = LaunchConfig {
18949 grid_dim: (n_head as u32, t as u32, 1),
18950 block_dim: (128, 1, 1),
18951 shared_mem_bytes: (t_kv * 4) as u32,
18952 };
18953 let (hd, nh, nhkv, ti, tkvi, cz) = (
18954 head_dim as i32,
18955 n_head as i32,
18956 n_head_kv as i32,
18957 t as i32,
18958 t_kv as i32,
18959 causal as i32,
18960 );
18961 let __s_b = self.gpu.stream();
18962 let mut b = __s_b.launch_builder(&f);
18963 b.arg(q)
18964 .arg(k)
18965 .arg(v)
18966 .arg(o)
18967 .arg(&hd)
18968 .arg(&nh)
18969 .arg(&nhkv)
18970 .arg(&ti)
18971 .arg(&tkvi)
18972 .arg(&scale)
18973 .arg(&cz);
18974 unsafe {
18975 b.launch(cfg)?;
18976 }
18977 Ok(())
18978 }
18979
18980 #[allow(clippy::too_many_arguments)]
18989 pub fn sdpa_naive_gmem(
18990 &self,
18991 q: &CudaSlice<f32>,
18992 k: &CudaSlice<f32>,
18993 v: &CudaSlice<f32>,
18994 o: &mut CudaSlice<f32>,
18995 head_dim: usize,
18996 n_head: usize,
18997 n_head_kv: usize,
18998 t: usize,
18999 t_kv: usize,
19000 scale: f32,
19001 causal: bool,
19002 ) -> Result<(), Box<dyn std::error::Error>> {
19003 let ws_len = n_head
19004 .checked_mul(t)
19005 .and_then(|x| x.checked_mul(t_kv))
19006 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19007 let ws_bytes = ws_len
19008 .checked_mul(std::mem::size_of::<f32>())
19009 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19010 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19011 return Err(format!(
19012 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19013 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19014 needs a tiled/flash kernel, not the naive oracle"
19015 )
19016 .into());
19017 }
19018 let mut scores = self.uninit(ws_len)?;
19019 let f = self.func("sdpa_naive_gmem_f32");
19020 let cfg = LaunchConfig {
19021 grid_dim: (n_head as u32, t as u32, 1),
19022 block_dim: (128, 1, 1),
19023 shared_mem_bytes: 0,
19024 };
19025 let (hd, nh, nhkv, ti, tkvi, cz) = (
19026 head_dim as i32,
19027 n_head as i32,
19028 n_head_kv as i32,
19029 t as i32,
19030 t_kv as i32,
19031 causal as i32,
19032 );
19033 let __s_b = self.gpu.stream();
19034 let mut b = __s_b.launch_builder(&f);
19035 b.arg(q)
19036 .arg(k)
19037 .arg(v)
19038 .arg(o)
19039 .arg(&mut scores)
19040 .arg(&hd)
19041 .arg(&nh)
19042 .arg(&nhkv)
19043 .arg(&ti)
19044 .arg(&tkvi)
19045 .arg(&scale)
19046 .arg(&cz);
19047 unsafe {
19048 b.launch(cfg)?;
19049 }
19050 Ok(())
19051 }
19052
19053 #[allow(clippy::too_many_arguments)]
19058 pub fn sdpa_naive_island(
19059 &self,
19060 q: &CudaSlice<f32>,
19061 k: &CudaSlice<f32>,
19062 v: &CudaSlice<f32>,
19063 o: &mut CudaSlice<f32>,
19064 span_id: &CudaSlice<i32>,
19065 head_dim: usize,
19066 n_head: usize,
19067 n_head_kv: usize,
19068 t: usize,
19069 t_kv: usize,
19070 scale: f32,
19071 window: usize,
19072 ) -> Result<(), Box<dyn std::error::Error>> {
19073 let f = self.func("sdpa_naive_island_f32");
19074 let cfg = LaunchConfig {
19075 grid_dim: (n_head as u32, t as u32, 1),
19076 block_dim: (128, 1, 1),
19077 shared_mem_bytes: (t_kv * 4) as u32,
19078 };
19079 let (hd, nh, nhkv, ti, tkvi, wi) = (
19080 head_dim as i32,
19081 n_head as i32,
19082 n_head_kv as i32,
19083 t as i32,
19084 t_kv as i32,
19085 window as i32,
19086 );
19087 let __s_b = self.gpu.stream();
19088 let mut b = __s_b.launch_builder(&f);
19089 b.arg(q)
19090 .arg(k)
19091 .arg(v)
19092 .arg(o)
19093 .arg(span_id)
19094 .arg(&hd)
19095 .arg(&nh)
19096 .arg(&nhkv)
19097 .arg(&ti)
19098 .arg(&tkvi)
19099 .arg(&scale)
19100 .arg(&wi);
19101 unsafe {
19102 b.launch(cfg)?;
19103 }
19104 Ok(())
19105 }
19106
19107 #[allow(clippy::too_many_arguments)]
19109 pub fn sdpa_naive_w(
19110 &self,
19111 q: &CudaSlice<f32>,
19112 k: &CudaSlice<f32>,
19113 v: &CudaSlice<f32>,
19114 o: &mut CudaSlice<f32>,
19115 head_dim: usize,
19116 n_head: usize,
19117 n_head_kv: usize,
19118 t: usize,
19119 t_kv: usize,
19120 scale: f32,
19121 causal: bool,
19122 window: usize,
19123 ) -> Result<(), Box<dyn std::error::Error>> {
19124 let f = self.func("sdpa_naive_w_f32");
19125 let cfg = LaunchConfig {
19126 grid_dim: (n_head as u32, t as u32, 1),
19127 block_dim: (128, 1, 1),
19128 shared_mem_bytes: (t_kv * 4) as u32,
19129 };
19130 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19131 head_dim as i32,
19132 n_head as i32,
19133 n_head_kv as i32,
19134 t as i32,
19135 t_kv as i32,
19136 causal as i32,
19137 window as i32,
19138 );
19139 let __s_b = self.gpu.stream();
19140 let mut b = __s_b.launch_builder(&f);
19141 b.arg(q)
19142 .arg(k)
19143 .arg(v)
19144 .arg(o)
19145 .arg(&hd)
19146 .arg(&nh)
19147 .arg(&nhkv)
19148 .arg(&ti)
19149 .arg(&tkvi)
19150 .arg(&scale)
19151 .arg(&cz)
19152 .arg(&wi);
19153 unsafe {
19154 b.launch(cfg)?;
19155 }
19156 Ok(())
19157 }
19158
19159 #[allow(clippy::too_many_arguments)]
19169 pub fn sdpa_naive_w_lo(
19170 &self,
19171 q: &CudaSlice<f32>,
19172 k: &CudaSlice<f32>,
19173 v: &CudaSlice<f32>,
19174 o: &mut CudaSlice<f32>,
19175 head_dim: usize,
19176 n_head: usize,
19177 n_head_kv: usize,
19178 t: usize,
19179 t_kv: usize,
19180 scale: f32,
19181 causal: bool,
19182 window: usize,
19183 ) -> Result<(), Box<dyn std::error::Error>> {
19184 let kv_lo = if window > 0 {
19185 (t_kv - t + 1).saturating_sub(window)
19186 } else {
19187 0
19188 };
19189 let smem = (t_kv - kv_lo) * 4;
19190 if smem > 48 * 1024 {
19191 return Err(format!(
19192 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19193 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19194 a window this wide needs the multi-pass long-ctx kernel"
19195 )
19196 .into());
19197 }
19198 let f = self.func("sdpa_naive_w_lo_f32");
19199 let cfg = LaunchConfig {
19200 grid_dim: (n_head as u32, t as u32, 1),
19201 block_dim: (128, 1, 1),
19202 shared_mem_bytes: smem as u32,
19203 };
19204 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19205 head_dim as i32,
19206 n_head as i32,
19207 n_head_kv as i32,
19208 t as i32,
19209 t_kv as i32,
19210 causal as i32,
19211 window as i32,
19212 kv_lo as i32,
19213 );
19214 let __s_b = self.gpu.stream();
19215 let mut b = __s_b.launch_builder(&f);
19216 b.arg(q)
19217 .arg(k)
19218 .arg(v)
19219 .arg(o)
19220 .arg(&hd)
19221 .arg(&nh)
19222 .arg(&nhkv)
19223 .arg(&ti)
19224 .arg(&tkvi)
19225 .arg(&scale)
19226 .arg(&cz)
19227 .arg(&wi)
19228 .arg(&lo);
19229 unsafe {
19230 b.launch(cfg)?;
19231 }
19232 Ok(())
19233 }
19234
19235 pub fn sdpa_naive_view(
19237 &self,
19238 q: &CudaSlice<f32>,
19239 k: &cudarc::driver::CudaView<f32>,
19240 v: &cudarc::driver::CudaView<f32>,
19241 o: &mut CudaSlice<f32>,
19242 head_dim: usize,
19243 n_head: usize,
19244 n_head_kv: usize,
19245 t: usize,
19246 t_kv: usize,
19247 scale: f32,
19248 causal: bool,
19249 ) -> Result<(), Box<dyn std::error::Error>> {
19250 let f = self.func("sdpa_naive_f32");
19251 let cfg = LaunchConfig {
19252 grid_dim: (n_head as u32, t as u32, 1),
19253 block_dim: (128, 1, 1),
19254 shared_mem_bytes: (t_kv * 4) as u32,
19255 };
19256 let (hd, nh, nhkv, ti, tkvi, cz) = (
19257 head_dim as i32,
19258 n_head as i32,
19259 n_head_kv as i32,
19260 t as i32,
19261 t_kv as i32,
19262 causal as i32,
19263 );
19264 let __s_b = self.gpu.stream();
19265 let mut b = __s_b.launch_builder(&f);
19266 b.arg(q)
19267 .arg(k)
19268 .arg(v)
19269 .arg(o)
19270 .arg(&hd)
19271 .arg(&nh)
19272 .arg(&nhkv)
19273 .arg(&ti)
19274 .arg(&tkvi)
19275 .arg(&scale)
19276 .arg(&cz);
19277 unsafe {
19278 b.launch(cfg)?;
19279 }
19280 Ok(())
19281 }
19282
19283 #[allow(clippy::too_many_arguments)]
19291 pub fn fa_dequant_kv_view_f32(
19292 &self,
19293 k: &cudarc::driver::CudaView<u8>,
19294 v: &cudarc::driver::CudaView<u8>,
19295 kf: &mut CudaSlice<f32>,
19296 vf: &mut CudaSlice<f32>,
19297 kv_dim_k: usize,
19298 kv_dim_v: usize,
19299 t_kv: usize,
19300 k_tok_bytes: usize,
19301 v_tok_bytes: usize,
19302 g: bool,
19303 ) -> Result<(), Box<dyn std::error::Error>> {
19304 let f = if g {
19305 self.func_g("fa_dequant_kv_ws_f32")
19306 } else {
19307 self.func("fa_dequant_kv_ws_f32")
19308 };
19309 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19310 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19311 let cfg = LaunchConfig {
19312 grid_dim: (nblk.max(1), 1, 1),
19313 block_dim: (256, 1, 1),
19314 shared_mem_bytes: 0,
19315 };
19316 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19317 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19318 let __s_b = self.gpu.stream();
19319 let mut b = __s_b.launch_builder(&f);
19320 b.arg(k)
19321 .arg(v)
19322 .arg(&mut *kf)
19323 .arg(&mut *vf)
19324 .arg(&kdk)
19325 .arg(&kdv)
19326 .arg(&tkvi)
19327 .arg(&ktb)
19328 .arg(&vtb);
19329 unsafe {
19330 b.launch(cfg)?;
19331 }
19332 Ok(())
19333 }
19334
19335 #[allow(clippy::too_many_arguments)]
19336 pub fn sdpa_naive_quantized_view(
19337 &self,
19338 q: &CudaSlice<f32>,
19339 k: &cudarc::driver::CudaView<u8>,
19340 v: &cudarc::driver::CudaView<u8>,
19341 o: &mut CudaSlice<f32>,
19342 head_dim: usize,
19343 n_head: usize,
19344 n_head_kv: usize,
19345 t: usize,
19346 t_kv: usize,
19347 scale: f32,
19348 causal: bool,
19349 k_tok_bytes: usize,
19350 v_tok_bytes: usize,
19351 ) -> Result<(), Box<dyn std::error::Error>> {
19352 let kv_dim = n_head_kv * head_dim;
19353 let mut kf = self.uninit(t_kv * kv_dim)?;
19354 let mut vf = self.uninit(t_kv * kv_dim)?;
19355 let f = self.func("fa_dequant_kv_ws_f32");
19356 let total = (2 * t_kv * kv_dim) as u64;
19357 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19358 let cfg = LaunchConfig {
19359 grid_dim: (nblk.max(1), 1, 1),
19360 block_dim: (256, 1, 1),
19361 shared_mem_bytes: 0,
19362 };
19363 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19364 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19365 let __s_b = self.gpu.stream();
19366 let mut b = __s_b.launch_builder(&f);
19367 b.arg(k)
19368 .arg(v)
19369 .arg(&mut kf)
19370 .arg(&mut vf)
19371 .arg(&kv_dim_i)
19372 .arg(&kv_dim_i)
19373 .arg(&t_kv_i)
19374 .arg(&k_tok_bytes_i)
19375 .arg(&v_tok_bytes_i);
19376 unsafe { b.launch(cfg)? };
19377 self.sdpa_naive(
19378 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19379 )
19380 }
19381
19382 #[allow(clippy::too_many_arguments)]
19394 pub fn sdpa_naive_w_quantized_view(
19395 &self,
19396 q: &CudaSlice<f32>,
19397 k: &cudarc::driver::CudaView<u8>,
19398 v: &cudarc::driver::CudaView<u8>,
19399 o: &mut CudaSlice<f32>,
19400 head_dim: usize,
19401 n_head: usize,
19402 n_head_kv: usize,
19403 t: usize,
19404 t_kv: usize,
19405 scale: f32,
19406 causal: bool,
19407 window: usize,
19408 k_tok_bytes: usize,
19409 v_tok_bytes: usize,
19410 ) -> Result<(), Box<dyn std::error::Error>> {
19411 let kv_dim = n_head_kv * head_dim;
19412 let mut kf = self.uninit(t_kv * kv_dim)?;
19413 let mut vf = self.uninit(t_kv * kv_dim)?;
19414 let f = self.func("fa_dequant_kv_ws_f32");
19415 let total = (2 * t_kv * kv_dim) as u64;
19416 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19417 let cfg = LaunchConfig {
19418 grid_dim: (nblk.max(1), 1, 1),
19419 block_dim: (256, 1, 1),
19420 shared_mem_bytes: 0,
19421 };
19422 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19423 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19424 let __s_b = self.gpu.stream();
19425 let mut b = __s_b.launch_builder(&f);
19426 b.arg(k)
19427 .arg(v)
19428 .arg(&mut kf)
19429 .arg(&mut vf)
19430 .arg(&kv_dim_i)
19431 .arg(&kv_dim_i)
19432 .arg(&t_kv_i)
19433 .arg(&k_tok_bytes_i)
19434 .arg(&v_tok_bytes_i);
19435 unsafe { b.launch(cfg)? };
19436 self.sdpa_naive_w(
19437 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19438 )
19439 }
19440
19441 pub fn fa_prefill(
19445 &self,
19446 q: &CudaSlice<f32>,
19447 k: &CudaSlice<f32>,
19448 v: &CudaSlice<f32>,
19449 o: &mut CudaSlice<f32>,
19450 head_dim: usize,
19451 n_head: usize,
19452 n_head_kv: usize,
19453 t: usize,
19454 t_kv: usize,
19455 scale: f32,
19456 causal: bool,
19457 ) -> Result<(), Box<dyn std::error::Error>> {
19458 if portable_mma_gated() {
19459 return self.sdpa_naive(
19460 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19461 );
19462 }
19463 let fa3_on = head_dim == 256
19471 && causal
19472 && t == t_kv
19473 && match std::env::var("MEMRA_FA3").as_deref() {
19474 Ok("0") => false,
19475 Ok("1") => {
19479 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19480 true
19481 }
19482 _ => cfg!(memra_hopper_mma),
19483 };
19484 if fa3_on {
19485 let n = t * n_head * head_dim;
19486 let nkv = t * n_head_kv * head_dim;
19487 let mut q16 = self.alloc_u8_uninit(n * 2)?;
19488 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19489 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19490 self.f32_to_bf16_into(q, &mut q16, n)?;
19491 self.f32_to_bf16_into(k, &mut k16, nkv)?;
19492 self.f32_to_bf16_into(v, &mut v16, nkv)?;
19493 let rc = {
19494 use cudarc::driver::{DevicePtr, DevicePtrMut};
19495 let stream = self.gpu.stream();
19496 let (qp, _g1) = q16.device_ptr(&stream);
19497 let (kp, _g2) = k16.device_ptr(&stream);
19498 let (vp, _g3) = v16.device_ptr(&stream);
19499 let (op, _g4) = o.device_ptr_mut(&stream);
19500 unsafe {
19501 memra_fa3_prefill(
19502 qp as *const core::ffi::c_void,
19503 kp as *const core::ffi::c_void,
19504 vp as *const core::ffi::c_void,
19505 op as *mut f32,
19506 t as i32,
19507 n_head as i32,
19508 n_head_kv as i32,
19509 head_dim as i32,
19510 scale,
19511 stream.cu_stream() as *mut core::ffi::c_void,
19512 )
19513 }
19514 };
19515 if rc != 0 {
19516 return Err(format!("memra_fa3_prefill rc={rc}").into());
19517 }
19518 return Ok(());
19519 }
19520 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19525 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19526 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19527 const BLOCK_Q: usize = 64;
19528 const BKX: usize = 32;
19529 let f = self.func("fa_prefill_bf16_p1");
19530 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19531 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19532 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19533 f.set_attribute(
19534 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19535 shmem as i32,
19536 )?;
19537 let cfg = LaunchConfig {
19538 grid_dim: (
19539 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19540 n_head as u32,
19541 1,
19542 ),
19543 block_dim: (32, 4, 1),
19544 shared_mem_bytes: shmem,
19545 };
19546 let (hd, nh, nhkv, ti, tkvi, cz) = (
19547 head_dim as i32,
19548 n_head as i32,
19549 n_head_kv as i32,
19550 t as i32,
19551 t_kv as i32,
19552 causal as i32,
19553 );
19554 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19555 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19556 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19557 let __s_b = self.gpu.stream();
19558 let mut b = __s_b.launch_builder(&f);
19559 b.arg(&qb)
19560 .arg(&kb)
19561 .arg(&vb)
19562 .arg(o)
19563 .arg(&hd)
19564 .arg(&nh)
19565 .arg(&nhkv)
19566 .arg(&ti)
19567 .arg(&tkvi)
19568 .arg(&scale)
19569 .arg(&cz);
19570 unsafe {
19571 b.launch(cfg)?;
19572 }
19573 return Ok(());
19574 }
19575 const BK: usize = 32;
19581 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19584 let (block_q, warps, w2_sfx): (usize, u32, &str) =
19585 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19586 let hd_sfx = fa_hd_suffix(head_dim)?;
19590 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19591 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
19596 let (kb16, vb16) = if bf16kv {
19597 let n = t_kv * n_head_kv * head_dim;
19598 let mut kb = self.alloc_u8_uninit(n * 2)?;
19599 let mut vb = self.alloc_u8_uninit(n * 2)?;
19600 let fcv = self.func("f32_to_bf16_bulk");
19601 let ni = n as i64;
19602 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
19603 let __s_b = self.gpu.stream();
19604 let mut b = __s_b.launch_builder(&fcv);
19605 b.arg(k).arg(&mut kb).arg(&ni);
19606 unsafe {
19607 b.launch(cfgc)?;
19608 }
19609 let __s_b = self.gpu.stream();
19610 let mut b = __s_b.launch_builder(&fcv);
19611 b.arg(v).arg(&mut vb).arg(&ni);
19612 unsafe {
19613 b.launch(cfgc)?;
19614 }
19615 (Some(kb), Some(vb))
19616 } else {
19617 (None, None)
19618 };
19619 let f = self.func(&if bf16kv {
19620 format!("fa_prefill_bf16kv_pp{hd_sfx}")
19621 } else {
19622 format!(
19623 "fa_prefill_f32{}{}{hd_sfx}",
19624 if floor { "" } else { "_pp" },
19625 if floor { "" } else { w2_sfx }
19626 )
19627 });
19628 let kv_stages = if bf16kv { 2 } else { 1 };
19631 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19632 + 4 * (block_q * BK + 2 * block_q)) as u32;
19633 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19634 f.set_attribute(
19635 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19636 shmem as i32,
19637 )?;
19638 let cfg = LaunchConfig {
19639 grid_dim: (
19640 (t as u32 + block_q as u32 - 1) / block_q as u32,
19641 n_head as u32,
19642 1,
19643 ),
19644 block_dim: (32, warps, 1),
19645 shared_mem_bytes: shmem,
19646 };
19647 let (hd, nh, nhkv, ti, tkvi, cz) = (
19648 head_dim as i32,
19649 n_head as i32,
19650 n_head_kv as i32,
19651 t as i32,
19652 t_kv as i32,
19653 causal as i32,
19654 );
19655 let __s_b = self.gpu.stream();
19656 let mut b = __s_b.launch_builder(&f);
19657 b.arg(q);
19658 match (&kb16, &vb16) {
19659 (Some(kb), Some(vb)) => {
19660 b.arg(kb).arg(vb);
19661 }
19662 _ => {
19663 b.arg(k).arg(v);
19664 }
19665 }
19666 b.arg(o)
19667 .arg(&hd)
19668 .arg(&nh)
19669 .arg(&nhkv)
19670 .arg(&ti)
19671 .arg(&tkvi)
19672 .arg(&scale)
19673 .arg(&cz);
19674 unsafe {
19675 b.launch(cfg)?;
19676 }
19677 Ok(())
19678 }
19679
19680 #[allow(clippy::too_many_arguments)]
19684 pub fn fa_prefill_w(
19685 &self,
19686 q: &CudaSlice<f32>,
19687 k: &CudaSlice<f32>,
19688 v: &CudaSlice<f32>,
19689 o: &mut CudaSlice<f32>,
19690 head_dim: usize,
19691 n_head: usize,
19692 n_head_kv: usize,
19693 t: usize,
19694 t_kv: usize,
19695 scale: f32,
19696 causal: bool,
19697 window: usize,
19698 ) -> Result<(), Box<dyn std::error::Error>> {
19699 if portable_mma_gated() {
19702 return self.sdpa_naive_w(
19703 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19704 );
19705 }
19706 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19710 let faw_f32 =
19711 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
19712 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19713 self.fa_prefill_w_arm(
19714 q,
19715 k,
19716 v,
19717 o,
19718 head_dim,
19719 n_head,
19720 n_head_kv,
19721 t,
19722 t_kv,
19723 scale,
19724 causal,
19725 window,
19726 floor || faw_f32,
19727 floor,
19728 )
19729 }
19730
19731 #[allow(clippy::too_many_arguments)]
19734 pub fn fa_prefill_w_pre(
19735 &self,
19736 qb: &CudaSlice<u8>,
19737 kb: &CudaSlice<u8>,
19738 vb: &CudaSlice<u8>,
19739 o: &mut CudaSlice<f32>,
19740 head_dim: usize,
19741 n_head: usize,
19742 n_head_kv: usize,
19743 t: usize,
19744 t_kv: usize,
19745 scale: f32,
19746 causal: bool,
19747 window: usize,
19748 v_f16: bool,
19749 ) -> Result<(), Box<dyn std::error::Error>> {
19750 const BLOCK_Q: usize = 64;
19751 const BK: usize = 32;
19752 debug_assert_eq!(head_dim, 256);
19753 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19754 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19755 if hp {
19756 const BLOCK_QH: usize = 32;
19757 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19760 let vh: &CudaSlice<u8> = if v_f16 {
19761 vb
19762 } else {
19763 let n = t_kv * n_head_kv * head_dim;
19764 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19765 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19766 }
19767 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19768 vguard.as_ref().unwrap()
19769 };
19770 let f = self.func("fa_prefill_w_bf16_p1h2");
19771 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19772 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19773 f.set_attribute(
19774 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19775 shmem as i32,
19776 )?;
19777 let cfg = LaunchConfig {
19778 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19779 block_dim: (32, 4, 1),
19780 shared_mem_bytes: shmem,
19781 };
19782 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19783 head_dim as i32,
19784 n_head as i32,
19785 n_head_kv as i32,
19786 t as i32,
19787 t_kv as i32,
19788 causal as i32,
19789 window as i32,
19790 );
19791 let __s_b = self.gpu.stream();
19792 let mut b = __s_b.launch_builder(&f);
19793 b.arg(qb)
19794 .arg(kb)
19795 .arg(vh)
19796 .arg(o)
19797 .arg(&hd)
19798 .arg(&nh)
19799 .arg(&nhkv)
19800 .arg(&ti)
19801 .arg(&tkvi)
19802 .arg(&scale)
19803 .arg(&cz)
19804 .arg(&wi);
19805 unsafe {
19806 b.launch(cfg)?;
19807 }
19808 return Ok(());
19809 }
19810 let f = self.func("fa_prefill_w_bf16_p1");
19811 let shmem =
19812 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19813 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19814 f.set_attribute(
19815 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19816 shmem as i32,
19817 )?;
19818 let cfg = LaunchConfig {
19819 grid_dim: (
19820 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19821 n_head as u32,
19822 1,
19823 ),
19824 block_dim: (32, 4, 1),
19825 shared_mem_bytes: shmem,
19826 };
19827 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19828 head_dim as i32,
19829 n_head as i32,
19830 n_head_kv as i32,
19831 t as i32,
19832 t_kv as i32,
19833 causal as i32,
19834 window as i32,
19835 );
19836 let __s_b = self.gpu.stream();
19837 let mut b = __s_b.launch_builder(&f);
19838 b.arg(qb)
19839 .arg(kb)
19840 .arg(vb)
19841 .arg(o)
19842 .arg(&hd)
19843 .arg(&nh)
19844 .arg(&nhkv)
19845 .arg(&ti)
19846 .arg(&tkvi)
19847 .arg(&scale)
19848 .arg(&cz)
19849 .arg(&wi);
19850 unsafe {
19851 b.launch(cfg)?;
19852 }
19853 Ok(())
19854 }
19855
19856 #[allow(clippy::too_many_arguments)]
19858 pub fn fa_prefill_w_arm(
19859 &self,
19860 q: &CudaSlice<f32>,
19861 k: &CudaSlice<f32>,
19862 v: &CudaSlice<f32>,
19863 o: &mut CudaSlice<f32>,
19864 head_dim: usize,
19865 n_head: usize,
19866 n_head_kv: usize,
19867 t: usize,
19868 t_kv: usize,
19869 scale: f32,
19870 causal: bool,
19871 window: usize,
19872 f32_stage: bool,
19873 floor: bool,
19874 ) -> Result<(), Box<dyn std::error::Error>> {
19875 const BLOCK_Q: usize = 64;
19876 const BK: usize = 32;
19877 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19878 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19882 let p1 = !floor
19883 && !f32_stage
19884 && *P1_ON.get_or_init(|| {
19885 std::env::var("MEMRA_FAW_P1")
19886 .map(|v| v != "0")
19887 .unwrap_or(true)
19888 });
19889 let hp =
19890 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19891 if hp {
19892 const BLOCK_QH: usize = 32;
19893 let f = self.func("fa_prefill_w_bf16_p1h2");
19894 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19895 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19896 f.set_attribute(
19897 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19898 shmem as i32,
19899 )?;
19900 let cfg = LaunchConfig {
19901 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19902 block_dim: (32, 4, 1),
19903 shared_mem_bytes: shmem,
19904 };
19905 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19906 head_dim as i32,
19907 n_head as i32,
19908 n_head_kv as i32,
19909 t as i32,
19910 t_kv as i32,
19911 causal as i32,
19912 window as i32,
19913 );
19914 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19915 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19916 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19917 let __s_b = self.gpu.stream();
19918 let mut b = __s_b.launch_builder(&f);
19919 b.arg(&qb)
19920 .arg(&kb)
19921 .arg(&vh)
19922 .arg(o)
19923 .arg(&hd)
19924 .arg(&nh)
19925 .arg(&nhkv)
19926 .arg(&ti)
19927 .arg(&tkvi)
19928 .arg(&scale)
19929 .arg(&cz)
19930 .arg(&wi);
19931 unsafe {
19932 b.launch(cfg)?;
19933 }
19934 return Ok(());
19935 }
19936 if p1 {
19937 let f = self.func("fa_prefill_w_bf16_p1");
19938 let shmem =
19939 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19940 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19941 f.set_attribute(
19942 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19943 shmem as i32,
19944 )?;
19945 let cfg = LaunchConfig {
19946 grid_dim: (
19947 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19948 n_head as u32,
19949 1,
19950 ),
19951 block_dim: (32, 4, 1),
19952 shared_mem_bytes: shmem,
19953 };
19954 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19955 head_dim as i32,
19956 n_head as i32,
19957 n_head_kv as i32,
19958 t as i32,
19959 t_kv as i32,
19960 causal as i32,
19961 window as i32,
19962 );
19963 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19964 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19965 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19966 let __s_b = self.gpu.stream();
19967 let mut b = __s_b.launch_builder(&f);
19968 b.arg(&qb)
19969 .arg(&kb)
19970 .arg(&vb)
19971 .arg(o)
19972 .arg(&hd)
19973 .arg(&nh)
19974 .arg(&nhkv)
19975 .arg(&ti)
19976 .arg(&tkvi)
19977 .arg(&scale)
19978 .arg(&cz)
19979 .arg(&wi);
19980 unsafe {
19981 b.launch(cfg)?;
19982 }
19983 return Ok(());
19984 }
19985 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19988 let g4 = !floor
19989 && !f32_stage
19990 && n_head_kv == 1
19991 && n_head % 4 == 0
19992 && *G4_ON.get_or_init(|| {
19993 std::env::var("MEMRA_FAW_G4")
19994 .map(|v| v != "0")
19995 .unwrap_or(true)
19996 });
19997 if g4 {
19998 const SP_M: usize = 16;
19999 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20002 let o2 = *O2_ON.get_or_init(|| {
20003 std::env::var("MEMRA_FAW_O2")
20004 .map(|v| v != "0")
20005 .unwrap_or(true)
20006 });
20007 let f = self.func(if o2 {
20008 "fa_prefill_w_bf16_g4o2"
20009 } else {
20010 "fa_prefill_w_bf16_g4"
20011 });
20012 let shmem = if o2 {
20013 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20014 } else {
20015 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20016 as u32
20017 };
20018 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20019 f.set_attribute(
20020 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20021 shmem as i32,
20022 )?;
20023 let cfg = LaunchConfig {
20024 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20025 block_dim: (32, 4, 1),
20026 shared_mem_bytes: shmem,
20027 };
20028 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20029 head_dim as i32,
20030 n_head as i32,
20031 n_head_kv as i32,
20032 t as i32,
20033 t_kv as i32,
20034 causal as i32,
20035 window as i32,
20036 );
20037 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20038 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20039 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20040 let __s_b = self.gpu.stream();
20041 let mut b = __s_b.launch_builder(&f);
20042 b.arg(&qb)
20043 .arg(&kb)
20044 .arg(&vb)
20045 .arg(o)
20046 .arg(&hd)
20047 .arg(&nh)
20048 .arg(&nhkv)
20049 .arg(&ti)
20050 .arg(&tkvi)
20051 .arg(&scale)
20052 .arg(&cz)
20053 .arg(&wi);
20054 unsafe {
20055 b.launch(cfg)?;
20056 }
20057 return Ok(());
20058 }
20059 let f = self.func(if floor {
20060 "fa_prefill_w_f32"
20061 } else if f32_stage {
20062 "fa_prefill_w_f32_pp"
20063 } else {
20064 "fa_prefill_w_bf16_pp"
20065 });
20066 let shmem =
20067 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20068 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20069 f.set_attribute(
20070 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20071 shmem as i32,
20072 )?;
20073 let cfg = LaunchConfig {
20074 grid_dim: (
20075 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20076 n_head as u32,
20077 1,
20078 ),
20079 block_dim: (32, 4, 1),
20080 shared_mem_bytes: shmem,
20081 };
20082 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20083 head_dim as i32,
20084 n_head as i32,
20085 n_head_kv as i32,
20086 t as i32,
20087 t_kv as i32,
20088 causal as i32,
20089 window as i32,
20090 );
20091 if f32_stage {
20092 let __s_b = self.gpu.stream();
20093 let mut b = __s_b.launch_builder(&f);
20094 b.arg(q)
20095 .arg(k)
20096 .arg(v)
20097 .arg(o)
20098 .arg(&hd)
20099 .arg(&nh)
20100 .arg(&nhkv)
20101 .arg(&ti)
20102 .arg(&tkvi)
20103 .arg(&scale)
20104 .arg(&cz)
20105 .arg(&wi);
20106 unsafe {
20107 b.launch(cfg)?;
20108 }
20109 } else {
20110 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20111 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20112 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20113 let __s_b = self.gpu.stream();
20114 let mut b = __s_b.launch_builder(&f);
20115 b.arg(&qb)
20116 .arg(&kb)
20117 .arg(&vb)
20118 .arg(o)
20119 .arg(&hd)
20120 .arg(&nh)
20121 .arg(&nhkv)
20122 .arg(&ti)
20123 .arg(&tkvi)
20124 .arg(&scale)
20125 .arg(&cz)
20126 .arg(&wi);
20127 unsafe {
20128 b.launch(cfg)?;
20129 }
20130 }
20131 Ok(())
20132 }
20133
20134 #[allow(clippy::too_many_arguments)]
20138 pub fn fa_prefill_hd512(
20139 &self,
20140 q: &CudaSlice<f32>,
20141 k: &CudaSlice<f32>,
20142 v: &CudaSlice<f32>,
20143 o: &mut CudaSlice<f32>,
20144 head_dim: usize,
20145 n_head: usize,
20146 n_head_kv: usize,
20147 t: usize,
20148 t_kv: usize,
20149 scale: f32,
20150 causal: bool,
20151 ) -> Result<(), Box<dyn std::error::Error>> {
20152 if portable_mma_gated() {
20154 return self.sdpa_naive(
20155 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20156 );
20157 }
20158 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20164 let f32_stage =
20165 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20166 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20170 let sp = !f32_stage
20171 && *SP_ON.get_or_init(|| {
20172 std::env::var("MEMRA_FA512_SP")
20173 .map(|v| v != "0")
20174 .unwrap_or(true)
20175 });
20176 self.fa_prefill_hd512_arm(
20177 q,
20178 k,
20179 v,
20180 o,
20181 head_dim,
20182 n_head,
20183 n_head_kv,
20184 t,
20185 t_kv,
20186 scale,
20187 causal,
20188 f32_stage,
20189 sp,
20190 sp && fa_f16pv_on(),
20191 )
20192 }
20193
20194 #[allow(clippy::too_many_arguments)]
20196 pub fn fa_prefill_hd512_pre(
20197 &self,
20198 qb: &CudaSlice<u8>,
20199 kb: &CudaSlice<u8>,
20200 vb: &CudaSlice<u8>,
20201 o: &mut CudaSlice<f32>,
20202 head_dim: usize,
20203 n_head: usize,
20204 n_head_kv: usize,
20205 t: usize,
20206 t_kv: usize,
20207 scale: f32,
20208 causal: bool,
20209 v_f16: bool,
20210 ) -> Result<(), Box<dyn std::error::Error>> {
20211 debug_assert_eq!(head_dim, 512);
20212 const SP_M: usize = 16;
20213 const BKS: usize = 32;
20214 let f16pv = fa_f16pv_on();
20218 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20219 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20220 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20221 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20222 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20223 let n = t_kv * n_head_kv * head_dim;
20225 let need = n * 2;
20226 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20227 *vguard = Some(self.alloc_uninit::<u8>(need)?);
20228 }
20229 let dst = vguard.as_mut().unwrap();
20230 self.bf16_to_f16_into(vb, n, dst)?;
20231 vguard.as_ref().unwrap()
20232 } else {
20233 vb
20234 };
20235 let f = self.func(if hp {
20236 "fa_prefill_bf16_hd512_sp16h2"
20237 } else {
20238 match (f16pv, nw) {
20239 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20240 (true, _) => "fa_prefill_bf16_hd512_sp16",
20241 _ => "fa_prefill_bf16_hd512_sp",
20242 }
20243 });
20244 let (nwarp, npart) = if hp {
20245 (4usize, 4usize)
20246 } else if nw > 2 {
20247 (nw, nw)
20248 } else {
20249 (2, 1)
20250 };
20251 let shmem = if hp {
20253 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20254 as u32
20255 } else {
20256 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20257 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20258 };
20259 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20260 f.set_attribute(
20261 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20262 shmem as i32,
20263 )?;
20264 let grid_y = if hp {
20265 (n_head / 2) as u32
20266 } else {
20267 n_head as u32
20268 };
20269 let cfg = LaunchConfig {
20270 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20271 block_dim: (32, nwarp as u32, 1),
20272 shared_mem_bytes: shmem,
20273 };
20274 let (hd, nh, nhkv, ti, tkvi, cz) = (
20275 head_dim as i32,
20276 n_head as i32,
20277 n_head_kv as i32,
20278 t as i32,
20279 t_kv as i32,
20280 causal as i32,
20281 );
20282 let __s_b = self.gpu.stream();
20283 let mut b = __s_b.launch_builder(&f);
20284 b.arg(qb)
20285 .arg(kb)
20286 .arg(vref)
20287 .arg(o)
20288 .arg(&hd)
20289 .arg(&nh)
20290 .arg(&nhkv)
20291 .arg(&ti)
20292 .arg(&tkvi)
20293 .arg(&scale)
20294 .arg(&cz);
20295 unsafe {
20296 b.launch(cfg)?;
20297 }
20298 Ok(())
20299 }
20300
20301 #[allow(clippy::too_many_arguments)]
20304 pub fn fa_prefill_hd512_arm(
20305 &self,
20306 q: &CudaSlice<f32>,
20307 k: &CudaSlice<f32>,
20308 v: &CudaSlice<f32>,
20309 o: &mut CudaSlice<f32>,
20310 head_dim: usize,
20311 n_head: usize,
20312 n_head_kv: usize,
20313 t: usize,
20314 t_kv: usize,
20315 scale: f32,
20316 causal: bool,
20317 f32_stage: bool,
20318 sp: bool,
20319 f16pv: bool,
20320 ) -> Result<(), Box<dyn std::error::Error>> {
20321 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20322 if sp && !f32_stage {
20323 const SP_M: usize = 16;
20327 const BKS: usize = 32;
20328 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20329 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20330 let f = self.func(if hp {
20331 "fa_prefill_bf16_hd512_sp16h2"
20332 } else {
20333 match (f16pv, nw) {
20334 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20335 (true, _) => "fa_prefill_bf16_hd512_sp16",
20336 _ => "fa_prefill_bf16_hd512_sp",
20337 }
20338 });
20339 let (nwarp, npart) = if hp {
20340 (4usize, 4usize)
20341 } else if nw > 2 {
20342 (nw, nw)
20343 } else {
20344 (2, 1)
20345 };
20346 let shmem = if hp {
20347 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20348 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20349 } else {
20350 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20351 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20352 };
20353 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20354 f.set_attribute(
20355 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20356 shmem as i32,
20357 )?;
20358 let grid_y = if hp {
20359 (n_head / 2) as u32
20360 } else {
20361 n_head as u32
20362 };
20363 let cfg = LaunchConfig {
20364 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20365 block_dim: (32, nwarp as u32, 1),
20366 shared_mem_bytes: shmem,
20367 };
20368 let (hd, nh, nhkv, ti, tkvi, cz) = (
20369 head_dim as i32,
20370 n_head as i32,
20371 n_head_kv as i32,
20372 t as i32,
20373 t_kv as i32,
20374 causal as i32,
20375 );
20376 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20377 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20378 let vb = if f16pv {
20379 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20380 } else {
20381 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20382 };
20383 let __s_b = self.gpu.stream();
20384 let mut b = __s_b.launch_builder(&f);
20385 b.arg(&qb)
20386 .arg(&kb)
20387 .arg(&vb)
20388 .arg(o)
20389 .arg(&hd)
20390 .arg(&nh)
20391 .arg(&nhkv)
20392 .arg(&ti)
20393 .arg(&tkvi)
20394 .arg(&scale)
20395 .arg(&cz);
20396 unsafe {
20397 b.launch(cfg)?;
20398 }
20399 return Ok(());
20400 }
20401 const BLOCK_Q: usize = 32;
20402 const BK: usize = 32;
20403 const HALF: usize = 256;
20404 let f = self.func(if f32_stage {
20405 "fa_prefill_f32_hd512"
20406 } else {
20407 "fa_prefill_bf16_hd512"
20408 });
20409 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20411 + 4 * BLOCK_Q) as u32;
20412 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20413 f.set_attribute(
20414 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20415 shmem as i32,
20416 )?;
20417 let cfg = LaunchConfig {
20418 grid_dim: (
20419 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20420 n_head as u32,
20421 2,
20422 ),
20423 block_dim: (32, 2, 1),
20424 shared_mem_bytes: shmem,
20425 };
20426 let (hd, nh, nhkv, ti, tkvi, cz) = (
20427 head_dim as i32,
20428 n_head as i32,
20429 n_head_kv as i32,
20430 t as i32,
20431 t_kv as i32,
20432 causal as i32,
20433 );
20434 if f32_stage {
20435 let __s_b = self.gpu.stream();
20436 let mut b = __s_b.launch_builder(&f);
20437 b.arg(q)
20438 .arg(k)
20439 .arg(v)
20440 .arg(o)
20441 .arg(&hd)
20442 .arg(&nh)
20443 .arg(&nhkv)
20444 .arg(&ti)
20445 .arg(&tkvi)
20446 .arg(&scale)
20447 .arg(&cz);
20448 unsafe {
20449 b.launch(cfg)?;
20450 }
20451 } else {
20452 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20453 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20454 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20455 let __s_b = self.gpu.stream();
20456 let mut b = __s_b.launch_builder(&f);
20457 b.arg(&qb)
20458 .arg(&kb)
20459 .arg(&vb)
20460 .arg(o)
20461 .arg(&hd)
20462 .arg(&nh)
20463 .arg(&nhkv)
20464 .arg(&ti)
20465 .arg(&tkvi)
20466 .arg(&scale)
20467 .arg(&cz);
20468 unsafe {
20469 b.launch(cfg)?;
20470 }
20471 }
20472 Ok(())
20473 }
20474
20475 #[allow(clippy::too_many_arguments)]
20479 pub fn rope_neox2_bf16e(
20480 &self,
20481 q: &mut CudaSlice<f32>,
20482 k: &mut CudaSlice<f32>,
20483 qb: &mut CudaSlice<u8>,
20484 kb: &mut CudaSlice<u8>,
20485 pos: &CudaSlice<i32>,
20486 head_dim: usize,
20487 n_dims: usize,
20488 nh_q: usize,
20489 nh_k: usize,
20490 n_tokens: usize,
20491 base: f32,
20492 freq_scale: f32,
20493 ff: Option<&CudaSlice<f32>>,
20494 ) -> Result<(), Box<dyn std::error::Error>> {
20495 let f = self.func("rope_neox2_bf16e_f32");
20496 let rows = ((nh_q + nh_k) * n_tokens) as u32;
20497 let cfg = LaunchConfig {
20498 grid_dim: (rows, 1, 1),
20499 block_dim: ((head_dim / 2) as u32, 1, 1),
20500 shared_mem_bytes: 0,
20501 };
20502 let theta_scale = base.powf(-2.0 / n_dims as f32);
20503 let (hd, nd, nhq, nhk, nt) = (
20504 head_dim as i32,
20505 n_dims as i32,
20506 nh_q as i32,
20507 nh_k as i32,
20508 n_tokens as i32,
20509 );
20510 let __s_b = self.gpu.stream();
20511 let mut b = __s_b.launch_builder(&f);
20512 match ff {
20513 Some(t) => {
20514 b.arg(&mut *q)
20515 .arg(&mut *k)
20516 .arg(&mut *qb)
20517 .arg(&mut *kb)
20518 .arg(pos)
20519 .arg(&hd)
20520 .arg(&nd)
20521 .arg(&nhq)
20522 .arg(&nhk)
20523 .arg(&nt)
20524 .arg(&theta_scale)
20525 .arg(&freq_scale)
20526 .arg(t);
20527 unsafe {
20528 b.launch(cfg)?;
20529 }
20530 }
20531 None => {
20532 let null: u64 = 0;
20533 b.arg(&mut *q)
20534 .arg(&mut *k)
20535 .arg(&mut *qb)
20536 .arg(&mut *kb)
20537 .arg(pos)
20538 .arg(&hd)
20539 .arg(&nd)
20540 .arg(&nhq)
20541 .arg(&nhk)
20542 .arg(&nt)
20543 .arg(&theta_scale)
20544 .arg(&freq_scale)
20545 .arg(&null);
20546 unsafe {
20547 b.launch(cfg)?;
20548 }
20549 }
20550 }
20551 Ok(())
20552 }
20553
20554 pub fn f32_to_bf16(
20557 &self,
20558 x: &CudaSlice<f32>,
20559 n: usize,
20560 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20561 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
20562 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20563 let f = self.func("f32_to_bf16_flat");
20564 let n_i = n as i64;
20565 let cfg = LaunchConfig {
20566 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20567 block_dim: (256, 1, 1),
20568 shared_mem_bytes: 0,
20569 };
20570 let __s_b = self.gpu.stream();
20571 let mut b = __s_b.launch_builder(&f);
20572 b.arg(x).arg(&mut y).arg(&n_i);
20573 unsafe {
20574 b.launch(cfg)?;
20575 }
20576 Ok(y)
20577 }
20578
20579 pub fn f32_to_f16(
20580 &self,
20581 x: &CudaSlice<f32>,
20582 n: usize,
20583 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20584 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20585 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20586 let f = self.func("f32_to_f16_flat");
20587 let n_i = n as i64;
20588 let cfg = LaunchConfig {
20589 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20590 block_dim: (256, 1, 1),
20591 shared_mem_bytes: 0,
20592 };
20593 let __s_b = self.gpu.stream();
20594 let mut b = __s_b.launch_builder(&f);
20595 b.arg(x).arg(&mut y).arg(&n_i);
20596 unsafe {
20597 b.launch(cfg)?;
20598 }
20599 Ok(y)
20600 }
20601
20602 pub fn bf16_to_f16(
20604 &self,
20605 xb: &CudaSlice<u8>,
20606 n: usize,
20607 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20608 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20609 self.bf16_to_f16_into(xb, n, &mut y)?;
20610 Ok(y)
20611 }
20612
20613 pub fn bf16_to_f16_into(
20615 &self,
20616 xb: &CudaSlice<u8>,
20617 n: usize,
20618 y: &mut CudaSlice<u8>,
20619 ) -> Result<(), Box<dyn std::error::Error>> {
20620 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
20621 assert!(y.len() >= n * 2);
20622 let f = self.func("bf16_to_f16_flat");
20623 let n2 = (n / 2) as i64;
20624 let cfg = LaunchConfig {
20625 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
20626 block_dim: (256, 1, 1),
20627 shared_mem_bytes: 0,
20628 };
20629 let __s_b = self.gpu.stream();
20630 let mut b = __s_b.launch_builder(&f);
20631 b.arg(xb).arg(y).arg(&n2);
20632 unsafe {
20633 b.launch(cfg)?;
20634 }
20635 Ok(())
20636 }
20637
20638 #[allow(clippy::too_many_arguments)]
20643 pub fn fa_prefill_vl8(
20644 &self,
20645 seqs: &[FaSeqVl],
20646 head_dim: usize,
20647 n_head: usize,
20648 n_head_kv: usize,
20649 scale: f32,
20650 ) -> Result<(), Box<dyn std::error::Error>> {
20651 const BK: usize = 32;
20652 let b = seqs.len();
20653 assert!(b >= 1 && b <= 8);
20654 let mut packed = [FaSeqVl::default(); 8];
20655 packed[..b].copy_from_slice(seqs);
20656 let v = FaVl8(packed);
20657 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20658 let ept = (n_head_kv * head_dim) as i32;
20659 {
20660 let f = self.func("fa_mirror_vl");
20661 let max_n = (max_t as i64) * ept as i64;
20662 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
20663 for which in 0..2i32 {
20664 let cfg = LaunchConfig {
20665 grid_dim: (blocks, 1, b as u32),
20666 block_dim: (256, 1, 1),
20667 shared_mem_bytes: 0,
20668 };
20669 let __s_lb = self.gpu.stream();
20670 let mut lb = __s_lb.launch_builder(&f);
20671 lb.arg(&v).arg(&ept).arg(&which);
20672 unsafe {
20673 lb.launch(cfg)?;
20674 }
20675 }
20676 }
20677 let hd_sfx = fa_hd_suffix(head_dim)?;
20678 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
20679 let block_q = 64usize;
20680 let kv_stages = 2usize;
20681 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20682 + 4 * (block_q * BK + 2 * block_q)) as u32;
20683 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20684 f.set_attribute(
20685 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20686 shmem as i32,
20687 )?;
20688 let cfg = LaunchConfig {
20689 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
20690 block_dim: (32, 4, 1),
20691 shared_mem_bytes: shmem,
20692 };
20693 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20694 let __s_lb = self.gpu.stream();
20695 let mut lb = __s_lb.launch_builder(&f);
20696 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
20697 unsafe {
20698 lb.launch(cfg)?;
20699 }
20700 Ok(())
20701 }
20702
20703 #[allow(clippy::too_many_arguments)]
20707 pub fn attn_pre_vl8(
20708 &self,
20709 seqs: &[AttnPreVl],
20710 wq: &CudaSlice<f32>,
20711 wk: &CudaSlice<f32>,
20712 head_dim: usize,
20713 rope_dims: usize,
20714 n_head: usize,
20715 n_head_kv: usize,
20716 eps: f32,
20717 freq_base: f32,
20718 freq_scale: f32,
20719 kv_dim_k: usize,
20720 kv_dim_v: usize,
20721 k_tok_bytes: usize,
20722 v_tok_bytes: usize,
20723 ) -> Result<(), Box<dyn std::error::Error>> {
20724 let b = seqs.len();
20725 assert!(b >= 1 && b <= 8);
20726 let mut packed = [AttnPreVl::default(); 8];
20727 packed[..b].copy_from_slice(seqs);
20728 let v = AttnPreVl8(packed);
20729 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20730 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20731 {
20732 let f = self.func("q_gate_split_vl");
20733 let n = max_t * (n_head * head_dim) as u32;
20734 let cfg = LaunchConfig {
20735 grid_dim: (n.div_ceil(256), 1, b as u32),
20736 block_dim: (256, 1, 1),
20737 shared_mem_bytes: 0,
20738 };
20739 let __s_lb = self.gpu.stream();
20740 let mut lb = __s_lb.launch_builder(&f);
20741 lb.arg(&v).arg(&hd).arg(&nh);
20742 unsafe {
20743 lb.launch(cfg)?;
20744 }
20745 }
20746 {
20747 let f = self.func("attn_rms_vl");
20748 let cfg = LaunchConfig {
20749 grid_dim: (max_t * n_head as u32, 2, b as u32),
20750 block_dim: (rms_block(), 1, 1),
20751 shared_mem_bytes: 0,
20752 };
20753 let __s_lb = self.gpu.stream();
20754 let mut lb = __s_lb.launch_builder(&f);
20755 lb.arg(&v)
20756 .arg(wq)
20757 .arg(wk)
20758 .arg(&hd)
20759 .arg(&nh)
20760 .arg(&nhkv)
20761 .arg(&eps);
20762 unsafe {
20763 lb.launch(cfg)?;
20764 }
20765 }
20766 {
20767 let f = self.func("attn_rope_vl");
20768 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20769 let nd = rope_dims as i32;
20770 let cfg = LaunchConfig {
20771 grid_dim: (max_t * n_head as u32, 2, b as u32),
20772 block_dim: ((head_dim / 2) as u32, 1, 1),
20773 shared_mem_bytes: 0,
20774 };
20775 let __s_lb = self.gpu.stream();
20776 let mut lb = __s_lb.launch_builder(&f);
20777 lb.arg(&v)
20778 .arg(&hd)
20779 .arg(&nd)
20780 .arg(&nh)
20781 .arg(&nhkv)
20782 .arg(&theta_scale)
20783 .arg(&freq_scale);
20784 unsafe {
20785 lb.launch(cfg)?;
20786 }
20787 }
20788 {
20789 let f = self.func("append_kv_vl");
20790 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20791 let cfg = LaunchConfig {
20792 grid_dim: (nblk, max_t, b as u32),
20793 block_dim: (32, 1, 1),
20794 shared_mem_bytes: 0,
20795 };
20796 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20797 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20798 let __s_lb = self.gpu.stream();
20799 let mut lb = __s_lb.launch_builder(&f);
20800 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20801 unsafe {
20802 lb.launch(cfg)?;
20803 }
20804 }
20805 Ok(())
20806 }
20807
20808 pub fn fa_prefill_view(
20813 &self,
20814 q: &CudaSlice<f32>,
20815 k: &cudarc::driver::CudaView<u8>,
20816 v: &cudarc::driver::CudaView<u8>,
20817 o: &mut CudaSlice<f32>,
20818 head_dim: usize,
20819 n_head: usize,
20820 n_head_kv: usize,
20821 t: usize,
20822 t_kv: usize,
20823 scale: f32,
20824 causal: bool,
20825 k_tok_bytes: usize,
20826 v_tok_bytes: usize,
20827 g: bool,
20828 ) -> Result<(), Box<dyn std::error::Error>> {
20829 if portable_mma_gated() {
20830 return self.sdpa_naive_quantized_view(
20831 q,
20832 k,
20833 v,
20834 o,
20835 head_dim,
20836 n_head,
20837 n_head_kv,
20838 t,
20839 t_kv,
20840 scale,
20841 causal,
20842 k_tok_bytes,
20843 v_tok_bytes,
20844 );
20845 }
20846 const BLOCK_Q: usize = 64;
20847 const BK: usize = 32;
20848 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20851 let f = if g {
20852 self.func_g(&name)
20853 } else {
20854 self.func(&name)
20855 };
20856 let shmem =
20857 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20858 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20859 f.set_attribute(
20860 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20861 shmem as i32,
20862 )?;
20863 let cfg = LaunchConfig {
20864 grid_dim: (
20865 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20866 n_head as u32,
20867 1,
20868 ),
20869 block_dim: (32, 4, 1),
20870 shared_mem_bytes: shmem,
20871 };
20872 let (hd, nh, nhkv, ti, tkvi, cz) = (
20873 head_dim as i32,
20874 n_head as i32,
20875 n_head_kv as i32,
20876 t as i32,
20877 t_kv as i32,
20878 causal as i32,
20879 );
20880 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20881 let __s_b = self.gpu.stream();
20882 let mut b = __s_b.launch_builder(&f);
20883 b.arg(q)
20884 .arg(k)
20885 .arg(v)
20886 .arg(o)
20887 .arg(&hd)
20888 .arg(&nh)
20889 .arg(&nhkv)
20890 .arg(&ti)
20891 .arg(&tkvi)
20892 .arg(&scale)
20893 .arg(&cz)
20894 .arg(&ktb)
20895 .arg(&vtb);
20896 unsafe {
20897 b.launch(cfg)?;
20898 }
20899 Ok(())
20900 }
20901
20902 #[allow(clippy::too_many_arguments)]
20912 pub fn fa_prefill_view_ws(
20913 &self,
20914 q: &CudaSlice<f32>,
20915 k: &cudarc::driver::CudaView<u8>,
20916 v: &cudarc::driver::CudaView<u8>,
20917 o: &mut CudaSlice<f32>,
20918 head_dim: usize,
20919 n_head: usize,
20920 n_head_kv: usize,
20921 t: usize,
20922 t_kv: usize,
20923 scale: f32,
20924 causal: bool,
20925 k_tok_bytes: usize,
20926 v_tok_bytes: usize,
20927 g: bool,
20928 ) -> Result<(), Box<dyn std::error::Error>> {
20929 if portable_mma_gated() {
20930 return self.sdpa_naive_quantized_view(
20931 q,
20932 k,
20933 v,
20934 o,
20935 head_dim,
20936 n_head,
20937 n_head_kv,
20938 t,
20939 t_kv,
20940 scale,
20941 causal,
20942 k_tok_bytes,
20943 v_tok_bytes,
20944 );
20945 }
20946 const BLOCK_Q: usize = 64;
20947 const BK: usize = 32;
20948 let kv_dim_k = n_head_kv * head_dim;
20949 let kv_dim_v = n_head_kv * head_dim;
20950 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20952 let mut guard = self.prime_deqw_ws.lock().unwrap();
20954 let need_grow = match guard.as_ref() {
20955 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20956 None => true,
20957 };
20958 if need_grow {
20959 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20960 let (ck, cv) = guard
20961 .as_ref()
20962 .map(|(a, b)| (a.len(), b.len()))
20963 .unwrap_or((0, 0));
20964 *guard = Some((
20965 self.alloc_u8(grow(ck, k_ws_bytes))?,
20966 self.alloc_u8(grow(cv, v_ws_bytes))?,
20967 ));
20968 }
20969 let (kw, vw) = guard.as_mut().unwrap();
20970 {
20972 let f = if g {
20974 self.func_g("fa_dequant_kv_ws_bf16")
20975 } else {
20976 self.func("fa_dequant_kv_ws_bf16")
20977 };
20978 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20979 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20980 let cfg = LaunchConfig {
20981 grid_dim: (nblk.max(1), 1, 1),
20982 block_dim: (256, 1, 1),
20983 shared_mem_bytes: 0,
20984 };
20985 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20986 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20987 let __s_b = self.gpu.stream();
20988 let mut b = __s_b.launch_builder(&f);
20989 b.arg(k)
20990 .arg(v)
20991 .arg(&mut *kw)
20992 .arg(&mut *vw)
20993 .arg(&kdk)
20994 .arg(&kdv)
20995 .arg(&tkvi)
20996 .arg(&ktb)
20997 .arg(&vtb);
20998 unsafe {
20999 b.launch(cfg)?;
21000 }
21001 }
21002 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21010 .map(|v| v != "0")
21011 .unwrap_or(true);
21012 {
21013 let hd_sfx = fa_hd_suffix(head_dim)?;
21014 let f = self.func(&format!(
21015 "fa_prefill_qw{}{hd_sfx}",
21016 if db { "_db" } else { "" }
21017 ));
21018 let shmem = if db {
21019 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21021 } else {
21022 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21023 };
21024 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21025 f.set_attribute(
21026 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21027 shmem as i32,
21028 )?;
21029 let cfg = LaunchConfig {
21030 grid_dim: (
21031 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21032 n_head as u32,
21033 1,
21034 ),
21035 block_dim: (32, 4, 1),
21036 shared_mem_bytes: shmem,
21037 };
21038 let (hd, nh, nhkv, ti, tkvi, cz) = (
21039 head_dim as i32,
21040 n_head as i32,
21041 n_head_kv as i32,
21042 t as i32,
21043 t_kv as i32,
21044 causal as i32,
21045 );
21046 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21047 let __s_b = self.gpu.stream();
21048 let mut b = __s_b.launch_builder(&f);
21049 b.arg(q)
21050 .arg(&*kw)
21051 .arg(&*vw)
21052 .arg(o)
21053 .arg(&hd)
21054 .arg(&nh)
21055 .arg(&nhkv)
21056 .arg(&ti)
21057 .arg(&tkvi)
21058 .arg(&scale)
21059 .arg(&cz)
21060 .arg(&kdk)
21061 .arg(&kdv);
21062 unsafe {
21063 b.launch(cfg)?;
21064 }
21065 }
21066 Ok(())
21067 }
21068
21069 #[allow(clippy::too_many_arguments)]
21085 pub fn fa_prefill_view_ws_w_hd128(
21086 &self,
21087 q: &CudaSlice<f32>,
21088 k: &cudarc::driver::CudaView<u8>,
21089 v: &cudarc::driver::CudaView<u8>,
21090 o: &mut CudaSlice<f32>,
21091 head_dim: usize,
21092 n_head: usize,
21093 n_head_kv: usize,
21094 t: usize,
21095 t_kv: usize,
21096 scale: f32,
21097 causal: bool,
21098 window: usize,
21099 k_tok_bytes: usize,
21100 v_tok_bytes: usize,
21101 ) -> Result<(), Box<dyn std::error::Error>> {
21102 assert_eq!(
21103 head_dim, 128,
21104 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21105 );
21106 if portable_mma_gated() {
21107 return self.sdpa_naive_w_quantized_view(
21108 q,
21109 k,
21110 v,
21111 o,
21112 head_dim,
21113 n_head,
21114 n_head_kv,
21115 t,
21116 t_kv,
21117 scale,
21118 causal,
21119 window,
21120 k_tok_bytes,
21121 v_tok_bytes,
21122 );
21123 }
21124 const BLOCK_Q: usize = 64;
21125 const BK: usize = 32;
21126 let kv_dim_k = n_head_kv * head_dim;
21127 let kv_dim_v = n_head_kv * head_dim;
21128 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21130 let mut guard = self.prime_deqw_ws.lock().unwrap();
21131 let need_grow = match guard.as_ref() {
21132 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21133 None => true,
21134 };
21135 if need_grow {
21136 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21137 let (ck, cv) = guard
21138 .as_ref()
21139 .map(|(a, b)| (a.len(), b.len()))
21140 .unwrap_or((0, 0));
21141 *guard = Some((
21142 self.alloc_u8(grow(ck, k_ws_bytes))?,
21143 self.alloc_u8(grow(cv, v_ws_bytes))?,
21144 ));
21145 }
21146 let (kw, vw) = guard.as_mut().unwrap();
21147 {
21150 let f = self.func("fa_dequant_kv_ws_bf16");
21151 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21152 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21153 let cfg = LaunchConfig {
21154 grid_dim: (nblk.max(1), 1, 1),
21155 block_dim: (256, 1, 1),
21156 shared_mem_bytes: 0,
21157 };
21158 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21159 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21160 let __s_b = self.gpu.stream();
21161 let mut b = __s_b.launch_builder(&f);
21162 b.arg(k)
21163 .arg(v)
21164 .arg(&mut *kw)
21165 .arg(&mut *vw)
21166 .arg(&kdk)
21167 .arg(&kdv)
21168 .arg(&tkvi)
21169 .arg(&ktb)
21170 .arg(&vtb);
21171 unsafe {
21172 b.launch(cfg)?;
21173 }
21174 }
21175 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21177 .map(|v| v != "0")
21178 .unwrap_or(true);
21179 {
21180 let f = self.func(if db {
21181 "fa_prefill_qw_db_w_hd128"
21182 } else {
21183 "fa_prefill_qw_w_hd128"
21184 });
21185 let shmem = if db {
21186 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21187 } else {
21188 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21189 };
21190 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21191 f.set_attribute(
21192 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21193 shmem as i32,
21194 )?;
21195 let cfg = LaunchConfig {
21196 grid_dim: (
21197 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21198 n_head as u32,
21199 1,
21200 ),
21201 block_dim: (32, 4, 1),
21202 shared_mem_bytes: shmem,
21203 };
21204 let (hd, nh, nhkv, ti, tkvi, cz) = (
21205 head_dim as i32,
21206 n_head as i32,
21207 n_head_kv as i32,
21208 t as i32,
21209 t_kv as i32,
21210 causal as i32,
21211 );
21212 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21213 let __s_b = self.gpu.stream();
21214 let mut b = __s_b.launch_builder(&f);
21215 b.arg(q)
21216 .arg(&*kw)
21217 .arg(&*vw)
21218 .arg(o)
21219 .arg(&hd)
21220 .arg(&nh)
21221 .arg(&nhkv)
21222 .arg(&ti)
21223 .arg(&tkvi)
21224 .arg(&scale)
21225 .arg(&cz)
21226 .arg(&kdk)
21227 .arg(&kdv)
21228 .arg(&wnd);
21229 unsafe {
21230 b.launch(cfg)?;
21231 }
21232 }
21233 Ok(())
21234 }
21235
21236 pub fn fa_decode(
21240 &self,
21241 q: &CudaSlice<f32>,
21242 k: &cudarc::driver::CudaView<u8>,
21243 v: &cudarc::driver::CudaView<u8>,
21244 o: &mut CudaSlice<f32>,
21245 head_dim: usize,
21246 n_head: usize,
21247 n_head_kv: usize,
21248 t_kv: usize,
21249 scale: f32,
21250 k_tok_bytes: usize,
21251 v_tok_bytes: usize,
21252 ) -> Result<(), Box<dyn std::error::Error>> {
21253 self.fa_decode_kvmod(
21254 q,
21255 k,
21256 v,
21257 o,
21258 head_dim,
21259 n_head,
21260 n_head_kv,
21261 t_kv,
21262 scale,
21263 k_tok_bytes,
21264 v_tok_bytes,
21265 false,
21266 )
21267 }
21268
21269 #[allow(clippy::too_many_arguments)]
21273 #[allow(clippy::too_many_arguments)]
21277 #[allow(clippy::too_many_arguments)]
21278 fn fa_decode_scalar_unified(
21279 &self,
21280 q: &cudarc::driver::CudaView<f32>,
21281 k: &cudarc::driver::CudaView<u8>,
21282 v: &cudarc::driver::CudaView<u8>,
21283 o: &mut cudarc::driver::CudaViewMut<f32>,
21284 head_dim: usize,
21285 n_head: usize,
21286 n_head_kv: usize,
21287 t_kv_host: usize,
21288 t_kv_dev: Option<&CudaSlice<i32>>,
21289 scale: f32,
21290 n_splits: usize,
21291 split_keys: usize,
21292 k_tok_bytes: usize,
21293 v_tok_bytes: usize,
21294 g: bool,
21295 part_o: &mut CudaSlice<f32>,
21296 part_m: &mut CudaSlice<f32>,
21297 part_l: &mut CudaSlice<f32>,
21298 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21299 ) -> Result<(), Box<dyn std::error::Error>> {
21300 let f = if g {
21301 self.func_g("fa_decode_f32")
21302 } else {
21303 self.fa_func("fa_decode_f32", head_dim)
21304 };
21305 let cfg = LaunchConfig {
21306 grid_dim: (n_head as u32, n_splits as u32, 1),
21307 block_dim: (head_dim as u32, 1, 1),
21308 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21309 };
21310 let (hd, nh, nhkv, nsp) = (
21311 head_dim as i32,
21312 n_head as i32,
21313 n_head_kv as i32,
21314 n_splits as i32,
21315 );
21316 let (ktb, vtb, tkvi, ski) = (
21317 k_tok_bytes as i64,
21318 v_tok_bytes as i64,
21319 t_kv_host as i32,
21320 split_keys as i32,
21321 );
21322 let __s_b = self.gpu.stream();
21323 let mut b = __s_b.launch_builder(&f);
21324 match t_kv_dev {
21325 Some(d) => {
21326 b.arg(q)
21327 .arg(k)
21328 .arg(v)
21329 .arg(&mut *part_o)
21330 .arg(&mut *part_m)
21331 .arg(&mut *part_l)
21332 .arg(&hd)
21333 .arg(&nh)
21334 .arg(&nhkv)
21335 .arg(&tkvi)
21336 .arg(d)
21337 .arg(&scale)
21338 .arg(&nsp)
21339 .arg(&ski)
21340 .arg(&ktb)
21341 .arg(&vtb);
21342 unsafe {
21343 b.launch(cfg)?;
21344 }
21345 }
21346 None => {
21347 let null: u64 = 0;
21348 b.arg(q)
21349 .arg(k)
21350 .arg(v)
21351 .arg(&mut *part_o)
21352 .arg(&mut *part_m)
21353 .arg(&mut *part_l)
21354 .arg(&hd)
21355 .arg(&nh)
21356 .arg(&nhkv)
21357 .arg(&tkvi)
21358 .arg(&null)
21359 .arg(&scale)
21360 .arg(&nsp)
21361 .arg(&ski)
21362 .arg(&ktb)
21363 .arg(&vtb);
21364 unsafe {
21365 b.launch(cfg)?;
21366 }
21367 }
21368 }
21369 let cfg2 = LaunchConfig {
21370 grid_dim: (n_head as u32, 1, 1),
21371 block_dim: (head_dim as u32, 1, 1),
21372 shared_mem_bytes: 0,
21373 };
21374 if let Some((oq, od)) = q8_out {
21375 let fc = if g {
21377 self.func_g("fa_decode_combine_q8_1")
21378 } else {
21379 self.fa_func("fa_decode_combine_q8_1", head_dim)
21380 };
21381 let __s_b2 = self.gpu.stream();
21382 let mut b2 = __s_b2.launch_builder(&fc);
21383 b2.arg(&*part_o)
21384 .arg(&*part_m)
21385 .arg(&*part_l)
21386 .arg(oq)
21387 .arg(od)
21388 .arg(&hd)
21389 .arg(&nh)
21390 .arg(&nsp);
21391 unsafe {
21392 b2.launch(cfg2)?;
21393 }
21394 return Ok(());
21395 }
21396 let fc = if g {
21397 self.func_g("fa_decode_combine_f32")
21398 } else {
21399 self.fa_func("fa_decode_combine_f32", head_dim)
21400 };
21401 let __s_b2 = self.gpu.stream();
21402 let mut b2 = __s_b2.launch_builder(&fc);
21403 b2.arg(&*part_o)
21404 .arg(&*part_m)
21405 .arg(&*part_l)
21406 .arg(o)
21407 .arg(&hd)
21408 .arg(&nh)
21409 .arg(&nsp);
21410 unsafe {
21411 b2.launch(cfg2)?;
21412 }
21413 Ok(())
21414 }
21415
21416 pub fn fa_decode_kvmod(
21417 &self,
21418 q: &CudaSlice<f32>,
21419 k: &cudarc::driver::CudaView<u8>,
21420 v: &cudarc::driver::CudaView<u8>,
21421 o: &mut CudaSlice<f32>,
21422 head_dim: usize,
21423 n_head: usize,
21424 n_head_kv: usize,
21425 t_kv: usize,
21426 scale: f32,
21427 k_tok_bytes: usize,
21428 v_tok_bytes: usize,
21429 g: bool,
21430 ) -> Result<(), Box<dyn std::error::Error>> {
21431 let q_view = q.as_view();
21432 let mut o_view = o.as_view_mut();
21433 self.fa_decode_kvmod_view(
21434 &q_view,
21435 k,
21436 v,
21437 &mut o_view,
21438 head_dim,
21439 n_head,
21440 n_head_kv,
21441 t_kv,
21442 scale,
21443 k_tok_bytes,
21444 v_tok_bytes,
21445 g,
21446 )
21447 }
21448
21449 #[allow(clippy::too_many_arguments)]
21454 pub fn fa_decode_kvmod_view(
21455 &self,
21456 q: &cudarc::driver::CudaView<f32>,
21457 k: &cudarc::driver::CudaView<u8>,
21458 v: &cudarc::driver::CudaView<u8>,
21459 o: &mut cudarc::driver::CudaViewMut<f32>,
21460 head_dim: usize,
21461 n_head: usize,
21462 n_head_kv: usize,
21463 t_kv: usize,
21464 scale: f32,
21465 k_tok_bytes: usize,
21466 v_tok_bytes: usize,
21467 g: bool,
21468 ) -> Result<(), Box<dyn std::error::Error>> {
21469 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21490 if g && head_dim == 256 && !fa_v4_at(t_kv) {
21494 fa_vec = false;
21495 }
21496 let sp = fa_split_keys(t_kv, n_head_kv);
21497 let n_splits = if fa_vec {
21498 ((t_kv + sp - 1) / sp).max(1)
21499 } else {
21500 ((t_kv + 255) / 256).max(1)
21501 };
21502 let o_len = n_head * n_splits * head_dim;
21503 let ml_len = n_head * n_splits;
21504 let mut part_guard = self.fa_part_pool.lock().unwrap();
21505 if part_guard
21506 .as_ref()
21507 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21508 .unwrap_or(true)
21509 {
21510 let old = part_guard.take();
21521 let (co, cm) = old
21522 .as_ref()
21523 .map(|pp| (pp.0.len(), pp.1.len()))
21524 .unwrap_or((0, 0));
21525 if let Some(old) = old {
21526 self.fa_part_retired.lock().unwrap().push(old);
21527 }
21528 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21529 eprintln!(
21530 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21531 co, o_len, cm, ml_len
21532 );
21533 }
21534 *part_guard = Some((
21535 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21536 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21537 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21538 ));
21539 }
21540 let pg = part_guard.as_mut().unwrap();
21541 self.gpu
21542 .stream()
21543 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21544 self.gpu
21545 .stream()
21546 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21547 self.gpu
21548 .stream()
21549 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21550 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21551 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21552 let (hd, nh, nhkv, tkvi, nsp) = (
21553 head_dim as i32,
21554 n_head as i32,
21555 n_head_kv as i32,
21556 t_kv as i32,
21557 n_splits as i32,
21558 );
21559 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21560 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
21564 let fa512_min = fa512_min_tkv();
21569 let deep = fa_vec
21572 && head_dim == 256
21573 && fa_v4_at(t_kv)
21574 && !g
21575 && fa_deep_at(t_kv)
21576 && !matches!(fa_v4_mode(), "noB3" | "stage");
21577 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21578 let gqa = (n_head / n_head_kv).max(1) as u32;
21581 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21582 (
21583 fv,
21584 LaunchConfig {
21585 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21586 block_dim: (32, gqa, 1),
21587 shared_mem_bytes: 0,
21588 },
21589 )
21590 } else if fa_vec && head_dim <= 256 {
21591 let gqa = (n_head / n_head_kv).max(1) as u32;
21592 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21603 let smem_tkv = *SMEM_TKV.get_or_init(|| {
21604 std::env::var("MEMRA_FA_SMEM_TKV")
21605 .ok()
21606 .and_then(|v| v.parse().ok())
21607 .unwrap_or_else(|| {
21608 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21609 })
21610 });
21611 if fa_v4_at(t_kv) && head_dim == 256 {
21612 let v4name = match fa_v4_mode() {
21616 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
21619 _ => "fa_decode_vec_q_v4",
21620 };
21621 let fv = if g {
21622 self.func_g(v4name)
21623 } else {
21624 self.func(v4name)
21625 };
21626 let shmem = (if deep { 12160 } else { 11520 }
21629 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
21630 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21631 fv.set_attribute(
21632 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21633 shmem as i32,
21634 )?;
21635 (
21636 fv,
21637 LaunchConfig {
21638 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21639 block_dim: (32, gqa, 1),
21640 shared_mem_bytes: shmem,
21641 },
21642 )
21643 } else if fa_v3_active(head_dim) {
21644 let fv = if g {
21647 self.func_g("fa_decode_vec_q_v3")
21648 } else {
21649 self.func("fa_decode_vec_q_v3")
21650 };
21651 let shmem = (32 * head_dim * 2) as u32; (
21653 fv,
21654 LaunchConfig {
21655 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21656 block_dim: (32, gqa, 1),
21657 shared_mem_bytes: shmem,
21658 },
21659 )
21660 } else if fa_v2_on() {
21661 let fv = if g {
21665 self.func_g("fa_decode_vec_q_v2")
21666 } else {
21667 self.func("fa_decode_vec_q_v2")
21668 };
21669 let shmem = (2 * 32 * head_dim * 2) as u32; (
21671 fv,
21672 LaunchConfig {
21673 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21674 block_dim: (32, gqa, 1),
21675 shared_mem_bytes: shmem,
21676 },
21677 )
21678 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
21679 {
21680 let fv = if g {
21684 self.func_g("fa_decode_vec_q_smem")
21685 } else {
21686 self.func("fa_decode_vec_q_smem")
21687 };
21688 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
21690 fv.set_attribute(
21691 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21692 shmem as i32,
21693 )?;
21694 (
21695 fv,
21696 LaunchConfig {
21697 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21698 block_dim: (32, gqa, 1),
21699 shared_mem_bytes: shmem,
21700 },
21701 )
21702 } else {
21703 let fv = if g {
21706 self.func_g("fa_decode_vec_q")
21707 } else {
21708 self.func("fa_decode_vec_q")
21709 };
21710 (
21711 fv,
21712 LaunchConfig {
21713 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21714 block_dim: (32, gqa, 1),
21715 shared_mem_bytes: 0,
21716 },
21717 )
21718 }
21719 } else {
21720 return self.fa_decode_scalar_unified(
21723 q,
21724 k,
21725 v,
21726 o,
21727 head_dim,
21728 n_head,
21729 n_head_kv,
21730 t_kv,
21731 None,
21732 scale,
21733 n_splits,
21734 if fa_vec { sp } else { 256 },
21735 k_tok_bytes,
21736 v_tok_bytes,
21737 g,
21738 part_o,
21739 part_m,
21740 part_l,
21741 None,
21742 );
21743 };
21744 let __s_b = self.gpu.stream();
21745 let mut b = __s_b.launch_builder(&f);
21746 b.arg(q)
21747 .arg(k)
21748 .arg(v)
21749 .arg(&mut *part_o)
21750 .arg(&mut *part_m)
21751 .arg(&mut *part_l)
21752 .arg(&hd)
21753 .arg(&nh)
21754 .arg(&nhkv)
21755 .arg(&tkvi)
21756 .arg(&scale)
21757 .arg(&nsp)
21758 .arg(&ktb)
21759 .arg(&vtb);
21760 unsafe {
21761 b.launch(cfg)?;
21762 }
21763 let (fc, cfg2) = (
21766 if g {
21767 self.func_g("fa_decode_combine_f32")
21768 } else {
21769 self.fa_func("fa_decode_combine_f32", head_dim)
21770 },
21771 LaunchConfig {
21772 grid_dim: (n_head as u32, 1, 1),
21773 block_dim: (head_dim as u32, 1, 1),
21774 shared_mem_bytes: 0,
21775 },
21776 );
21777 let __s_b2 = self.gpu.stream();
21778 let mut b2 = __s_b2.launch_builder(&fc);
21779 b2.arg(&*part_o)
21780 .arg(&*part_m)
21781 .arg(&*part_l)
21782 .arg(o)
21783 .arg(&hd)
21784 .arg(&nh)
21785 .arg(&nsp);
21786 unsafe {
21787 b2.launch(cfg2)?;
21788 }
21789 Ok(())
21790 }
21791
21792 #[allow(clippy::too_many_arguments)]
21803 pub fn fa_decode_batch_seqs_v4(
21804 &self,
21805 q: &CudaSlice<f32>,
21806 kv_ptrs: &cudarc::driver::CudaView<u64>,
21807 pos_seq: &CudaSlice<i32>,
21808 o: &mut CudaSlice<f32>,
21809 head_dim: usize,
21810 n_head: usize,
21811 n_head_kv: usize,
21812 b_n: usize,
21813 t_kv_max: usize,
21814 scale: f32,
21815 split_keys: usize,
21816 k_tok_bytes: usize,
21817 v_tok_bytes: usize,
21818 ) -> Result<(), Box<dyn std::error::Error>> {
21819 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21820 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21821 let o_len = b_n * n_head * n_splits_max * head_dim;
21822 let ml_len = b_n * n_head * n_splits_max;
21823 let mut part_guard = self.fa_part_pool.lock().unwrap();
21824 if part_guard
21825 .as_ref()
21826 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21827 .unwrap_or(true)
21828 {
21829 let old = part_guard.take();
21840 let (co, cm) = old
21841 .as_ref()
21842 .map(|pp| (pp.0.len(), pp.1.len()))
21843 .unwrap_or((0, 0));
21844 if let Some(old) = old {
21845 self.fa_part_retired.lock().unwrap().push(old);
21846 }
21847 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21848 eprintln!(
21849 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21850 co, o_len, cm, ml_len
21851 );
21852 }
21853 *part_guard = Some((
21854 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21855 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21856 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21857 ));
21858 }
21859 let pg = part_guard.as_mut().unwrap();
21860 self.gpu
21861 .stream()
21862 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21863 self.gpu
21864 .stream()
21865 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21866 self.gpu
21867 .stream()
21868 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21869 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21870 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21871 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21872 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21873 let gqa = (n_head / n_head_kv).max(1) as u32;
21874 let f = self.func("fa_decode_vec_q_seqs_v4");
21875 let shmem = (11520 + 32 * head_dim * 2) as u32;
21877 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21878 f.set_attribute(
21879 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21880 shmem as i32,
21881 )?;
21882 let cfg = LaunchConfig {
21883 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21884 block_dim: (32, gqa, 1),
21885 shared_mem_bytes: shmem,
21886 };
21887 {
21888 let __s_b = self.gpu.stream();
21889 let mut b = __s_b.launch_builder(&f);
21890 b.arg(q)
21891 .arg(kv_ptrs)
21892 .arg(pos_seq)
21893 .arg(&mut *part_o)
21894 .arg(&mut *part_m)
21895 .arg(&mut *part_l)
21896 .arg(&hd)
21897 .arg(&nh)
21898 .arg(&nhkv)
21899 .arg(&scale)
21900 .arg(&nspm)
21901 .arg(&spk)
21902 .arg(&ktb)
21903 .arg(&vtb);
21904 unsafe {
21905 b.launch(cfg)?;
21906 }
21907 }
21908 let fc = self.func("fa_decode_combine_seqs");
21909 let cfg2 = LaunchConfig {
21910 grid_dim: (n_head as u32, b_n as u32, 1),
21911 block_dim: (head_dim as u32, 1, 1),
21912 shared_mem_bytes: 0,
21913 };
21914 let __s_b2 = self.gpu.stream();
21915 let mut b2 = __s_b2.launch_builder(&fc);
21916 b2.arg(&*part_o)
21917 .arg(&*part_m)
21918 .arg(&*part_l)
21919 .arg(o)
21920 .arg(&hd)
21921 .arg(&nh)
21922 .arg(pos_seq)
21923 .arg(&nspm)
21924 .arg(&spk);
21925 unsafe {
21926 b2.launch(cfg2)?;
21927 }
21928 Ok(())
21929 }
21930
21931 #[allow(clippy::too_many_arguments)]
21938 pub fn append_kv_quantized_seqs(
21939 &self,
21940 k_rows: &CudaSlice<f32>,
21941 v_rows: &CudaSlice<f32>,
21942 kv_ptrs: &cudarc::driver::CudaView<u64>,
21943 pos_seq: &CudaSlice<i32>,
21944 b_n: usize,
21945 kv_dim_k: usize,
21946 kv_dim_v: usize,
21947 k_tok_bytes: usize,
21948 v_tok_bytes: usize,
21949 ) -> Result<(), Box<dyn std::error::Error>> {
21950 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21951 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21952 let cfg = LaunchConfig {
21953 grid_dim: (nblk, b_n as u32, 1),
21954 block_dim: (32, 1, 1),
21955 shared_mem_bytes: 0,
21956 };
21957 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21958 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21959 let __s_b = self.gpu.stream();
21960 let mut b = __s_b.launch_builder(&f);
21961 b.arg(k_rows)
21962 .arg(v_rows)
21963 .arg(kv_ptrs)
21964 .arg(pos_seq)
21965 .arg(&kdk)
21966 .arg(&kdv)
21967 .arg(&ktb)
21968 .arg(&vtb);
21969 unsafe {
21970 b.launch(cfg)?;
21971 }
21972 Ok(())
21973 }
21974
21975 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21981 std::env::var("MEMRA_NO_FA_VEC").is_err()
21982 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21983 && base_len + 1 >= fa_vec_min_tkv()
21984 && head_dim <= 256
21985 && head_dim % 32 == 0
21986 }
21987
21988 #[allow(clippy::too_many_arguments)]
21997 pub fn fa_decode_rows(
21998 &self,
21999 q: &CudaSlice<f32>,
22000 k: &cudarc::driver::CudaView<u8>,
22001 v: &cudarc::driver::CudaView<u8>,
22002 o: &mut CudaSlice<f32>,
22003 head_dim: usize,
22004 n_head: usize,
22005 n_head_kv: usize,
22006 base_len: usize,
22007 t: usize,
22008 scale: f32,
22009 k_tok_bytes: usize,
22010 v_tok_bytes: usize,
22011 base_dev: Option<(&CudaSlice<i32>, i32)>,
22015 kv_shared: bool,
22018 g: bool,
22022 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22025 ) -> Result<(), Box<dyn std::error::Error>> {
22026 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22027 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
22034 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22035 let v = *SP512.get_or_init(|| {
22038 std::env::var("MEMRA_FA_SP512")
22039 .ok()
22040 .and_then(|x| x.parse().ok())
22041 .unwrap_or(0)
22042 });
22043 sp = if v >= 8 {
22044 v
22045 } else {
22046 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22047 };
22048 }
22049 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22050 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22051 let gqa = (n_head / n_head_kv).max(1) as u32;
22052 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22063 groups.push((0, t, sp));
22064 } else {
22065 let mut r0 = 0usize;
22066 while r0 < t {
22067 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22068 let mut r1 = r0 + 1;
22069 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22070 r1 += 1;
22071 }
22072 groups.push((r0, r1 - r0, sp_g));
22073 r0 = r1;
22074 }
22075 }
22076 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22080 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22081 std::env::var("MEMRA_FA_SMEM_TKV")
22082 .ok()
22083 .and_then(|v| v.parse().ok())
22084 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22085 });
22086 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22087 let v3 = fa_v3_active(head_dim);
22088 let smem_rows =
22089 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22090 let _ = kv_shared;
22095 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22098 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22112 let tb512 = head_dim == 512
22114 && sp <= 32
22115 && n_head / n_head_kv.max(1) <= 16
22116 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22117 let fname = if tb512 {
22118 "fa_decode_vec_q_rows_v4_512_tb"
22119 } else if i2 {
22120 "fa_decode_vec_q_rows_dpl16_i2"
22121 } else if head_dim == 512 {
22122 "fa_decode_vec_q_rows_dpl16"
22123 }
22124 else if v4 {
22126 "fa_decode_vec_q_rows_v4"
22127 } else if v3 {
22128 "fa_decode_vec_q_rows_v3"
22129 } else if fa_v2_on() {
22130 "fa_decode_vec_q_rows_v2"
22131 } else if smem_rows {
22132 "fa_decode_vec_q_rows_smem"
22133 } else {
22134 "fa_decode_vec_q_rows"
22135 };
22136 let f = if head_dim == 512 {
22137 self.fa_func(fname, head_dim)
22138 } else if g {
22139 self.func_g(if smem_rows {
22147 "fa_decode_vec_q_rows"
22148 } else {
22149 fname
22150 })
22151 } else {
22152 self.func(fname)
22153 };
22154 let shmem = if tb512 {
22155 let gk = Self::gkv_on();
22157 let sh =
22158 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22159 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22160 f.set_attribute(
22161 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22162 sh as i32,
22163 )?;
22164 sh
22165 } else if v4 || v3 || smem_rows || fa_v2_on() {
22166 let sh = (if v4 {
22168 11520 + 32 * head_dim * if g { 1 } else { 2 }
22169 } else if v3 {
22170 32 * head_dim * 2
22171 } else {
22172 2 * 32 * head_dim * 2
22173 }) as u32;
22174 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22175 f.set_attribute(
22176 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22177 sh as i32,
22178 )?;
22179 sh
22180 } else {
22181 0
22182 };
22183 for &(r0, t_g, sp_g) in &groups {
22187 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22188 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22189 let base_i = (base_len + r0) as i32;
22190 let o_len = t_g * n_head * n_splits_g * head_dim;
22191 let ml_len = t_g * n_head * n_splits_g;
22192 let mut part_guard = self.fa_part_pool.lock().unwrap();
22193 if part_guard
22194 .as_ref()
22195 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22196 .unwrap_or(true)
22197 {
22198 let old = part_guard.take();
22209 let (co, cm) = old
22210 .as_ref()
22211 .map(|pp| (pp.0.len(), pp.1.len()))
22212 .unwrap_or((0, 0));
22213 if let Some(old) = old {
22214 self.fa_part_retired.lock().unwrap().push(old);
22215 }
22216 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22217 eprintln!(
22218 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22219 co, o_len, cm, ml_len
22220 );
22221 }
22222 *part_guard = Some((
22223 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22224 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22225 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22226 ));
22227 }
22228 let pg = part_guard.as_mut().unwrap();
22229 self.gpu
22230 .stream()
22231 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22232 self.gpu
22233 .stream()
22234 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22235 self.gpu
22236 .stream()
22237 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22238 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22239 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22240 let qv = self.view(q, t * n_head * head_dim);
22241 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22242 let cfg = LaunchConfig {
22243 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22244 block_dim: (32, gqa, 1),
22245 shared_mem_bytes: shmem,
22246 };
22247 {
22248 let __s_b = self.gpu.stream();
22249 let mut b = __s_b.launch_builder(&f);
22250 if tb512 {
22251 let (bd, plus) =
22253 base_dev.expect("hd512 rows twin requires a device base counter");
22254 let plus_g = plus + r0 as i32;
22255 let nr = t_g as i32;
22256 if Self::pdl_on() && Self::pdl_wb_on() {
22257 use cudarc::driver::{DevicePtr, DevicePtrMut};
22259 let s = &self.gpu.stream();
22260 let (pq, _b0) = q_g.device_ptr(s);
22261 let (pk, _b1) = k.device_ptr(s);
22262 let (pv, _b2) = v.device_ptr(s);
22263 let (po, _b3) = part_o.device_ptr_mut(s);
22264 let (pm, _b4) = part_m.device_ptr_mut(s);
22265 let (pl, _b5) = part_l.device_ptr_mut(s);
22266 let (pb, _b6) = bd.device_ptr(s);
22267 let mut ps = [
22268 &pq as *const _ as *mut std::ffi::c_void,
22269 &pk as *const _ as *mut _,
22270 &pv as *const _ as *mut _,
22271 &po as *const _ as *mut _,
22272 &pm as *const _ as *mut _,
22273 &pl as *const _ as *mut _,
22274 &hd as *const _ as *mut _,
22275 &nh as *const _ as *mut _,
22276 &nhkv as *const _ as *mut _,
22277 &pb as *const _ as *mut _,
22278 &plus_g as *const _ as *mut _,
22279 &scale as *const _ as *mut _,
22280 &nspm as *const _ as *mut _,
22281 &spk as *const _ as *mut _,
22282 &ktb as *const _ as *mut _,
22283 &vtb as *const _ as *mut _,
22284 &nr as *const _ as *mut _,
22285 ];
22286 unsafe {
22287 self.launch_pdl_flash(
22288 Self::gkv_on(),
22289 "fa_decode_vec_q_rows_v4_512_tb",
22290 (n_head_kv as u32, n_splits_g as u32, 1),
22291 (32, gqa, 1),
22292 shmem,
22293 &mut ps,
22294 )?;
22295 }
22296 } else {
22297 let cfg_tb = LaunchConfig {
22298 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22299 block_dim: (32, gqa, 1),
22300 shared_mem_bytes: shmem,
22301 };
22302 b.arg(&q_g)
22303 .arg(k)
22304 .arg(v)
22305 .arg(&mut *part_o)
22306 .arg(&mut *part_m)
22307 .arg(&mut *part_l)
22308 .arg(&hd)
22309 .arg(&nh)
22310 .arg(&nhkv)
22311 .arg(bd)
22312 .arg(&plus_g)
22313 .arg(&scale)
22314 .arg(&nspm)
22315 .arg(&spk)
22316 .arg(&ktb)
22317 .arg(&vtb)
22318 .arg(&nr);
22319 unsafe {
22320 b.launch(cfg_tb)?;
22321 }
22322 }
22323 } else if head_dim == 512 {
22324 let (bd, plus) =
22325 base_dev.expect("hd512 rows twin requires a device base counter");
22326 let plus_g = plus + r0 as i32;
22327 b.arg(&q_g)
22328 .arg(k)
22329 .arg(v)
22330 .arg(&mut *part_o)
22331 .arg(&mut *part_m)
22332 .arg(&mut *part_l)
22333 .arg(&hd)
22334 .arg(&nh)
22335 .arg(&nhkv)
22336 .arg(bd)
22337 .arg(&plus_g)
22338 .arg(&scale)
22339 .arg(&nspm)
22340 .arg(&spk)
22341 .arg(&ktb)
22342 .arg(&vtb);
22343 unsafe {
22344 b.launch(cfg)?;
22345 }
22346 } else {
22347 b.arg(&q_g)
22348 .arg(k)
22349 .arg(v)
22350 .arg(&mut *part_o)
22351 .arg(&mut *part_m)
22352 .arg(&mut *part_l)
22353 .arg(&hd)
22354 .arg(&nh)
22355 .arg(&nhkv)
22356 .arg(&base_i)
22357 .arg(&scale)
22358 .arg(&nspm)
22359 .arg(&spk)
22360 .arg(&ktb)
22361 .arg(&vtb);
22362 unsafe {
22363 b.launch(cfg)?;
22364 }
22365 }
22366 }
22367 let cfg2 = LaunchConfig {
22368 grid_dim: (n_head as u32, t_g as u32, 1),
22369 block_dim: (head_dim as u32, 1, 1),
22370 shared_mem_bytes: 0,
22371 };
22372 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22373 if head_dim == 512 {
22374 let (bd, plus) = base_dev.unwrap();
22377 let plus_g = plus + r0 as i32;
22378 if let Some((oq, od)) = q8_out.as_mut() {
22379 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22381 if Self::pdl_on() && Self::pdl_wb_on() {
22382 use cudarc::driver::{DevicePtr, DevicePtrMut};
22384 let s = &self.gpu.stream();
22385 let (po, _g0) = part_o.device_ptr(s);
22386 let (pm, _g1) = part_m.device_ptr(s);
22387 let (pl, _g2) = part_l.device_ptr(s);
22388 let (pq, _g3) = oq.device_ptr_mut(s);
22389 let (pd, _g4) = od.device_ptr_mut(s);
22390 let (pb, _g5) = bd.device_ptr(s);
22391 let mut ps = [
22392 &po as *const _ as *mut std::ffi::c_void,
22393 &pm as *const _ as *mut _,
22394 &pl as *const _ as *mut _,
22395 &pq as *const _ as *mut _,
22396 &pd as *const _ as *mut _,
22397 &hd as *const _ as *mut _,
22398 &nh as *const _ as *mut _,
22399 &pb as *const _ as *mut _,
22400 &plus_g as *const _ as *mut _,
22401 &nspm as *const _ as *mut _,
22402 &spk as *const _ as *mut _,
22403 ];
22404 unsafe {
22405 self.launch_pdl_flash(
22406 Self::gkv_on(),
22407 "fa_decode_combine_rows_dc_q8_1",
22408 cfg2.grid_dim,
22409 cfg2.block_dim,
22410 0,
22411 &mut ps,
22412 )?;
22413 }
22414 continue;
22415 }
22416 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22417 let __s_b2 = self.gpu.stream();
22418 let mut b2 = __s_b2.launch_builder(&fc);
22419 b2.arg(&*part_o)
22420 .arg(&*part_m)
22421 .arg(&*part_l)
22422 .arg(&mut **oq)
22423 .arg(&mut **od)
22424 .arg(&hd)
22425 .arg(&nh)
22426 .arg(bd)
22427 .arg(&plus_g)
22428 .arg(&nspm)
22429 .arg(&spk);
22430 unsafe {
22431 b2.launch(cfg2)?;
22432 }
22433 continue;
22434 }
22435 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22436 let __s_b2 = self.gpu.stream();
22437 let mut b2 = __s_b2.launch_builder(&fc);
22438 b2.arg(&*part_o)
22439 .arg(&*part_m)
22440 .arg(&*part_l)
22441 .arg(&mut o_g)
22442 .arg(&hd)
22443 .arg(&nh)
22444 .arg(bd)
22445 .arg(&plus_g)
22446 .arg(&nspm)
22447 .arg(&spk);
22448 unsafe {
22449 b2.launch(cfg2)?;
22450 }
22451 } else {
22452 assert!(
22455 q8_out.is_none(),
22456 "rows q8 emit requires the hd512 dc combine"
22457 );
22458 let fc = self.func("fa_decode_combine_rows");
22459 let __s_b2 = self.gpu.stream();
22460 let mut b2 = __s_b2.launch_builder(&fc);
22461 b2.arg(&*part_o)
22462 .arg(&*part_m)
22463 .arg(&*part_l)
22464 .arg(&mut o_g)
22465 .arg(&hd)
22466 .arg(&nh)
22467 .arg(&base_i)
22468 .arg(&nspm)
22469 .arg(&spk);
22470 unsafe {
22471 b2.launch(cfg2)?;
22472 }
22473 }
22474 }
22475 Ok(())
22476 }
22477
22478 #[allow(clippy::too_many_arguments)]
22482 pub fn fa_decode_rows_w(
22483 &self,
22484 q: &CudaSlice<f32>,
22485 k: &cudarc::driver::CudaView<u8>,
22486 v: &cudarc::driver::CudaView<u8>,
22487 o: &mut CudaSlice<f32>,
22488 head_dim: usize,
22489 n_head: usize,
22490 n_head_kv: usize,
22491 base_dev: &CudaSlice<i32>,
22492 base_plus: i32,
22493 t: usize,
22494 scale: f32,
22495 window: usize,
22496 k_tok_bytes: usize,
22497 v_tok_bytes: usize,
22498 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22499 ) -> Result<(), Box<dyn std::error::Error>> {
22500 debug_assert!(head_dim == 256);
22505 let sp = {
22513 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22514 let v = *SPW.get_or_init(|| {
22515 std::env::var("MEMRA_FA_SPW")
22516 .ok()
22517 .and_then(|x| x.parse().ok())
22518 .unwrap_or(0)
22519 });
22520 if v >= 8 {
22521 v
22522 } else {
22523 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22524 }
22525 };
22526 let n_splits_max = (window + sp - 1) / sp;
22527 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22528 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22529 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22530 let gqa = (n_head / n_head_kv).max(1) as u32;
22531 let o_len = t * n_head * n_splits_max * head_dim;
22532 let ml_len = t * n_head * n_splits_max;
22533 let mut part_guard = self.fa_part_pool.lock().unwrap();
22534 if part_guard
22535 .as_ref()
22536 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22537 .unwrap_or(true)
22538 {
22539 let old = part_guard.take();
22550 let (co, cm) = old
22551 .as_ref()
22552 .map(|pp| (pp.0.len(), pp.1.len()))
22553 .unwrap_or((0, 0));
22554 if let Some(old) = old {
22555 self.fa_part_retired.lock().unwrap().push(old);
22556 }
22557 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22558 eprintln!(
22559 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22560 co, o_len, cm, ml_len
22561 );
22562 }
22563 *part_guard = Some((
22564 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22565 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22566 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22567 ));
22568 }
22569 let pg = part_guard.as_mut().unwrap();
22570 self.gpu
22571 .stream()
22572 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22573 self.gpu
22574 .stream()
22575 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22576 self.gpu
22577 .stream()
22578 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22579 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22580 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22586 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22587 std::env::var("MEMRA_FA_SMEM_TKV")
22588 .ok()
22589 .and_then(|v| v.parse().ok())
22590 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22591 });
22592 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22598 let wg = Self::wkv_on();
22603 let sp2 =
22606 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
22607 if sp2 {
22608 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22609 if Self::pdl_on() && Self::pdl_wb_on() {
22610 use cudarc::driver::{DevicePtr, DevicePtrMut};
22612 let s = &self.gpu.stream();
22613 let (pq, _b0) = q.device_ptr(s);
22614 let (pk, _b1) = k.device_ptr(s);
22615 let (pv, _b2) = v.device_ptr(s);
22616 let (po, _b3) = part_o.device_ptr_mut(s);
22617 let (pm, _b4) = part_m.device_ptr_mut(s);
22618 let (pl, _b5) = part_l.device_ptr_mut(s);
22619 let (pb, _b6) = base_dev.device_ptr(s);
22620 let mut ps = [
22621 &pq as *const _ as *mut std::ffi::c_void,
22622 &pk as *const _ as *mut _,
22623 &pv as *const _ as *mut _,
22624 &po as *const _ as *mut _,
22625 &pm as *const _ as *mut _,
22626 &pl as *const _ as *mut _,
22627 &hd as *const _ as *mut _,
22628 &nh as *const _ as *mut _,
22629 &nhkv as *const _ as *mut _,
22630 &pb as *const _ as *mut _,
22631 &base_plus as *const _ as *mut _,
22632 &scale as *const _ as *mut _,
22633 &nspm as *const _ as *mut _,
22634 &spk as *const _ as *mut _,
22635 &ktb as *const _ as *mut _,
22636 &vtb as *const _ as *mut _,
22637 &wini as *const _ as *mut _,
22638 ];
22639 unsafe {
22640 self.launch_pdl_flash(
22641 wg,
22642 "fa_decode_vec_q_rows_v4_w_sp",
22643 (n_head_kv as u32, n_splits_max as u32, t as u32),
22644 (32, gqa + 1, 1),
22645 sh,
22646 &mut ps,
22647 )?;
22648 }
22649 } else {
22650 let f = if wg {
22651 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
22652 } else {
22653 self.func("fa_decode_vec_q_rows_v4_w_sp")
22654 };
22655 f.set_attribute(
22656 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22657 sh as i32,
22658 )?;
22659 let cfg = LaunchConfig {
22660 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22661 block_dim: (32, gqa + 1, 1),
22662 shared_mem_bytes: sh,
22663 };
22664 let __s_b = self.gpu.stream();
22665 let mut b = __s_b.launch_builder(&f);
22666 b.arg(q)
22667 .arg(k)
22668 .arg(v)
22669 .arg(&mut *part_o)
22670 .arg(&mut *part_m)
22671 .arg(&mut *part_l)
22672 .arg(&hd)
22673 .arg(&nh)
22674 .arg(&nhkv)
22675 .arg(base_dev)
22676 .arg(&base_plus)
22677 .arg(&scale)
22678 .arg(&nspm)
22679 .arg(&spk)
22680 .arg(&ktb)
22681 .arg(&vtb)
22682 .arg(&wini);
22683 unsafe {
22684 b.launch(cfg)?;
22685 }
22686 }
22687 } else {
22688 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
22689 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22691 use cudarc::driver::{DevicePtr, DevicePtrMut};
22692 let s = &self.gpu.stream();
22693 let (pq, _b0) = q.device_ptr(s);
22694 let (pk, _b1) = k.device_ptr(s);
22695 let (pv, _b2) = v.device_ptr(s);
22696 let (po, _b3) = part_o.device_ptr_mut(s);
22697 let (pm, _b4) = part_m.device_ptr_mut(s);
22698 let (pl, _b5) = part_l.device_ptr_mut(s);
22699 let (pb, _b6) = base_dev.device_ptr(s);
22700 let mut ps = [
22701 &pq as *const _ as *mut std::ffi::c_void,
22702 &pk as *const _ as *mut _,
22703 &pv as *const _ as *mut _,
22704 &po as *const _ as *mut _,
22705 &pm as *const _ as *mut _,
22706 &pl as *const _ as *mut _,
22707 &hd as *const _ as *mut _,
22708 &nh as *const _ as *mut _,
22709 &nhkv as *const _ as *mut _,
22710 &pb as *const _ as *mut _,
22711 &base_plus as *const _ as *mut _,
22712 &scale as *const _ as *mut _,
22713 &nspm as *const _ as *mut _,
22714 &spk as *const _ as *mut _,
22715 &ktb as *const _ as *mut _,
22716 &vtb as *const _ as *mut _,
22717 &wini as *const _ as *mut _,
22718 ];
22719 unsafe {
22720 self.launch_pdl_flash(
22721 wg,
22722 "fa_decode_vec_q_rows_v4_w",
22723 (n_head_kv as u32, n_splits_max as u32, t as u32),
22724 (32, gqa, 1),
22725 sh,
22726 &mut ps,
22727 )?;
22728 }
22729 } else {
22730 let pick = |name: &str| {
22731 if wg {
22732 self.func_g(name)
22733 } else {
22734 self.func(name)
22735 }
22736 };
22737 let (f, sh) = if fa_v4_at(window) {
22738 let f = pick("fa_decode_vec_q_rows_v4_w");
22739 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22740 } else if smem_tkv > 0 && window >= smem_tkv {
22741 (
22744 pick("fa_decode_vec_q_rows_smem_w"),
22745 (2 * 32 * head_dim * 2) as u32,
22746 )
22747 } else {
22748 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22749 };
22750 f.set_attribute(
22751 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22752 sh as i32,
22753 )?;
22754 let cfg = LaunchConfig {
22755 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22756 block_dim: (32, gqa, 1),
22757 shared_mem_bytes: sh,
22758 };
22759 let __s_b = self.gpu.stream();
22760 let mut b = __s_b.launch_builder(&f);
22761 b.arg(q)
22762 .arg(k)
22763 .arg(v)
22764 .arg(&mut *part_o)
22765 .arg(&mut *part_m)
22766 .arg(&mut *part_l)
22767 .arg(&hd)
22768 .arg(&nh)
22769 .arg(&nhkv)
22770 .arg(base_dev)
22771 .arg(&base_plus)
22772 .arg(&scale)
22773 .arg(&nspm)
22774 .arg(&spk)
22775 .arg(&ktb)
22776 .arg(&vtb)
22777 .arg(&wini);
22778 unsafe {
22779 b.launch(cfg)?;
22780 }
22781 }
22782 }
22783 let cfg2 = LaunchConfig {
22784 grid_dim: (n_head as u32, t as u32, 1),
22785 block_dim: (head_dim as u32, 1, 1),
22786 shared_mem_bytes: 0,
22787 };
22788 if let Some((oq, od)) = q8_out {
22789 if Self::pdl_on() && Self::pdl_wb_on() {
22792 use cudarc::driver::{DevicePtr, DevicePtrMut};
22794 let s = &self.gpu.stream();
22795 let (po, _g0) = part_o.device_ptr(s);
22796 let (pm, _g1) = part_m.device_ptr(s);
22797 let (pl, _g2) = part_l.device_ptr(s);
22798 let (pq, _g3) = oq.device_ptr_mut(s);
22799 let (pd, _g4) = od.device_ptr_mut(s);
22800 let mut ps = [
22801 &po as *const _ as *mut std::ffi::c_void,
22802 &pm as *const _ as *mut _,
22803 &pl as *const _ as *mut _,
22804 &pq as *const _ as *mut _,
22805 &pd as *const _ as *mut _,
22806 &hd as *const _ as *mut _,
22807 &nh as *const _ as *mut _,
22808 &nspm as *const _ as *mut _,
22809 &spk as *const _ as *mut _,
22810 &wini as *const _ as *mut _,
22811 ];
22812 unsafe {
22813 self.launch_pdl_flash(
22814 wg,
22815 "fa_decode_combine_rows_w_q8_1",
22816 cfg2.grid_dim,
22817 cfg2.block_dim,
22818 0,
22819 &mut ps,
22820 )?;
22821 }
22822 return Ok(());
22823 }
22824 let fc = if wg {
22825 self.func_g("fa_decode_combine_rows_w_q8_1")
22826 } else {
22827 self.func("fa_decode_combine_rows_w_q8_1")
22828 };
22829 let __s_b2 = self.gpu.stream();
22830 let mut b2 = __s_b2.launch_builder(&fc);
22831 b2.arg(&*part_o)
22832 .arg(&*part_m)
22833 .arg(&*part_l)
22834 .arg(oq)
22835 .arg(od)
22836 .arg(&hd)
22837 .arg(&nh)
22838 .arg(&nspm)
22839 .arg(&spk)
22840 .arg(&wini);
22841 unsafe {
22842 b2.launch(cfg2)?;
22843 }
22844 return Ok(());
22845 }
22846 let fc = if wg {
22847 self.func_g("fa_decode_combine_rows_w")
22848 } else {
22849 self.func("fa_decode_combine_rows_w")
22850 };
22851 let __s_b2 = self.gpu.stream();
22852 let mut b2 = __s_b2.launch_builder(&fc);
22853 b2.arg(&*part_o)
22854 .arg(&*part_m)
22855 .arg(&*part_l)
22856 .arg(o)
22857 .arg(&hd)
22858 .arg(&nh)
22859 .arg(&nspm)
22860 .arg(&spk)
22861 .arg(&wini);
22862 unsafe {
22863 b2.launch(cfg2)?;
22864 }
22865 Ok(())
22866 }
22867
22868 #[allow(clippy::too_many_arguments)]
22874 pub fn fa_decode_rows_dc(
22875 &self,
22876 q: &CudaSlice<f32>,
22877 k: &cudarc::driver::CudaView<u8>,
22878 v: &cudarc::driver::CudaView<u8>,
22879 o: &mut CudaSlice<f32>,
22880 head_dim: usize,
22881 n_head: usize,
22882 n_head_kv: usize,
22883 base_dev: &CudaSlice<i32>,
22884 t_kv_upper: usize,
22885 t: usize,
22886 scale: f32,
22887 k_tok_bytes: usize,
22888 v_tok_bytes: usize,
22889 base_plus: i32,
22890 g: bool,
22891 ) -> Result<(), Box<dyn std::error::Error>> {
22892 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22893 assert!(
22894 v4 || fa_v3_active(head_dim),
22895 "stream fa rows requires the v3 or v4 lane"
22896 );
22897 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22898 if v4 {
22899 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22900 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22901 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22902 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22903 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22904 let gqa = (n_head / n_head_kv).max(1) as u32;
22905 let o_len = t * n_head * n_splits_max * head_dim;
22906 let ml_len = t * n_head * n_splits_max;
22907 let mut part_guard = self.fa_part_pool.lock().unwrap();
22908 if part_guard
22909 .as_ref()
22910 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22911 .unwrap_or(true)
22912 {
22913 let old = part_guard.take();
22924 let (co, cm) = old
22925 .as_ref()
22926 .map(|pp| (pp.0.len(), pp.1.len()))
22927 .unwrap_or((0, 0));
22928 if let Some(old) = old {
22929 self.fa_part_retired.lock().unwrap().push(old);
22930 }
22931 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22932 eprintln!(
22933 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22934 co, o_len, cm, ml_len
22935 );
22936 }
22937 *part_guard = Some((
22938 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22939 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22940 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22941 ));
22942 }
22943 let pg = part_guard.as_mut().unwrap();
22944 self.gpu
22945 .stream()
22946 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22947 self.gpu
22948 .stream()
22949 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22950 self.gpu
22951 .stream()
22952 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22953 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22954 let f = if g {
22955 self.func_g("fa_decode_vec_q_rows_v4_dc")
22956 } else {
22957 self.func("fa_decode_vec_q_rows_v4_dc")
22958 };
22959 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22960 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22961 f.set_attribute(
22962 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22963 sh as i32,
22964 )?;
22965 let cfg = LaunchConfig {
22966 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22967 block_dim: (32, gqa, 1),
22968 shared_mem_bytes: sh,
22969 };
22970 let __s_b = self.gpu.stream();
22971 let mut b = __s_b.launch_builder(&f);
22972 b.arg(q)
22973 .arg(k)
22974 .arg(v)
22975 .arg(&mut *part_o)
22976 .arg(&mut *part_m)
22977 .arg(&mut *part_l)
22978 .arg(&hd)
22979 .arg(&nh)
22980 .arg(&nhkv)
22981 .arg(base_dev)
22982 .arg(&base_plus)
22983 .arg(&scale)
22984 .arg(&nspm)
22985 .arg(&spk)
22986 .arg(&ktb)
22987 .arg(&vtb);
22988 unsafe {
22989 b.launch(cfg)?;
22990 }
22991 let fc = self.func("fa_decode_combine_rows_dc");
22992 let cfg2 = LaunchConfig {
22993 grid_dim: (n_head as u32, t as u32, 1),
22994 block_dim: (head_dim as u32, 1, 1),
22995 shared_mem_bytes: 0,
22996 };
22997 let __s_b2 = self.gpu.stream();
22998 let mut b2 = __s_b2.launch_builder(&fc);
22999 b2.arg(&*part_o)
23000 .arg(&*part_m)
23001 .arg(&*part_l)
23002 .arg(o)
23003 .arg(&hd)
23004 .arg(&nh)
23005 .arg(base_dev)
23006 .arg(&base_plus)
23007 .arg(&nspm)
23008 .arg(&spk);
23009 unsafe {
23010 b2.launch(cfg2)?;
23011 }
23012 return Ok(());
23013 }
23014 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23015 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23016 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23017 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23018 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23019 let gqa = (n_head / n_head_kv).max(1) as u32;
23020 let o_len = t * n_head * n_splits_max * head_dim;
23021 let ml_len = t * n_head * n_splits_max;
23022 let mut part_guard = self.fa_part_pool.lock().unwrap();
23023 if part_guard
23024 .as_ref()
23025 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23026 .unwrap_or(true)
23027 {
23028 let old = part_guard.take();
23039 let (co, cm) = old
23040 .as_ref()
23041 .map(|pp| (pp.0.len(), pp.1.len()))
23042 .unwrap_or((0, 0));
23043 if let Some(old) = old {
23044 self.fa_part_retired.lock().unwrap().push(old);
23045 }
23046 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23047 eprintln!(
23048 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23049 co, o_len, cm, ml_len
23050 );
23051 }
23052 *part_guard = Some((
23053 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23054 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23055 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23056 ));
23057 }
23058 let pg = part_guard.as_mut().unwrap();
23059 self.gpu
23060 .stream()
23061 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23062 self.gpu
23063 .stream()
23064 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23065 self.gpu
23066 .stream()
23067 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23068 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23069 let f = self.func("fa_decode_vec_q_rows_v3_dc");
23070 let sh = (32 * head_dim * 2) as u32;
23071 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23072 f.set_attribute(
23073 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23074 sh as i32,
23075 )?;
23076 let cfg = LaunchConfig {
23077 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23078 block_dim: (32, gqa, 1),
23079 shared_mem_bytes: sh,
23080 };
23081 let __s_b = self.gpu.stream();
23082 let mut b = __s_b.launch_builder(&f);
23083 b.arg(q)
23084 .arg(k)
23085 .arg(v)
23086 .arg(&mut *part_o)
23087 .arg(&mut *part_m)
23088 .arg(&mut *part_l)
23089 .arg(&hd)
23090 .arg(&nh)
23091 .arg(&nhkv)
23092 .arg(base_dev)
23093 .arg(&scale)
23094 .arg(&nspm)
23095 .arg(&spk)
23096 .arg(&ktb)
23097 .arg(&vtb);
23098 unsafe {
23099 b.launch(cfg)?;
23100 }
23101 let fc = self.func("fa_decode_combine_rows_dc");
23102 let cfg2 = LaunchConfig {
23103 grid_dim: (n_head as u32, t as u32, 1),
23104 block_dim: (head_dim as u32, 1, 1),
23105 shared_mem_bytes: 0,
23106 };
23107 let plus0 = 0i32;
23108 let __s_b2 = self.gpu.stream();
23109 let mut b2 = __s_b2.launch_builder(&fc);
23110 b2.arg(&*part_o)
23111 .arg(&*part_m)
23112 .arg(&*part_l)
23113 .arg(o)
23114 .arg(&hd)
23115 .arg(&nh)
23116 .arg(base_dev)
23117 .arg(&plus0)
23118 .arg(&nspm)
23119 .arg(&spk);
23120 unsafe {
23121 b2.launch(cfg2)?;
23122 }
23123 Ok(())
23124 }
23125
23126 pub fn fa_decode_dc(
23137 &self,
23138 q: &CudaSlice<f32>,
23139 k: &cudarc::driver::CudaView<u8>,
23140 v: &cudarc::driver::CudaView<u8>,
23141 o: &mut CudaSlice<f32>,
23142 head_dim: usize,
23143 n_head: usize,
23144 n_head_kv: usize,
23145 t_kv_dev: &CudaSlice<i32>,
23146 bucket_max: usize,
23147 scale: f32,
23148 k_tok_bytes: usize,
23149 v_tok_bytes: usize,
23150 g: bool,
23151 ) -> Result<(), Box<dyn std::error::Error>> {
23152 self.fa_decode_dc_q8(
23153 q,
23154 k,
23155 v,
23156 o,
23157 head_dim,
23158 n_head,
23159 n_head_kv,
23160 t_kv_dev,
23161 bucket_max,
23162 scale,
23163 k_tok_bytes,
23164 v_tok_bytes,
23165 g,
23166 None,
23167 )
23168 }
23169
23170 #[allow(clippy::too_many_arguments)]
23173 pub fn fa_decode_dc_q8(
23174 &self,
23175 q: &CudaSlice<f32>,
23176 k: &cudarc::driver::CudaView<u8>,
23177 v: &cudarc::driver::CudaView<u8>,
23178 o: &mut CudaSlice<f32>,
23179 head_dim: usize,
23180 n_head: usize,
23181 n_head_kv: usize,
23182 t_kv_dev: &CudaSlice<i32>,
23183 bucket_max: usize,
23184 scale: f32,
23185 k_tok_bytes: usize,
23186 v_tok_bytes: usize,
23187 g: bool,
23188 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23189 ) -> Result<(), Box<dyn std::error::Error>> {
23190 let mut fa_vec =
23198 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23199 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23200 fa_vec = false;
23201 } let sp = fa_split_keys(bucket_max, n_head_kv);
23203 let n_splits = if fa_vec {
23204 ((bucket_max + sp - 1) / sp).max(1)
23205 } else {
23206 ((bucket_max + 255) / 256).max(1)
23207 };
23208 let o_len = n_head * n_splits * head_dim;
23209 let ml_len = n_head * n_splits;
23210 let mut part_guard = self.fa_part_pool.lock().unwrap();
23211 if part_guard
23212 .as_ref()
23213 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23214 .unwrap_or(true)
23215 {
23216 let old = part_guard.take();
23227 let (co, cm) = old
23228 .as_ref()
23229 .map(|pp| (pp.0.len(), pp.1.len()))
23230 .unwrap_or((0, 0));
23231 if let Some(old) = old {
23232 self.fa_part_retired.lock().unwrap().push(old);
23233 }
23234 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23235 eprintln!(
23236 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23237 co, o_len, cm, ml_len
23238 );
23239 }
23240 *part_guard = Some((
23241 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23242 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23243 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23244 ));
23245 }
23246 let pg = part_guard.as_mut().unwrap();
23247 self.gpu
23248 .stream()
23249 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23250 self.gpu
23251 .stream()
23252 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23253 self.gpu
23254 .stream()
23255 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23256 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23257 let (hd, nh, nhkv, nsp) = (
23258 head_dim as i32,
23259 n_head as i32,
23260 n_head_kv as i32,
23261 n_splits as i32,
23262 );
23263 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23264 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23265 let deep = fa_vec
23268 && head_dim == 256
23269 && fa_v4_at(bucket_max)
23270 && !g
23271 && fa_deep_at(bucket_max)
23272 && !matches!(fa_v4_mode(), "noB3" | "stage");
23273 let (f, cfg) = if fa_vec
23274 && head_dim == 512
23275 && bucket_max >= {
23276 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23277 *FA512_MIN_DC.get_or_init(|| {
23278 std::env::var("MEMRA_FA512_MIN")
23279 .ok()
23280 .and_then(|v| v.parse().ok())
23281 .unwrap_or(512)
23282 })
23283 } {
23284 let gqa = (n_head / n_head_kv).max(1) as u32;
23286 (
23287 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23288 LaunchConfig {
23289 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23290 block_dim: (32, gqa, 1),
23291 shared_mem_bytes: 0,
23292 },
23293 )
23294 } else if fa_vec && head_dim == 512 {
23295 let q_view = q.as_view();
23298 let mut o_view = o.as_view_mut();
23299 return self.fa_decode_scalar_unified(
23300 &q_view,
23301 k,
23302 v,
23303 &mut o_view,
23304 head_dim,
23305 n_head,
23306 n_head_kv,
23307 0,
23308 Some(t_kv_dev),
23309 scale,
23310 n_splits,
23311 sp,
23312 k_tok_bytes,
23313 v_tok_bytes,
23314 g,
23315 &mut *part_o,
23316 &mut *part_m,
23317 &mut *part_l,
23318 q8_out,
23319 );
23320 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23321 let gqa = (n_head / n_head_kv).max(1) as u32;
23324 let fv = if g {
23325 self.func_g("fa_decode_vec_q_v4_dc")
23326 } else if deep {
23327 self.func("fa_decode_vec_q_v4_deep_dc")
23328 } else {
23329 self.func("fa_decode_vec_q_v4_dc")
23330 };
23331 let shmem =
23332 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23333 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23334 fv.set_attribute(
23335 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23336 shmem as i32,
23337 )?;
23338 (
23339 fv,
23340 LaunchConfig {
23341 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23342 block_dim: (32, gqa, 1),
23343 shared_mem_bytes: shmem,
23344 },
23345 )
23346 } else if fa_vec && fa_v3_active(head_dim) {
23347 let gqa = (n_head / n_head_kv).max(1) as u32;
23350 let fv = if g {
23351 self.func_g("fa_decode_vec_q_v3_dc")
23352 } else {
23353 self.func("fa_decode_vec_q_v3_dc")
23354 };
23355 let shmem = (32 * head_dim * 2) as u32; (
23357 fv,
23358 LaunchConfig {
23359 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23360 block_dim: (32, gqa, 1),
23361 shared_mem_bytes: shmem,
23362 },
23363 )
23364 } else if fa_vec && fa_v2_on() {
23365 let gqa = (n_head / n_head_kv).max(1) as u32;
23369 let fv = if g {
23370 self.func_g("fa_decode_vec_q_v2_dc")
23371 } else {
23372 self.func("fa_decode_vec_q_v2_dc")
23373 };
23374 let shmem = (2 * 32 * head_dim * 2) as u32; (
23376 fv,
23377 LaunchConfig {
23378 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23379 block_dim: (32, gqa, 1),
23380 shared_mem_bytes: shmem,
23381 },
23382 )
23383 } else if fa_vec {
23384 let gqa = (n_head / n_head_kv).max(1) as u32;
23385 let fv = if g {
23387 self.func_g("fa_decode_vec_q_dc")
23388 } else {
23389 self.func("fa_decode_vec_q_dc")
23390 };
23391 (
23392 fv,
23393 LaunchConfig {
23394 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23395 block_dim: (32, gqa, 1),
23396 shared_mem_bytes: 0,
23397 },
23398 )
23399 } else {
23400 let q_view = q.as_view();
23401 let mut o_view = o.as_view_mut();
23402 return self.fa_decode_scalar_unified(
23403 &q_view,
23404 k,
23405 v,
23406 &mut o_view,
23407 head_dim,
23408 n_head,
23409 n_head_kv,
23410 0,
23411 Some(t_kv_dev),
23412 scale,
23413 n_splits,
23414 if fa_vec { sp } else { 256 },
23415 k_tok_bytes,
23416 v_tok_bytes,
23417 g,
23418 &mut *part_o,
23419 &mut *part_m,
23420 &mut *part_l,
23421 q8_out,
23422 );
23423 };
23424 let ski = sp as i32; let __s_b = self.gpu.stream();
23426 let mut b = __s_b.launch_builder(&f);
23427 b.arg(q)
23428 .arg(k)
23429 .arg(v)
23430 .arg(&mut *part_o)
23431 .arg(&mut *part_m)
23432 .arg(&mut *part_l)
23433 .arg(&hd)
23434 .arg(&nh)
23435 .arg(&nhkv)
23436 .arg(t_kv_dev)
23437 .arg(&scale)
23438 .arg(&nsp)
23439 .arg(&ski)
23440 .arg(&ktb)
23441 .arg(&vtb);
23442 unsafe {
23443 b.launch(cfg)?;
23444 }
23445 let cfg2 = LaunchConfig {
23446 grid_dim: (n_head as u32, 1, 1),
23447 block_dim: (head_dim as u32, 1, 1),
23448 shared_mem_bytes: 0,
23449 };
23450 if let Some((oq, od)) = q8_out {
23451 let fc = if g {
23452 self.func_g("fa_decode_combine_q8_1")
23453 } else {
23454 self.fa_func("fa_decode_combine_q8_1", head_dim)
23455 };
23456 let __s_b2 = self.gpu.stream();
23457 let mut b2 = __s_b2.launch_builder(&fc);
23458 b2.arg(&*part_o)
23459 .arg(&*part_m)
23460 .arg(&*part_l)
23461 .arg(oq)
23462 .arg(od)
23463 .arg(&hd)
23464 .arg(&nh)
23465 .arg(&nsp);
23466 unsafe {
23467 b2.launch(cfg2)?;
23468 }
23469 return Ok(());
23470 }
23471 let fc = if g {
23472 self.func_g("fa_decode_combine_f32")
23473 } else {
23474 self.fa_func("fa_decode_combine_f32", head_dim)
23475 };
23476 let __s_b2 = self.gpu.stream();
23477 let mut b2 = __s_b2.launch_builder(&fc);
23478 b2.arg(&*part_o)
23479 .arg(&*part_m)
23480 .arg(&*part_l)
23481 .arg(o)
23482 .arg(&hd)
23483 .arg(&nh)
23484 .arg(&nsp);
23485 unsafe {
23486 b2.launch(cfg2)?;
23487 }
23488 Ok(())
23489 }
23490
23491 #[allow(clippy::too_many_arguments)]
23495 pub fn append_kv_quantized_dcw(
23496 &self,
23497 k_row: &CudaSlice<f32>,
23498 v_row: &CudaSlice<f32>,
23499 kc: &mut CudaSlice<u8>,
23500 vc: &mut CudaSlice<u8>,
23501 len_dev: &CudaSlice<i32>,
23502 base_dev: Option<&CudaSlice<i32>>,
23503 kv_dim_k: usize,
23504 kv_dim_v: usize,
23505 k_tok_bytes: usize,
23506 v_tok_bytes: usize,
23507 ) -> Result<(), Box<dyn std::error::Error>> {
23508 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23509 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23510 let cfg = LaunchConfig {
23511 grid_dim: (nblk, 1, 1),
23512 block_dim: (32, 1, 1),
23513 shared_mem_bytes: 0,
23514 };
23515 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23516 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23517 let null: u64 = 0;
23518 let __s_b = self.gpu.stream();
23519 let mut b = __s_b.launch_builder(&f);
23520 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23521 match base_dev {
23522 Some(base) => {
23523 b.arg(base);
23524 }
23525 None => {
23526 b.arg(&null);
23527 }
23528 }
23529 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23530 unsafe {
23531 b.launch(cfg)?;
23532 }
23533 Ok(())
23534 }
23535
23536 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23538 let f = self.func("inc_i32");
23539 let cfg = LaunchConfig {
23540 grid_dim: (1, 1, 1),
23541 block_dim: (1, 1, 1),
23542 shared_mem_bytes: 0,
23543 };
23544 let __s_b = self.gpu.stream();
23545 let mut b = __s_b.launch_builder(&f);
23546 b.arg(counter);
23547 unsafe {
23548 b.launch(cfg)?;
23549 }
23550 Ok(())
23551 }
23552
23553 #[allow(clippy::too_many_arguments)]
23562 fn fa_part_pool_grow(
23566 &self,
23567 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
23568 o_len: usize,
23569 ml_len: usize,
23570 ) -> Result<(), Box<dyn std::error::Error>> {
23571 if part_guard
23572 .as_ref()
23573 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23574 .unwrap_or(true)
23575 {
23576 let old = part_guard.take();
23577 let (co, cm) = old
23578 .as_ref()
23579 .map(|pp| (pp.0.len(), pp.1.len()))
23580 .unwrap_or((0, 0));
23581 if let Some(old) = old {
23582 self.fa_part_retired.lock().unwrap().push(old);
23583 }
23584 *part_guard = Some((
23585 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23586 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23587 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23588 ));
23589 }
23590 Ok(())
23591 }
23592
23593 pub fn fa_dcw_pool_ensure(
23596 &self,
23597 head_dim: usize,
23598 n_head: usize,
23599 n_head_kv: usize,
23600 bucket_max: usize,
23601 ) -> Result<(), Box<dyn std::error::Error>> {
23602 let sp = fa_split_keys(bucket_max, n_head_kv);
23603 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23604 let o_len = n_head * n_splits * head_dim;
23605 let ml_len = n_head * n_splits;
23606 let mut part_guard = self.fa_part_pool.lock().unwrap();
23607 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
23608 }
23609
23610 #[allow(clippy::too_many_arguments)]
23618 pub fn fa_decode_dcw2(
23619 &self,
23620 q2: &CudaSlice<f32>,
23621 k_ring: &cudarc::driver::CudaView<u8>,
23622 v_ring: &cudarc::driver::CudaView<u8>,
23623 o2: &mut CudaSlice<f32>,
23624 head_dim: usize,
23625 n_head: usize,
23626 n_head_kv: usize,
23627 len_dev: &CudaSlice<i32>,
23628 base_dev: Option<&CudaSlice<i32>>,
23629 window: usize,
23630 bucket_max: usize,
23631 scale: f32,
23632 k_tok_bytes: usize,
23633 v_tok_bytes: usize,
23634 gate2: &CudaSlice<f32>,
23635 ) -> Result<(), Box<dyn std::error::Error>> {
23636 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23637 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23638 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
23639 }
23640 let sp = fa_split_keys(bucket_max, n_head_kv);
23641 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23642 let o_len = 2 * n_head * n_splits * head_dim;
23644 let ml_len = 2 * n_head * n_splits;
23645 let mut part_guard = self.fa_part_pool.lock().unwrap();
23646 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23647 let pg = part_guard.as_mut().unwrap();
23648 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23649 let (hd, nh, nhkv, nsp) = (
23650 head_dim as i32,
23651 n_head as i32,
23652 n_head_kv as i32,
23653 n_splits as i32,
23654 );
23655 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23656 let (ski, win) = (sp as i32, window as i32);
23657 let gqa = (n_head / n_head_kv).max(1) as u32;
23658 let smem = (32 * head_dim * 2) as u32;
23659 let f = self.func("fa_decode_vec_q_v3_dcw2");
23660 let cfg = LaunchConfig {
23661 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23662 block_dim: (32, gqa, 1),
23663 shared_mem_bytes: smem,
23664 };
23665 let null: u64 = 0;
23666 {
23667 let __s_b = self.gpu.stream();
23668 let mut b = __s_b.launch_builder(&f);
23669 b.arg(q2)
23670 .arg(k_ring)
23671 .arg(v_ring)
23672 .arg(&mut *part_o)
23673 .arg(&mut *part_m)
23674 .arg(&mut *part_l)
23675 .arg(&hd)
23676 .arg(&nh)
23677 .arg(&nhkv)
23678 .arg(len_dev);
23679 match base_dev {
23680 Some(base) => {
23681 b.arg(base);
23682 }
23683 None => {
23684 b.arg(&null);
23685 }
23686 }
23687 b.arg(&win)
23688 .arg(&scale)
23689 .arg(&nsp)
23690 .arg(&ski)
23691 .arg(&ktb)
23692 .arg(&vtb);
23693 unsafe {
23694 b.launch(cfg)?;
23695 }
23696 }
23697 let fc = {
23701 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23702 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23703 self.func("fa_decode_combine_gate_f32_s")
23704 } else {
23705 self.func("fa_decode_combine_gate_f32")
23706 }
23707 };
23708 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23709 let nh2 = (2 * n_head) as i32;
23710 let cfg2 = LaunchConfig {
23711 grid_dim: ((2 * n_head) as u32, 1, 1),
23712 block_dim: (head_dim as u32, 1, 1),
23713 shared_mem_bytes: if combine_shared {
23714 (2 * n_splits * 4) as u32
23715 } else {
23716 0
23717 },
23718 };
23719 let __s_b2 = self.gpu.stream();
23720 let mut b2 = __s_b2.launch_builder(&fc);
23721 b2.arg(&*part_o)
23722 .arg(&*part_m)
23723 .arg(&*part_l)
23724 .arg(gate2)
23725 .arg(o2)
23726 .arg(&hd)
23727 .arg(&nh2)
23728 .arg(&nsp);
23729 unsafe {
23730 b2.launch(cfg2)?;
23731 }
23732 Ok(())
23733 }
23734
23735 #[allow(clippy::too_many_arguments)]
23744 pub fn fa_decode_dcw_rows(
23745 &self,
23746 q_rows: &CudaSlice<f32>,
23747 tab: &CudaSlice<u64>,
23748 o_rows: &mut CudaSlice<f32>,
23749 t: usize,
23750 head_dim: usize,
23751 n_head: usize,
23752 n_head_kv: usize,
23753 window: usize,
23754 max_ns: usize,
23755 scale: f32,
23756 k_tok_bytes: usize,
23757 v_tok_bytes: usize,
23758 gate_rows: &CudaSlice<f32>,
23759 ) -> Result<(), Box<dyn std::error::Error>> {
23760 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
23761 || head_dim > 256
23762 || head_dim % 32 != 0
23763 || !fa_v3_on()
23764 {
23765 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
23766 }
23767 if fa_sm_count() < 128
23768 || std::env::var("MEMRA_FA_SPLIT").is_ok()
23769 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
23770 || std::env::var("MEMRA_FA_SP16").is_ok()
23771 {
23772 return Err(
23773 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
23774 (or a <128-SM rig) keep the per-row path"
23775 .into(),
23776 );
23777 }
23778 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
23779 return Err("fa_decode_dcw_rows geometry".into());
23780 }
23781 let o_len = t * n_head * max_ns * head_dim;
23782 let ml_len = t * n_head * max_ns;
23783 let mut part_guard = self.fa_part_pool.lock().unwrap();
23784 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23785 let pg = part_guard.as_mut().unwrap();
23786 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23787 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23788 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23789 let (win, mns) = (window as i32, max_ns as i32);
23790 let gqa = (n_head / n_head_kv).max(1) as u32;
23791 let smem = (32 * head_dim * 2) as u32;
23792 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
23793 let cfg = LaunchConfig {
23794 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
23795 block_dim: (32, gqa, 1),
23796 shared_mem_bytes: smem,
23797 };
23798 {
23799 let __s_b = self.gpu.stream();
23800 let mut b = __s_b.launch_builder(&f);
23801 b.arg(q_rows)
23802 .arg(tab)
23803 .arg(&mut *part_o)
23804 .arg(&mut *part_m)
23805 .arg(&mut *part_l)
23806 .arg(&hd)
23807 .arg(&nh)
23808 .arg(&nhkv)
23809 .arg(&win)
23810 .arg(&scale)
23811 .arg(&mns)
23812 .arg(&ktb)
23813 .arg(&vtb);
23814 unsafe {
23815 b.launch(cfg)?;
23816 }
23817 }
23818 let fc = {
23822 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23823 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23824 self.func("fa_decode_combine_gate_f32_s")
23825 } else {
23826 self.func("fa_decode_combine_gate_f32")
23827 }
23828 };
23829 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23830 let nht = (t * n_head) as i32;
23831 let cfg2 = LaunchConfig {
23832 grid_dim: ((t * n_head) as u32, 1, 1),
23833 block_dim: (head_dim as u32, 1, 1),
23834 shared_mem_bytes: if combine_shared {
23835 (2 * max_ns * 4) as u32
23836 } else {
23837 0
23838 },
23839 };
23840 let __s_b2 = self.gpu.stream();
23841 let mut b2 = __s_b2.launch_builder(&fc);
23842 b2.arg(&*part_o)
23843 .arg(&*part_m)
23844 .arg(&*part_l)
23845 .arg(gate_rows)
23846 .arg(o_rows)
23847 .arg(&hd)
23848 .arg(&nht)
23849 .arg(&mns);
23850 unsafe {
23851 b2.launch(cfg2)?;
23852 }
23853 Ok(())
23854 }
23855
23856 pub fn fa_decode_dcw(
23857 &self,
23858 q: &CudaSlice<f32>,
23859 k_ring: &cudarc::driver::CudaView<u8>,
23860 v_ring: &cudarc::driver::CudaView<u8>,
23861 o: &mut CudaSlice<f32>,
23862 head_dim: usize,
23863 n_head: usize,
23864 n_head_kv: usize,
23865 len_dev: &CudaSlice<i32>,
23866 base_dev: Option<&CudaSlice<i32>>,
23867 window: usize,
23868 bucket_max: usize,
23869 scale: f32,
23870 k_tok_bytes: usize,
23871 v_tok_bytes: usize,
23872 fused_gate: Option<&CudaSlice<f32>>,
23876 ) -> Result<(), Box<dyn std::error::Error>> {
23877 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23878 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23879 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"
23880 .into());
23881 }
23882 let sp = fa_split_keys(bucket_max, n_head_kv);
23883 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23884 let o_len = n_head * n_splits * head_dim;
23885 let ml_len = n_head * n_splits;
23886 let mut part_guard = self.fa_part_pool.lock().unwrap();
23887 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23888 let pg = part_guard.as_mut().unwrap();
23889 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23894 let memset_on = *MEMSET_ON
23899 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23900 || crate::tp::token_graph_building();
23901 if memset_on {
23902 self.gpu
23903 .stream()
23904 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23905 self.gpu
23906 .stream()
23907 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23908 self.gpu
23909 .stream()
23910 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23911 }
23912 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23913 let (hd, nh, nhkv, nsp) = (
23914 head_dim as i32,
23915 n_head as i32,
23916 n_head_kv as i32,
23917 n_splits as i32,
23918 );
23919 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23920 let (ski, win) = (sp as i32, window as i32);
23921 let gqa = (n_head / n_head_kv).max(1) as u32;
23922 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23926 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23927 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23928 Ok("2") => 2,
23929 Ok("1") => 1,
23930 _ => 0,
23931 });
23932 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23937 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23938 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23939 std::sync::Mutex::new(None);
23940 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23944 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23945 && (n_head / n_head_kv) % 2 == 0
23946 && (n_head / n_head_kv) >= 2;
23947 let f = if fprof {
23948 self.func("fa_decode_vec_q_v3_dcw_prof")
23949 } else if hs2 {
23950 self.func("fa_decode_vec_q_v3_dcw_hs2")
23951 } else if hoist == 2 {
23952 self.func("fa_decode_vec_q_v3_dcw_hc")
23954 } else if hoist == 1 {
23955 self.func("fa_decode_vec_q_v3_dcw_h")
23957 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23958 self.func("fa_decode_vec_q_v3_dcw_u8")
23959 } else {
23960 self.func("fa_decode_vec_q_v3_dcw")
23961 };
23962 let cfg = LaunchConfig {
23963 grid_dim: if hs2 {
23964 ((2 * n_head_kv) as u32, n_splits as u32, 1)
23965 } else {
23966 (n_head_kv as u32, n_splits as u32, 1)
23967 },
23968 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23969 shared_mem_bytes: smem,
23970 };
23971 let null: u64 = 0;
23972 let __s_b = self.gpu.stream();
23973 let mut b = __s_b.launch_builder(&f);
23974 b.arg(q)
23975 .arg(k_ring)
23976 .arg(v_ring)
23977 .arg(&mut *part_o)
23978 .arg(&mut *part_m)
23979 .arg(&mut *part_l)
23980 .arg(&hd)
23981 .arg(&nh)
23982 .arg(&nhkv)
23983 .arg(len_dev);
23984 match base_dev {
23985 Some(base) => {
23986 b.arg(base);
23987 }
23988 None => {
23989 b.arg(&null);
23990 }
23991 }
23992 b.arg(&win)
23993 .arg(&scale)
23994 .arg(&nsp)
23995 .arg(&ski)
23996 .arg(&ktb)
23997 .arg(&vtb);
23998 if fprof {
23999 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24000 if guard
24001 .as_ref()
24002 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24003 {
24004 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24005 }
24006 let (_, buf) = guard.as_mut().expect("armed above");
24007 b.arg(&*buf);
24008 unsafe {
24009 b.launch(cfg)?;
24010 }
24011 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24012 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24013 if n % 430 == 0 {
24014 self.stream().synchronize()?;
24015 let h = self.dtoh_u64(buf)?;
24016 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24017 let tot: u64 = h[..6].iter().sum();
24018 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24019 for (i, name) in phases.iter().enumerate() {
24020 let pct = if tot > 0 {
24021 h[i] as f64 / tot as f64 * 100.0
24022 } else {
24023 0.0
24024 };
24025 line.push_str(&format!(" {name}={pct:.1}%"));
24026 }
24027 if h[6] > 0 {
24028 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24029 }
24030 eprintln!("{line}");
24031 }
24032 } else {
24033 unsafe {
24034 b.launch(cfg)?;
24035 }
24036 }
24037 let mut combine_shared = false;
24038 let fc = if fused_gate.is_some() {
24039 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24042 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24043 combine_shared = true;
24044 self.func("fa_decode_combine_gate_f32_s")
24045 } else {
24046 self.func("fa_decode_combine_gate_f32")
24047 }
24048 } else {
24049 self.fa_func("fa_decode_combine_f32", head_dim)
24050 };
24051 let cfg2 = LaunchConfig {
24052 grid_dim: (n_head as u32, 1, 1),
24053 block_dim: (head_dim as u32, 1, 1),
24054 shared_mem_bytes: if combine_shared {
24055 (2 * n_splits * 4) as u32
24056 } else {
24057 0
24058 },
24059 };
24060 let __s_b2 = self.gpu.stream();
24061 let mut b2 = __s_b2.launch_builder(&fc);
24062 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24063 if let Some(gate_row) = fused_gate {
24064 b2.arg(gate_row);
24065 }
24066 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24067 unsafe {
24068 b2.launch(cfg2)?;
24069 }
24070 Ok(())
24071 }
24072
24073 pub fn fa_geom_eager(
24079 &self,
24080 t_kv: usize,
24081 head_dim: usize,
24082 n_head_kv: usize,
24083 g: bool,
24084 ) -> (bool, usize) {
24085 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24089 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24095 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24096 if g && head_dim == 256 && !fa_v4_at(t_kv) {
24102 fa_vec = false;
24103 }
24104 let sp = fa_split_keys(t_kv, n_head_kv);
24105 let n_splits = if fa_vec {
24106 ((t_kv + sp - 1) / sp).max(1)
24107 } else {
24108 ((t_kv + 255) / 256).max(1)
24109 };
24110 (fa_vec, n_splits)
24111 }
24112
24113 pub fn fa_bucket_key(
24119 &self,
24120 t_kv: usize,
24121 head_dim: usize,
24122 n_head_kv: usize,
24123 g: bool,
24124 ) -> (bool, usize) {
24125 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24126 }
24127
24128 pub fn capture_graph_retained<F>(
24140 &self,
24141 step: F,
24142 ) -> Result<
24143 (
24144 cudarc::driver::CudaGraph,
24145 Vec<Box<dyn std::any::Any + Send>>,
24146 ),
24147 Box<dyn std::error::Error>,
24148 >
24149 where
24150 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24151 {
24152 use cudarc::driver::sys::CUgraphInstantiate_flags;
24153 self.capture_graph_retained_flags(
24154 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24155 step,
24156 )
24157 }
24158
24159 pub fn capture_graph_retained_flags<F>(
24164 &self,
24165 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24166 mut step: F,
24167 ) -> Result<
24168 (
24169 cudarc::driver::CudaGraph,
24170 Vec<Box<dyn std::any::Any + Send>>,
24171 ),
24172 Box<dyn std::error::Error>,
24173 >
24174 where
24175 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24176 {
24177 use cudarc::driver::sys::CUstreamCaptureMode;
24178 self.capture_keep.lock().unwrap().clear();
24186 let was_tracking = self.gpu.ctx.is_event_tracking();
24187 if was_tracking {
24188 unsafe {
24189 self.gpu.ctx.disable_event_tracking();
24190 }
24191 }
24192 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24193 self.capture_keep_on
24194 .store(true, std::sync::atomic::Ordering::Relaxed);
24195 let w = (|| {
24196 step(self)?;
24197 step(self)
24198 })();
24199 self.capture_keep_on
24200 .store(false, std::sync::atomic::Ordering::Relaxed);
24201 w?;
24202 self.gpu.stream().synchronize()?;
24203 self.gpu
24204 .stream()
24205 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24206 let r = step(self);
24207 let g = self.gpu.stream().end_capture(flags);
24208 r?;
24209 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24210 graph.upload()?;
24211 Ok(graph)
24212 };
24213 let result = run();
24214 self.capture_keep_on
24215 .store(false, std::sync::atomic::Ordering::Relaxed);
24216 if was_tracking {
24217 unsafe {
24218 self.gpu.ctx.enable_event_tracking();
24219 }
24220 }
24221 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24222 Ok((result?, keeper))
24223 }
24224
24225 pub fn capture_graph_retained_nowarm<F>(
24231 &self,
24232 mut step: F,
24233 ) -> Result<
24234 (
24235 cudarc::driver::CudaGraph,
24236 Vec<Box<dyn std::any::Any + Send>>,
24237 ),
24238 Box<dyn std::error::Error>,
24239 >
24240 where
24241 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24242 {
24243 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24244 let was_tracking = self.gpu.ctx.is_event_tracking();
24245 if was_tracking {
24246 unsafe {
24247 self.gpu.ctx.disable_event_tracking();
24248 }
24249 }
24250 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24251 self.gpu.stream().synchronize()?;
24252 self.gpu
24253 .stream()
24254 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24255 let r = step(self);
24256 let g = self.gpu.stream().end_capture(
24257 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24258 );
24259 r?;
24260 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24261 graph.upload()?;
24262 Ok(graph)
24263 };
24264 let result = run();
24265 if was_tracking {
24266 unsafe {
24267 self.gpu.ctx.enable_event_tracking();
24268 }
24269 }
24270 Ok((result?, Vec::new()))
24271 }
24272
24273 pub fn capture_graph<F>(
24274 &self,
24275 mut step: F,
24276 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24277 where
24278 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24279 {
24280 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24281 let was_tracking = self.gpu.ctx.is_event_tracking();
24289 if was_tracking {
24290 unsafe {
24291 self.gpu.ctx.disable_event_tracking();
24292 }
24293 }
24294 let iflag = {
24301 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24302 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24303 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24306 Ok("priority") => {
24307 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24308 }
24309 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24310 })
24311 };
24312 let ct = {
24319 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24320 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24321 };
24322 let warmups = {
24345 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24346 *W.get_or_init(|| {
24347 std::env::var("MEMRA_GRAPH_WARMUPS")
24348 .ok()
24349 .and_then(|v| v.parse().ok())
24350 .filter(|n| *n >= 1)
24351 .unwrap_or(1)
24352 })
24353 };
24354 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24355 let t_w = std::time::Instant::now();
24356 for _ in 0..warmups {
24358 step(self)?;
24359 }
24360 self.gpu.stream().synchronize()?;
24361 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24362 let t_c = std::time::Instant::now();
24364 self.gpu
24365 .stream()
24366 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24367 let r = step(self);
24370 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24371 let t_i = std::time::Instant::now();
24372 let g = self.gpu.stream().end_capture(iflag);
24373 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24374 r?;
24375 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24376 let t_u = std::time::Instant::now();
24377 graph.upload()?;
24378 if ct {
24379 println!(
24380 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
24381 instantiate {ms_inst:.2} ms upload {:.2} ms",
24382 t_u.elapsed().as_secs_f64() * 1e3
24383 );
24384 }
24385 Ok(graph)
24386 };
24387 let result = run();
24388 if was_tracking {
24389 unsafe {
24390 self.gpu.ctx.enable_event_tracking();
24391 }
24392 }
24393 result
24394 }
24395
24396 pub fn gdn_scan_s128_view(
24398 &self,
24399 q: &CudaSlice<f32>,
24400 k: &CudaSlice<f32>,
24401 v: &CudaSlice<f32>,
24402 g: &CudaSlice<f32>,
24403 beta: &CudaSlice<f32>,
24404 state_in: &cudarc::driver::CudaView<f32>,
24405 state_out: &mut cudarc::driver::CudaViewMut<f32>,
24406 o: &mut CudaSlice<f32>,
24407 n_head: usize,
24408 t: usize,
24409 scale: f32,
24410 ) -> Result<(), Box<dyn std::error::Error>> {
24411 let f = self.func("gdn_scan_s128");
24412 const S_V: u32 = 128;
24413 const WARP: u32 = 32;
24414 const COLS: u32 = 4;
24415 let cfg = LaunchConfig {
24416 grid_dim: (n_head as u32, 1, S_V / COLS),
24417 block_dim: (WARP, COLS, 1),
24418 shared_mem_bytes: 0,
24419 };
24420 let (h, ti) = (n_head as i32, t as i32);
24421 let __s_b = self.gpu.stream();
24422 let mut b = __s_b.launch_builder(&f);
24423 b.arg(q)
24424 .arg(k)
24425 .arg(v)
24426 .arg(g)
24427 .arg(beta)
24428 .arg(state_in)
24429 .arg(state_out)
24430 .arg(o)
24431 .arg(&h)
24432 .arg(&ti)
24433 .arg(&scale);
24434 unsafe {
24435 b.launch(cfg)?;
24436 }
24437 Ok(())
24438 }
24439
24440 pub fn ssm_conv1d_view(
24442 &self,
24443 x: &cudarc::driver::CudaView<f32>,
24444 w: &CudaSlice<f32>,
24445 y: &mut CudaSlice<f32>,
24446 conv_dim: usize,
24447 t: usize,
24448 d_conv: usize,
24449 silu: bool,
24450 ) -> Result<(), Box<dyn std::error::Error>> {
24451 let f = self.func("ssm_conv1d_silu_f32");
24452 let cfg = LaunchConfig {
24454 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24455 block_dim: (256, 1, 1),
24456 shared_mem_bytes: 0,
24457 };
24458 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24459 let __s_b = self.gpu.stream();
24460 let mut b = __s_b.launch_builder(&f);
24461 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24462 unsafe {
24463 b.launch(cfg)?;
24464 }
24465 Ok(())
24466 }
24467
24468 pub fn ssm_conv1d_tm(
24475 &self,
24476 qkv_tm: &CudaSlice<f32>,
24477 w: &CudaSlice<f32>,
24478 y: &mut CudaSlice<f32>,
24479 conv_dim: usize,
24480 t: usize,
24481 d_conv: usize,
24482 ) -> Result<(), Box<dyn std::error::Error>> {
24483 let f = self.func("ssm_conv1d_tm_f32");
24484 let cfg = LaunchConfig {
24485 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24486 block_dim: (256, 1, 1),
24487 shared_mem_bytes: 0,
24488 };
24489 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24490 let __s_b = self.gpu.stream();
24491 let mut b = __s_b.launch_builder(&f);
24492 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24493 unsafe {
24494 b.launch(cfg)?;
24495 }
24496 Ok(())
24497 }
24498
24499 pub fn ssm_conv1d_tm_state(
24507 &self,
24508 qkv_tm: &CudaSlice<f32>,
24509 conv_state: &mut CudaSlice<f32>,
24510 w: &CudaSlice<f32>,
24511 y: &mut CudaSlice<f32>,
24512 conv_dim: usize,
24513 t: usize,
24514 d_conv: usize,
24515 ) -> Result<(), Box<dyn std::error::Error>> {
24516 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24517 }
24518
24519 #[allow(clippy::too_many_arguments)]
24522 pub fn ssm_conv1d_tm_state_pad(
24523 &self,
24524 qkv_tm: &CudaSlice<f32>,
24525 conv_state: &mut CudaSlice<f32>,
24526 w: &CudaSlice<f32>,
24527 y: &mut CudaSlice<f32>,
24528 conv_dim: usize,
24529 t: usize,
24530 d_conv: usize,
24531 pad_len: Option<&CudaSlice<i32>>,
24532 ) -> Result<(), Box<dyn std::error::Error>> {
24533 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24534 let ring_old = if t < d_conv - 1 {
24538 Some(self.clone_dtod(conv_state)?)
24539 } else {
24540 None
24541 };
24542 {
24543 let f = self.func("ssm_conv1d_tm_state_f32");
24544 let cfg = LaunchConfig {
24545 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24546 block_dim: (256, 1, 1),
24547 shared_mem_bytes: 0,
24548 };
24549 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24550 let __s_b = self.gpu.stream();
24551 let mut b = __s_b.launch_builder(&f);
24552 b.arg(qkv_tm)
24553 .arg(&*conv_state)
24554 .arg(w)
24555 .arg(y)
24556 .arg(&cd)
24557 .arg(&ti)
24558 .arg(&dc);
24559 unsafe {
24560 b.launch(cfg)?;
24561 }
24562 }
24563 match (ring_old, pad_len) {
24564 (None, Some(len_d)) => {
24565 let f = self.func("ssm_conv_ring_update_dev_f32");
24566 let n = conv_dim * (d_conv - 1);
24567 let cfg = LaunchConfig::for_num_elems(n as u32);
24568 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24569 let __s_b = self.gpu.stream();
24570 let mut b = __s_b.launch_builder(&f);
24571 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24572 unsafe {
24573 b.launch(cfg)?;
24574 }
24575 }
24576 (None, None) => {
24577 let f = self.func("ssm_conv_ring_update_f32");
24578 let n = conv_dim * (d_conv - 1);
24579 let cfg = LaunchConfig::for_num_elems(n as u32);
24580 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24581 let __s_b = self.gpu.stream();
24582 let mut b = __s_b.launch_builder(&f);
24583 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24584 unsafe {
24585 b.launch(cfg)?;
24586 }
24587 }
24588 (Some(old), _) => {
24589 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
24590 }
24591 }
24592 Ok(())
24593 }
24594
24595 pub fn ssm_conv1d_tm_state_pad_v(
24597 &self,
24598 qkv_tm: &cudarc::driver::CudaView<f32>,
24599 conv_state: &mut CudaSlice<f32>,
24600 w: &CudaSlice<f32>,
24601 y: &mut CudaSlice<f32>,
24602 conv_dim: usize,
24603 t: usize,
24604 d_conv: usize,
24605 pad_len: Option<&CudaSlice<i32>>,
24606 ) -> Result<(), Box<dyn std::error::Error>> {
24607 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24608 let ring_old = if t < d_conv - 1 {
24612 Some(self.clone_dtod(conv_state)?)
24613 } else {
24614 None
24615 };
24616 {
24617 let f = self.func("ssm_conv1d_tm_state_f32");
24618 let cfg = LaunchConfig {
24619 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24620 block_dim: (256, 1, 1),
24621 shared_mem_bytes: 0,
24622 };
24623 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24624 let __s_b = self.gpu.stream();
24625 let mut b = __s_b.launch_builder(&f);
24626 b.arg(qkv_tm)
24627 .arg(&*conv_state)
24628 .arg(w)
24629 .arg(y)
24630 .arg(&cd)
24631 .arg(&ti)
24632 .arg(&dc);
24633 unsafe {
24634 b.launch(cfg)?;
24635 }
24636 }
24637 match (ring_old, pad_len) {
24638 (None, Some(len_d)) => {
24639 let f = self.func("ssm_conv_ring_update_dev_f32");
24640 let n = conv_dim * (d_conv - 1);
24641 let cfg = LaunchConfig::for_num_elems(n as u32);
24642 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24643 let __s_b = self.gpu.stream();
24644 let mut b = __s_b.launch_builder(&f);
24645 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24646 unsafe {
24647 b.launch(cfg)?;
24648 }
24649 }
24650 (None, None) => {
24651 let f = self.func("ssm_conv_ring_update_f32");
24652 let n = conv_dim * (d_conv - 1);
24653 let cfg = LaunchConfig::for_num_elems(n as u32);
24654 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24655 let __s_b = self.gpu.stream();
24656 let mut b = __s_b.launch_builder(&f);
24657 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24658 unsafe {
24659 b.launch(cfg)?;
24660 }
24661 }
24662 (Some(_), _) => unreachable!(
24663 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
24664 ),
24665 }
24666 Ok(())
24667 }
24668
24669 pub fn ssm_conv_ring_rebuild(
24674 &self,
24675 qkv_tm: &CudaSlice<f32>,
24676 ring_old: &CudaSlice<f32>,
24677 conv_state: &mut CudaSlice<f32>,
24678 conv_dim: usize,
24679 tc: usize,
24680 d_conv: usize,
24681 ) -> Result<(), Box<dyn std::error::Error>> {
24682 let f = self.func("ssm_conv_ring_rebuild_f32");
24683 let n = conv_dim * (d_conv - 1);
24684 let cfg = LaunchConfig::for_num_elems(n as u32);
24685 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
24686 let __s_b = self.gpu.stream();
24687 let mut b = __s_b.launch_builder(&f);
24688 b.arg(qkv_tm)
24689 .arg(ring_old)
24690 .arg(conv_state)
24691 .arg(&cd)
24692 .arg(&ti)
24693 .arg(&dc);
24694 unsafe {
24695 b.launch(cfg)?;
24696 }
24697 Ok(())
24698 }
24699
24700 #[allow(clippy::too_many_arguments)]
24705 pub fn gdn_prep_decode(
24706 &self,
24707 conv_out: &CudaSlice<f32>,
24708 beta_raw: &CudaSlice<f32>,
24709 alpha: &CudaSlice<f32>,
24710 dt_bias: &CudaSlice<f32>,
24711 a: &CudaSlice<f32>,
24712 q_l2: &mut CudaSlice<f32>,
24713 k_l2: &mut CudaSlice<f32>,
24714 v_g: &mut CudaSlice<f32>,
24715 beta: &mut CudaSlice<f32>,
24716 g_log: &mut CudaSlice<f32>,
24717 d_state: usize,
24718 num_v: usize,
24719 num_k: usize,
24720 key_dim: usize,
24721 eps: f32,
24722 ) -> Result<(), Box<dyn std::error::Error>> {
24723 let f = self.func("gdn_prep_decode_f32");
24724 let cfg = LaunchConfig {
24725 grid_dim: (num_v as u32, 1, 1),
24726 block_dim: (32, 4, 1),
24727 shared_mem_bytes: 0,
24728 };
24729 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24730 let __s_b = self.gpu.stream();
24731 let mut b = __s_b.launch_builder(&f);
24732 b.arg(conv_out)
24733 .arg(beta_raw)
24734 .arg(alpha)
24735 .arg(dt_bias)
24736 .arg(a)
24737 .arg(q_l2)
24738 .arg(k_l2)
24739 .arg(v_g)
24740 .arg(beta)
24741 .arg(g_log)
24742 .arg(&ds)
24743 .arg(&nv)
24744 .arg(&nk)
24745 .arg(&kd)
24746 .arg(&eps);
24747 unsafe {
24748 b.launch(cfg)?;
24749 }
24750 Ok(())
24751 }
24752
24753 #[allow(clippy::too_many_arguments)]
24757 pub fn ssm_conv1d_gdn(
24758 &self,
24759 qkv_tm: &CudaSlice<f32>,
24760 w: &CudaSlice<f32>,
24761 q_g: &mut CudaSlice<f32>,
24762 k_g: &mut CudaSlice<f32>,
24763 v_g: &mut CudaSlice<f32>,
24764 conv_dim: usize,
24765 t: usize,
24766 d_conv: usize,
24767 d_state: usize,
24768 num_v: usize,
24769 num_k: usize,
24770 key_dim: usize,
24771 ) -> Result<(), Box<dyn std::error::Error>> {
24772 let f = self.func("ssm_conv1d_gdn_f32");
24773 let cfg = LaunchConfig {
24774 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24775 block_dim: (256, 1, 1),
24776 shared_mem_bytes: 0,
24777 };
24778 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24779 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24780 let __s_b = self.gpu.stream();
24781 let mut b = __s_b.launch_builder(&f);
24782 b.arg(qkv_tm)
24783 .arg(w)
24784 .arg(q_g)
24785 .arg(k_g)
24786 .arg(v_g)
24787 .arg(&cd)
24788 .arg(&ti)
24789 .arg(&dc)
24790 .arg(&ds)
24791 .arg(&nv)
24792 .arg(&nk)
24793 .arg(&kd);
24794 unsafe {
24795 b.launch(cfg)?;
24796 }
24797 Ok(())
24798 }
24799
24800 pub fn ssm_conv1d(
24801 &self,
24802 x: &CudaSlice<f32>,
24803 w: &CudaSlice<f32>,
24804 y: &mut CudaSlice<f32>,
24805 conv_dim: usize,
24806 t: usize,
24807 d_conv: usize,
24808 silu: bool,
24809 ) -> Result<(), Box<dyn std::error::Error>> {
24810 let f = self.func("ssm_conv1d_silu_f32");
24811 let cfg = LaunchConfig {
24812 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24813 block_dim: (256, 1, 1),
24814 shared_mem_bytes: 0,
24815 };
24816 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24817 let __s_b = self.gpu.stream();
24818 let mut b = __s_b.launch_builder(&f);
24819 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24820 unsafe {
24821 b.launch(cfg)?;
24822 }
24823 Ok(())
24824 }
24825
24826 pub fn gdn_scan_s128(
24829 &self,
24830 q: &CudaSlice<f32>,
24831 k: &CudaSlice<f32>,
24832 v: &CudaSlice<f32>,
24833 g: &CudaSlice<f32>,
24834 beta: &CudaSlice<f32>,
24835 state_in: &CudaSlice<f32>,
24836 state_out: &mut CudaSlice<f32>,
24837 o: &mut CudaSlice<f32>,
24838 n_head: usize,
24839 t: usize,
24840 scale: f32,
24841 ) -> Result<(), Box<dyn std::error::Error>> {
24842 let f = self.func("gdn_scan_s128");
24843 const S_V: u32 = 128;
24844 const WARP: u32 = 32;
24845 const COLS_PER_BLOCK: u32 = 4;
24846 let cfg = LaunchConfig {
24847 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
24848 block_dim: (WARP, COLS_PER_BLOCK, 1),
24849 shared_mem_bytes: 0,
24850 };
24851 let (h, ti) = (n_head as i32, t as i32);
24852 let __s_b = self.gpu.stream();
24853 let mut b = __s_b.launch_builder(&f);
24854 b.arg(q)
24855 .arg(k)
24856 .arg(v)
24857 .arg(g)
24858 .arg(beta)
24859 .arg(state_in)
24860 .arg(state_out)
24861 .arg(o)
24862 .arg(&h)
24863 .arg(&ti)
24864 .arg(&scale);
24865 unsafe {
24866 b.launch(cfg)?;
24867 }
24868 Ok(())
24869 }
24870
24871 #[allow(clippy::too_many_arguments)]
24876 pub fn ssm_conv1d_fused_decode_b(
24877 &self,
24878 qkv_cols: &CudaSlice<f32>,
24879 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24880 w: &CudaSlice<f32>,
24881 conv_outs: &mut CudaSlice<f32>,
24882 conv_dim: usize,
24883 d_conv: usize,
24884 b_n: usize,
24885 ) -> Result<(), Box<dyn std::error::Error>> {
24886 let f = self.func("ssm_conv1d_fused_decode_b_f32");
24887 let cfg = LaunchConfig {
24888 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24889 block_dim: (256, 1, 1),
24890 shared_mem_bytes: 0,
24891 };
24892 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24893 let __s_b = self.gpu.stream();
24894 let mut b = __s_b.launch_builder(&f);
24895 b.arg(qkv_cols)
24896 .arg(conv_state_ptrs)
24897 .arg(w)
24898 .arg(conv_outs)
24899 .arg(&cd)
24900 .arg(&dc);
24901 unsafe {
24902 b.launch(cfg)?;
24903 }
24904 Ok(())
24905 }
24906
24907 #[allow(clippy::too_many_arguments)]
24908 pub fn gdn_prep_decode_b(
24909 &self,
24910 conv_outs: &CudaSlice<f32>,
24911 beta_raws: &CudaSlice<f32>,
24912 alphas: &CudaSlice<f32>,
24913 dt_bias: &CudaSlice<f32>,
24914 a: &CudaSlice<f32>,
24915 q_l2: &mut CudaSlice<f32>,
24916 k_l2: &mut CudaSlice<f32>,
24917 v_g: &mut CudaSlice<f32>,
24918 beta: &mut CudaSlice<f32>,
24919 g_log: &mut CudaSlice<f32>,
24920 d_state: usize,
24921 num_v: usize,
24922 num_k: usize,
24923 key_dim: usize,
24924 eps: f32,
24925 conv_dim: usize,
24926 b_n: usize,
24927 ) -> Result<(), Box<dyn std::error::Error>> {
24928 let f = self.func("gdn_prep_decode_b_f32");
24929 let cfg = LaunchConfig {
24930 grid_dim: (num_v as u32, 1, b_n as u32),
24931 block_dim: (32, 4, 1),
24932 shared_mem_bytes: 0,
24933 };
24934 let (ds, nv, nk, kd, cd) = (
24935 d_state as i32,
24936 num_v as i32,
24937 num_k as i32,
24938 key_dim as i32,
24939 conv_dim as i32,
24940 );
24941 let __s_b = self.gpu.stream();
24942 let mut b = __s_b.launch_builder(&f);
24943 b.arg(conv_outs)
24944 .arg(beta_raws)
24945 .arg(alphas)
24946 .arg(dt_bias)
24947 .arg(a)
24948 .arg(q_l2)
24949 .arg(k_l2)
24950 .arg(v_g)
24951 .arg(beta)
24952 .arg(g_log)
24953 .arg(&ds)
24954 .arg(&nv)
24955 .arg(&nk)
24956 .arg(&kd)
24957 .arg(&eps)
24958 .arg(&cd);
24959 unsafe {
24960 b.launch(cfg)?;
24961 }
24962 Ok(())
24963 }
24964
24965 #[allow(clippy::too_many_arguments)]
24966 pub fn gdn_scan_s128_batched(
24967 &self,
24968 q: &CudaSlice<f32>,
24969 k: &CudaSlice<f32>,
24970 v: &CudaSlice<f32>,
24971 g: &CudaSlice<f32>,
24972 beta: &CudaSlice<f32>,
24973 state_in_ptrs: &cudarc::driver::CudaView<u64>,
24974 state_out_ptrs: &cudarc::driver::CudaView<u64>,
24975 o: &mut CudaSlice<f32>,
24976 n_head: usize,
24977 b_n: usize,
24978 scale: f32,
24979 ) -> Result<(), Box<dyn std::error::Error>> {
24980 let f = self.func("gdn_scan_s128_b");
24981 const S_V: u32 = 128;
24982 const WARP: u32 = 32;
24983 const COLS_PER_BLOCK: u32 = 4;
24984 let cfg = LaunchConfig {
24985 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24986 block_dim: (WARP, COLS_PER_BLOCK, 1),
24987 shared_mem_bytes: 0,
24988 };
24989 let h = n_head as i32;
24990 let __s_b = self.gpu.stream();
24991 let mut b = __s_b.launch_builder(&f);
24992 b.arg(q)
24993 .arg(k)
24994 .arg(v)
24995 .arg(g)
24996 .arg(beta)
24997 .arg(state_in_ptrs)
24998 .arg(state_out_ptrs)
24999 .arg(o)
25000 .arg(&h)
25001 .arg(&scale);
25002 unsafe {
25003 b.launch(cfg)?;
25004 }
25005 Ok(())
25006 }
25007
25008 #[allow(clippy::too_many_arguments)]
25014 pub fn ssm_conv1d_fused_decode_b_view(
25015 &self,
25016 qkv_cols: &cudarc::driver::CudaView<f32>,
25017 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25018 w: &CudaSlice<f32>,
25019 conv_outs: &mut CudaSlice<f32>,
25020 conv_dim: usize,
25021 d_conv: usize,
25022 b_n: usize,
25023 ) -> Result<(), Box<dyn std::error::Error>> {
25024 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25025 let cfg = LaunchConfig {
25026 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25027 block_dim: (256, 1, 1),
25028 shared_mem_bytes: 0,
25029 };
25030 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25031 let __s_b = self.gpu.stream();
25032 let mut b = __s_b.launch_builder(&f);
25033 b.arg(qkv_cols)
25034 .arg(conv_state_ptrs)
25035 .arg(w)
25036 .arg(conv_outs)
25037 .arg(&cd)
25038 .arg(&dc);
25039 unsafe {
25040 b.launch(cfg)?;
25041 }
25042 Ok(())
25043 }
25044
25045 #[allow(clippy::too_many_arguments)]
25046 pub fn gdn_prep_decode_b_view(
25047 &self,
25048 conv_outs: &CudaSlice<f32>,
25049 beta_raws: &cudarc::driver::CudaView<f32>,
25050 alphas: &cudarc::driver::CudaView<f32>,
25051 dt_bias: &CudaSlice<f32>,
25052 a: &CudaSlice<f32>,
25053 q_l2: &mut CudaSlice<f32>,
25054 k_l2: &mut CudaSlice<f32>,
25055 v_g: &mut CudaSlice<f32>,
25056 beta: &mut CudaSlice<f32>,
25057 g_log: &mut CudaSlice<f32>,
25058 d_state: usize,
25059 num_v: usize,
25060 num_k: usize,
25061 key_dim: usize,
25062 eps: f32,
25063 conv_dim: usize,
25064 b_n: usize,
25065 ) -> Result<(), Box<dyn std::error::Error>> {
25066 let f = self.func("gdn_prep_decode_b_f32");
25067 let cfg = LaunchConfig {
25068 grid_dim: (num_v as u32, 1, b_n as u32),
25069 block_dim: (32, 4, 1),
25070 shared_mem_bytes: 0,
25071 };
25072 let (ds, nv, nk, kd, cd) = (
25073 d_state as i32,
25074 num_v as i32,
25075 num_k as i32,
25076 key_dim as i32,
25077 conv_dim as i32,
25078 );
25079 let __s_b = self.gpu.stream();
25080 let mut b = __s_b.launch_builder(&f);
25081 b.arg(conv_outs)
25082 .arg(beta_raws)
25083 .arg(alphas)
25084 .arg(dt_bias)
25085 .arg(a)
25086 .arg(q_l2)
25087 .arg(k_l2)
25088 .arg(v_g)
25089 .arg(beta)
25090 .arg(g_log)
25091 .arg(&ds)
25092 .arg(&nv)
25093 .arg(&nk)
25094 .arg(&kd)
25095 .arg(&eps)
25096 .arg(&cd);
25097 unsafe {
25098 b.launch(cfg)?;
25099 }
25100 Ok(())
25101 }
25102
25103 #[allow(clippy::too_many_arguments)]
25104 pub fn gdn_scan_s128_batched_view(
25105 &self,
25106 q: &CudaSlice<f32>,
25107 k: &CudaSlice<f32>,
25108 v: &CudaSlice<f32>,
25109 g: &CudaSlice<f32>,
25110 beta: &CudaSlice<f32>,
25111 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25112 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25113 o: &mut cudarc::driver::CudaViewMut<f32>,
25114 n_head: usize,
25115 b_n: usize,
25116 scale: f32,
25117 ) -> Result<(), Box<dyn std::error::Error>> {
25118 let f = self.func("gdn_scan_s128_b");
25119 const S_V: u32 = 128;
25120 const WARP: u32 = 32;
25121 const COLS_PER_BLOCK: u32 = 4;
25122 let cfg = LaunchConfig {
25123 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25124 block_dim: (WARP, COLS_PER_BLOCK, 1),
25125 shared_mem_bytes: 0,
25126 };
25127 let h = n_head as i32;
25128 let __s_b = self.gpu.stream();
25129 let mut b = __s_b.launch_builder(&f);
25130 b.arg(q)
25131 .arg(k)
25132 .arg(v)
25133 .arg(g)
25134 .arg(beta)
25135 .arg(state_in_ptrs)
25136 .arg(state_out_ptrs)
25137 .arg(o)
25138 .arg(&h)
25139 .arg(&scale);
25140 unsafe {
25141 b.launch(cfg)?;
25142 }
25143 Ok(())
25144 }
25145
25146 pub fn gdn_chunked_enabled() -> bool {
25155 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25156 *E.get_or_init(|| {
25157 std::env::var("MEMRA_GDN_CHUNKED")
25158 .map(|v| v != "0")
25159 .unwrap_or(true)
25160 })
25161 }
25162
25163 pub fn gdn_chunk_size() -> usize {
25168 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25169 *C.get_or_init(|| {
25170 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25171 .ok()
25172 .and_then(|v| v.parse().ok())
25173 .unwrap_or(32);
25174 c.clamp(32, 128) / 32 * 32
25175 })
25176 }
25177
25178 #[allow(clippy::too_many_arguments)]
25183 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25186 #[allow(clippy::too_many_arguments)]
25187 pub fn gdn_chunk_k123(
25188 &self,
25189 q: &CudaSlice<f32>,
25190 k: &CudaSlice<f32>,
25191 v: &CudaSlice<f32>,
25192 g: &CudaSlice<f32>,
25193 beta: &CudaSlice<f32>,
25194 wb16: Option<&mut CudaSlice<u8>>,
25195 n_head: usize,
25196 t: usize,
25197 c: usize,
25198 hk: usize,
25199 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25200 ) -> Result<
25201 (
25202 CudaSlice<f32>,
25203 CudaSlice<f32>,
25204 CudaSlice<f32>,
25205 CudaSlice<f32>,
25206 ),
25207 Box<dyn std::error::Error>,
25208 > {
25209 const D: usize = 128;
25210 let h = n_head;
25211 let nc = (t + c - 1) / c;
25212 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25213 let mut gcum = self.uninit(t * h)?;
25214 let mut a = self.uninit(nc * h * c * c)?;
25215 let mut p = self.uninit(nc * h * c * c)?;
25216 let mut u = self.uninit(nc * h * c * D)?;
25217 let mut w = self.uninit(nc * h * c * D)?;
25218 {
25219 let f = self.func("gdn_chunk_cumgate_f32");
25221 let cfg = LaunchConfig {
25222 grid_dim: (nc as u32, h as u32, 1),
25223 block_dim: (32, 1, 1),
25224 shared_mem_bytes: 0,
25225 };
25226 let __s_b = self.gpu.stream();
25227 let mut b = __s_b.launch_builder(&f);
25228 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25229 unsafe {
25230 b.launch(cfg)?;
25231 }
25232 }
25233 if let Some((qb, kb, pb)) = k2w {
25234 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25237 let f = self.func("gdn_k2_wgmma");
25238 let cfg = LaunchConfig {
25239 grid_dim: (nc as u32, h as u32, 1),
25240 block_dim: (128, 1, 1),
25241 shared_mem_bytes: 0,
25242 };
25243 let hki = hk as i32;
25244 let __s_b = self.gpu.stream();
25245 let mut b = __s_b.launch_builder(&f);
25246 b.arg(qb)
25247 .arg(kb)
25248 .arg(&gcum)
25249 .arg(beta)
25250 .arg(&mut a)
25251 .arg(&mut *pb)
25252 .arg(&hi)
25253 .arg(&ti)
25254 .arg(&ci)
25255 .arg(&hki);
25256 unsafe {
25257 b.launch(cfg)?;
25258 }
25259 } else if c <= 64 && !portable_mma_gated() {
25260 let f = self.func("gdn_chunk_attn_f32");
25262 f.set_attribute(
25263 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25264 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25265 )?;
25266 let jt = ((c + 31) / 32) as u32;
25267 let cfg = LaunchConfig {
25268 grid_dim: (nc as u32, h as u32, jt),
25269 block_dim: (256, 1, 1),
25270 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25271 };
25272 let hki = hk as i32;
25273 let __s_b = self.gpu.stream();
25274 let mut b = __s_b.launch_builder(&f);
25275 b.arg(q)
25276 .arg(k)
25277 .arg(&gcum)
25278 .arg(beta)
25279 .arg(&mut a)
25280 .arg(&mut p)
25281 .arg(&hi)
25282 .arg(&ti)
25283 .arg(&ci)
25284 .arg(&hki);
25285 unsafe {
25286 b.launch(cfg)?;
25287 }
25288 } else {
25289 assert!(
25291 hk == h,
25292 "generic K2 is broadcast-only (de-broadcast rides C==32)"
25293 );
25294 let f = self.func("gdn_chunk_attn_g_f32");
25295 let cfg = LaunchConfig {
25296 grid_dim: (nc as u32, h as u32, 1),
25297 block_dim: (32, 8, 1),
25298 shared_mem_bytes: 0,
25299 };
25300 let __s_b = self.gpu.stream();
25301 let mut b = __s_b.launch_builder(&f);
25302 b.arg(q)
25303 .arg(k)
25304 .arg(&gcum)
25305 .arg(beta)
25306 .arg(&mut a)
25307 .arg(&mut p)
25308 .arg(&hi)
25309 .arg(&ti)
25310 .arg(&ci);
25311 unsafe {
25312 b.launch(cfg)?;
25313 }
25314 }
25315 {
25316 let cfg = LaunchConfig {
25318 grid_dim: (nc as u32, h as u32, 1),
25319 block_dim: (256, 1, 1),
25320 shared_mem_bytes: 0,
25321 };
25322 match c {
25323 32 | 64 => {
25324 let f = self.func(if c == 32 {
25325 "gdn_chunk_solve32_f32"
25326 } else {
25327 "gdn_chunk_solve64_f32"
25328 });
25329 let wb: u64 = match wb16 {
25331 Some(d) => self.addr_u8(d),
25332 None => 0,
25333 };
25334 let hki = hk as i32;
25335 let __s_b = self.gpu.stream();
25336 let mut b = __s_b.launch_builder(&f);
25337 b.arg(v)
25338 .arg(k)
25339 .arg(&a)
25340 .arg(&gcum)
25341 .arg(&mut u)
25342 .arg(&mut w)
25343 .arg(&wb)
25344 .arg(&hi)
25345 .arg(&ti)
25346 .arg(&hki);
25347 unsafe {
25348 b.launch(cfg)?;
25349 }
25350 }
25351 _ => {
25352 assert!(hk == h, "generic K3 is broadcast-only");
25353 let f = self.func("gdn_chunk_solve_f32");
25354 let __s_b = self.gpu.stream();
25355 let mut b = __s_b.launch_builder(&f);
25356 b.arg(v)
25357 .arg(k)
25358 .arg(&a)
25359 .arg(&gcum)
25360 .arg(&mut u)
25361 .arg(&mut w)
25362 .arg(&hi)
25363 .arg(&ti)
25364 .arg(&ci);
25365 unsafe {
25366 b.launch(cfg)?;
25367 }
25368 }
25369 }
25370 }
25371 Ok((gcum, p, u, w))
25372 }
25373
25374 pub fn gdn_db_on() -> bool {
25378 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25379 }
25380
25381 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25390 !portable_mma_gated()
25391 && c == 32
25392 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25393 Ok("1") => true,
25394 Ok("0") => false,
25395 _ => gdn_mma_default_on(),
25396 }
25397 }
25398
25399 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25406 cfg!(memra_hopper_mma)
25407 && self.gdn_mma_enabled(c)
25408 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25409 }
25410
25411 #[allow(clippy::too_many_arguments)]
25416 pub fn ssm_conv1d_gdn_state_pad(
25417 &self,
25418 qkv_tm: &cudarc::driver::CudaView<f32>,
25419 conv_state: &mut CudaSlice<f32>,
25420 w: &CudaSlice<f32>,
25421 q_g: &mut CudaSlice<f32>,
25422 k_g: &mut CudaSlice<f32>,
25423 v_g: &mut CudaSlice<f32>,
25424 conv_dim: usize,
25425 t: usize,
25426 d_conv: usize,
25427 d_state: usize,
25428 num_v: usize,
25429 num_k: usize,
25430 key_dim: usize,
25431 hk: usize,
25432 pad_len: Option<&CudaSlice<i32>>,
25433 ) -> Result<(), Box<dyn std::error::Error>> {
25434 assert!(
25435 t >= d_conv - 1,
25436 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25437 );
25438 {
25439 let f = self.func("ssm_conv1d_gdn_state_f32");
25440 let cfg = LaunchConfig {
25441 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25442 block_dim: (256, 1, 1),
25443 shared_mem_bytes: 0,
25444 };
25445 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25446 let (ds, nv, nk, kd, hki) = (
25447 d_state as i32,
25448 num_v as i32,
25449 num_k as i32,
25450 key_dim as i32,
25451 hk as i32,
25452 );
25453 let __s_b = self.gpu.stream();
25454 let mut b = __s_b.launch_builder(&f);
25455 b.arg(qkv_tm)
25456 .arg(&*conv_state)
25457 .arg(w)
25458 .arg(q_g)
25459 .arg(k_g)
25460 .arg(v_g)
25461 .arg(&cd)
25462 .arg(&ti)
25463 .arg(&dc)
25464 .arg(&ds)
25465 .arg(&nv)
25466 .arg(&nk)
25467 .arg(&kd)
25468 .arg(&hki);
25469 unsafe {
25470 b.launch(cfg)?;
25471 }
25472 }
25473 match pad_len {
25474 Some(len_d) => {
25475 let f = self.func("ssm_conv_ring_update_dev_f32");
25476 let n = conv_dim * (d_conv - 1);
25477 let cfg = LaunchConfig::for_num_elems(n as u32);
25478 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25479 let __s_b = self.gpu.stream();
25480 let mut b = __s_b.launch_builder(&f);
25481 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25482 unsafe {
25483 b.launch(cfg)?;
25484 }
25485 }
25486 None => {
25487 let f = self.func("ssm_conv_ring_update_f32");
25488 let n = conv_dim * (d_conv - 1);
25489 let cfg = LaunchConfig::for_num_elems(n as u32);
25490 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25491 let __s_b = self.gpu.stream();
25492 let mut b = __s_b.launch_builder(&f);
25493 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25494 unsafe {
25495 b.launch(cfg)?;
25496 }
25497 }
25498 }
25499 Ok(())
25500 }
25501
25502 pub fn gdn_chunk_alloc(
25506 &self,
25507 n_head: usize,
25508 t: usize,
25509 c: usize,
25510 hk: usize,
25511 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25512 const D: usize = 128;
25513 assert!(
25514 c == 32,
25515 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25516 );
25517 let h = n_head;
25518 let nc = (t + c - 1) / c;
25519 Ok(GdnChunkBufs {
25520 gcum: self.uninit(t * h)?,
25521 a: self.uninit(nc * h * c * c)?,
25522 p: self.uninit(nc * h * c * c)?,
25523 u: self.uninit(nc * h * c * D)?,
25524 w: self.uninit(nc * h * c * D)?,
25525 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25526 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25527 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25528 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25529 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25530 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25531 o: self.uninit(D * h * t)?,
25532 t,
25533 nc,
25534 })
25535 }
25536
25537 pub fn f32_to_bf16_v(
25539 &self,
25540 x: &cudarc::driver::CudaView<f32>,
25541 dst: &mut CudaSlice<u8>,
25542 n: usize,
25543 ) -> Result<(), Box<dyn std::error::Error>> {
25544 let f = self.func("f32_to_bf16_bulk");
25545 let ni = n as i64;
25546 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25547 let __s_b = self.gpu.stream();
25548 let mut b = __s_b.launch_builder(&f);
25549 b.arg(x).arg(dst).arg(&ni);
25550 unsafe {
25551 b.launch(cfg)?;
25552 }
25553 Ok(())
25554 }
25555
25556 pub fn f32_to_bf16_into(
25558 &self,
25559 x: &CudaSlice<f32>,
25560 dst: &mut CudaSlice<u8>,
25561 n: usize,
25562 ) -> Result<(), Box<dyn std::error::Error>> {
25563 let f = self.func("f32_to_bf16_bulk");
25564 let ni = n as i64;
25565 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25566 let __s_b = self.gpu.stream();
25567 let mut b = __s_b.launch_builder(&f);
25568 b.arg(x).arg(dst).arg(&ni);
25569 unsafe {
25570 b.launch(cfg)?;
25571 }
25572 Ok(())
25573 }
25574
25575 pub fn gdn_chunk_k123_vl8(
25578 &self,
25579 seqs: &[GdnSeqVl],
25580 n_head: usize,
25581 hk: usize,
25582 wq: Option<&GdnWVl8>,
25583 ) -> Result<(), Box<dyn std::error::Error>> {
25584 let b = seqs.len();
25585 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
25586 let mut packed = [GdnSeqVl::default(); 8];
25587 packed[..b].copy_from_slice(seqs);
25588 let v = GdnVl8(packed);
25589 let (hi, ci) = (n_head as i32, 32i32);
25590 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25591 {
25592 let f = self.func("gdn_chunk_cumgate_vl");
25593 let cfg = LaunchConfig {
25594 grid_dim: (max_nc, n_head as u32, b as u32),
25595 block_dim: (32, 1, 1),
25596 shared_mem_bytes: 0,
25597 };
25598 let __s_lb = self.gpu.stream();
25599 let mut lb = __s_lb.launch_builder(&f);
25600 lb.arg(&v).arg(&hi).arg(&ci);
25601 unsafe {
25602 lb.launch(cfg)?;
25603 }
25604 }
25605 let hki = hk as i32;
25606 if let Some(w) = wq {
25607 let f = self.func("gdn_k2_wgmma_vl");
25609 let cfg = LaunchConfig {
25610 grid_dim: (max_nc, n_head as u32, b as u32),
25611 block_dim: (128, 1, 1),
25612 shared_mem_bytes: 0,
25613 };
25614 let __s_lb = self.gpu.stream();
25615 let mut lb = __s_lb.launch_builder(&f);
25616 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
25617 unsafe {
25618 lb.launch(cfg)?;
25619 }
25620 } else {
25621 let f = self.func("gdn_chunk_attn_vl");
25622 f.set_attribute(
25623 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25624 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25625 )?;
25626 let cfg = LaunchConfig {
25627 grid_dim: (max_nc, n_head as u32, b as u32),
25628 block_dim: (256, 1, 1),
25629 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25630 };
25631 let __s_lb = self.gpu.stream();
25632 let mut lb = __s_lb.launch_builder(&f);
25633 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25634 unsafe {
25635 lb.launch(cfg)?;
25636 }
25637 }
25638 {
25639 let f = self.func("gdn_chunk_solve32_vl");
25640 let cfg = LaunchConfig {
25641 grid_dim: (max_nc, n_head as u32, b as u32),
25642 block_dim: (256, 1, 1),
25643 shared_mem_bytes: 0,
25644 };
25645 let __s_lb = self.gpu.stream();
25646 let mut lb = __s_lb.launch_builder(&f);
25647 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25648 unsafe {
25649 lb.launch(cfg)?;
25650 }
25651 }
25652 Ok(())
25653 }
25654
25655 #[allow(clippy::too_many_arguments)]
25659 pub fn gdn_prep_vl8(
25660 &self,
25661 seqs: &[GdnPrepVl],
25662 conv_w: &CudaSlice<f32>,
25663 dt_bias: &CudaSlice<f32>,
25664 a: &CudaSlice<f32>,
25665 conv_dim: usize,
25666 d_conv: usize,
25667 d_state: usize,
25668 num_v: usize,
25669 num_k: usize,
25670 key_dim: usize,
25671 hk: usize,
25672 eps: f32,
25673 ) -> Result<(), Box<dyn std::error::Error>> {
25674 let b = seqs.len();
25675 assert!(b >= 1 && b <= 8);
25676 let mut packed = [GdnPrepVl::default(); 8];
25677 packed[..b].copy_from_slice(seqs);
25678 let v = GdnPrepVl8(packed);
25679 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25680 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
25681 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
25682 assert!(
25683 conv_fuse || hk == num_v,
25684 "de-broadcast requires the fused conv"
25685 );
25686 if conv_fuse {
25687 let f = self.func("ssm_conv1d_gdn_state_vl");
25688 let cfg = LaunchConfig {
25689 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25690 block_dim: (256, 1, 1),
25691 shared_mem_bytes: 0,
25692 };
25693 let (dsi, nvi, nki, kdi, hki) = (
25694 d_state as i32,
25695 num_v as i32,
25696 num_k as i32,
25697 key_dim as i32,
25698 hk as i32,
25699 );
25700 let __s_lb = self.gpu.stream();
25701 let mut lb = __s_lb.launch_builder(&f);
25702 lb.arg(&v)
25703 .arg(conv_w)
25704 .arg(&cdi)
25705 .arg(&dci)
25706 .arg(&dsi)
25707 .arg(&nvi)
25708 .arg(&nki)
25709 .arg(&kdi)
25710 .arg(&hki);
25711 unsafe {
25712 lb.launch(cfg)?;
25713 }
25714 } else {
25715 let f = self.func("ssm_conv1d_tm_state_vl");
25716 let cfg = LaunchConfig {
25717 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25718 block_dim: (256, 1, 1),
25719 shared_mem_bytes: 0,
25720 };
25721 let __s_lb = self.gpu.stream();
25722 let mut lb = __s_lb.launch_builder(&f);
25723 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
25724 unsafe {
25725 lb.launch(cfg)?;
25726 }
25727 }
25728 {
25729 let f = self.func("ssm_conv_ring_update_vl");
25730 let n = (conv_dim * (d_conv - 1)) as u32;
25731 let cfg = LaunchConfig {
25732 grid_dim: (n.div_ceil(256), 1, b as u32),
25733 block_dim: (256, 1, 1),
25734 shared_mem_bytes: 0,
25735 };
25736 let __s_lb = self.gpu.stream();
25737 let mut lb = __s_lb.launch_builder(&f);
25738 lb.arg(&v).arg(&cdi).arg(&dci);
25739 unsafe {
25740 lb.launch(cfg)?;
25741 }
25742 }
25743 if !conv_fuse {
25744 let f = self.func("qkv_to_gdn_repack_vl");
25745 let n = max_t * (num_v * d_state) as u32;
25746 let cfg = LaunchConfig {
25747 grid_dim: (n.div_ceil(256), 1, b as u32),
25748 block_dim: (256, 1, 1),
25749 shared_mem_bytes: 0,
25750 };
25751 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25752 let __s_lb = self.gpu.stream();
25753 let mut lb = __s_lb.launch_builder(&f);
25754 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
25755 unsafe {
25756 lb.launch(cfg)?;
25757 }
25758 }
25759 if Self::l2_v2_on(d_state) {
25760 let f = self.func("gdn_l2_v2_vl");
25761 let cfg = LaunchConfig {
25762 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
25763 block_dim: (256, 1, 1),
25764 shared_mem_bytes: 0,
25765 };
25766 let (dsi, nvi) = (d_state as i32, hk as i32);
25767 let __s_lb = self.gpu.stream();
25768 let mut lb = __s_lb.launch_builder(&f);
25769 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25770 unsafe {
25771 lb.launch(cfg)?;
25772 }
25773 } else {
25774 let f = self.func("gdn_l2_vl");
25775 let cfg = LaunchConfig {
25776 grid_dim: (max_t * hk as u32, 2, b as u32),
25777 block_dim: (256, 1, 1),
25778 shared_mem_bytes: 0,
25779 };
25780 let (dsi, nvi) = (d_state as i32, hk as i32);
25781 let __s_lb = self.gpu.stream();
25782 let mut lb = __s_lb.launch_builder(&f);
25783 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25784 unsafe {
25785 lb.launch(cfg)?;
25786 }
25787 }
25788 {
25789 let f = self.func("gdn_gate_prep_vl");
25790 let n = max_t * num_v as u32;
25791 let cfg = LaunchConfig {
25792 grid_dim: (n.div_ceil(256), 1, b as u32),
25793 block_dim: (256, 1, 1),
25794 shared_mem_bytes: 0,
25795 };
25796 let nvi = num_v as i32;
25797 let __s_lb = self.gpu.stream();
25798 let mut lb = __s_lb.launch_builder(&f);
25799 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
25800 unsafe {
25801 lb.launch(cfg)?;
25802 }
25803 }
25804 Ok(())
25805 }
25806
25807 pub fn gdn_mirror_vl8(
25809 &self,
25810 seqs: &[GdnSeqVl],
25811 n_head: usize,
25812 which: i32,
25813 hk: usize,
25814 ) -> Result<(), Box<dyn std::error::Error>> {
25815 let b = seqs.len();
25816 assert!(b >= 1 && b <= 8);
25817 let mut packed = [GdnSeqVl::default(); 8];
25818 packed[..b].copy_from_slice(seqs);
25819 let v = GdnVl8(packed);
25820 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
25821 let max_n = seqs
25822 .iter()
25823 .map(|s| {
25824 if which == 0 {
25825 s.t as i64 * ept as i64
25826 } else {
25827 s.nc as i64 * ept as i64 * 32
25828 }
25829 })
25830 .max()
25831 .unwrap();
25832 let f = self.func("gdn_mirror_vl");
25833 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25834 let cfg = LaunchConfig {
25835 grid_dim: (blocks, 1, b as u32),
25836 block_dim: (256, 1, 1),
25837 shared_mem_bytes: 0,
25838 };
25839 let __s_lb = self.gpu.stream();
25840 let mut lb = __s_lb.launch_builder(&f);
25841 lb.arg(&v).arg(&ept).arg(&which);
25842 unsafe {
25843 lb.launch(cfg)?;
25844 }
25845 Ok(())
25846 }
25847
25848 pub fn gdn_tail_vl8(
25850 &self,
25851 seqs: &[GdnPrepVl],
25852 norm_w: &CudaSlice<f32>,
25853 d_state: usize,
25854 num_v: usize,
25855 eps: f32,
25856 ) -> Result<(), Box<dyn std::error::Error>> {
25857 let b = seqs.len();
25858 assert!(b >= 1 && b <= 8);
25859 let mut packed = [GdnPrepVl::default(); 8];
25860 packed[..b].copy_from_slice(seqs);
25861 let v = GdnPrepVl8(packed);
25862 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25863 let f = self.func("gated_rmsnorm_f16out_vl");
25864 let cfg = LaunchConfig {
25866 grid_dim: (max_t * num_v as u32, 1, b as u32),
25867 block_dim: (128, 1, 1),
25868 shared_mem_bytes: 0,
25869 };
25870 let (dsi, nvi) = (d_state as i32, num_v as i32);
25871 let __s_lb = self.gpu.stream();
25872 let mut lb = __s_lb.launch_builder(&f);
25873 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25874 unsafe {
25875 lb.launch(cfg)?;
25876 }
25877 Ok(())
25878 }
25879
25880 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25883 use cudarc::driver::DevicePtr;
25884 let s = self.gpu.stream();
25885 let (p, _g) = x.device_ptr(&s);
25886 p as u64
25887 }
25888 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25889 use cudarc::driver::DevicePtrMut;
25890 let s = self.gpu.stream();
25891 let (p, _g) = x.device_ptr_mut(&s);
25892 p as u64
25893 }
25894 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25895 use cudarc::driver::DevicePtr;
25896 let s = self.gpu.stream();
25897 let (p, _g) = x.device_ptr(&s);
25898 p as u64
25899 }
25900 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25901 use cudarc::driver::DevicePtr;
25902 let s = self.gpu.stream();
25903 let (p, _g) = x.device_ptr(&s);
25904 p as u64
25905 }
25906
25907 pub fn gdn_chunk_vl8(
25911 &self,
25912 seqs: &[GdnSeqVl],
25913 n_head: usize,
25914 scale: f32,
25915 hk: usize,
25916 wq: Option<&GdnWVl8>,
25917 ) -> Result<(), Box<dyn std::error::Error>> {
25918 const NSPLIT: u32 = 4;
25919 let b = seqs.len();
25920 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25921 let mut packed = [GdnSeqVl::default(); 8];
25922 packed[..b].copy_from_slice(seqs);
25923 let v = GdnVl8(packed);
25924 let (hi, ci) = (n_head as i32, 32i32);
25925 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25926 let hki = hk as i32;
25927 if let Some(w) = wq {
25928 let f = self.func("gdn_k45_wgmma_vl");
25930 let cfg = LaunchConfig {
25931 grid_dim: (n_head as u32, NSPLIT, b as u32),
25932 block_dim: (256, 1, 1),
25933 shared_mem_bytes: 0,
25934 };
25935 let __s_lb = self.gpu.stream();
25936 let mut lb = __s_lb.launch_builder(&f);
25937 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25938 unsafe {
25939 lb.launch(cfg)?;
25940 }
25941 let _ = max_nc;
25942 return Ok(());
25943 }
25944 {
25945 let f = self.func("gdn_chunk_state_mma_vl");
25946 let cfg = LaunchConfig {
25947 grid_dim: (n_head as u32, NSPLIT, b as u32),
25948 block_dim: (256, 1, 1),
25949 shared_mem_bytes: 0,
25950 };
25951 let __s_lb = self.gpu.stream();
25952 let mut lb = __s_lb.launch_builder(&f);
25953 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25954 unsafe {
25955 lb.launch(cfg)?;
25956 }
25957 }
25958 {
25959 let f = self.func("gdn_chunk_output_mma_vl");
25960 let cfg = LaunchConfig {
25961 grid_dim: (max_nc, n_head as u32, b as u32),
25962 block_dim: (256, 1, 1),
25963 shared_mem_bytes: 0,
25964 };
25965 let __s_lb = self.gpu.stream();
25966 let mut lb = __s_lb.launch_builder(&f);
25967 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25968 unsafe {
25969 lb.launch(cfg)?;
25970 }
25971 }
25972 Ok(())
25973 }
25974 pub fn gdn_scan_chunked(
25975 &self,
25976 q: &CudaSlice<f32>,
25977 k: &CudaSlice<f32>,
25978 v: &CudaSlice<f32>,
25979 g: &CudaSlice<f32>,
25980 beta: &CudaSlice<f32>,
25981 kb16_pre: Option<&CudaSlice<u8>>,
25982 qb16_pre: Option<&CudaSlice<u8>>,
25983 state_in: &CudaSlice<f32>,
25984 state_out: &mut CudaSlice<f32>,
25985 o: &mut CudaSlice<f32>,
25986 n_head: usize,
25987 t: usize,
25988 scale: f32,
25989 c: usize,
25990 hk: usize,
25991 ) -> Result<(), Box<dyn std::error::Error>> {
25992 const D: usize = 128;
25993 const NSPLIT: u32 = 4;
25994 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25995 let h = n_head;
25996 let nc = (t + c - 1) / c;
25997 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25998 let gdn_mma_pre = !portable_mma_gated()
26003 && c == 32
26004 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26005 Ok("1") => true,
26006 Ok("0") => false,
26007 _ => gdn_mma_default_on(),
26008 };
26009 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26010 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26011 } else {
26012 None
26013 };
26014 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26019 && gdn_mma_pre
26020 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26021 let nk = t * hk * D;
26022 let mut kb16_local: Option<CudaSlice<u8>> = None;
26023 if gdn_mma_pre && kb16_pre.is_none() {
26024 let mut kb = self.alloc_u8_uninit(nk * 2)?;
26025 let f = self.func("f32_to_bf16_bulk");
26026 let n2 = nk as i64;
26027 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26028 let __s_b = self.gpu.stream();
26029 let mut b = __s_b.launch_builder(&f);
26030 b.arg(k).arg(&mut kb).arg(&n2);
26031 unsafe {
26032 b.launch(cfg2)?;
26033 }
26034 kb16_local = Some(kb);
26035 }
26036 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26037 if let Some(kb) = kb16_pre {
26038 assert!(kb.len() >= nk * 2, "kb16_pre too small");
26039 }
26040 let mut qb16: Option<CudaSlice<u8>> = None;
26041 let mut pb16: Option<CudaSlice<u8>> = None;
26042 if gdn_wgmma_pre {
26043 if qb16_pre.is_none() {
26046 let mut qb = self.alloc_u8_uninit(nk * 2)?;
26047 let f = self.func("f32_to_bf16_bulk");
26048 let n2 = nk as i64;
26049 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26050 let __s_b = self.gpu.stream();
26051 let mut b = __s_b.launch_builder(&f);
26052 b.arg(q).arg(&mut qb).arg(&n2);
26053 unsafe {
26054 b.launch(cfg2)?;
26055 }
26056 qb16 = Some(qb);
26057 } else if let Some(qb) = qb16_pre {
26058 assert!(qb.len() >= nk * 2, "qb16_pre too small");
26059 }
26060 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26061 }
26062 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26063 let k2w = if gdn_wgmma_pre {
26064 Some((
26065 *qb16_ref0.as_ref().unwrap(),
26066 *kb16_ref0.as_ref().unwrap(),
26067 pb16.as_mut().unwrap(),
26068 ))
26069 } else {
26070 None
26071 };
26072 let (gcum, p, u, w) =
26073 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26074 let _ = &w;
26075 let mut y = self.uninit(nc * h * c * D)?;
26076 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
26090 && c == 32
26091 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26092 Ok("1") => true,
26093 Ok("0") => false,
26094 _ => gdn_mma_default_on(),
26095 };
26096 if gdn_mma {
26097 let wb16 = wb16_pre
26098 .take()
26099 .expect("mma path pre-allocates wb16 (K3 store fold)");
26100 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26101 if gdn_wgmma_pre {
26113 let qb16 = qb16_ref0.unwrap();
26115 let pb16 = pb16.as_ref().unwrap();
26116 {
26117 let f = self.func("gdn_k45_wgmma");
26118 let cfg = LaunchConfig {
26119 grid_dim: (h as u32, 4, 1),
26120 block_dim: (256, 1, 1),
26121 shared_mem_bytes: 0,
26122 };
26123 let hki = hk as i32;
26124 let __s_b = self.gpu.stream();
26125 let mut b = __s_b.launch_builder(&f);
26126 b.arg(kb16_ref)
26127 .arg(&gcum)
26128 .arg(beta)
26129 .arg(&u)
26130 .arg(&wb16)
26131 .arg(qb16)
26132 .arg(pb16)
26133 .arg(o)
26134 .arg(&scale)
26135 .arg(state_in)
26136 .arg(&mut *state_out)
26137 .arg(&hi)
26138 .arg(&ti)
26139 .arg(&ci)
26140 .arg(&hki);
26141 unsafe {
26142 b.launch(cfg)?;
26143 }
26144 }
26145 return Ok(());
26146 }
26147 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26151 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26152 {
26153 let f = self.func("gdn_chunk_state_mma");
26154 let cfg = LaunchConfig {
26155 grid_dim: (h as u32, NSPLIT, 1),
26156 block_dim: (256, 1, 1),
26157 shared_mem_bytes: 0,
26158 };
26159 let hki = hk as i32;
26160 let __s_b = self.gpu.stream();
26161 let mut b = __s_b.launch_builder(&f);
26162 b.arg(kb16_ref)
26163 .arg(&gcum)
26164 .arg(beta)
26165 .arg(&u)
26166 .arg(&wb16)
26167 .arg(&mut y16)
26168 .arg(&mut ssnap16)
26169 .arg(state_in)
26170 .arg(&mut *state_out)
26171 .arg(&hi)
26172 .arg(&ti)
26173 .arg(&ci)
26174 .arg(&hki);
26175 unsafe {
26176 b.launch(cfg)?;
26177 }
26178 }
26179 {
26180 let f = self.func("gdn_chunk_output_mma");
26182 let jt = ((c + 31) / 32) as u32;
26183 let cfg = LaunchConfig {
26184 grid_dim: (nc as u32, h as u32, jt),
26185 block_dim: (256, 1, 1),
26186 shared_mem_bytes: 0,
26187 };
26188 let hki = hk as i32;
26189 let __s_b = self.gpu.stream();
26190 let mut b = __s_b.launch_builder(&f);
26191 b.arg(q)
26192 .arg(&gcum)
26193 .arg(&p)
26194 .arg(&y16)
26195 .arg(&ssnap16)
26196 .arg(o)
26197 .arg(&hi)
26198 .arg(&ti)
26199 .arg(&ci)
26200 .arg(&scale)
26201 .arg(&hki);
26202 unsafe {
26203 b.launch(cfg)?;
26204 }
26205 }
26206 return Ok(());
26207 }
26208 {
26209 let f = self.func("gdn_chunk_state_f32");
26211 let cfg = LaunchConfig {
26212 grid_dim: (h as u32, NSPLIT, 1),
26213 block_dim: (256, 1, 1),
26214 shared_mem_bytes: 0,
26215 };
26216 let __s_b = self.gpu.stream();
26217 let mut b = __s_b.launch_builder(&f);
26218 b.arg(k)
26219 .arg(&gcum)
26220 .arg(beta)
26221 .arg(&u)
26222 .arg(&w)
26223 .arg(&mut y)
26224 .arg(&mut ssnap)
26225 .arg(state_in)
26226 .arg(&mut *state_out)
26227 .arg(&hi)
26228 .arg(&ti)
26229 .arg(&ci);
26230 unsafe {
26231 b.launch(cfg)?;
26232 }
26233 }
26234 {
26235 let f = self.func("gdn_chunk_output_f32");
26237 let jt = ((c + 31) / 32) as u32;
26238 let cfg = LaunchConfig {
26239 grid_dim: (nc as u32, h as u32, jt),
26240 block_dim: (256, 1, 1),
26241 shared_mem_bytes: 0,
26242 };
26243 let __s_b = self.gpu.stream();
26244 let mut b = __s_b.launch_builder(&f);
26245 b.arg(q)
26246 .arg(&gcum)
26247 .arg(&p)
26248 .arg(&y)
26249 .arg(&ssnap)
26250 .arg(o)
26251 .arg(&hi)
26252 .arg(&ti)
26253 .arg(&ci)
26254 .arg(&scale);
26255 unsafe {
26256 b.launch(cfg)?;
26257 }
26258 }
26259 Ok(())
26260 }
26261
26262 #[allow(clippy::too_many_arguments)]
26271 #[allow(clippy::too_many_arguments)]
26272 pub fn gdn_scan_prefill(
26273 &self,
26274 q: &CudaSlice<f32>,
26275 k: &CudaSlice<f32>,
26276 v: &CudaSlice<f32>,
26277 g: &CudaSlice<f32>,
26278 beta: &CudaSlice<f32>,
26279 kb16_pre: Option<&CudaSlice<u8>>,
26280 qb16_pre: Option<&CudaSlice<u8>>,
26281 state_in: &CudaSlice<f32>,
26282 state_out: &mut CudaSlice<f32>,
26283 o: &mut CudaSlice<f32>,
26284 n_head: usize,
26285 t: usize,
26286 scale: f32,
26287 hk: usize,
26288 ) -> Result<(), Box<dyn std::error::Error>> {
26289 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26290 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26291 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26292 }
26293 if Self::gdn_chunked_enabled() && t >= 16 {
26294 self.gdn_scan_chunked(
26295 q,
26296 k,
26297 v,
26298 g,
26299 beta,
26300 kb16_pre,
26301 qb16_pre,
26302 state_in,
26303 state_out,
26304 o,
26305 n_head,
26306 t,
26307 scale,
26308 Self::gdn_chunk_size(),
26309 hk,
26310 )
26311 } else {
26312 assert!(
26313 hk == n_head,
26314 "s128 scan is broadcast-only (prep guarantees by predicate)"
26315 );
26316 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26317 }
26318 }
26319
26320 #[allow(clippy::too_many_arguments)]
26322 fn gdn_scan_diff(
26323 &self,
26324 q: &CudaSlice<f32>,
26325 k: &CudaSlice<f32>,
26326 v: &CudaSlice<f32>,
26327 g: &CudaSlice<f32>,
26328 beta: &CudaSlice<f32>,
26329 state_in: &CudaSlice<f32>,
26330 state_out: &mut CudaSlice<f32>,
26331 o: &mut CudaSlice<f32>,
26332 n_head: usize,
26333 t: usize,
26334 scale: f32,
26335 ) -> Result<(), Box<dyn std::error::Error>> {
26336 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26337 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26338 let mut o_c = self.uninit(o.len())?;
26339 let mut st_c = self.uninit(state_out.len())?;
26340 self.gdn_scan_chunked(
26341 q,
26342 k,
26343 v,
26344 g,
26345 beta,
26346 None,
26347 None,
26348 state_in,
26349 &mut st_c,
26350 &mut o_c,
26351 n_head,
26352 t,
26353 scale,
26354 Self::gdn_chunk_size(),
26355 n_head,
26356 )?;
26357 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26358 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26359 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26360 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26361 let mut max_abs = 0f32;
26362 let mut max_rel = 0f32;
26363 let mut sum_rel = 0f64;
26364 for (x, y) in a.iter().zip(b) {
26365 let ad = (x - y).abs();
26366 let rel = ad / x.abs().max(y.abs()).max(1e-3);
26367 if ad > max_abs {
26368 max_abs = ad;
26369 }
26370 if rel > max_rel {
26371 max_rel = rel;
26372 }
26373 sum_rel += rel as f64;
26374 }
26375 (max_abs, max_rel, sum_rel / a.len() as f64)
26376 };
26377 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26378 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26379 println!(
26380 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26381 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26382 Self::gdn_chunk_size()
26383 );
26384 Ok(())
26385 }
26386
26387 pub fn gdn_glog(
26389 &self,
26390 alpha: &CudaSlice<f32>,
26391 dt_bias: &CudaSlice<f32>,
26392 a: &CudaSlice<f32>,
26393 g_log: &mut CudaSlice<f32>,
26394 n_head: usize,
26395 t: usize,
26396 ) -> Result<(), Box<dyn std::error::Error>> {
26397 let f = self.func("gdn_glog_f32");
26398 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26399 let (h, ti) = (n_head as i32, t as i32);
26400 let __s_b = self.gpu.stream();
26401 let mut b = __s_b.launch_builder(&f);
26402 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26403 unsafe {
26404 b.launch(cfg)?;
26405 }
26406 Ok(())
26407 }
26408
26409 pub fn sigmoid_v(
26412 &self,
26413 x: &cudarc::driver::CudaView<f32>,
26414 y: &mut CudaSlice<f32>,
26415 n: usize,
26416 ) -> Result<(), Box<dyn std::error::Error>> {
26417 let f = self.func("sigmoid_f32");
26418 let cfg = LaunchConfig::for_num_elems(n as u32);
26419 let ni = n as i32;
26420 let __s_b = self.gpu.stream();
26421 let mut b = __s_b.launch_builder(&f);
26422 b.arg(x).arg(y).arg(&ni);
26423 unsafe {
26424 b.launch(cfg)?;
26425 }
26426 Ok(())
26427 }
26428
26429 pub fn gdn_glog_v(
26430 &self,
26431 alpha: &cudarc::driver::CudaView<f32>,
26432 dt_bias: &CudaSlice<f32>,
26433 a: &CudaSlice<f32>,
26434 g_log: &mut CudaSlice<f32>,
26435 n_head: usize,
26436 t: usize,
26437 ) -> Result<(), Box<dyn std::error::Error>> {
26438 let f = self.func("gdn_glog_f32");
26439 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26440 let (h, ti) = (n_head as i32, t as i32);
26441 let __s_b = self.gpu.stream();
26442 let mut b = __s_b.launch_builder(&f);
26443 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26444 unsafe {
26445 b.launch(cfg)?;
26446 }
26447 Ok(())
26448 }
26449
26450 pub fn sigmoid(
26451 &self,
26452 x: &CudaSlice<f32>,
26453 y: &mut CudaSlice<f32>,
26454 n: usize,
26455 ) -> Result<(), Box<dyn std::error::Error>> {
26456 let f = self.func("sigmoid_f32");
26457 let cfg = LaunchConfig::for_num_elems(n as u32);
26458 let ni = n as i32;
26459 let __s_b = self.gpu.stream();
26460 let mut b = __s_b.launch_builder(&f);
26461 b.arg(x).arg(y).arg(&ni);
26462 unsafe {
26463 b.launch(cfg)?;
26464 }
26465 Ok(())
26466 }
26467
26468 pub fn sig_mul_f16out(
26471 &self,
26472 a: &CudaSlice<f32>,
26473 g: &CudaSlice<f32>,
26474 dst: &mut CudaSlice<f32>,
26475 dst16: &mut CudaSlice<u8>,
26476 n: usize,
26477 ) -> Result<(), Box<dyn std::error::Error>> {
26478 let f = self.func("sig_mul_f16out_f32");
26479 let cfg = LaunchConfig::for_num_elems(n as u32);
26480 let ni = n as i32;
26481 let __s_b = self.gpu.stream();
26482 let mut b = __s_b.launch_builder(&f);
26483 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26484 unsafe {
26485 b.launch(cfg)?;
26486 }
26487 Ok(())
26488 }
26489
26490 #[allow(clippy::too_many_arguments)]
26499 pub fn attn_head_gate(
26500 &self,
26501 a: &CudaSlice<f32>,
26502 g: &CudaSlice<f32>,
26503 dst: &mut CudaSlice<f32>,
26504 dst16: Option<&mut CudaSlice<u8>>,
26505 head_dim: usize,
26506 n_head: usize,
26507 t: usize,
26508 ) -> Result<(), Box<dyn std::error::Error>> {
26509 let f = self.func("attn_head_gate_f32");
26510 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26511 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26512 let d16: u64 = match dst16 {
26514 Some(d) => self.addr_u8(d),
26515 None => 0,
26516 };
26517 let __s_b = self.gpu.stream();
26518 let mut b = __s_b.launch_builder(&f);
26519 b.arg(a)
26520 .arg(g)
26521 .arg(dst)
26522 .arg(&d16)
26523 .arg(&hd)
26524 .arg(&nh)
26525 .arg(&ti);
26526 unsafe {
26527 b.launch(cfg)?;
26528 }
26529 Ok(())
26530 }
26531
26532 #[allow(clippy::too_many_arguments)]
26541 pub fn swiglu_clamped_mul_scaled(
26542 &self,
26543 gate: &CudaSlice<f32>,
26544 up: &CudaSlice<f32>,
26545 gs: f32,
26546 us: f32,
26547 limit: f32,
26548 dst: &mut CudaSlice<f32>,
26549 n: usize,
26550 ) -> Result<(), Box<dyn std::error::Error>> {
26551 debug_assert!(
26552 limit > 1e-6,
26553 "swiglu_clamped needs a live limit; use silu_mul_scaled"
26554 );
26555 let f = self.func("swiglu_clamped_mul_scaled_f32");
26556 let cfg = LaunchConfig::for_num_elems(n as u32);
26557 let ni = n as i32;
26558 let __s_b = self.gpu.stream();
26559 let mut b = __s_b.launch_builder(&f);
26560 b.arg(gate)
26561 .arg(up)
26562 .arg(&gs)
26563 .arg(&us)
26564 .arg(&limit)
26565 .arg(dst)
26566 .arg(&ni);
26567 unsafe {
26568 b.launch(cfg)?;
26569 }
26570 Ok(())
26571 }
26572
26573 pub fn gated_rmsnorm(
26575 &self,
26576 o: &CudaSlice<f32>,
26577 w: &CudaSlice<f32>,
26578 z: &CudaSlice<f32>,
26579 dst: &mut CudaSlice<f32>,
26580 ncols: usize,
26581 nrows: usize,
26582 eps: f32,
26583 ) -> Result<(), Box<dyn std::error::Error>> {
26584 let f = self.func("gated_rmsnorm_f32");
26585 let cfg = LaunchConfig {
26586 grid_dim: (nrows as u32, 1, 1),
26587 block_dim: (128, 1, 1),
26588 shared_mem_bytes: 0,
26589 };
26590 let (nc, e) = (ncols as i32, eps);
26591 let __s_b = self.gpu.stream();
26592 let mut b = __s_b.launch_builder(&f);
26593 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26594 unsafe {
26595 b.launch(cfg)?;
26596 }
26597 Ok(())
26598 }
26599
26600 pub fn gated_rmsnorm_f16out(
26603 &self,
26604 o: &CudaSlice<f32>,
26605 w: &CudaSlice<f32>,
26606 z: &CudaSlice<f32>,
26607 dst: &mut CudaSlice<f32>,
26608 dst16: &mut CudaSlice<u8>,
26609 ncols: usize,
26610 nrows: usize,
26611 eps: f32,
26612 ) -> Result<(), Box<dyn std::error::Error>> {
26613 let f = self.func("gated_rmsnorm_f16out_f32");
26614 let cfg = LaunchConfig {
26616 grid_dim: (nrows as u32, 1, 1),
26617 block_dim: (128, 1, 1),
26618 shared_mem_bytes: 0,
26619 };
26620 let (nc, e) = (ncols as i32, eps);
26621 let __s_b = self.gpu.stream();
26622 let mut b = __s_b.launch_builder(&f);
26623 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26624 unsafe {
26625 b.launch(cfg)?;
26626 }
26627 Ok(())
26628 }
26629
26630 #[allow(clippy::too_many_arguments)]
26634 pub fn add_rms_norm_zq8(
26635 &self,
26636 a: &CudaSlice<f32>,
26637 b_in: &CudaSlice<f32>,
26638 w: &CudaSlice<f32>,
26639 res: &mut CudaSlice<f32>,
26640 z: &mut CudaSlice<f32>,
26641 ncols: usize,
26642 nrows: usize,
26643 eps: f32,
26644 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26645 assert!(ncols % 32 == 0);
26646 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
26647 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26648 let f = self.func("add_rms_norm_zq8");
26649 let cfg = LaunchConfig {
26650 grid_dim: (nrows as u32, 1, 1),
26651 block_dim: (1024, 1, 1),
26652 shared_mem_bytes: 0,
26653 };
26654 let (nc, ep) = (ncols as i32, eps);
26655 let __s_b = self.gpu.stream();
26656 let mut b = __s_b.launch_builder(&f);
26657 b.arg(a)
26658 .arg(b_in)
26659 .arg(w)
26660 .arg(res)
26661 .arg(z)
26662 .arg(&mut q)
26663 .arg(&mut d)
26664 .arg(&nc)
26665 .arg(&ep);
26666 unsafe {
26667 b.launch(cfg)?;
26668 }
26669 Ok((q, d))
26670 }
26671
26672 pub fn gated_rmsnorm_zv(
26677 &self,
26678 o: &CudaSlice<f32>,
26679 w: &CudaSlice<f32>,
26680 z: &cudarc::driver::CudaView<f32>,
26681 dst: &mut CudaSlice<f32>,
26682 ncols: usize,
26683 nrows: usize,
26684 eps: f32,
26685 ) -> Result<(), Box<dyn std::error::Error>> {
26686 let f = self.func("gated_rmsnorm_f32");
26687 let cfg = LaunchConfig {
26688 grid_dim: (nrows as u32, 1, 1),
26689 block_dim: (128, 1, 1),
26690 shared_mem_bytes: 0,
26691 };
26692 let (nc, e) = (ncols as i32, eps);
26693 let __s_b = self.gpu.stream();
26694 let mut b = __s_b.launch_builder(&f);
26695 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26696 unsafe {
26697 b.launch(cfg)?;
26698 }
26699 Ok(())
26700 }
26701
26702 pub fn gated_rmsnorm_f16out_zv(
26703 &self,
26704 o: &CudaSlice<f32>,
26705 w: &CudaSlice<f32>,
26706 z: &cudarc::driver::CudaView<f32>,
26707 dst: &mut CudaSlice<f32>,
26708 dst16: &mut CudaSlice<u8>,
26709 ncols: usize,
26710 nrows: usize,
26711 eps: f32,
26712 ) -> Result<(), Box<dyn std::error::Error>> {
26713 let f = self.func("gated_rmsnorm_f16out_f32");
26714 let cfg = LaunchConfig {
26716 grid_dim: (nrows as u32, 1, 1),
26717 block_dim: (128, 1, 1),
26718 shared_mem_bytes: 0,
26719 };
26720 let (nc, e) = (ncols as i32, eps);
26721 let __s_b = self.gpu.stream();
26722 let mut b = __s_b.launch_builder(&f);
26723 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26724 unsafe {
26725 b.launch(cfg)?;
26726 }
26727 Ok(())
26728 }
26729
26730 pub fn gated_rmsnorm_q8_1(
26731 &self,
26732 o: &CudaSlice<f32>,
26733 w: &CudaSlice<f32>,
26734 z: &CudaSlice<f32>,
26735 ncols: usize,
26736 nrows: usize,
26737 eps: f32,
26738 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26739 assert!(ncols % 32 == 0);
26740 let f = self.func("gated_rmsnorm_q8_1");
26741 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
26742 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26743 let cfg = LaunchConfig {
26744 grid_dim: (nrows as u32, 1, 1),
26745 block_dim: (128, 1, 1),
26746 shared_mem_bytes: 0,
26747 };
26748 let (nc, ep) = (ncols as i32, eps);
26749 let __s_b = self.gpu.stream();
26750 let mut b = __s_b.launch_builder(&f);
26751 b.arg(o)
26752 .arg(w)
26753 .arg(z)
26754 .arg(&mut out_q)
26755 .arg(&mut out_d)
26756 .arg(&nc)
26757 .arg(&ep);
26758 unsafe {
26759 b.launch(cfg)?;
26760 }
26761 Ok((out_q, out_d))
26762 }
26763
26764 pub fn transpose(
26766 &self,
26767 inp: &CudaSlice<f32>,
26768 rows: usize,
26769 cols: usize,
26770 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26771 let f = self.func("transpose_f32");
26772 let mut out = self.zeros(rows * cols)?;
26773 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
26774 let (r, c) = (rows as i32, cols as i32);
26775 let __s_b = self.gpu.stream();
26776 let mut b = __s_b.launch_builder(&f);
26777 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
26778 unsafe {
26779 b.launch(cfg)?;
26780 }
26781 Ok(out)
26782 }
26783
26784 pub fn repeat_heads(
26786 &self,
26787 inp: &CudaSlice<f32>,
26788 out: &mut CudaSlice<f32>,
26789 head_dim: usize,
26790 n_in: usize,
26791 n_out: usize,
26792 t: usize,
26793 ) -> Result<(), Box<dyn std::error::Error>> {
26794 let f = self.func("repeat_heads_f32");
26795 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
26796 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
26797 let __s_b = self.gpu.stream();
26798 let mut b = __s_b.launch_builder(&f);
26799 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
26800 unsafe {
26801 b.launch(cfg)?;
26802 }
26803 Ok(())
26804 }
26805
26806 pub fn q_gate_split(
26813 &self,
26814 qf: &CudaSlice<f32>,
26815 q_out: &mut CudaSlice<f32>,
26816 gate_out: &mut CudaSlice<f32>,
26817 head_dim: usize,
26818 n_head: usize,
26819 t: usize,
26820 ) -> Result<(), Box<dyn std::error::Error>> {
26821 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
26822 let out_need = head_dim * n_head * t;
26823 if q_out.len() < out_need || gate_out.len() < out_need {
26824 return Err(format!(
26825 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
26826 q_out.len(),
26827 gate_out.len()
26828 )
26829 .into());
26830 }
26831 let f = self.func("q_gate_split_f32");
26832 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26833 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26834 let __s_b = self.gpu.stream();
26835 let mut b = __s_b.launch_builder(&f);
26836 b.arg(qf)
26837 .arg(q_out)
26838 .arg(gate_out)
26839 .arg(&hd)
26840 .arg(&nh)
26841 .arg(&ti);
26842 unsafe {
26843 b.launch(cfg)?;
26844 }
26845 Ok(())
26846 }
26847
26848 pub fn qkv_to_gdn_repack(
26852 &self,
26853 conv_out: &CudaSlice<f32>,
26854 q_g: &mut CudaSlice<f32>,
26855 k_g: &mut CudaSlice<f32>,
26856 v_g: &mut CudaSlice<f32>,
26857 d_state: usize,
26858 num_v: usize,
26859 num_k: usize,
26860 key_dim: usize,
26861 t: usize,
26862 ) -> Result<(), Box<dyn std::error::Error>> {
26863 let f = self.func("qkv_to_gdn_repack_f32");
26864 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26865 let (ds, nv, nk, kd, ti) = (
26866 d_state as i32,
26867 num_v as i32,
26868 num_k as i32,
26869 key_dim as i32,
26870 t as i32,
26871 );
26872 let __s_b = self.gpu.stream();
26873 let mut b = __s_b.launch_builder(&f);
26874 b.arg(conv_out)
26875 .arg(q_g)
26876 .arg(k_g)
26877 .arg(v_g)
26878 .arg(&ds)
26879 .arg(&nv)
26880 .arg(&nk)
26881 .arg(&kd)
26882 .arg(&ti);
26883 unsafe {
26884 b.launch(cfg)?;
26885 }
26886 Ok(())
26887 }
26888
26889 pub fn conv_left_pad(
26892 &self,
26893 src: &CudaSlice<f32>,
26894 dst: &mut CudaSlice<f32>,
26895 conv_dim: usize,
26896 t: usize,
26897 pad: usize,
26898 ) -> Result<(), Box<dyn std::error::Error>> {
26899 let f = self.func("conv_left_pad_f32");
26900 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26901 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26902 let __s_b = self.gpu.stream();
26903 let mut b = __s_b.launch_builder(&f);
26904 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26905 unsafe {
26906 b.launch(cfg)?;
26907 }
26908 Ok(())
26909 }
26910
26911 pub fn conv_assemble_and_roll(
26915 &self,
26916 qkv_col: &CudaSlice<f32>,
26917 conv_state: &mut CudaSlice<f32>,
26918 conv_in: &mut CudaSlice<f32>,
26919 conv_dim: usize,
26920 pad: usize,
26921 ) -> Result<(), Box<dyn std::error::Error>> {
26922 let f = self.func("conv_assemble_and_roll_f32");
26923 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26924 let (cd, p) = (conv_dim as i32, pad as i32);
26925 let __s_b = self.gpu.stream();
26926 let mut b = __s_b.launch_builder(&f);
26927 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26928 unsafe {
26929 b.launch(cfg)?;
26930 }
26931 Ok(())
26932 }
26933
26934 pub fn ssm_conv1d_fused_decode(
26940 &self,
26941 qkv_col: &CudaSlice<f32>,
26942 conv_state: &mut CudaSlice<f32>,
26943 w: &CudaSlice<f32>,
26944 conv_out: &mut CudaSlice<f32>,
26945 conv_dim: usize,
26946 d_conv: usize,
26947 ) -> Result<(), Box<dyn std::error::Error>> {
26948 let f = self.func("ssm_conv1d_fused_decode_f32");
26949 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26950 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26951 let __s_b = self.gpu.stream();
26952 let mut b = __s_b.launch_builder(&f);
26953 b.arg(qkv_col)
26954 .arg(conv_state)
26955 .arg(w)
26956 .arg(conv_out)
26957 .arg(&cd)
26958 .arg(&dc);
26959 unsafe {
26960 b.launch(cfg)?;
26961 }
26962 Ok(())
26963 }
26964
26965 pub fn slice_range(
26968 &self,
26969 src: &CudaSlice<f32>,
26970 start: usize,
26971 len: usize,
26972 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26973 let host = self.gpu.stream().clone_dtoh(src)?;
26974 self.gpu.stream().synchronize()?;
26975 Ok(self.htod(&host[start..start + len])?)
26976 }
26977}
26978
26979#[cfg(test)]
26980mod target_dispatch_tests {
26981 use super::legacy_quant_gemm_allowed;
26982
26983 #[test]
26984 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26985 assert!(legacy_quant_gemm_allowed(false, false, false));
26987 assert!(!legacy_quant_gemm_allowed(false, false, true));
26988 assert!(!legacy_quant_gemm_allowed(true, false, false));
26990 assert!(!legacy_quant_gemm_allowed(true, false, true));
26991 assert!(legacy_quant_gemm_allowed(true, true, false));
26993 assert!(!legacy_quant_gemm_allowed(true, true, true));
26994 }
26995
26996 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26997 #[test]
26998 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26999 assert!(!legacy_quant_gemm_allowed(
27000 cfg!(memra_portable_cuda),
27001 cfg!(memra_hopper_mma),
27002 false
27003 ));
27004 }
27005
27006 #[cfg(memra_hopper_mma)]
27007 #[test]
27008 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27009 assert!(legacy_quant_gemm_allowed(
27010 cfg!(memra_portable_cuda),
27011 cfg!(memra_hopper_mma),
27012 false
27013 ));
27014 assert!(super::portable_mma_gated() == false);
27015 }
27016}
27017
27018impl memra_kv::KvDev for Engine {
27021 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27022 Engine::zeros(self, n)
27023 }
27024 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27025 Engine::uninit(self, n)
27026 }
27027 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27028 Engine::alloc_u8(self, n)
27029 }
27030 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27031 Engine::htod_i32(self, v)
27032 }
27033 fn clone_dtod(
27034 &self,
27035 src: &CudaSlice<f32>,
27036 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27037 Engine::clone_dtod(self, src)
27038 }
27039 fn copy_into(
27040 &self,
27041 dst: &mut CudaSlice<f32>,
27042 off: usize,
27043 src: &CudaSlice<f32>,
27044 len: usize,
27045 ) -> Result<(), Box<dyn std::error::Error>> {
27046 Engine::copy_into(self, dst, off, src, len)
27047 }
27048 fn set_i32_one(
27049 &self,
27050 d: &mut CudaSlice<i32>,
27051 v: i32,
27052 ) -> Result<(), Box<dyn std::error::Error>> {
27053 Engine::set_i32_one(self, d, v)
27054 }
27055}
27056
27057#[cfg(test)]
27058mod fused_gate_bounds_tests {
27059 use super::*;
27060
27061 #[test]
27074 #[ignore = "requires a CUDA GPU"]
27075 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27076 let e = Engine::new(0).unwrap();
27077 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27078 let fused = 2 * head_dim * n_head * t;
27079 let out_n = head_dim * n_head * t;
27080
27081 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27083 let mut q = e.uninit(out_n).unwrap();
27084 let mut gate = e.uninit(out_n).unwrap();
27085 let err = e
27086 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27087 .expect_err("half-width wq must be refused, not read past")
27088 .to_string();
27089 assert!(err.contains("NO fused gate"), "{err}");
27090 assert!(err.contains(&format!("{fused}")), "{err}");
27091
27092 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27095 let wide = e.htod(&host).unwrap();
27096 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27097 .expect("full-width wq splits");
27098 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27099 for tok in 0..t {
27100 for hh in 0..n_head {
27101 for d in 0..head_dim {
27102 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27103 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27104 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27105 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27106 }
27107 }
27108 }
27109
27110 let mut small = e.uninit(out_n - 1).unwrap();
27112 assert!(
27113 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27114 .is_err()
27115 );
27116 }
27117}
27118
27119#[cfg(test)]
27123mod fused_rope_width_tests {
27124 use super::Engine;
27125
27126 #[test]
27129 fn full_width_is_accepted() {
27130 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27131 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27132 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27133 }
27134
27135 #[test]
27148 fn gemma4_official_artifact_widths_pass() {
27149 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27150 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27151 }
27152
27153 #[test]
27156 fn partial_rotary_is_refused_with_the_geometry_named() {
27157 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27159 .expect_err("partial rotary must refuse");
27160 let msg = err.to_string();
27161 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27162 assert!(msg.contains("n_rot 64"), "{msg}");
27163 assert!(msg.contains("head_dim 256"), "{msg}");
27164 assert!(
27165 msg.contains("64..256"),
27166 "names the band it would corrupt: {msg}"
27167 );
27168 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27170 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27172 }
27173}