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 step_tp_w8_on() -> bool {
599 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
600 *ON.get_or_init(|| std::env::var("MEMRA_STEP_TP_W8").as_deref() == Ok("1"))
601}
602
603pub(crate) fn sig_expf_dev_on() -> bool {
608 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
609 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
610}
611
612pub(crate) fn topk_fast_on() -> bool {
613 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
614 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
615}
616
617pub(crate) fn rms_block() -> u32 {
618 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
619 *V.get_or_init(|| {
620 std::env::var("MEMRA_RMS_BLOCK")
621 .ok()
622 .and_then(|v| v.parse().ok())
623 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
624 })
625}
626
627pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
628 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
629 if let Some(forced) = *S.get_or_init(|| {
630 std::env::var("MEMRA_FA_SPLIT")
631 .ok()
632 .and_then(|v| v.parse().ok())
633 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
634 }) {
635 return forced;
636 }
637 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
655 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
656 {
657 return if t_kv <= 8192 {
658 16
659 } else if t_kv <= 16384 {
660 64
661 } else {
662 128
663 };
664 }
665 let big_rig = fa_sm_count() >= 128;
666 if big_rig {
667 let _ = n_head_kv;
668 if t_kv <= 2048 {
669 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
678 if let Some(sp) = *SHORT.get_or_init(|| {
679 std::env::var("MEMRA_FA_SP_SHORT")
680 .ok()
681 .and_then(|v| v.parse().ok())
682 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
683 }) {
684 return sp;
685 }
686 16
687 } else if t_kv <= 16384 {
688 64
689 } else {
690 128
691 }
692 } else if n_head_kv <= 4 {
693 if t_kv <= 512 {
714 8
715 } else if t_kv <= 16384 {
716 64
717 } else {
718 128
719 }
720 } else {
721 if t_kv <= 8192 {
722 32
723 } else if t_kv <= 16384 {
724 64
725 } else {
726 128
727 }
728 }
729}
730
731pub(crate) fn fa_sm_count() -> i32 {
734 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
735 *N.get_or_init(|| {
736 cudarc::driver::result::init().ok();
737 cudarc::driver::result::device::get(0)
738 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
739 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
740 .unwrap_or(82)
741 })
742}
743
744fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
748 match head_dim {
749 256 => Ok(""),
750 128 => Ok("_hd128"),
751 d => Err(format!(
752 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
753 callers must gate to sdpa_naive"
754 )
755 .into()),
756 }
757}
758
759pub const QT_Q8_0: i32 = 0;
761pub const QT_Q4_K: i32 = 1;
762pub const QT_Q6_K: i32 = 2;
763pub const QT_Q5_K: i32 = 3;
764pub const QT_Q3_K: i32 = 4;
765pub const QT_IQ4_XS: i32 = 5;
766pub const QT_IQ3_S: i32 = 6;
767pub const QT_NVFP4: i32 = 7;
768pub const QT_F8_E4M3: i32 = 10;
774pub const QT_NVFP4_RP: i32 = 9;
777pub const QT_F32: i32 = 8;
779pub const QT_BF16: i32 = 11;
780pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
785pub const QT_F8_E4M3_BLK: i32 = 14;
801
802pub struct Engine {
804 pub gpu: memra_runtime::Gpu,
805 module: Arc<CudaModule>,
806 hybrid: Arc<CudaModule>,
807 qmatvec: Arc<CudaModule>,
808 flash: Arc<CudaModule>,
809 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
813 gemm: Arc<CudaModule>,
814 router: Arc<CudaModule>,
815 sample: Arc<CudaModule>,
817 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
821 w8_mirrors: Mutex<std::collections::HashMap<u64, CudaSlice<u8>>>,
826 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
829 moe_cache_layout: Mutex<Option<Vec<usize>>>,
833 capture_keep_on: std::sync::atomic::AtomicBool,
839 verify_exact: std::sync::atomic::AtomicBool,
844 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
845 pub copy_stream: Arc<CudaStream>,
847 #[cfg(memra_cutlass)]
854 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
855 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
859 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
862 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
866 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
870 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
872 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
873 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
878 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
883 router_stage: Mutex<Option<PinnedStage>>,
887}
888
889fn fa_v2_on() -> bool {
899 std::env::var("MEMRA_FA_V2")
905 .map(|v| v != "0")
906 .unwrap_or(true)
907}
908
909pub(crate) fn fa_v3_on() -> bool {
917 std::env::var("MEMRA_FA_V3")
921 .map(|v| v != "0")
922 .unwrap_or(true)
923}
924
925fn fa_v4_mode() -> &'static str {
930 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
931 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
932}
933fn fa_v4_on() -> bool {
934 fa_v4_mode() != "0"
935} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
944 std::sync::atomic::AtomicUsize::new(1024);
945pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
946 std::sync::atomic::AtomicUsize::new(usize::MAX);
947pub fn fa_v4_at_pub(t_kv: usize) -> bool {
948 fa_v4_at(t_kv)
949}
950fn fa_v4_at(t_kv: usize) -> bool {
951 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
952 let mx = *M.get_or_init(|| {
953 std::env::var("MEMRA_FA_V4_MAX")
954 .ok()
955 .and_then(|v| v.parse().ok())
956 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
957 });
958 fa_v4_on() && t_kv < mx
959}
960pub const FA_DEEP_MIN_DEFAULT: usize = 0;
974fn fa_deep_at(t_kv: usize) -> bool {
975 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
976 return false;
977 }
978 let min = std::env::var("MEMRA_FA_DEEP_MIN")
979 .ok()
980 .and_then(|v| v.parse().ok())
981 .unwrap_or(FA_DEEP_MIN_DEFAULT);
982 t_kv >= min
983}
984pub fn fa_deep_at_pub(t_kv: usize) -> bool {
986 fa_deep_at(t_kv)
987}
988
989fn fa_v3_active(head_dim: usize) -> bool {
990 fa_v3_on()
993 && head_dim % 128 == 0
994 && kv_cache_formats() == ("q8_0", "q5_1")
995 && !Engine::kv_fp8_on()
996}
997
998pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1006 std::env::var("MEMRA_NO_FA_VEC").is_err()
1007 && t_kv >= fa_vec_min_tkv()
1008 && head_dim == 256
1009 && fa_v4_at(t_kv)
1010 && !matches!(fa_v4_mode(), "noB3" | "stage")
1011 && !Engine::kv_fp8_on()
1012}
1013pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1015 fa_split_keys(t_kv, n_head_kv)
1016}
1017
1018struct PinnedStage {
1023 ptr: *mut u8,
1024 cap: usize,
1025}
1026unsafe impl Send for PinnedStage {}
1027impl PinnedStage {
1028 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1029 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1030 Ok(PinnedStage { ptr, cap })
1031 }
1032}
1033impl Drop for PinnedStage {
1034 fn drop(&mut self) {
1035 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1036 }
1037}
1038
1039pub const ARGMAX_NB: usize = 256;
1042
1043pub(crate) use memra_fa3_vl as fa3_vl_raw;
1045
1046unsafe extern "C" {
1047 fn memra_fa3_prefill(
1049 q16: *const core::ffi::c_void,
1050 k16: *const core::ffi::c_void,
1051 v16: *const core::ffi::c_void,
1052 o: *mut f32,
1053 t: i32,
1054 h: i32,
1055 hkv: i32,
1056 d: i32,
1057 scale: f32,
1058 stream: *mut core::ffi::c_void,
1059 ) -> i32;
1060 pub(crate) fn memra_fa3_vl(
1062 q16s: *const *const core::ffi::c_void,
1063 k16s: *const *const core::ffi::c_void,
1064 v16s: *const *const core::ffi::c_void,
1065 os: *const *mut f32,
1066 ts: *const i32,
1067 b: i32,
1068 h: i32,
1069 hkv: i32,
1070 d: i32,
1071 scale: f32,
1072 stream: *mut core::ffi::c_void,
1073 ) -> i32;
1074}
1075
1076#[repr(C)]
1081#[derive(Clone, Copy)]
1082pub struct WPtr8(pub [u64; 8]);
1083unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1084
1085#[repr(C)]
1090#[derive(Clone, Copy, Default)]
1091pub struct GdnSeqVl {
1092 pub kb16: u64,
1093 pub gcum: u64,
1094 pub beta: u64,
1095 pub u: u64,
1096 pub wb16: u64,
1097 pub y: u64,
1098 pub ssnap: u64,
1099 pub state_in: u64,
1100 pub state_out: u64,
1101 pub q: u64,
1102 pub p: u64,
1103 pub o: u64,
1104 pub k: u64,
1105 pub v: u64,
1106 pub g: u64,
1107 pub a: u64,
1108 pub w: u64,
1109 pub t: i32,
1110 pub nc: i32,
1111}
1112unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1113#[repr(C)]
1114#[derive(Clone, Copy)]
1115pub struct GdnVl8(pub [GdnSeqVl; 8]);
1116unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1117
1118#[repr(C)]
1121#[derive(Clone, Copy, Default)]
1122pub struct GdnWVl {
1123 pub qb16: u64,
1124 pub pb16: u64,
1125}
1126unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1127#[repr(C)]
1128#[derive(Clone, Copy)]
1129pub struct GdnWVl8(pub [GdnWVl; 8]);
1130unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1131
1132#[repr(C)]
1134#[derive(Clone, Copy, Default)]
1135pub struct GdnPrepVl {
1136 pub qkv: u64,
1137 pub conv_state: u64,
1138 pub conv_out: u64,
1139 pub q_g: u64,
1140 pub k_g: u64,
1141 pub v_g: u64,
1142 pub q_l2: u64,
1143 pub k_l2: u64,
1144 pub beta_raw: u64,
1145 pub alpha: u64,
1146 pub beta: u64,
1147 pub g_log: u64,
1148 pub o: u64,
1149 pub z: u64,
1150 pub gn: u64,
1151 pub gn16: u64,
1152 pub kb16: u64,
1153 pub qb16: u64,
1154 pub t: i32,
1155 pub pad: i32,
1156}
1157unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1158#[repr(C)]
1159#[derive(Clone, Copy)]
1160pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1161unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1162
1163#[repr(C)]
1165#[derive(Clone, Copy, Default)]
1166pub struct FaSeqVl {
1167 pub q: u64,
1168 pub k16: u64,
1169 pub v16: u64,
1170 pub o: u64,
1171 pub kf: u64,
1172 pub vf: u64,
1173 pub t: i32,
1174 pub pad: i32,
1175}
1176unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1177#[repr(C)]
1178#[derive(Clone, Copy)]
1179pub struct FaVl8(pub [FaSeqVl; 8]);
1180unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1181
1182#[repr(C)]
1184#[derive(Clone, Copy, Default)]
1185pub struct AttnPreVl {
1186 pub qf: u64,
1187 pub kf: u64,
1188 pub vf: u64,
1189 pub q: u64,
1190 pub gate: u64,
1191 pub qn: u64,
1192 pub kn: u64,
1193 pub kc: u64,
1194 pub vc: u64,
1195 pub t: i32,
1196 pub pad: i32,
1197}
1198unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1199#[repr(C)]
1200#[derive(Clone, Copy)]
1201pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1202unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1203
1204pub struct GdnChunkBufs {
1207 pub gcum: CudaSlice<f32>,
1208 pub a: CudaSlice<f32>,
1209 pub p: CudaSlice<f32>,
1210 pub u: CudaSlice<f32>,
1211 pub w: CudaSlice<f32>,
1212 pub kb16: CudaSlice<u8>,
1213 pub wb16: CudaSlice<u8>,
1214 pub y16: CudaSlice<u8>,
1215 pub ssnap16: CudaSlice<u8>,
1216 pub qb16: CudaSlice<u8>,
1217 pub pb16: CudaSlice<u8>,
1218 pub o: CudaSlice<f32>,
1219 pub t: usize,
1220 pub nc: usize,
1221}
1222
1223#[repr(C)]
1225#[derive(Clone, Copy)]
1226pub struct F32x8(pub [f32; 8]);
1227unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1228
1229pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1233
1234impl Engine {
1235 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1236 let gpu = memra_runtime::Gpu::new(ordinal)?;
1237 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1241 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1242 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1243 .and_then(|d| unsafe {
1244 Ok((
1245 cudarc::driver::result::device::get_attribute(
1246 d,
1247 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1248 )?,
1249 cudarc::driver::result::device::get_attribute(
1250 d,
1251 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1252 )?,
1253 ))
1254 })
1255 .unwrap_or((0, 0));
1256 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1257 let ok = matches!(
1258 (built, maj, min),
1259 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1260 );
1261 if !ok {
1262 return Err(format!(
1263 "memra was built for sm_{built} but device {ordinal} reports compute \
1264 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1265 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1266 )
1267 .into());
1268 }
1269 }
1270 unsafe {
1275 use cudarc::driver::sys;
1276 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1277 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1278 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1279 let mut thresh: u64 = u64::MAX;
1280 let _ = sys::cuMemPoolSetAttribute(
1281 pool,
1282 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1283 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1284 );
1285 }
1286 }
1287 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1288 let hybrid = gpu
1289 .ctx
1290 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1291 let qmatvec = gpu
1292 .ctx
1293 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1294 let flash = gpu
1295 .ctx
1296 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1297 let gemm = gpu
1298 .ctx
1299 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1300 let router = gpu
1301 .ctx
1302 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1303 let sample = gpu
1304 .ctx
1305 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1306 let copy_stream = gpu.ctx.new_stream()?;
1307 if std::env::var("MEMRA_EVT")
1323 .map(|v| v == "1")
1324 .unwrap_or(false)
1325 {
1326 } else {
1328 unsafe {
1329 gpu.ctx.disable_event_tracking();
1330 }
1331 }
1332 Ok(Self {
1333 gpu,
1334 module,
1335 hybrid,
1336 qmatvec,
1337 flash,
1338 flash_g: std::sync::OnceLock::new(),
1339 gemm,
1340 router,
1341 sample,
1342 moe_cache: Mutex::new(None),
1343 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
1344 w8_act: Mutex::new(std::collections::HashMap::new()),
1345 moe_cache_layout: Mutex::new(None),
1346 copy_stream,
1347 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1348 verify_exact: std::sync::atomic::AtomicBool::new(false),
1349 capture_keep: Mutex::new(Vec::new()),
1350 argmax_partials: Mutex::new(None),
1351 prime_deqw_ws: Mutex::new(None),
1352 router_stage: Mutex::new(None),
1353 fp8_scratch: Mutex::new(None),
1354 fa_vf16_scratch: Mutex::new(None),
1355 fa_part_pool: Mutex::new(None),
1356 fa_part_retired: Mutex::new(Vec::new()),
1357 fn_cache: Mutex::new(Default::default()),
1358 f16_scratch: Mutex::new(None),
1359 #[cfg(memra_cutlass)]
1360 cutlass_scratch: Mutex::new(None),
1361 })
1362 }
1363
1364 pub fn ctx(&self) -> &Arc<CudaContext> {
1365 &self.gpu.ctx
1366 }
1367
1368 pub fn pool_cached_bytes(&self) -> usize {
1386 let (reserved, used) = self.pool_reserved_used();
1387 reserved.saturating_sub(used)
1388 }
1389
1390 pub fn device_graph_mem_reserved(&self) -> usize {
1400 use cudarc::driver::sys as cus;
1401 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1402 return 0;
1403 };
1404 let mut bytes: u64 = 0;
1405 let rc = unsafe {
1406 cus::cuDeviceGetGraphMemAttribute(
1407 dev,
1408 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1409 &mut bytes as *mut u64 as *mut std::ffi::c_void,
1410 )
1411 };
1412 if rc == cus::cudaError_enum::CUDA_SUCCESS {
1413 bytes as usize
1414 } else {
1415 0
1416 }
1417 }
1418
1419 pub fn pool_reserved_used(&self) -> (usize, usize) {
1426 use cudarc::driver::sys;
1427 unsafe {
1428 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1429 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1430 != sys::CUresult::CUDA_SUCCESS
1431 {
1432 return (0, 0);
1433 }
1434 let (mut reserved, mut used) = (0u64, 0u64);
1435 if sys::cuMemPoolGetAttribute(
1436 pool,
1437 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1438 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1439 ) != sys::CUresult::CUDA_SUCCESS
1440 {
1441 return (0, 0);
1442 }
1443 if sys::cuMemPoolGetAttribute(
1444 pool,
1445 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1446 &mut used as *mut u64 as *mut core::ffi::c_void,
1447 ) != sys::CUresult::CUDA_SUCCESS
1448 {
1449 return (0, 0);
1450 }
1451 (reserved as usize, used as usize)
1452 }
1453 }
1454
1455 pub fn stream(&self) -> Arc<CudaStream> {
1458 self.gpu.stream()
1459 }
1460 pub fn gkv_on() -> bool {
1463 memra_kv::gkv_on()
1464 }
1465
1466 pub fn wkv_on() -> bool {
1478 memra_kv::wkv_on()
1479 }
1480
1481 pub fn kv_fp8_on() -> bool {
1487 memra_kv::kv_fp8_on()
1488 }
1489
1490 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1493 if head_dim == 512 && Self::gkv_on() {
1494 self.func_g(name)
1495 } else {
1496 self.func(name)
1497 }
1498 }
1499
1500 fn func_g(&self, name: &str) -> CudaFunction {
1504 let m = self.flash_g.get_or_init(|| {
1505 self.gpu
1506 .ctx
1507 .load_module(cudarc::nvrtc::Ptx::from_binary(
1508 FLASH_FATBIN_KF8VF8.to_vec(),
1509 ))
1510 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1511 });
1512 let key = format!("g:{name}");
1513 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1514 return f.clone();
1515 }
1516 let f = match m.load_function(name) {
1517 Ok(f) => f,
1518 Err(_) => self.func(name),
1519 };
1520 self.fn_cache.lock().unwrap().insert(key, f.clone());
1521 f
1522 }
1523
1524 fn func(&self, name: &str) -> CudaFunction {
1525 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1528 return f.clone();
1529 }
1530 let f = self
1531 .module
1532 .load_function(name)
1533 .or_else(|_| self.hybrid.load_function(name))
1534 .or_else(|_| self.qmatvec.load_function(name))
1535 .or_else(|_| self.flash.load_function(name))
1536 .or_else(|_| self.gemm.load_function(name))
1537 .or_else(|_| self.router.load_function(name))
1538 .or_else(|_| self.sample.load_function(name))
1539 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1540 self.fn_cache
1541 .lock()
1542 .unwrap()
1543 .insert(name.to_string(), f.clone());
1544 f
1545 }
1546
1547 pub fn scatter_trim_logits(
1550 &self,
1551 src: &CudaSlice<f32>,
1552 d2t: &CudaSlice<u32>,
1553 dst: &mut CudaSlice<f32>,
1554 d_vocab: usize,
1555 n_vocab: usize,
1556 ) -> Result<(), Box<dyn std::error::Error>> {
1557 let f1 = self.func("scatter_trim_logits_f32");
1558 let f2 = self.func("scatter_trim_logits_pass2_f32");
1559 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1560 let cfg1 = LaunchConfig {
1561 grid_dim: (256, 1, 1),
1562 block_dim: (256, 1, 1),
1563 shared_mem_bytes: 0,
1564 };
1565 let __s_b1 = self.gpu.stream();
1566 let mut b1 = __s_b1.launch_builder(&f1);
1567 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1568 unsafe {
1569 b1.launch(cfg1)?;
1570 }
1571 let cfg2 = LaunchConfig {
1572 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1573 block_dim: (256, 1, 1),
1574 shared_mem_bytes: 0,
1575 };
1576 let __s_b2 = self.gpu.stream();
1577 let mut b2 = __s_b2.launch_builder(&f2);
1578 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1579 unsafe {
1580 b2.launch(cfg2)?;
1581 }
1582 Ok(())
1583 }
1584
1585 #[allow(clippy::too_many_arguments)]
1591 pub fn filter_stats(
1592 &self,
1593 x: &CudaSlice<f32>,
1594 row_stride: usize,
1595 rows: &CudaSlice<i32>,
1596 out_th: &mut CudaSlice<f32>,
1597 out_z: &mut CudaSlice<f32>,
1598 out_max: &mut CudaSlice<f32>,
1599 n: usize,
1600 nrow: usize,
1601 temp: f32,
1602 top_k: i32,
1603 top_p: f32,
1604 min_p: f32,
1605 ) -> Result<(), Box<dyn std::error::Error>> {
1606 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1621 let coop_on =
1622 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1623 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1624 if coop_on && 16 * nrow <= self.sm_count() as usize {
1625 let f = self.func("filter_stats_coop_f32");
1626 let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1627 let cfg = LaunchConfig {
1628 grid_dim: (16, nrow as u32, 1),
1629 block_dim: (512, 1, 1),
1630 shared_mem_bytes: 0,
1631 };
1632 let __s_b = self.gpu.stream();
1633 let mut b = __s_b.launch_builder(&f);
1634 b.arg(x)
1635 .arg(&rs)
1636 .arg(rows)
1637 .arg(&mut *out_th)
1638 .arg(&mut *out_z)
1639 .arg(&mut *out_max)
1640 .arg(&mut ws)
1641 .arg(&ni)
1642 .arg(&nr)
1643 .arg(&temp)
1644 .arg(&top_k)
1645 .arg(&top_p)
1646 .arg(&min_p);
1647 unsafe {
1648 b.launch_cooperative(cfg)?;
1649 }
1650 return Ok(());
1651 }
1652 let f = self.func("filter_stats_f32");
1653 let cfg = LaunchConfig {
1654 grid_dim: (nrow as u32, 1, 1),
1655 block_dim: (1024, 1, 1),
1656 shared_mem_bytes: 0,
1657 };
1658 let __s_b = self.gpu.stream();
1659 let mut b = __s_b.launch_builder(&f);
1660 b.arg(x)
1661 .arg(&rs)
1662 .arg(rows)
1663 .arg(&mut *out_th)
1664 .arg(&mut *out_z)
1665 .arg(&mut *out_max)
1666 .arg(&ni)
1667 .arg(&nr)
1668 .arg(&temp)
1669 .arg(&top_k)
1670 .arg(&top_p)
1671 .arg(&min_p);
1672 unsafe {
1673 b.launch(cfg)?;
1674 }
1675 Ok(())
1676 }
1677
1678 #[allow(clippy::too_many_arguments)]
1680 pub fn softmax_gather_filtered(
1681 &self,
1682 x: &CudaSlice<f32>,
1683 row_stride: usize,
1684 ids: &CudaSlice<u32>,
1685 rows: &CudaSlice<i32>,
1686 th: &CudaSlice<f32>,
1687 z: &CudaSlice<f32>,
1688 out: &mut CudaSlice<f32>,
1689 n: usize,
1690 npair: usize,
1691 temp: f32,
1692 ) -> Result<(), Box<dyn std::error::Error>> {
1693 let f = self.func("softmax_gather_filtered_f32");
1694 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1695 let cfg = LaunchConfig {
1696 grid_dim: (npair as u32, 1, 1),
1697 block_dim: (256, 1, 1),
1698 shared_mem_bytes: 0,
1699 };
1700 let __s_b = self.gpu.stream();
1701 let mut b = __s_b.launch_builder(&f);
1702 b.arg(x)
1703 .arg(&rs)
1704 .arg(ids)
1705 .arg(rows)
1706 .arg(th)
1707 .arg(z)
1708 .arg(&mut *out)
1709 .arg(&ni)
1710 .arg(&np)
1711 .arg(&temp);
1712 unsafe {
1713 b.launch(cfg)?;
1714 }
1715 Ok(())
1716 }
1717
1718 #[allow(clippy::too_many_arguments)]
1720 pub fn residual_sample_filtered(
1721 &self,
1722 p: &CudaSlice<f32>,
1723 q: Option<&CudaSlice<f32>>,
1724 n: usize,
1725 temp: f32,
1726 seed: u64,
1727 stream_pos: u32,
1728 p_stats: (f32, f32, f32),
1729 q_stats: (f32, f32, f32),
1730 out_tok: &mut CudaSlice<u32>,
1731 ) -> Result<(), Box<dyn std::error::Error>> {
1732 let f = self.func("residual_sample_filtered_f32");
1733 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1734 let has_q: i32 = q.is_some() as i32;
1735 let qbuf = q.unwrap_or(p);
1736 let (pm, pth, pz) = p_stats;
1737 let (qm, qth, qz) = q_stats;
1738 let cfg = LaunchConfig {
1739 grid_dim: (1, 1, 1),
1740 block_dim: (1024, 1, 1),
1741 shared_mem_bytes: 0,
1742 };
1743 let __s_b = self.gpu.stream();
1744 let mut b = __s_b.launch_builder(&f);
1745 b.arg(p)
1746 .arg(qbuf)
1747 .arg(&has_q)
1748 .arg(&ni)
1749 .arg(&temp)
1750 .arg(&slo)
1751 .arg(&shi)
1752 .arg(&stream_pos)
1753 .arg(&pm)
1754 .arg(&pth)
1755 .arg(&pz)
1756 .arg(&qm)
1757 .arg(&qth)
1758 .arg(&qz)
1759 .arg(&mut *out_tok);
1760 unsafe {
1761 b.launch(cfg)?;
1762 }
1763 Ok(())
1764 }
1765
1766 #[allow(clippy::too_many_arguments)]
1772 pub fn residual_sample_sparse_q(
1773 &self,
1774 p: &CudaSlice<f32>,
1775 cand_ids: &CudaSlice<u32>,
1776 q_probs: &CudaSlice<f32>,
1777 n_cand: usize,
1778 n: usize,
1779 temp: f32,
1780 seed: u64,
1781 stream_pos: u32,
1782 p_stats: (f32, f32, f32),
1783 out_tok: &mut CudaSlice<u32>,
1784 ) -> Result<(), Box<dyn std::error::Error>> {
1785 assert!(
1786 n_cand >= 1 && n_cand <= 32,
1787 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1788 );
1789 let f = self.func("residual_sample_sparse_q_f32");
1790 let (ni, nc) = (n as i32, n_cand as i32);
1791 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1792 let (pm, pth, pz) = p_stats;
1793 let cfg = LaunchConfig {
1794 grid_dim: (1, 1, 1),
1795 block_dim: (1024, 1, 1),
1796 shared_mem_bytes: 0,
1797 };
1798 let __s_b = self.gpu.stream();
1799 let mut b = __s_b.launch_builder(&f);
1800 b.arg(p)
1801 .arg(cand_ids)
1802 .arg(q_probs)
1803 .arg(&nc)
1804 .arg(&ni)
1805 .arg(&temp)
1806 .arg(&slo)
1807 .arg(&shi)
1808 .arg(&stream_pos)
1809 .arg(&pm)
1810 .arg(&pth)
1811 .arg(&pz)
1812 .arg(&mut *out_tok);
1813 unsafe {
1814 b.launch(cfg)?;
1815 }
1816 Ok(())
1817 }
1818
1819 #[allow(clippy::too_many_arguments)]
1821 pub fn gumbel_perturb_filtered(
1822 &self,
1823 x: &CudaSlice<f32>,
1824 y: &mut CudaSlice<f32>,
1825 n: usize,
1826 seed: u64,
1827 stream_pos: u32,
1828 temp: f32,
1829 row_max: f32,
1830 th: f32,
1831 ) -> Result<(), Box<dyn std::error::Error>> {
1832 let f = self.func("gumbel_perturb_filtered_f32");
1833 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1834 let cfg = LaunchConfig {
1835 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1836 block_dim: (256, 1, 1),
1837 shared_mem_bytes: 0,
1838 };
1839 let __s_b = self.gpu.stream();
1840 let mut b = __s_b.launch_builder(&f);
1841 b.arg(x)
1842 .arg(&mut *y)
1843 .arg(&ni)
1844 .arg(&slo)
1845 .arg(&shi)
1846 .arg(&stream_pos)
1847 .arg(&temp)
1848 .arg(&row_max)
1849 .arg(&th);
1850 unsafe {
1851 b.launch(cfg)?;
1852 }
1853 Ok(())
1854 }
1855
1856 #[allow(clippy::too_many_arguments)]
1860 pub fn penalize_logits(
1861 &self,
1862 x: &mut CudaSlice<f32>,
1863 hist: &CudaSlice<u32>,
1864 n_hist: usize,
1865 rep: f32,
1866 freq: f32,
1867 present: f32,
1868 n: usize,
1869 ) -> Result<(), Box<dyn std::error::Error>> {
1870 if n_hist == 0 {
1871 return Ok(());
1872 }
1873 let f = self.func("penalize_logits_f32");
1874 let (nh, ni) = (n_hist as i32, n as i32);
1875 let cfg = LaunchConfig {
1876 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1877 block_dim: (128, 1, 1),
1878 shared_mem_bytes: 0,
1879 };
1880 let __s_b = self.gpu.stream();
1881 let mut b = __s_b.launch_builder(&f);
1882 b.arg(&mut *x)
1883 .arg(hist)
1884 .arg(&nh)
1885 .arg(&rep)
1886 .arg(&freq)
1887 .arg(&present)
1888 .arg(&ni);
1889 unsafe {
1890 b.launch(cfg)?;
1891 }
1892 Ok(())
1893 }
1894
1895 #[allow(clippy::too_many_arguments)]
1897 pub fn penalize_logits_rows(
1898 &self,
1899 x: &mut CudaSlice<f32>,
1900 hist: &CudaSlice<u32>,
1901 n_hist: usize,
1902 rep: f32,
1903 freq: f32,
1904 present: f32,
1905 n: usize,
1906 nrow: usize,
1907 ) -> Result<(), Box<dyn std::error::Error>> {
1908 if n_hist == 0 || nrow == 0 {
1909 return Ok(());
1910 }
1911 let f = self.func("penalize_logits_rows_f32");
1912 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1913 let cfg = LaunchConfig {
1914 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1915 block_dim: (128, 1, 1),
1916 shared_mem_bytes: 0,
1917 };
1918 let __s_b = self.gpu.stream();
1919 let mut b = __s_b.launch_builder(&f);
1920 b.arg(&mut *x)
1921 .arg(hist)
1922 .arg(&nh)
1923 .arg(&rep)
1924 .arg(&freq)
1925 .arg(&present)
1926 .arg(&ni)
1927 .arg(&nr);
1928 unsafe {
1929 b.launch(cfg)?;
1930 }
1931 Ok(())
1932 }
1933
1934 #[allow(clippy::too_many_arguments)]
1940 pub fn penalize_logits_sparse_rows(
1941 &self,
1942 x: &mut CudaSlice<f32>,
1943 ids: &[u32],
1944 counts: &[u32],
1945 offsets: &[i32],
1946 rows: &[i32],
1947 reps: &[f32],
1948 freqs: &[f32],
1949 presents: &[f32],
1950 n: usize,
1951 ) -> Result<(), Box<dyn std::error::Error>> {
1952 let nrow = rows.len();
1953 if nrow == 0 {
1954 return Ok(());
1955 }
1956 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
1957 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
1958 let entry_count =
1959 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
1960 if ids.len() != counts.len()
1961 || offsets.len() != nrow + 1
1962 || reps.len() != nrow
1963 || freqs.len() != nrow
1964 || presents.len() != nrow
1965 || offsets.first().copied() != Some(0)
1966 || offsets.last().copied() != Some(entry_count)
1967 {
1968 return Err("sparse penalty row metadata shape mismatch".into());
1969 }
1970 if counts.contains(&0) {
1971 return Err("sparse penalty counts must be positive".into());
1972 }
1973 let mut max_len = 0usize;
1974 for pair in offsets.windows(2) {
1975 if pair[0] < 0 || pair[1] < pair[0] {
1976 return Err("sparse penalty offsets must be monotonic".into());
1977 }
1978 max_len = max_len.max((pair[1] - pair[0]) as usize);
1979 }
1980 if max_len == 0 {
1981 return Ok(());
1982 }
1983
1984 let mut seen = std::collections::HashSet::with_capacity(ids.len());
1985 for (r, &row) in rows.iter().enumerate() {
1986 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
1987 return Err("sparse penalty row index exceeds logits shape".into());
1988 }
1989 let begin = offsets[r] as usize;
1990 let end = offsets[r + 1] as usize;
1991 for &id in &ids[begin..end] {
1992 if id as usize >= n {
1993 return Err("sparse penalty token id exceeds logits row".into());
1994 }
1995 if !seen.insert((row, id)) {
1996 return Err("sparse penalty entries must be unique per logits row".into());
1997 }
1998 }
1999 }
2000
2001 unsafe {
2003 self.penalize_logits_sparse_rows_unchecked(
2004 x, ids, counts, offsets, rows, reps, freqs, presents, n,
2005 )
2006 }
2007 }
2008
2009 #[allow(clippy::too_many_arguments)]
2018 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2019 &self,
2020 x: &mut CudaSlice<f32>,
2021 ids: &[u32],
2022 counts: &[u32],
2023 offsets: &[i32],
2024 rows: &[i32],
2025 reps: &[f32],
2026 freqs: &[f32],
2027 presents: &[f32],
2028 n: usize,
2029 ) -> Result<(), Box<dyn std::error::Error>> {
2030 let nrow = rows.len();
2031 if nrow == 0 {
2032 return Ok(());
2033 }
2034 let max_len = offsets
2035 .windows(2)
2036 .map(|pair| (pair[1] - pair[0]) as usize)
2037 .max()
2038 .unwrap_or(0);
2039 if max_len == 0 {
2040 return Ok(());
2041 }
2042 let ids_d = self.htod_u32_v(ids)?;
2043 let counts_d = self.htod_u32_v(counts)?;
2044 let offsets_d = self.htod_i32(offsets)?;
2045 let rows_d = self.htod_i32(rows)?;
2046 let reps_d = self.htod(reps)?;
2047 let freqs_d = self.htod(freqs)?;
2048 let presents_d = self.htod(presents)?;
2049 let f = self.func("penalize_logits_sparse_rows_f32");
2050 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2051 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2052 let cfg = LaunchConfig {
2053 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2054 block_dim: (128, 1, 1),
2055 shared_mem_bytes: 0,
2056 };
2057 let __s_b = self.gpu.stream();
2058 let mut b = __s_b.launch_builder(&f);
2059 b.arg(&mut *x)
2060 .arg(&ids_d)
2061 .arg(&counts_d)
2062 .arg(&offsets_d)
2063 .arg(&rows_d)
2064 .arg(&reps_d)
2065 .arg(&freqs_d)
2066 .arg(&presents_d)
2067 .arg(&ni)
2068 .arg(&nr);
2069 unsafe {
2070 b.launch(cfg)?;
2071 }
2072 Ok(())
2073 }
2074
2075 #[allow(clippy::too_many_arguments)]
2083 pub fn penalize_logits_rows_inc(
2084 &self,
2085 x: &mut CudaSlice<f32>,
2086 hist: &CudaSlice<u32>,
2087 n_hist0: usize,
2088 rep: f32,
2089 freq: f32,
2090 present: f32,
2091 n: usize,
2092 nrow: usize,
2093 win: usize,
2094 ) -> Result<(), Box<dyn std::error::Error>> {
2095 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2096 return Ok(());
2097 }
2098 debug_assert!(
2099 hist.len() >= n_hist0 + nrow - 1,
2100 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2101 );
2102 let f = self.func("penalize_logits_rows_inc_f32");
2103 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2104 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2105 let cfg = LaunchConfig {
2106 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2107 block_dim: (128, 1, 1),
2108 shared_mem_bytes: 0,
2109 };
2110 let __s_b = self.gpu.stream();
2111 let mut b = __s_b.launch_builder(&f);
2112 b.arg(&mut *x)
2113 .arg(hist)
2114 .arg(&nh)
2115 .arg(&rep)
2116 .arg(&freq)
2117 .arg(&present)
2118 .arg(&ni)
2119 .arg(&nr)
2120 .arg(&wi);
2121 unsafe {
2122 b.launch(cfg)?;
2123 }
2124 Ok(())
2125 }
2126
2127 pub fn wpf_level() -> u32 {
2135 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2136 *ON.get_or_init(|| {
2137 std::env::var("MEMRA_WPF")
2138 .ok()
2139 .and_then(|v| v.parse().ok())
2140 .unwrap_or(1)
2141 })
2142 }
2143
2144 pub fn set_verify_exact(&self, on: bool) {
2156 self.verify_exact
2157 .store(on, std::sync::atomic::Ordering::Relaxed);
2158 }
2159 pub(crate) fn verify_exact_on(&self) -> bool {
2160 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2161 }
2162
2163 pub fn qkv_append_on() -> bool {
2166 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2167 *ON.get_or_init(|| {
2168 std::env::var("MEMRA_QKV_APPEND")
2169 .map(|v| v != "0")
2170 .unwrap_or(true)
2171 })
2172 }
2173
2174 pub fn pdl_wb_on() -> bool {
2177 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2178 *ON.get_or_init(|| {
2179 std::env::var("MEMRA_PDL_WB")
2180 .map(|v| v != "0")
2181 .unwrap_or(true)
2182 })
2183 }
2184
2185 pub fn norm_ilp_on() -> bool {
2193 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2194 *ON.get_or_init(|| {
2195 std::env::var("MEMRA_NORM_ILP")
2196 .map(|v| v != "0")
2197 .unwrap_or(true)
2198 })
2199 }
2200
2201 pub fn tk_ffn_dual_on() -> bool {
2209 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2210 *ON.get_or_init(|| {
2211 std::env::var("MEMRA_TK_FFN_DUAL")
2212 .map(|v| v != "0")
2213 .unwrap_or(true)
2214 })
2215 }
2216
2217 pub fn pdl_mmvq_on() -> bool {
2221 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2222 *ON.get_or_init(|| {
2223 std::env::var("MEMRA_PDL_MMVQ")
2224 .map(|v| v != "0")
2225 .unwrap_or(true)
2226 })
2227 }
2228
2229 pub fn pdl_on() -> bool {
2230 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2231 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2232 }
2233
2234 pub fn pdl_nvfp4q8_on() -> bool {
2240 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2241 *ON.get_or_init(|| {
2242 std::env::var("MEMRA_PDL_NVFP4")
2243 .map(|v| v != "0")
2244 .unwrap_or(true)
2245 })
2246 }
2247
2248 fn q40_mr1_on() -> bool {
2254 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2255 match *Q40MR.get_or_init(|| {
2256 std::env::var("MEMRA_Q40_MR")
2257 .ok()
2258 .and_then(|v| v.parse().ok())
2259 }) {
2260 Some(v) => v == 1,
2261 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2262 }
2263 }
2264
2265 fn pdl_func_flash(
2270 &self,
2271 g: bool,
2272 name: &'static str,
2273 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2274 use cudarc::driver::sys as cu;
2275 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2282 std::sync::Mutex::new(None);
2283 static FNS: std::sync::Mutex<
2284 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2285 > = std::sync::Mutex::new(None);
2286 let ctx_key = self.ctx().cu_ctx() as usize;
2287 if let Some(&f) = FNS
2288 .lock()
2289 .unwrap()
2290 .get_or_insert_with(Default::default)
2291 .get(&(ctx_key, g, name))
2292 {
2293 return Ok(f as cu::CUfunction);
2294 }
2295 let module = {
2296 let mut mods = MODS.lock().unwrap();
2297 let map = mods.get_or_insert_with(Default::default);
2298 match map.get(&(ctx_key, g)) {
2299 Some(&m) => m,
2300 None => {
2301 let m = self.pdl_load_module_in_ctx(if g {
2302 FLASH_FATBIN_KF8VF8
2303 } else {
2304 FLASH_FATBIN
2305 })?;
2306 map.insert((ctx_key, g), m);
2307 m
2308 }
2309 }
2310 };
2311 let cname = std::ffi::CString::new(name)?;
2312 let mut f: cu::CUfunction = std::ptr::null_mut();
2313 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2314 if r != cu::CUresult::CUDA_SUCCESS {
2315 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2316 }
2317 FNS.lock()
2318 .unwrap()
2319 .get_or_insert_with(Default::default)
2320 .insert((ctx_key, g, name), f as usize);
2321 Ok(f)
2322 }
2323
2324 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2329 use cudarc::driver::sys as cu;
2330 let mut prev: cu::CUcontext = std::ptr::null_mut();
2331 unsafe {
2332 cu::cuCtxGetCurrent(&mut prev).result()?;
2333 }
2334 self.ctx().bind_to_thread()?;
2335 let mut m: cu::CUmodule = std::ptr::null_mut();
2336 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2337 let restore = if prev.is_null() {
2338 cu::CUresult::CUDA_SUCCESS
2339 } else {
2340 unsafe { cu::cuCtxSetCurrent(prev) }
2341 };
2342 if r != cu::CUresult::CUDA_SUCCESS {
2343 return Err(format!("pdl module load: {r:?}").into());
2344 }
2345 if restore != cu::CUresult::CUDA_SUCCESS {
2346 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2347 }
2348 Ok(m as usize)
2349 }
2350
2351 pub fn raw_kernel_function(
2354 &self,
2355 name: &'static str,
2356 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2357 self.pdl_func(name)
2358 }
2359
2360 fn pdl_func(
2361 &self,
2362 name: &'static str,
2363 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2364 use cudarc::driver::sys as cu;
2365 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2368 std::sync::Mutex::new(None);
2369 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2372 std::sync::Mutex::new(None);
2373 static FNS: std::sync::Mutex<
2374 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2375 > = std::sync::Mutex::new(None);
2376 let ctx_key = self.ctx().cu_ctx() as usize;
2377 if let Some(&f) = FNS
2378 .lock()
2379 .unwrap()
2380 .get_or_insert_with(Default::default)
2381 .get(&(ctx_key, name))
2382 {
2383 return Ok(f as cu::CUfunction);
2384 }
2385 let module = {
2386 let mut mods = MODULES.lock().unwrap();
2387 let map = mods.get_or_insert_with(Default::default);
2388 match map.get(&ctx_key) {
2389 Some(&m) => m,
2390 None => {
2391 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2392 map.insert(ctx_key, m);
2393 m
2394 }
2395 }
2396 };
2397 let cname = std::ffi::CString::new(name)?;
2398 let mut f: cu::CUfunction = std::ptr::null_mut();
2399 let mut r =
2400 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2401 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2402 let qmodule = {
2403 let mut mods = QMODULES.lock().unwrap();
2404 let map = mods.get_or_insert_with(Default::default);
2405 match map.get(&ctx_key) {
2406 Some(&m) => m,
2407 None => {
2408 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2409 map.insert(ctx_key, m);
2410 m
2411 }
2412 }
2413 };
2414 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2415 }
2416 if r != cu::CUresult::CUDA_SUCCESS {
2417 return Err(format!("pdl_func {name}: {r:?}").into());
2418 }
2419 FNS.lock()
2420 .unwrap()
2421 .get_or_insert_with(Default::default)
2422 .insert((ctx_key, name), f as usize);
2423 Ok(f)
2424 }
2425
2426 unsafe fn launch_pdl_flash(
2438 &self,
2439 g: bool,
2440 name: &'static str,
2441 grid: (u32, u32, u32),
2442 block: (u32, u32, u32),
2443 smem: u32,
2444 params: &mut [*mut std::ffi::c_void],
2445 ) -> Result<(), Box<dyn std::error::Error>> {
2446 use cudarc::driver::sys as cu;
2447 let f = self.pdl_func_flash(g, name)?;
2448 if smem > 0 {
2449 let r =
2451 unsafe {
2452 cu::cuFuncSetAttribute(f,
2453 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2454 smem as i32)
2455 };
2456 if r != cu::CUresult::CUDA_SUCCESS {
2457 return Err(format!("pdl smem attr {name}: {r:?}").into());
2458 }
2459 }
2460 let mut attr = cu::CUlaunchAttribute {
2461 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2462 pad: [0; 4],
2463 value: cu::CUlaunchAttributeValue {
2464 programmaticStreamSerializationAllowed: 1,
2465 },
2466 };
2467 let cfg = cu::CUlaunchConfig {
2468 gridDimX: grid.0,
2469 gridDimY: grid.1,
2470 gridDimZ: grid.2,
2471 blockDimX: block.0,
2472 blockDimY: block.1,
2473 blockDimZ: block.2,
2474 sharedMemBytes: smem,
2475 hStream: self.gpu.stream().cu_stream(),
2476 attrs: &mut attr,
2477 numAttrs: 1,
2478 };
2479 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2480 if r != cu::CUresult::CUDA_SUCCESS {
2481 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2482 }
2483 Ok(())
2484 }
2485
2486 unsafe fn launch_pdl(
2487 &self,
2488 name: &'static str,
2489 grid: (u32, u32, u32),
2490 block: (u32, u32, u32),
2491 params: &mut [*mut std::ffi::c_void],
2492 ) -> Result<(), Box<dyn std::error::Error>> {
2493 use cudarc::driver::sys as cu;
2494 let f = self.pdl_func(name)?;
2495 let mut attr = cu::CUlaunchAttribute {
2496 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2497 pad: [0; 4],
2498 value: cu::CUlaunchAttributeValue {
2499 programmaticStreamSerializationAllowed: 1,
2500 },
2501 };
2502 let cfg = cu::CUlaunchConfig {
2503 gridDimX: grid.0,
2504 gridDimY: grid.1,
2505 gridDimZ: grid.2,
2506 blockDimX: block.0,
2507 blockDimY: block.1,
2508 blockDimZ: block.2,
2509 sharedMemBytes: 0,
2510 hStream: self.gpu.stream().cu_stream(),
2511 attrs: &mut attr,
2512 numAttrs: 1,
2513 };
2514 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2515 if r != cu::CUresult::CUDA_SUCCESS {
2516 return Err(format!("launch_pdl {name}: {r:?}").into());
2517 }
2518 Ok(())
2519 }
2520
2521 pub fn prefetch_weight_l2(
2524 &self,
2525 w: &crate::model::GpuTensor,
2526 ) -> Result<(), Box<dyn std::error::Error>> {
2527 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2528 let p = rp4.as_ref().unwrap_or(bytes);
2529 self.prefetch_l2(p, p.len())?;
2530 }
2531 Ok(())
2532 }
2533
2534 pub fn gather_row_bf16(
2537 &self,
2538 table: &CudaSlice<u8>,
2539 tok: &CudaSlice<u32>,
2540 idx: usize,
2541 dst: &mut CudaSlice<f32>,
2542 ncols: usize,
2543 ) -> Result<(), Box<dyn std::error::Error>> {
2544 let f = self.func("gather_row_bf16_f32");
2545 let cfg = LaunchConfig {
2546 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2547 block_dim: (256, 1, 1),
2548 shared_mem_bytes: 0,
2549 };
2550 let (nc, ix) = (ncols as i32, idx as i32);
2551 let __s_b = self.gpu.stream();
2552 let mut b = __s_b.launch_builder(&f);
2553 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2554 unsafe {
2555 b.launch(cfg)?;
2556 }
2557 Ok(())
2558 }
2559
2560 #[allow(clippy::too_many_arguments)]
2566 pub fn dflash2_dynconv(
2567 &self,
2568 x: &CudaSlice<f32>,
2569 dyn_: &CudaSlice<f32>,
2570 base: &CudaSlice<f32>,
2571 out: &mut CudaSlice<f32>,
2572 rows: usize,
2573 hidden: usize,
2574 group_size: usize,
2575 ksize: usize,
2576 half: usize,
2577 ) -> Result<(), Box<dyn std::error::Error>> {
2578 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2579 let f = self.func("dflash2_dynconv_f32");
2580 let n = rows * hidden;
2581 let cfg = LaunchConfig {
2582 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2583 block_dim: (256, 1, 1),
2584 shared_mem_bytes: 0,
2585 };
2586 let (ri, hi, gi, ki, hf) = (
2587 rows as i32,
2588 hidden as i32,
2589 group_size as i32,
2590 ksize as i32,
2591 half as i32,
2592 );
2593 let __s_b = self.gpu.stream();
2594 let mut b = __s_b.launch_builder(&f);
2595 b.arg(x)
2596 .arg(dyn_)
2597 .arg(base)
2598 .arg(out)
2599 .arg(&ri)
2600 .arg(&hi)
2601 .arg(&gi)
2602 .arg(&ki)
2603 .arg(&hf);
2604 unsafe {
2605 b.launch(cfg)?;
2606 }
2607 Ok(())
2608 }
2609
2610 pub fn topk_rows(
2614 &self,
2615 logits: &CudaSlice<f32>,
2616 n_rows: usize,
2617 n_cols: usize,
2618 k: usize,
2619 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2620 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2621 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2622 let f = self.func("topk_rows_f32");
2623 let nth = 256usize;
2624 let mut vals = self.uninit(n_rows * k)?;
2625 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2626 let cfg = LaunchConfig {
2627 grid_dim: (n_rows as u32, 1, 1),
2628 block_dim: (nth as u32, 1, 1),
2629 shared_mem_bytes: (nth * k * 8) as u32,
2630 };
2631 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2632 let __s_b = self.gpu.stream();
2633 let mut b = __s_b.launch_builder(&f);
2634 b.arg(logits)
2635 .arg(&nr)
2636 .arg(&nc)
2637 .arg(&ki)
2638 .arg(&mut vals)
2639 .arg(&mut idxs);
2640 unsafe {
2641 b.launch(cfg)?;
2642 }
2643 Ok((vals, idxs))
2644 }
2645
2646 pub fn add_row_inplace(
2648 &self,
2649 logits: &mut CudaSlice<f32>,
2650 bias: &CudaSlice<f32>,
2651 n: usize,
2652 row_off: usize,
2653 ) -> Result<(), Box<dyn std::error::Error>> {
2654 let f = self.func("add_row_inplace_f32");
2655 let cfg = LaunchConfig {
2656 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2657 block_dim: (256, 1, 1),
2658 shared_mem_bytes: 0,
2659 };
2660 let (ni, off) = (n as i32, row_off as i64);
2661 let __s_b = self.gpu.stream();
2662 let mut b = __s_b.launch_builder(&f);
2663 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2664 unsafe {
2665 b.launch(cfg)?;
2666 }
2667 Ok(())
2668 }
2669
2670 pub fn prefetch_l2(
2672 &self,
2673 p: &CudaSlice<u8>,
2674 n: usize,
2675 ) -> Result<(), Box<dyn std::error::Error>> {
2676 let f = self.func("prefetch_l2_bytes");
2677 let lines = n.div_ceil(128);
2678 let ni = n as i64;
2679 let cfg = LaunchConfig {
2680 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2681 block_dim: (256, 1, 1),
2682 shared_mem_bytes: 0,
2683 };
2684 let __s_b = self.gpu.stream();
2685 let mut b = __s_b.launch_builder(&f);
2686 b.arg(p).arg(&ni);
2687 unsafe {
2688 b.launch(cfg)?;
2689 }
2690 Ok(())
2691 }
2692
2693 pub fn router_gemv(
2696 &self,
2697 w: &CudaSlice<f32>,
2698 x: &CudaSlice<f32>,
2699 n_embd: usize,
2700 n_experts: usize,
2701 t: usize,
2702 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2703 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2709 Ok("0") => false,
2710 Ok(_) => true,
2711 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2712 };
2713 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2723 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2724 }
2725
2726 pub fn router_gemv_form(
2729 &self,
2730 w: &CudaSlice<f32>,
2731 x: &CudaSlice<f32>,
2732 n_embd: usize,
2733 n_experts: usize,
2734 t: usize,
2735 w8: bool,
2736 batch: bool,
2737 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2738 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2739 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2740 let f = if batch {
2741 self.func("router_gemv_f32_w8_batch")
2742 } else if w8 {
2743 self.func("router_gemv_f32_w8")
2744 } else {
2745 self.func("router_gemv_f32")
2746 };
2747 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2748 let cfg = if batch {
2749 LaunchConfig {
2750 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2751 block_dim: (32, 8, 1),
2752 shared_mem_bytes: 0,
2753 }
2754 } else {
2755 LaunchConfig {
2756 grid_dim: (n_experts as u32, t as u32, 1),
2757 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2758 shared_mem_bytes: 0,
2759 }
2760 };
2761 let __s_b = self.gpu.stream();
2762 let mut b = __s_b.launch_builder(&f);
2763 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2764 unsafe {
2765 b.launch(cfg)?;
2766 }
2767 Ok(y)
2768 }
2769
2770 pub fn router_gemv_into(
2773 &self,
2774 w: &CudaSlice<f32>,
2775 x: &CudaSlice<f32>,
2776 y: &mut CudaSlice<f32>,
2777 n_embd: usize,
2778 n_experts: usize,
2779 t: usize,
2780 ) -> Result<(), Box<dyn std::error::Error>> {
2781 if y.len() < t * n_experts {
2782 return Err("router_gemv_into output too small".into());
2783 }
2784 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2785 Ok("0") => false,
2786 Ok(_) => true,
2787 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2788 };
2789 let f = if w8 {
2790 self.func("router_gemv_f32_w8")
2791 } else {
2792 self.func("router_gemv_f32")
2793 };
2794 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2795 let cfg = LaunchConfig {
2796 grid_dim: (n_experts as u32, t as u32, 1),
2797 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2798 shared_mem_bytes: 0,
2799 };
2800 let __s_b = self.gpu.stream();
2801 let mut b = __s_b.launch_builder(&f);
2802 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2803 unsafe {
2804 b.launch(cfg)?;
2805 }
2806 Ok(())
2807 }
2808
2809 pub fn rows_permute(
2811 &self,
2812 src: &CudaSlice<f32>,
2813 idx: &CudaSlice<i32>,
2814 nrows: usize,
2815 ncols: usize,
2816 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2817 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2818 let f = self.func("rows_permute_f32");
2819 let (nc, nr) = (ncols as i32, nrows as i32);
2820 let cfg = LaunchConfig {
2821 grid_dim: (nrows as u32, 1, 1),
2822 block_dim: (256, 1, 1),
2823 shared_mem_bytes: 0,
2824 };
2825 let __s_b = self.gpu.stream();
2826 let mut b = __s_b.launch_builder(&f);
2827 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2828 unsafe {
2829 b.launch(cfg)?;
2830 }
2831 Ok(dst)
2832 }
2833
2834 pub fn sigmoid_dot_rows(
2839 &self,
2840 x: &CudaSlice<f32>,
2841 w: &CudaSlice<f32>,
2842 n_embd: usize,
2843 t: usize,
2844 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2845 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2848 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2849 let gs = self.linear(x, w, t, n_embd, 1)?;
2850 let mut g = self.uninit(t)?;
2851 self.sigmoid(&gs, &mut g, t)?;
2852 return Ok(g);
2853 }
2854 let mut g = self.alloc_uninit::<f32>(t)?;
2860 let f = self.func("sigmoid_dot_rows_f32");
2861 let (ne, ti) = (n_embd as i32, t as i32);
2862 let cfg = LaunchConfig {
2863 grid_dim: (t as u32, 1, 1),
2864 block_dim: (32, 8, 1),
2865 shared_mem_bytes: 0,
2866 };
2867 let __s_b = self.gpu.stream();
2868 let mut b = __s_b.launch_builder(&f);
2869 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2870 unsafe {
2871 b.launch(cfg)?;
2872 }
2873 Ok(g)
2874 }
2875
2876 pub fn sigmoid_dot_rows_into(
2878 &self,
2879 x: &CudaSlice<f32>,
2880 w: &CudaSlice<f32>,
2881 g: &mut CudaSlice<f32>,
2882 n_embd: usize,
2883 t: usize,
2884 ) -> Result<(), Box<dyn std::error::Error>> {
2885 if g.len() < t {
2886 return Err("sigmoid_dot_rows_into output too small".into());
2887 }
2888 let f = self.func("sigmoid_dot_rows_f32");
2889 let (ne, ti) = (n_embd as i32, t as i32);
2890 let cfg = LaunchConfig {
2891 grid_dim: (t as u32, 1, 1),
2892 block_dim: (32, 8, 1),
2893 shared_mem_bytes: 0,
2894 };
2895 let __s_b = self.gpu.stream();
2896 let mut b = __s_b.launch_builder(&f);
2897 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2898 unsafe {
2899 b.launch(cfg)?;
2900 }
2901 Ok(())
2902 }
2903
2904 pub fn spec_rollback_stream(
2906 &self,
2907 len_ptrs: &CudaSlice<u64>,
2908 pos_start: &CudaSlice<i32>,
2909 acc: &CudaSlice<u32>,
2910 base: usize,
2911 n_rows: usize,
2912 ) -> Result<(), Box<dyn std::error::Error>> {
2913 let f = self.func("spec_rollback_stream");
2914 let (b, nr) = (base as i32, n_rows as i32);
2915 let cfg = LaunchConfig {
2916 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2917 block_dim: (64, 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(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2923 unsafe {
2924 bl.launch(cfg)?;
2925 }
2926 Ok(())
2927 }
2928
2929 pub fn plain_tok_ring(
2931 &self,
2932 vam: &CudaSlice<u32>,
2933 pos_start: &CudaSlice<i32>,
2934 base: usize,
2935 ring: &mut CudaSlice<u32>,
2936 ) -> Result<(), Box<dyn std::error::Error>> {
2937 let f = self.func("plain_tok_ring");
2938 let (b, cap) = (base as i32, ring.len() as i32);
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_bl = self.gpu.stream();
2945 let mut bl = __s_bl.launch_builder(&f);
2946 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2947 unsafe {
2948 bl.launch(cfg)?;
2949 }
2950 Ok(())
2951 }
2952
2953 pub fn spec_ring_commit(
2955 &self,
2956 vtok: &CudaSlice<u32>,
2957 acc: &CudaSlice<u32>,
2958 brk: &CudaSlice<u32>,
2959 ring: &mut CudaSlice<u32>,
2960 pend: &mut CudaSlice<u32>,
2961 ) -> Result<(), Box<dyn std::error::Error>> {
2962 let f = self.func("spec_ring_commit");
2963 let cfg = LaunchConfig {
2964 grid_dim: (1, 1, 1),
2965 block_dim: (32, 1, 1),
2966 shared_mem_bytes: 0,
2967 };
2968 let __s_b = self.gpu.stream();
2969 let mut b = __s_b.launch_builder(&f);
2970 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2971 unsafe {
2972 b.launch(cfg)?;
2973 }
2974 Ok(())
2975 }
2976 pub fn i32_copy_add(
2977 &self,
2978 src: &CudaSlice<i32>,
2979 dst: &mut CudaSlice<i32>,
2980 delta: i32,
2981 ) -> Result<(), Box<dyn std::error::Error>> {
2982 let f = self.func("i32_copy_add");
2983 let cfg = LaunchConfig {
2984 grid_dim: (1, 1, 1),
2985 block_dim: (32, 1, 1),
2986 shared_mem_bytes: 0,
2987 };
2988 let __s_b = self.gpu.stream();
2989 let mut b = __s_b.launch_builder(&f);
2990 b.arg(src).arg(dst).arg(&delta);
2991 unsafe {
2992 b.launch(cfg)?;
2993 }
2994 Ok(())
2995 }
2996 pub fn u32_copy(
2997 &self,
2998 src: &CudaSlice<u32>,
2999 dst: &mut CudaSlice<u32>,
3000 ) -> Result<(), Box<dyn std::error::Error>> {
3001 let f = self.func("u32_copy");
3002 let cfg = LaunchConfig {
3003 grid_dim: (1, 1, 1),
3004 block_dim: (32, 1, 1),
3005 shared_mem_bytes: 0,
3006 };
3007 let __s_b = self.gpu.stream();
3008 let mut b = __s_b.launch_builder(&f);
3009 b.arg(src).arg(dst);
3010 unsafe {
3011 b.launch(cfg)?;
3012 }
3013 Ok(())
3014 }
3015
3016 pub fn spec_adapt_k(
3020 &self,
3021 acc: &CudaSlice<u32>,
3022 brk: &mut CudaSlice<u32>,
3023 floor: usize,
3024 cap: usize,
3025 ) -> Result<(), Box<dyn std::error::Error>> {
3026 let f = self.func("spec_adapt_k");
3027 let (fl, cp) = (floor as i32, cap as i32);
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(acc).arg(brk).arg(&fl).arg(&cp);
3036 unsafe {
3037 b.launch(cfg)?;
3038 }
3039 Ok(())
3040 }
3041
3042 pub fn spec_accept_greedy_dc(
3044 &self,
3045 preds: &CudaSlice<u32>,
3046 vtok: &CudaSlice<u32>,
3047 last_pred: &CudaSlice<u32>,
3048 brk: &CudaSlice<u32>,
3049 out: &mut CudaSlice<u32>,
3050 ) -> Result<(), Box<dyn std::error::Error>> {
3051 let f = self.func("spec_accept_greedy_dc");
3052 let cfg = LaunchConfig {
3053 grid_dim: (1, 1, 1),
3054 block_dim: (32, 1, 1),
3055 shared_mem_bytes: 0,
3056 };
3057 let __s_b = self.gpu.stream();
3058 let mut b = __s_b.launch_builder(&f);
3059 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3060 unsafe {
3061 b.launch(cfg)?;
3062 }
3063 Ok(())
3064 }
3065
3066 pub fn pos_iota(
3068 &self,
3069 pos0: &CudaSlice<i32>,
3070 out: &mut CudaSlice<i32>,
3071 t: usize,
3072 ) -> Result<(), Box<dyn std::error::Error>> {
3073 let f = self.func("pos_iota_i32");
3074 let ti = t as i32;
3075 let cfg = LaunchConfig {
3076 grid_dim: (1, 1, 1),
3077 block_dim: (t.max(1) as u32, 1, 1),
3078 shared_mem_bytes: 0,
3079 };
3080 let __s_b = self.gpu.stream();
3081 let mut b = __s_b.launch_builder(&f);
3082 b.arg(pos0).arg(out).arg(&ti);
3083 unsafe {
3084 b.launch(cfg)?;
3085 }
3086 Ok(())
3087 }
3088 #[allow(clippy::too_many_arguments)]
3089 pub fn append_kv_quantized_rows_dc(
3090 &self,
3091 k_rows: &CudaSlice<f32>,
3092 v_rows: &CudaSlice<f32>,
3093 kc: &mut CudaSlice<u8>,
3094 vc: &mut CudaSlice<u8>,
3095 t0_dev: &CudaSlice<i32>,
3096 t: usize,
3097 kv_dim_k: usize,
3098 kv_dim_v: usize,
3099 k_tok_bytes: usize,
3100 v_tok_bytes: usize,
3101 g: bool,
3102 ) -> Result<(), Box<dyn std::error::Error>> {
3103 let f = if g {
3104 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3105 } else {
3106 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3107 };
3108 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3109 let cfg = LaunchConfig {
3110 grid_dim: (nblk, t as u32, 1),
3111 block_dim: (32, 1, 1),
3112 shared_mem_bytes: 0,
3113 };
3114 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3115 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3116 let __s_b = self.gpu.stream();
3117 let mut b = __s_b.launch_builder(&f);
3118 b.arg(k_rows)
3119 .arg(v_rows)
3120 .arg(kc)
3121 .arg(vc)
3122 .arg(t0_dev)
3123 .arg(&kdk)
3124 .arg(&kdv)
3125 .arg(&ktb)
3126 .arg(&vtb);
3127 unsafe {
3128 b.launch(cfg)?;
3129 }
3130 Ok(())
3131 }
3132
3133 #[allow(clippy::too_many_arguments)]
3136 pub fn append_kv_quantized_row_dc_inc(
3137 &self,
3138 k_row: &CudaSlice<f32>,
3139 v_row: &CudaSlice<f32>,
3140 kc: &mut CudaSlice<u8>,
3141 vc: &mut CudaSlice<u8>,
3142 t0_dev: &mut CudaSlice<i32>,
3143 kv_dim_k: usize,
3144 kv_dim_v: usize,
3145 k_tok_bytes: usize,
3146 v_tok_bytes: usize,
3147 g: bool,
3148 ) -> Result<(), Box<dyn std::error::Error>> {
3149 let f = if g {
3150 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3151 } else {
3152 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3153 };
3154 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3155 let cfg = LaunchConfig {
3156 grid_dim: (1, 1, 1),
3157 block_dim: (nthreads, 1, 1),
3158 shared_mem_bytes: 0,
3159 };
3160 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3161 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3162 let __s_b = self.gpu.stream();
3163 let mut b = __s_b.launch_builder(&f);
3164 b.arg(k_row)
3165 .arg(v_row)
3166 .arg(kc)
3167 .arg(vc)
3168 .arg(t0_dev)
3169 .arg(&kdk)
3170 .arg(&kdv)
3171 .arg(&ktb)
3172 .arg(&vtb);
3173 unsafe {
3174 b.launch(cfg)?;
3175 }
3176 Ok(())
3177 }
3178
3179 pub fn pack_tok_p(
3181 &self,
3182 tok: &CudaSlice<u32>,
3183 p: &CudaSlice<f32>,
3184 out: &mut CudaSlice<u32>,
3185 slot: usize,
3186 ) -> Result<(), Box<dyn std::error::Error>> {
3187 let f = self.func("pack_tok_p");
3188 let sl = slot as i32;
3189 let cfg = LaunchConfig {
3190 grid_dim: (1, 1, 1),
3191 block_dim: (32, 1, 1),
3192 shared_mem_bytes: 0,
3193 };
3194 let __s_b = self.gpu.stream();
3195 let mut b = __s_b.launch_builder(&f);
3196 b.arg(tok).arg(p).arg(out).arg(&sl);
3197 unsafe {
3198 b.launch(cfg)?;
3199 }
3200 Ok(())
3201 }
3202 pub fn tok_map_u32(
3203 &self,
3204 tok: &mut CudaSlice<u32>,
3205 map: &CudaSlice<u32>,
3206 ) -> Result<(), Box<dyn std::error::Error>> {
3207 let f = self.func("tok_map_u32");
3208 let cfg = LaunchConfig {
3209 grid_dim: (1, 1, 1),
3210 block_dim: (32, 1, 1),
3211 shared_mem_bytes: 0,
3212 };
3213 let __s_b = self.gpu.stream();
3214 let mut b = __s_b.launch_builder(&f);
3215 b.arg(tok).arg(map);
3216 unsafe {
3217 b.launch(cfg)?;
3218 }
3219 Ok(())
3220 }
3221
3222 #[allow(clippy::too_many_arguments)]
3224 pub fn spec_assemble_verify(
3225 &self,
3226 tokp: &CudaSlice<u32>,
3227 pend: &CudaSlice<u32>,
3228 d2t: Option<&CudaSlice<u32>>,
3229 vtok: &mut CudaSlice<u32>,
3230 brk: &mut CudaSlice<u32>,
3231 p_min: f32,
3232 k: usize,
3233 pmin0: bool,
3234 ) -> Result<(), Box<dyn std::error::Error>> {
3235 let f = self.func("spec_assemble_verify");
3236 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3237 let cfg = LaunchConfig {
3238 grid_dim: (1, 1, 1),
3239 block_dim: (32, 1, 1),
3240 shared_mem_bytes: 0,
3241 };
3242 let __s_b = self.gpu.stream();
3243 let mut b = __s_b.launch_builder(&f);
3244 match d2t {
3245 Some(m) => {
3246 b.arg(tokp)
3247 .arg(pend)
3248 .arg(m)
3249 .arg(vtok)
3250 .arg(brk)
3251 .arg(&p_min)
3252 .arg(&ki)
3253 .arg(&pm);
3254 unsafe {
3255 b.launch(cfg)?;
3256 }
3257 }
3258 None => {
3259 let null: u64 = 0;
3260 b.arg(tokp)
3261 .arg(pend)
3262 .arg(&null)
3263 .arg(vtok)
3264 .arg(brk)
3265 .arg(&p_min)
3266 .arg(&ki)
3267 .arg(&pm);
3268 unsafe {
3269 b.launch(cfg)?;
3270 }
3271 }
3272 }
3273 Ok(())
3274 }
3275
3276 #[allow(clippy::too_many_arguments)]
3278 pub fn ssm_conv_ring_rebuild_dc(
3279 &self,
3280 qkv_tm: &CudaSlice<f32>,
3281 ring_old: &CudaSlice<f32>,
3282 conv_state: &mut CudaSlice<f32>,
3283 conv_dim: usize,
3284 acc: &CudaSlice<u32>,
3285 base: usize,
3286 t_v: usize,
3287 d_conv: usize,
3288 ) -> Result<(), Box<dyn std::error::Error>> {
3289 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3290 let n = conv_dim * (d_conv - 1);
3291 let cfg = LaunchConfig::for_num_elems(n as u32);
3292 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3293 let __s_b = self.gpu.stream();
3294 let mut b = __s_b.launch_builder(&f);
3295 b.arg(qkv_tm)
3296 .arg(ring_old)
3297 .arg(conv_state)
3298 .arg(&cd)
3299 .arg(acc)
3300 .arg(&b0)
3301 .arg(&tv)
3302 .arg(&dc);
3303 unsafe {
3304 b.launch(cfg)?;
3305 }
3306 Ok(())
3307 }
3308 #[allow(clippy::too_many_arguments)]
3309 pub fn gdn_scan_s128_dc(
3310 &self,
3311 q: &CudaSlice<f32>,
3312 k: &CudaSlice<f32>,
3313 v: &CudaSlice<f32>,
3314 g: &CudaSlice<f32>,
3315 beta: &CudaSlice<f32>,
3316 state_in: &CudaSlice<f32>,
3317 state_out: &mut CudaSlice<f32>,
3318 o: &mut CudaSlice<f32>,
3319 n_head: usize,
3320 acc: &CudaSlice<u32>,
3321 base: usize,
3322 t_v: usize,
3323 scale: f32,
3324 ) -> Result<(), Box<dyn std::error::Error>> {
3325 let f = self.func("gdn_scan_s128_dc");
3326 const S_V: u32 = 128;
3327 const WARP: u32 = 32;
3328 const COLS_PER_BLOCK: u32 = 4;
3329 let cfg = LaunchConfig {
3330 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3331 block_dim: (WARP, COLS_PER_BLOCK, 1),
3332 shared_mem_bytes: 0,
3333 };
3334 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3335 let __s_b = self.gpu.stream();
3336 let mut b = __s_b.launch_builder(&f);
3337 b.arg(q)
3338 .arg(k)
3339 .arg(v)
3340 .arg(g)
3341 .arg(beta)
3342 .arg(state_in)
3343 .arg(state_out)
3344 .arg(o)
3345 .arg(&h)
3346 .arg(acc)
3347 .arg(&b0)
3348 .arg(&tv)
3349 .arg(&scale);
3350 unsafe {
3351 b.launch(cfg)?;
3352 }
3353 Ok(())
3354 }
3355
3356 pub fn spec_rollback_kv(
3358 &self,
3359 len_ptrs: &CudaSlice<u64>,
3360 saved: &CudaSlice<i32>,
3361 acc: &CudaSlice<u32>,
3362 base: usize,
3363 n_layer: usize,
3364 ) -> Result<(), Box<dyn std::error::Error>> {
3365 let f = self.func("spec_rollback_kv");
3366 let (b, nl) = (base as i32, n_layer as i32);
3367 let cfg = LaunchConfig {
3368 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3369 block_dim: (64, 1, 1),
3370 shared_mem_bytes: 0,
3371 };
3372 let __s_bl = self.gpu.stream();
3373 let mut bl = __s_bl.launch_builder(&f);
3374 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3375 unsafe {
3376 bl.launch(cfg)?;
3377 }
3378 Ok(())
3379 }
3380
3381 pub fn spec_fork_valid(
3383 &self,
3384 acc: &CudaSlice<u32>,
3385 optimistic_pending: u32,
3386 valid: &mut CudaSlice<u32>,
3387 ) -> Result<(), Box<dyn std::error::Error>> {
3388 let f = self.func("spec_fork_valid");
3389 let cfg = LaunchConfig {
3390 grid_dim: (1, 1, 1),
3391 block_dim: (1, 1, 1),
3392 shared_mem_bytes: 0,
3393 };
3394 let __s_bl = self.gpu.stream();
3395 let mut bl = __s_bl.launch_builder(&f);
3396 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3397 unsafe {
3398 bl.launch(cfg)?;
3399 }
3400 Ok(())
3401 }
3402
3403 pub fn spec_fork_reconcile_kv(
3405 &self,
3406 len_ptrs: &CudaSlice<u64>,
3407 saved: &CudaSlice<i32>,
3408 acc: &CudaSlice<u32>,
3409 valid: &CudaSlice<u32>,
3410 base: usize,
3411 n_layer: usize,
3412 ) -> Result<(), Box<dyn std::error::Error>> {
3413 let f = self.func("spec_fork_reconcile_kv");
3414 let (b, nl) = (base as i32, n_layer as i32);
3415 let cfg = LaunchConfig {
3416 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3417 block_dim: (64, 1, 1),
3418 shared_mem_bytes: 0,
3419 };
3420 let __s_bl = self.gpu.stream();
3421 let mut bl = __s_bl.launch_builder(&f);
3422 bl.arg(len_ptrs)
3423 .arg(saved)
3424 .arg(acc)
3425 .arg(valid)
3426 .arg(&b)
3427 .arg(&nl);
3428 unsafe {
3429 bl.launch(cfg)?;
3430 }
3431 Ok(())
3432 }
3433
3434 pub fn spec_fork_restore_f32(
3436 &self,
3437 snapshot: &CudaSlice<f32>,
3438 state: &mut CudaSlice<f32>,
3439 valid: &CudaSlice<u32>,
3440 ) -> Result<(), Box<dyn std::error::Error>> {
3441 assert_eq!(
3442 snapshot.len(),
3443 state.len(),
3444 "fork recurrent snapshot shape mismatch"
3445 );
3446 let f = self.func("spec_fork_restore_f32");
3447 let n = state.len() as i32;
3448 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3449 let cfg = LaunchConfig {
3450 grid_dim: (blocks, 1, 1),
3451 block_dim: (256, 1, 1),
3452 shared_mem_bytes: 0,
3453 };
3454 let __s_bl = self.gpu.stream();
3455 let mut bl = __s_bl.launch_builder(&f);
3456 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3457 unsafe {
3458 bl.launch(cfg)?;
3459 }
3460 Ok(())
3461 }
3462
3463 pub fn spec_seed_gather(
3466 &self,
3467 vx: &CudaSlice<f32>,
3468 fill_prev: &CudaSlice<f32>,
3469 acc: &CudaSlice<u32>,
3470 h_seed: &mut CudaSlice<f32>,
3471 base: usize,
3472 n_embd: usize,
3473 ) -> Result<(), Box<dyn std::error::Error>> {
3474 let f = self.func("spec_seed_gather");
3475 let (b, ne) = (base as i32, n_embd as i32);
3476 let cfg = LaunchConfig {
3477 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3478 block_dim: (256, 1, 1),
3479 shared_mem_bytes: 0,
3480 };
3481 let __s_bl = self.gpu.stream();
3482 let mut bl = __s_bl.launch_builder(&f);
3483 bl.arg(vx)
3484 .arg(fill_prev)
3485 .arg(acc)
3486 .arg(h_seed)
3487 .arg(&b)
3488 .arg(&ne);
3489 unsafe {
3490 bl.launch(cfg)?;
3491 }
3492 Ok(())
3493 }
3494
3495 pub fn spec_accept_greedy(
3497 &self,
3498 preds: &CudaSlice<u32>,
3499 draft: &CudaSlice<u32>,
3500 last_pred: u32,
3501 base: usize,
3502 k_round: usize,
3503 out: &mut CudaSlice<u32>,
3504 ) -> Result<(), Box<dyn std::error::Error>> {
3505 let f = self.func("spec_accept_greedy");
3506 let (b, k) = (base as i32, k_round as i32);
3507 let cfg = LaunchConfig {
3508 grid_dim: (1, 1, 1),
3509 block_dim: (32, 1, 1),
3510 shared_mem_bytes: 0,
3511 };
3512 let __s_bl = self.gpu.stream();
3513 let mut bl = __s_bl.launch_builder(&f);
3514 bl.arg(preds)
3515 .arg(draft)
3516 .arg(&last_pred)
3517 .arg(&b)
3518 .arg(&k)
3519 .arg(out);
3520 unsafe {
3521 bl.launch(cfg)?;
3522 }
3523 Ok(())
3524 }
3525
3526 pub fn gumbel_perturb(
3533 &self,
3534 x: &CudaSlice<f32>,
3535 y: &mut CudaSlice<f32>,
3536 n: usize,
3537 seed: u64,
3538 stream_pos: u32,
3539 temp: f32,
3540 ) -> Result<(), Box<dyn std::error::Error>> {
3541 let f = self.func("gumbel_perturb_f32");
3542 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3543 let cfg = LaunchConfig {
3544 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3545 block_dim: (256, 1, 1),
3546 shared_mem_bytes: 0,
3547 };
3548 let __s_b = self.gpu.stream();
3549 let mut b = __s_b.launch_builder(&f);
3550 b.arg(x)
3551 .arg(&mut *y)
3552 .arg(&ni)
3553 .arg(&slo)
3554 .arg(&shi)
3555 .arg(&stream_pos)
3556 .arg(&temp);
3557 unsafe {
3558 b.launch(cfg)?;
3559 }
3560 Ok(())
3561 }
3562
3563 pub fn mask_logits_col(
3571 &self,
3572 logits: &mut CudaSlice<f32>,
3573 mask: &CudaSlice<u32>,
3574 col: usize,
3575 n: usize,
3576 mask_words: usize,
3577 ) -> Result<(), Box<dyn std::error::Error>> {
3578 let f = self.func("mask_logits_f32");
3579 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3580 let cfg = LaunchConfig {
3581 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3582 block_dim: (256, 1, 1),
3583 shared_mem_bytes: 0,
3584 };
3585 let __s_b = self.gpu.stream();
3586 let mut b = __s_b.launch_builder(&f);
3587 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3588 unsafe {
3589 b.launch(cfg)?;
3590 }
3591 Ok(())
3592 }
3593
3594 pub fn gumbel_perturb_col(
3601 &self,
3602 x: &CudaSlice<f32>,
3603 col: usize,
3604 y: &mut CudaSlice<f32>,
3605 n: usize,
3606 seed: u64,
3607 stream_pos: u32,
3608 temp: f32,
3609 ) -> Result<(), Box<dyn std::error::Error>> {
3610 let f = self.func("gumbel_perturb_f32");
3611 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3612 let col_view = x.slice(col * n..(col + 1) * n);
3613 let cfg = LaunchConfig {
3614 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3615 block_dim: (256, 1, 1),
3616 shared_mem_bytes: 0,
3617 };
3618 let __s_b = self.gpu.stream();
3619 let mut b = __s_b.launch_builder(&f);
3620 b.arg(&col_view)
3621 .arg(&mut *y)
3622 .arg(&ni)
3623 .arg(&slo)
3624 .arg(&shi)
3625 .arg(&stream_pos)
3626 .arg(&temp);
3627 unsafe {
3628 b.launch(cfg)?;
3629 }
3630 Ok(())
3631 }
3632
3633 #[allow(clippy::too_many_arguments)]
3639 pub fn gumbel_perturb_filtered_col(
3640 &self,
3641 x: &CudaSlice<f32>,
3642 col: usize,
3643 y: &mut CudaSlice<f32>,
3644 n: usize,
3645 seed: u64,
3646 stream_pos: u32,
3647 temp: f32,
3648 stat_max: &CudaSlice<f32>,
3649 stat_th: &CudaSlice<f32>,
3650 stat_idx: usize,
3651 ) -> Result<(), Box<dyn std::error::Error>> {
3652 let f = self.func("gumbel_perturb_filtered_col_f32");
3653 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3654 let (ci, si) = (col as i32, stat_idx as i32);
3655 let cfg = LaunchConfig {
3656 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3657 block_dim: (256, 1, 1),
3658 shared_mem_bytes: 0,
3659 };
3660 let __s_b = self.gpu.stream();
3661 let mut b = __s_b.launch_builder(&f);
3662 b.arg(x)
3663 .arg(&ci)
3664 .arg(&mut *y)
3665 .arg(&ni)
3666 .arg(&slo)
3667 .arg(&shi)
3668 .arg(&stream_pos)
3669 .arg(&temp)
3670 .arg(stat_max)
3671 .arg(stat_th)
3672 .arg(&si);
3673 unsafe {
3674 b.launch(cfg)?;
3675 }
3676 Ok(())
3677 }
3678
3679 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3684 let f = self.func("memra_sctr_inc");
3685 let cfg = LaunchConfig {
3686 grid_dim: (1, 1, 1),
3687 block_dim: (1, 1, 1),
3688 shared_mem_bytes: 0,
3689 };
3690 let __s_b = self.gpu.stream();
3691 let mut b = __s_b.launch_builder(&f);
3692 b.arg(&mut *ctr);
3693 unsafe {
3694 b.launch(cfg)?;
3695 }
3696 Ok(())
3697 }
3698
3699 pub fn gumbel_perturb_ctr(
3704 &self,
3705 x: &CudaSlice<f32>,
3706 y: &mut CudaSlice<f32>,
3707 n: usize,
3708 seed: u64,
3709 ctr: &CudaSlice<u32>,
3710 temp: f32,
3711 ) -> Result<(), Box<dyn std::error::Error>> {
3712 let f = self.func("gumbel_perturb_ctr_f32");
3713 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3714 let cfg = LaunchConfig {
3715 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3716 block_dim: (256, 1, 1),
3717 shared_mem_bytes: 0,
3718 };
3719 let __s_b = self.gpu.stream();
3720 let mut b = __s_b.launch_builder(&f);
3721 b.arg(x)
3722 .arg(&mut *y)
3723 .arg(&ni)
3724 .arg(&slo)
3725 .arg(&shi)
3726 .arg(ctr)
3727 .arg(&temp);
3728 unsafe {
3729 b.launch(cfg)?;
3730 }
3731 Ok(())
3732 }
3733
3734 pub fn softmax_gather(
3738 &self,
3739 x: &CudaSlice<f32>,
3740 row_stride: usize,
3741 ids: &CudaSlice<u32>,
3742 rows: &CudaSlice<i32>,
3743 out: &mut CudaSlice<f32>,
3744 n: usize,
3745 npair: usize,
3746 temp: f32,
3747 ) -> Result<(), Box<dyn std::error::Error>> {
3748 let f = self.func("softmax_gather_f32");
3749 let (ni, rs) = (n as i32, row_stride as i64);
3750 let np = npair as i32;
3751 let cfg = LaunchConfig {
3752 grid_dim: (npair as u32, 1, 1),
3753 block_dim: (256, 1, 1),
3754 shared_mem_bytes: 0,
3755 };
3756 let __s_b = self.gpu.stream();
3757 let mut b = __s_b.launch_builder(&f);
3758 b.arg(x)
3759 .arg(&rs)
3760 .arg(ids)
3761 .arg(rows)
3762 .arg(&mut *out)
3763 .arg(&ni)
3764 .arg(&np)
3765 .arg(&temp);
3766 unsafe {
3767 b.launch(cfg)?;
3768 }
3769 Ok(())
3770 }
3771
3772 pub fn residual_sample(
3776 &self,
3777 p: &CudaSlice<f32>,
3778 q: Option<&CudaSlice<f32>>,
3779 n: usize,
3780 temp: f32,
3781 seed: u64,
3782 stream_pos: u32,
3783 out_tok: &mut CudaSlice<u32>,
3784 ) -> Result<(), Box<dyn std::error::Error>> {
3785 let f = self.func("residual_sample_f32");
3786 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3787 let nth = 1024u32;
3788 let cfg = LaunchConfig {
3789 grid_dim: (1, 1, 1),
3790 block_dim: (nth, 1, 1),
3791 shared_mem_bytes: 0,
3792 };
3793 let has_q: i32 = q.is_some() as i32;
3794 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3796 let mut b = __s_b.launch_builder(&f);
3797 b.arg(p)
3798 .arg(qbuf)
3799 .arg(&has_q)
3800 .arg(&ni)
3801 .arg(&temp)
3802 .arg(&slo)
3803 .arg(&shi)
3804 .arg(&stream_pos)
3805 .arg(&mut *out_tok);
3806 unsafe {
3807 b.launch(cfg)?;
3808 }
3809 Ok(())
3810 }
3811
3812 pub fn with_moe_cache<R>(
3817 &self,
3818 max_block_bytes: usize,
3819 f: impl FnOnce(
3820 &mut crate::moe_cache::MoeSlotCache,
3821 &Engine,
3822 ) -> Result<R, Box<dyn std::error::Error>>,
3823 ) -> Result<R, Box<dyn std::error::Error>> {
3824 let mut guard = self.moe_cache.lock().unwrap();
3825 if guard.is_none() {
3826 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3827 }
3828 let cache = guard.as_mut().unwrap();
3829 f(cache, self)
3830 }
3831
3832 pub fn freeze_moe_cache(&self) {
3835 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3836 cache.freeze();
3837 }
3838 }
3839
3840 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3843 self.moe_cache
3844 .lock()
3845 .unwrap()
3846 .as_ref()
3847 .map(crate::moe_cache::MoeSlotCache::export_residency)
3848 }
3849
3850 pub(crate) fn moe_cache_frozen(&self) -> bool {
3851 self.moe_cache
3852 .lock()
3853 .unwrap()
3854 .as_ref()
3855 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3856 }
3857
3858 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3865 crate::cpu_experts::configured()
3866 && self.moe_cache_frozen()
3867 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3868 }
3869
3870 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3872 assert!(
3873 self.moe_cache.lock().unwrap().is_none(),
3874 "MoE cache layout configured after cache construction"
3875 );
3876 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3877 }
3878
3879 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3880 self.moe_cache_layout.lock().unwrap().clone()
3881 }
3882
3883 pub fn moe_cache_enabled() -> bool {
3885 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3886 }
3887
3888 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3891 let guard = self.moe_cache.lock().unwrap();
3892 guard
3893 .as_ref()
3894 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3895 }
3896
3897 pub fn cpu_expert_stats(
3901 &self,
3902 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3903 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3904 }
3905
3906 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3909 crate::cpu_experts::predictor_stats()
3910 }
3911
3912 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3913 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3914 }
3915
3916 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3919 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3920 }
3921
3922 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3925 let guard = self.moe_cache.lock().unwrap();
3926 guard
3927 .as_ref()
3928 .and_then(|cache| cache.pread_stats())
3929 .map(|stats| {
3930 (
3931 stats.reads,
3932 stats.bytes,
3933 stats.read_errors,
3934 stats.short_reads,
3935 stats.fallbacks,
3936 stats.buffer_waits,
3937 stats.ring_full,
3938 )
3939 })
3940 }
3941
3942 pub fn spill_config_fallbacks(&self) -> u64 {
3944 crate::spill_pread::config_fallbacks()
3945 }
3946
3947 pub fn moe_cache_reset_counters(&self) {
3949 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3950 c.reset_counters();
3951 }
3952 }
3953
3954 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3955 Ok(self.gpu.stream().clone_htod(v)?)
3956 }
3957
3958 pub fn htod_bytes_padded(
3962 &self,
3963 v: &[u8],
3964 pad: usize,
3965 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3966 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3967 {
3968 let mut view = d.slice_mut(0..v.len());
3969 self.gpu.stream().memcpy_htod(v, &mut view)?;
3970 }
3971 Ok(d)
3972 }
3973
3974 pub fn copy_into(
3976 &self,
3977 dst: &mut CudaSlice<f32>,
3978 off: usize,
3979 src: &CudaSlice<f32>,
3980 len: usize,
3981 ) -> Result<(), Box<dyn std::error::Error>> {
3982 let mut view = dst.slice_mut(off..off + len);
3983 self.gpu
3984 .stream()
3985 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3986 Ok(())
3987 }
3988
3989 pub fn copy_u8_into(
3992 &self,
3993 dst: &mut CudaSlice<u8>,
3994 off: usize,
3995 src: &CudaSlice<u8>,
3996 len: usize,
3997 ) -> Result<(), Box<dyn std::error::Error>> {
3998 let mut view = dst.slice_mut(off..off + len);
3999 self.gpu
4000 .stream()
4001 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4002 Ok(())
4003 }
4004
4005 pub fn copy_u8_range_into(
4007 &self,
4008 dst: &mut CudaSlice<u8>,
4009 dst_off: usize,
4010 src: &CudaSlice<u8>,
4011 src_off: usize,
4012 len: usize,
4013 ) -> Result<(), Box<dyn std::error::Error>> {
4014 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
4015 self.gpu
4016 .stream()
4017 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4018 Ok(())
4019 }
4020
4021 pub fn prepare_kv_append(
4025 &self,
4026 kv: &mut crate::cache::KvLayer,
4027 retain_from: usize,
4028 append_rows: usize,
4029 ) -> Result<usize, Box<dyn std::error::Error>> {
4030 let Some(plan) = kv
4031 .ring
4032 .as_ref()
4033 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4034 .transpose()?
4035 else {
4036 return Ok(kv.len);
4037 };
4038 match plan {
4039 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4040 crate::cache::KvRingAppend::Rebase {
4041 src_row,
4042 keep_rows,
4043 new_base,
4044 write_row,
4045 } => {
4046 if keep_rows > 0 {
4047 let k_len = keep_rows * kv.k_tok_bytes;
4048 let v_len = keep_rows * kv.v_tok_bytes;
4049 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4050 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4051 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4052 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4053 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4054 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4055 }
4056 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4057 Ok(write_row)
4058 }
4059 }
4060 }
4061
4062 pub fn htod_u8_into(
4065 &self,
4066 dst: &mut CudaSlice<u8>,
4067 off: usize,
4068 src: &[u8],
4069 ) -> Result<(), Box<dyn std::error::Error>> {
4070 let mut view = dst.slice_mut(off..off + src.len());
4071 self.gpu.stream().memcpy_htod(src, &mut view)?;
4072 Ok(())
4073 }
4074
4075 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4076 b.slice(0..len)
4077 }
4078
4079 pub fn view_u8_range<'a>(
4082 &self,
4083 b: &'a CudaSlice<u8>,
4084 start: usize,
4085 end: usize,
4086 ) -> cudarc::driver::CudaView<'a, u8> {
4087 b.slice(start..end)
4088 }
4089 pub fn view_u8<'a>(
4090 &self,
4091 b: &'a CudaSlice<u8>,
4092 len: usize,
4093 ) -> cudarc::driver::CudaView<'a, u8> {
4094 b.slice(0..len)
4095 }
4096
4097 pub fn append_kv_quantized(
4101 &self,
4102 k_row: &CudaSlice<f32>,
4103 v_row: &CudaSlice<f32>,
4104 kc: &mut CudaSlice<u8>,
4105 vc: &mut CudaSlice<u8>,
4106 t: usize,
4107 kv_dim_k: usize,
4108 kv_dim_v: usize,
4109 k_tok_bytes: usize,
4110 v_tok_bytes: usize,
4111 g: bool,
4112 ) -> Result<(), Box<dyn std::error::Error>> {
4113 let f = if g {
4114 self.func_g("append_quantize_kv_q8_0_q5_1")
4115 } else {
4116 self.func("append_quantize_kv_q8_0_q5_1")
4117 };
4118 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4119 let cfg = LaunchConfig {
4120 grid_dim: (nblk, 1, 1),
4121 block_dim: (32, 1, 1),
4122 shared_mem_bytes: 0,
4123 };
4124 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4125 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4126 let __s_b = self.gpu.stream();
4127 let mut b = __s_b.launch_builder(&f);
4128 b.arg(k_row)
4129 .arg(v_row)
4130 .arg(kc)
4131 .arg(vc)
4132 .arg(&ti)
4133 .arg(&kdk)
4134 .arg(&kdv)
4135 .arg(&ktb)
4136 .arg(&vtb);
4137 unsafe {
4138 b.launch(cfg)?;
4139 }
4140 Ok(())
4141 }
4142
4143 pub fn append_kv_quantized_dc(
4147 &self,
4148 k_row: &CudaSlice<f32>,
4149 v_row: &CudaSlice<f32>,
4150 kc: &mut CudaSlice<u8>,
4151 vc: &mut CudaSlice<u8>,
4152 t_dev: &CudaSlice<i32>,
4153 kv_dim_k: usize,
4154 kv_dim_v: usize,
4155 k_tok_bytes: usize,
4156 v_tok_bytes: usize,
4157 g: bool,
4158 ) -> Result<(), Box<dyn std::error::Error>> {
4159 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4160 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4161 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4162 if Self::pdl_on() && Self::pdl_wb_on() {
4164 use cudarc::driver::{DevicePtr, DevicePtrMut};
4165 let s = &self.gpu.stream();
4166 let (pk, _g0) = k_row.device_ptr(s);
4167 let (pv, _g1) = v_row.device_ptr(s);
4168 let (pkc, _g2) = kc.device_ptr_mut(s);
4169 let (pvc, _g3) = vc.device_ptr_mut(s);
4170 let (pt, _g4) = t_dev.device_ptr(s);
4171 let mut ps = [
4172 &pk as *const _ as *mut std::ffi::c_void,
4173 &pv as *const _ as *mut _,
4174 &pkc as *const _ as *mut _,
4175 &pvc as *const _ as *mut _,
4176 &pt as *const _ as *mut _,
4177 &kdk as *const _ as *mut _,
4178 &kdv as *const _ as *mut _,
4179 &ktb as *const _ as *mut _,
4180 &vtb as *const _ as *mut _,
4181 ];
4182 unsafe {
4183 self.launch_pdl_flash(
4184 g,
4185 "append_quantize_kv_q8_0_q5_1_dc",
4186 (nblk, 1, 1),
4187 (32, 1, 1),
4188 0,
4189 &mut ps,
4190 )?;
4191 }
4192 return Ok(());
4193 }
4194 let f = if g {
4195 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4196 } else {
4197 self.func("append_quantize_kv_q8_0_q5_1_dc")
4198 };
4199 let cfg = LaunchConfig {
4200 grid_dim: (nblk, 1, 1),
4201 block_dim: (32, 1, 1),
4202 shared_mem_bytes: 0,
4203 };
4204 let __s_b = self.gpu.stream();
4205 let mut b = __s_b.launch_builder(&f);
4206 b.arg(k_row)
4207 .arg(v_row)
4208 .arg(kc)
4209 .arg(vc)
4210 .arg(t_dev)
4211 .arg(&kdk)
4212 .arg(&kdv)
4213 .arg(&ktb)
4214 .arg(&vtb);
4215 unsafe {
4216 b.launch(cfg)?;
4217 }
4218 Ok(())
4219 }
4220
4221 #[allow(clippy::too_many_arguments)]
4228 pub fn append_kv_quantized_rows(
4229 &self,
4230 k_rows: &CudaSlice<f32>,
4231 v_rows: &CudaSlice<f32>,
4232 kc: &mut CudaSlice<u8>,
4233 vc: &mut CudaSlice<u8>,
4234 t0: usize,
4235 t: usize,
4236 kv_dim_k: usize,
4237 kv_dim_v: usize,
4238 k_tok_bytes: usize,
4239 v_tok_bytes: usize,
4240 g: bool,
4241 ) -> Result<(), Box<dyn std::error::Error>> {
4242 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4243 for i in 0..t {
4244 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4245 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4246 self.append_kv_quantized_view(
4247 &k_row,
4248 &v_row,
4249 kc,
4250 vc,
4251 t0 + i,
4252 kv_dim_k,
4253 kv_dim_v,
4254 k_tok_bytes,
4255 v_tok_bytes,
4256 g,
4257 )?;
4258 }
4259 return Ok(());
4260 }
4261 let f = if g {
4262 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4263 } else {
4264 self.func("append_quantize_kv_q8_0_q5_1_rows")
4265 };
4266 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4267 let cfg = LaunchConfig {
4268 grid_dim: (nblk, t as u32, 1),
4269 block_dim: (32, 1, 1),
4270 shared_mem_bytes: 0,
4271 };
4272 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4273 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4274 let __s_b = self.gpu.stream();
4275 let mut b = __s_b.launch_builder(&f);
4276 b.arg(k_rows)
4277 .arg(v_rows)
4278 .arg(kc)
4279 .arg(vc)
4280 .arg(&t0i)
4281 .arg(&kdk)
4282 .arg(&kdv)
4283 .arg(&ktb)
4284 .arg(&vtb);
4285 unsafe {
4286 b.launch(cfg)?;
4287 }
4288 Ok(())
4289 }
4290
4291 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4295 let f = self.func("inc_i32");
4296 let cfg = LaunchConfig {
4297 grid_dim: (1, 1, 1),
4298 block_dim: (1, 1, 1),
4299 shared_mem_bytes: 0,
4300 };
4301 let __s_b = self.gpu.stream();
4302 let mut b = __s_b.launch_builder(&f);
4303 b.arg(p);
4304 unsafe {
4305 b.launch(cfg)?;
4306 }
4307 Ok(())
4308 }
4309
4310 pub fn append_kv_quantized_view(
4313 &self,
4314 k_row: &cudarc::driver::CudaView<f32>,
4315 v_row: &cudarc::driver::CudaView<f32>,
4316 kc: &mut CudaSlice<u8>,
4317 vc: &mut CudaSlice<u8>,
4318 t: usize,
4319 kv_dim_k: usize,
4320 kv_dim_v: usize,
4321 k_tok_bytes: usize,
4322 v_tok_bytes: usize,
4323 g: bool,
4324 ) -> Result<(), Box<dyn std::error::Error>> {
4325 let stream = self.gpu.stream();
4326 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4327 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4328 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4329 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4330 let f = if g {
4331 self.func_g("append_quantize_kv_q8_0_q5_1")
4332 } else {
4333 self.func("append_quantize_kv_q8_0_q5_1")
4334 };
4335 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4336 let cfg = LaunchConfig {
4337 grid_dim: (nblk, 1, 1),
4338 block_dim: (32, 1, 1),
4339 shared_mem_bytes: 0,
4340 };
4341 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4342 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4343 let mut b = stream.launch_builder(&f);
4344 b.arg(k_row)
4345 .arg(v_row)
4346 .arg(kc)
4347 .arg(vc)
4348 .arg(&ti)
4349 .arg(&kdk)
4350 .arg(&kdv)
4351 .arg(&ktb)
4352 .arg(&vtb);
4353 unsafe {
4354 b.launch(cfg)?;
4355 }
4356 Ok(())
4357 }
4358
4359 pub fn copy_view_into(
4362 &self,
4363 dst: &mut CudaSlice<f32>,
4364 off: usize,
4365 src: &cudarc::driver::CudaView<f32>,
4366 len: usize,
4367 ) -> Result<(), Box<dyn std::error::Error>> {
4368 let mut view = dst.slice_mut(off..off + len);
4369 self.gpu
4370 .stream()
4371 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4372 Ok(())
4373 }
4374
4375 pub fn clone_dtod(
4379 &self,
4380 src: &CudaSlice<f32>,
4381 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4382 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4383 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4384 Ok(dst)
4385 }
4386
4387 pub fn dtod_copy_view(
4390 &self,
4391 src: &cudarc::driver::CudaView<f32>,
4392 dst: &mut CudaSlice<f32>,
4393 ) -> Result<(), Box<dyn std::error::Error>> {
4394 self.gpu.stream().memcpy_dtod(src, dst)?;
4395 Ok(())
4396 }
4397
4398 pub fn dtod_copy_view_i8(
4400 &self,
4401 src: &cudarc::driver::CudaView<i8>,
4402 dst: &mut CudaSlice<i8>,
4403 ) -> Result<(), Box<dyn std::error::Error>> {
4404 self.gpu.stream().memcpy_dtod(src, dst)?;
4405 Ok(())
4406 }
4407
4408 pub fn dtod_copy_into(
4410 &self,
4411 src: &CudaSlice<f32>,
4412 dst: &mut CudaSlice<f32>,
4413 offset: usize,
4414 ) -> Result<(), Box<dyn std::error::Error>> {
4415 let n = src.len();
4416 let mut dv = dst.slice_mut(offset..offset + n);
4417 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4418 Ok(())
4419 }
4420
4421 pub fn copy_batch_uniform_f32(
4427 &self,
4428 table: &CudaSlice<u64>,
4429 n: usize,
4430 words: usize,
4431 ) -> Result<(), Box<dyn std::error::Error>> {
4432 if n == 0 || words == 0 {
4433 return Ok(());
4434 }
4435 debug_assert!(
4436 table.len() >= 2 * n,
4437 "pointer table must hold n srcs + n dsts"
4438 );
4439 let f = self.func("copy_batch_uniform_f32");
4440 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4443 let (ni, wi) = (n as i32, words as i32);
4444 let cfg = LaunchConfig {
4445 grid_dim: (chunks, n as u32, 1),
4446 block_dim: (256, 1, 1),
4447 shared_mem_bytes: 0,
4448 };
4449 let __s = self.gpu.stream();
4450 let mut b = __s.launch_builder(&f);
4451 b.arg(table).arg(&ni).arg(&wi);
4452 unsafe {
4453 b.launch(cfg)?;
4454 }
4455 Ok(())
4456 }
4457
4458 pub fn htod_u64_into(
4461 &self,
4462 v: &[u64],
4463 dst: &mut CudaSlice<u64>,
4464 ) -> Result<(), Box<dyn std::error::Error>> {
4465 let mut view = dst.slice_mut(0..v.len());
4466 self.gpu.stream().memcpy_htod(v, &mut view)?;
4467 Ok(())
4468 }
4469
4470 pub fn copy_indirect_src_f32(
4475 &self,
4476 src_entry: &cudarc::driver::CudaView<u64>,
4477 dst: &mut CudaSlice<f32>,
4478 dst_off: usize,
4479 words: usize,
4480 ) -> Result<(), Box<dyn std::error::Error>> {
4481 let f = self.func("copy_indirect_src_f32");
4482 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4483 let wi = words as i32;
4484 let cfg = LaunchConfig {
4485 grid_dim: (chunks, 1, 1),
4486 block_dim: (256, 1, 1),
4487 shared_mem_bytes: 0,
4488 };
4489 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4490 let __s = self.gpu.stream();
4491 let mut b = __s.launch_builder(&f);
4492 b.arg(src_entry).arg(&mut dv).arg(&wi);
4493 unsafe {
4494 b.launch(cfg)?;
4495 }
4496 Ok(())
4497 }
4498
4499 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4501 self.alloc_uninit::<i8>(n)
4502 }
4503
4504 pub fn qmatvec(
4506 &self,
4507 w: &CudaSlice<u8>,
4508 x: &CudaSlice<f32>,
4509 m: usize,
4510 in_f: usize,
4511 out_f: usize,
4512 qtype: i32,
4513 row_bytes: usize,
4514 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4515 let f = self.func("qmatvec_f32");
4516 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4518 grid_dim: (out_f as u32, m as u32, 1),
4519 block_dim: (256, 1, 1),
4520 shared_mem_bytes: 0,
4521 };
4522 let (inf, outf, mi, qt, rb) =
4523 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4524 let __s_b = self.gpu.stream();
4525 let mut b = __s_b.launch_builder(&f);
4526 b.arg(w)
4527 .arg(x)
4528 .arg(&mut y)
4529 .arg(&inf)
4530 .arg(&outf)
4531 .arg(&mi)
4532 .arg(&qt)
4533 .arg(&rb);
4534 unsafe {
4535 b.launch(cfg)?;
4536 }
4537 Ok(y)
4538 }
4539
4540 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4542 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4543 self.keep_if_capturing(&s);
4544 Ok(s)
4545 }
4546
4547 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4551 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4552 self.keep_if_capturing(&s);
4553 Ok(s)
4554 }
4555
4556 pub fn memset_zeros_view(
4559 &self,
4560 dst: &mut cudarc::driver::CudaViewMut<f32>,
4561 ) -> Result<(), Box<dyn std::error::Error>> {
4562 self.gpu.stream().memset_zeros(dst)?;
4563 Ok(())
4564 }
4565
4566 pub fn stage_expert(
4572 &self,
4573 host_bytes: &[u8],
4574 scratch: &mut CudaSlice<u8>,
4575 off: usize,
4576 ) -> Result<(), Box<dyn std::error::Error>> {
4577 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4580 }
4581
4582 pub fn moe_router_topk(
4588 &self,
4589 logits: &CudaSlice<f32>,
4590 t: usize,
4591 n_expert: usize,
4592 n_used: usize,
4593 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4594 let f = self.func("moe_router_topk_f32");
4595 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 {
4598 grid_dim: (t as u32, 1, 1),
4599 block_dim: (n_expert as u32, 1, 1),
4600 shared_mem_bytes: 0,
4601 };
4602 let (ne, nu) = (n_expert as i32, n_used as i32);
4603 let __s_b = self.gpu.stream();
4604 let mut b = __s_b.launch_builder(&f);
4605 b.arg(logits)
4606 .arg(&mut sel_idx)
4607 .arg(&mut sel_w)
4608 .arg(&ne)
4609 .arg(&nu);
4610 unsafe {
4611 b.launch(cfg)?;
4612 }
4613 Ok((sel_idx, sel_w))
4614 }
4615
4616 pub fn moe_router_topk_scaled(
4619 &self,
4620 logits: &CudaSlice<f32>,
4621 t: usize,
4622 n_expert: usize,
4623 n_used: usize,
4624 ex_scale: &CudaSlice<f32>,
4625 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4626 let f = self.func("moe_router_topk_scaled_f32");
4631 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4632 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4633 let cfg = LaunchConfig {
4634 grid_dim: (t as u32, 1, 1),
4635 block_dim: (n_expert as u32, 1, 1),
4636 shared_mem_bytes: 0,
4637 };
4638 let (ne, nu) = (n_expert as i32, n_used as i32);
4639 let __s_b = self.gpu.stream();
4640 let mut b = __s_b.launch_builder(&f);
4641 b.arg(logits)
4642 .arg(&mut sel_idx)
4643 .arg(&mut sel_w)
4644 .arg(&ne)
4645 .arg(&nu)
4646 .arg(ex_scale);
4647 unsafe {
4648 b.launch(cfg)?;
4649 }
4650 Ok((sel_idx, sel_w))
4651 }
4652
4653 pub fn moe_router_topk_host(
4661 &self,
4662 logits: &CudaSlice<f32>,
4663 t: usize,
4664 n_expert: usize,
4665 n_used: usize,
4666 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4667 let f = self.func("moe_router_topk_f32");
4668 let n = t * n_used;
4669 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4670 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4671 let cfg = LaunchConfig {
4672 grid_dim: (t as u32, 1, 1),
4673 block_dim: (n_expert as u32, 1, 1),
4674 shared_mem_bytes: 0,
4675 };
4676 let (ne, nu) = (n_expert as i32, n_used as i32);
4677 let __s_b = self.gpu.stream();
4678 let mut b = __s_b.launch_builder(&f);
4679 b.arg(logits)
4680 .arg(&mut sel_idx)
4681 .arg(&mut sel_w)
4682 .arg(&ne)
4683 .arg(&nu);
4684 unsafe {
4685 b.launch(cfg)?;
4686 }
4687 let bytes = n * 8;
4689 let mut guard = self.router_stage.lock().unwrap();
4690 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4691 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4692 }
4693 let stage = guard.as_mut().unwrap();
4694 let (si, sw) = unsafe {
4695 (
4696 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4697 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4698 )
4699 };
4700 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()))
4704 }
4705
4706 #[allow(clippy::too_many_arguments)]
4710 pub fn moe_router_sigmoid_topk(
4711 &self,
4712 logits: &CudaSlice<f32>,
4713 t: usize,
4714 n_expert: usize,
4715 n_used: usize,
4716 active_count: usize,
4717 correction_bias: &CudaSlice<f32>,
4718 active: &CudaSlice<u8>,
4719 scaling_factor: f32,
4720 route_norm: bool,
4721 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4722 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4723 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4724 return Err(format!(
4725 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4726 )
4727 .into());
4728 }
4729 if logits.len() < t * n_expert
4730 || correction_bias.len() != n_expert
4731 || active.len() != n_expert
4732 {
4733 return Err(format!(
4734 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4735 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4736 ).into());
4737 }
4738 let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4739 self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4742 } else if crate::sig_expf_dev_on() {
4743 self.func("moe_router_sigmoid_topk_f32_dexp")
4744 } else if crate::topk_fast_on() {
4745 self.func("moe_router_sigmoid_topk_f32_fast")
4746 } else {
4747 self.func("moe_router_sigmoid_topk_f32")
4748 };
4749 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4750 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4751 let threads = n_expert.div_ceil(32) * 32;
4752 let cfg = LaunchConfig {
4753 grid_dim: (t as u32, 1, 1),
4754 block_dim: (threads as u32, 1, 1),
4755 shared_mem_bytes: 0,
4756 };
4757 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4758 let __s_b = self.gpu.stream();
4759 let mut b = __s_b.launch_builder(&f);
4760 b.arg(logits)
4761 .arg(correction_bias)
4762 .arg(active)
4763 .arg(&mut sel_idx)
4764 .arg(&mut sel_w)
4765 .arg(&ne)
4766 .arg(&nu)
4767 .arg(&scaling_factor)
4768 .arg(&rn);
4769 unsafe {
4770 b.launch(cfg)?;
4771 }
4772 Ok((sel_idx, sel_w))
4773 }
4774
4775 #[allow(clippy::too_many_arguments)]
4778 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4782 if ptr == 0 {
4783 return Err("ring_flag_raw: unarmed flag".into());
4784 }
4785 let f = self.func("memra_ring_flag");
4786 let cfg = LaunchConfig {
4787 grid_dim: (1, 1, 1),
4788 block_dim: (32, 1, 1),
4789 shared_mem_bytes: 0,
4790 };
4791 let __s_b = self.gpu.stream();
4792 let mut b = __s_b.launch_builder(&f);
4793 b.arg(&ptr).arg(&value);
4794 unsafe {
4795 b.launch(cfg)?;
4796 }
4797 Ok(())
4798 }
4799
4800 pub fn moe_sel_w_mirror(
4803 &self,
4804 sel_src: &CudaSlice<i32>,
4805 w_src: &CudaSlice<f32>,
4806 sel_dst: &mut CudaSlice<i32>,
4807 w_dst: &mut CudaSlice<f32>,
4808 n: usize,
4809 ) -> Result<(), Box<dyn std::error::Error>> {
4810 if n == 0
4811 || n > 32
4812 || sel_src.len() < n
4813 || w_src.len() < n
4814 || sel_dst.len() < n
4815 || w_dst.len() < n
4816 {
4817 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4818 }
4819 let f = self.func("moe_sel_w_mirror");
4820 let cfg = LaunchConfig {
4821 grid_dim: (1, 1, 1),
4822 block_dim: (32, 1, 1),
4823 shared_mem_bytes: 0,
4824 };
4825 let ni = n as i32;
4826 let __s_b = self.gpu.stream();
4827 let mut b = __s_b.launch_builder(&f);
4828 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
4829 unsafe {
4830 b.launch(cfg)?;
4831 }
4832 Ok(())
4833 }
4834
4835 pub fn moe_router_sigmoid_topk_into(
4836 &self,
4837 logits: &CudaSlice<f32>,
4838 t: usize,
4839 n_expert: usize,
4840 n_used: usize,
4841 active_count: usize,
4842 correction_bias: &CudaSlice<f32>,
4843 active: &CudaSlice<u8>,
4844 scaling_factor: f32,
4845 route_norm: bool,
4846 sel_idx: &mut CudaSlice<i32>,
4847 sel_w: &mut CudaSlice<f32>,
4848 ) -> Result<(), Box<dyn std::error::Error>> {
4849 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4850 if n_expert == 0
4851 || n_expert > 1024
4852 || n_used == 0
4853 || n_used > 32 || n_used > n_expert
4855 || logits.len() < t * n_expert
4856 || correction_bias.len() != n_expert
4857 || active.len() != n_expert
4858 || sel_idx.len() < t * n_used
4859 || sel_w.len() < t * n_used
4860 {
4861 return Err("sigmoid router _into geometry mismatch".into());
4862 }
4863 let f = if crate::sig_expf_dev_on() {
4864 self.func("moe_router_sigmoid_topk_f32_dexp")
4865 } else if crate::topk_fast_on() {
4866 self.func("moe_router_sigmoid_topk_f32_fast")
4867 } else {
4868 self.func("moe_router_sigmoid_topk_f32")
4869 };
4870 let threads = n_expert.div_ceil(32) * 32;
4871 let cfg = LaunchConfig {
4872 grid_dim: (t as u32, 1, 1),
4873 block_dim: (threads as u32, 1, 1),
4874 shared_mem_bytes: 0,
4875 };
4876 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4877 let __s_b = self.gpu.stream();
4878 let mut b = __s_b.launch_builder(&f);
4879 b.arg(logits)
4880 .arg(correction_bias)
4881 .arg(active)
4882 .arg(&mut *sel_idx)
4883 .arg(&mut *sel_w)
4884 .arg(&ne)
4885 .arg(&nu)
4886 .arg(&scaling_factor)
4887 .arg(&rn);
4888 unsafe {
4889 b.launch(cfg)?;
4890 }
4891 Ok(())
4892 }
4893
4894 #[allow(clippy::too_many_arguments)]
4897 pub fn moe_router_sigmoid_topk_host(
4898 &self,
4899 logits: &CudaSlice<f32>,
4900 t: usize,
4901 n_expert: usize,
4902 n_used: usize,
4903 active_count: usize,
4904 correction_bias: &CudaSlice<f32>,
4905 active: &CudaSlice<u8>,
4906 scaling_factor: f32,
4907 route_norm: bool,
4908 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4909 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4910 logits,
4911 t,
4912 n_expert,
4913 n_used,
4914 active_count,
4915 correction_bias,
4916 active,
4917 scaling_factor,
4918 route_norm,
4919 )?;
4920 let n = t * n_used;
4921 let bytes = n * 8;
4922 let mut guard = self.router_stage.lock().unwrap();
4923 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4924 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4925 }
4926 let stage = guard.as_mut().unwrap();
4927 let (si, sw) = unsafe {
4928 (
4929 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4930 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4931 )
4932 };
4933 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4934 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4935 self.gpu.stream().synchronize()?;
4936 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4937 }
4938
4939 pub fn stage_expert_async(
4943 &self,
4944 host_bytes: &[u8],
4945 scratch: &mut CudaSlice<u8>,
4946 off: usize,
4947 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4948 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4949 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4950 Ok(self.copy_stream.record_event(None)?)
4951 }
4952
4953 pub fn compute_wait(
4955 &self,
4956 ev: &cudarc::driver::CudaEvent,
4957 ) -> Result<(), Box<dyn std::error::Error>> {
4958 self.gpu.stream().wait(ev)?;
4959 Ok(())
4960 }
4961
4962 pub fn qmatvec_view(
4967 &self,
4968 w: &CudaSlice<u8>,
4969 range: std::ops::Range<usize>,
4970 x: &cudarc::driver::CudaView<f32>,
4971 m: usize,
4972 in_f: usize,
4973 out_f: usize,
4974 qtype: i32,
4975 row_bytes: usize,
4976 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4977 let f = self.func("qmatvec_f32");
4978 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4981 grid_dim: (out_f as u32, m as u32, 1),
4982 block_dim: (256, 1, 1),
4983 shared_mem_bytes: 0,
4984 };
4985 let (inf, outf, mi, qt, rb) =
4986 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4987 let __s_b = self.gpu.stream();
4988 let mut b = __s_b.launch_builder(&f);
4989 b.arg(&wv)
4990 .arg(x)
4991 .arg(&mut y)
4992 .arg(&inf)
4993 .arg(&outf)
4994 .arg(&mi)
4995 .arg(&qt)
4996 .arg(&rb);
4997 unsafe {
4998 b.launch(cfg)?;
4999 }
5000 Ok(y)
5001 }
5002
5003 #[allow(clippy::too_many_arguments)]
5010 pub fn moe_gate_up_silu8_q8(
5014 &self,
5015 gp: WPtr8,
5016 up: WPtr8,
5017 aq: &CudaSlice<i8>,
5018 ad: &CudaSlice<f32>,
5019 in_f: usize,
5020 n_ff: usize,
5021 n_used: usize,
5022 qt_g: i32,
5023 qt_u: i32,
5024 rb_g: usize,
5025 rb_u: usize,
5026 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5027 let f = self.func("moe_gate_up_silu8_q8");
5028 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5029 let cfg = LaunchConfig {
5030 grid_dim: (n_ff as u32, n_used as u32, 1),
5031 block_dim: (32, 1, 1),
5032 shared_mem_bytes: 0,
5033 };
5034 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5035 let __s_b = self.gpu.stream();
5036 let mut b = __s_b.launch_builder(&f);
5037 b.arg(&gp)
5038 .arg(&up)
5039 .arg(aq)
5040 .arg(ad)
5041 .arg(&mut act)
5042 .arg(&inf)
5043 .arg(&nff)
5044 .arg(&qt_g)
5045 .arg(&qt_u)
5046 .arg(&rbg)
5047 .arg(&rbu);
5048 unsafe {
5049 b.launch(cfg)?;
5050 }
5051 Ok(act)
5052 }
5053
5054 #[allow(clippy::too_many_arguments)]
5055 pub fn moe_down8_fma_q8(
5056 &self,
5057 dp: WPtr8,
5058 w: F32x8,
5059 aq2: &CudaSlice<i8>,
5060 ad2: &CudaSlice<f32>,
5061 dst: &mut cudarc::driver::CudaViewMut<f32>,
5062 in_f: usize,
5063 out_f: usize,
5064 n_used: usize,
5065 qt: i32,
5066 rb: usize,
5067 ) -> Result<(), Box<dyn std::error::Error>> {
5068 let f = self.func("moe_down8_fma_q8");
5069 let cfg = LaunchConfig {
5070 grid_dim: (out_f as u32, 1, 1),
5071 block_dim: (32, 1, 1),
5072 shared_mem_bytes: 0,
5073 };
5074 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5075 let __s_b = self.gpu.stream();
5076 let mut b = __s_b.launch_builder(&f);
5077 b.arg(&dp)
5078 .arg(&w)
5079 .arg(aq2)
5080 .arg(ad2)
5081 .arg(dst)
5082 .arg(&inf)
5083 .arg(&outf)
5084 .arg(&nu)
5085 .arg(&qt)
5086 .arg(&rbi);
5087 unsafe {
5088 b.launch(cfg)?;
5089 }
5090 Ok(())
5091 }
5092
5093 pub fn qmatvec_expert_q8(
5095 &self,
5096 w: &CudaSlice<u8>,
5097 range: std::ops::Range<usize>,
5098 aq: &CudaSlice<i8>,
5099 ad: &CudaSlice<f32>,
5100 m: usize,
5101 in_f: usize,
5102 out_f: usize,
5103 qtype: i32,
5104 row_bytes: usize,
5105 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5106 let f = self.func("qmatvec_expert_q8");
5107 let wv = w.slice(range);
5108 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5109 const ROWS: u32 = 4; let cfg = LaunchConfig {
5111 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5112 block_dim: (32, ROWS, 1),
5113 shared_mem_bytes: 0,
5114 };
5115 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5116 let __s_b = self.gpu.stream();
5117 let mut b = __s_b.launch_builder(&f);
5118 b.arg(&wv)
5119 .arg(aq)
5120 .arg(ad)
5121 .arg(&mut y)
5122 .arg(&inf)
5123 .arg(&outf)
5124 .arg(&mi)
5125 .arg(&qtype)
5126 .arg(&rbi);
5127 unsafe {
5128 b.launch(cfg)?;
5129 }
5130 Ok(y)
5131 }
5132
5133 pub fn moe_gate_up_silu8(
5134 &self,
5135 gp: WPtr8,
5136 up: WPtr8,
5137 x: &cudarc::driver::CudaView<f32>,
5138 in_f: usize,
5139 n_ff: usize,
5140 n_used: usize,
5141 qt_g: i32,
5142 qt_u: i32,
5143 rb_g: usize,
5144 rb_u: usize,
5145 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5146 let f = self.func("moe_gate_up_silu8_f32");
5147 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5149 grid_dim: (n_ff as u32, n_used as u32, 1),
5150 block_dim: (256, 1, 1),
5151 shared_mem_bytes: 0,
5152 };
5153 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5154 let __s_b = self.gpu.stream();
5155 let mut b = __s_b.launch_builder(&f);
5156 b.arg(&gp)
5157 .arg(&up)
5158 .arg(x)
5159 .arg(&mut act)
5160 .arg(&inf)
5161 .arg(&nff)
5162 .arg(&qt_g)
5163 .arg(&qt_u)
5164 .arg(&rbg)
5165 .arg(&rbu);
5166 unsafe {
5167 b.launch(cfg)?;
5168 }
5169 Ok(act)
5170 }
5171
5172 #[allow(clippy::too_many_arguments)]
5178 pub fn moe_down8_fma_into(
5179 &self,
5180 dp: WPtr8,
5181 w: F32x8,
5182 act: &CudaSlice<f32>,
5183 dst: &mut cudarc::driver::CudaViewMut<f32>,
5184 in_f: usize,
5185 out_f: usize,
5186 n_used: usize,
5187 qt: i32,
5188 rb: usize,
5189 ) -> Result<(), Box<dyn std::error::Error>> {
5190 let f = self.func("moe_down8_fma_f32");
5191 let cfg = LaunchConfig {
5192 grid_dim: (out_f as u32, 1, 1),
5193 block_dim: (256, 1, 1),
5194 shared_mem_bytes: 0,
5195 };
5196 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5197 let __s_b = self.gpu.stream();
5198 let mut b = __s_b.launch_builder(&f);
5199 b.arg(&dp)
5200 .arg(&w)
5201 .arg(act)
5202 .arg(dst)
5203 .arg(&inf)
5204 .arg(&outf)
5205 .arg(&nu)
5206 .arg(&qt)
5207 .arg(&rbv);
5208 unsafe {
5209 b.launch(cfg)?;
5210 }
5211 Ok(())
5212 }
5213
5214 #[allow(clippy::too_many_arguments)]
5219 #[allow(clippy::too_many_arguments)]
5234 #[allow(clippy::too_many_arguments)]
5236 pub fn moe_pairs_matvec_q8(
5237 &self,
5238 table: &CudaSlice<u64>,
5239 proj: i32,
5240 pair_tok: &CudaSlice<i32>,
5241 pair_ex: &CudaSlice<i32>,
5242 aq: &CudaSlice<i8>,
5243 ad: &CudaSlice<f32>,
5244 in_f: usize,
5245 out_f: usize,
5246 n_expert: usize,
5247 n_pairs: usize,
5248 qtype: i32,
5249 row_bytes: usize,
5250 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5251 let f = self.func("moe_pairs_matvec_q8");
5252 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5253 const ROWS: u32 = 4;
5254 let cfg = LaunchConfig {
5255 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5256 block_dim: (32, ROWS, 1),
5257 shared_mem_bytes: 0,
5258 };
5259 let (inf, outf, ne, np, rbi) = (
5260 in_f as i32,
5261 out_f as i32,
5262 n_expert as i32,
5263 n_pairs as i32,
5264 row_bytes as i64,
5265 );
5266 let __s_b = self.gpu.stream();
5267 let mut b = __s_b.launch_builder(&f);
5268 b.arg(table)
5269 .arg(&proj)
5270 .arg(pair_tok)
5271 .arg(pair_ex)
5272 .arg(aq)
5273 .arg(ad)
5274 .arg(&mut y)
5275 .arg(&inf)
5276 .arg(&outf)
5277 .arg(&ne)
5278 .arg(&np)
5279 .arg(&qtype)
5280 .arg(&rbi);
5281 unsafe {
5282 b.launch(cfg)?;
5283 }
5284 Ok(y)
5285 }
5286
5287 #[allow(clippy::too_many_arguments)]
5289 pub fn moe_pairs_matvec_q8_em(
5290 &self,
5291 table: &CudaSlice<u64>,
5292 proj: i32,
5293 ex_ids: &CudaSlice<i32>,
5294 ex_off: &CudaSlice<i32>,
5295 ex_pairs: &CudaSlice<i32>,
5296 pair_tok: &CudaSlice<i32>,
5297 aq: &CudaSlice<i8>,
5298 ad: &CudaSlice<f32>,
5299 in_f: usize,
5300 out_f: usize,
5301 n_expert: usize,
5302 n_active: usize,
5303 n_pairs: usize,
5304 qtype: i32,
5305 row_bytes: usize,
5306 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5307 let f = self.func("moe_pairs_matvec_q8_em");
5308 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5309 const ROWS: u32 = 4;
5310 let cfg = LaunchConfig {
5311 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5312 block_dim: (32, ROWS, 1),
5313 shared_mem_bytes: 0,
5314 };
5315 let (inf, outf, ne, na, rbi) = (
5316 in_f as i32,
5317 out_f as i32,
5318 n_expert as i32,
5319 n_active as i32,
5320 row_bytes as i64,
5321 );
5322 let __s_b = self.gpu.stream();
5323 let mut b = __s_b.launch_builder(&f);
5324 b.arg(table)
5325 .arg(&proj)
5326 .arg(ex_ids)
5327 .arg(ex_off)
5328 .arg(ex_pairs)
5329 .arg(pair_tok)
5330 .arg(aq)
5331 .arg(ad)
5332 .arg(&mut y)
5333 .arg(&inf)
5334 .arg(&outf)
5335 .arg(&ne)
5336 .arg(&na)
5337 .arg(&qtype)
5338 .arg(&rbi);
5339 unsafe {
5340 b.launch(cfg)?;
5341 }
5342 Ok(y)
5343 }
5344
5345 #[allow(clippy::too_many_arguments)]
5348 pub fn moe_pairs_matvec_q8_dec(
5349 &self,
5350 table: &CudaSlice<u64>,
5351 proj: i32,
5352 ex_ids: &CudaSlice<i32>,
5353 ex_off: &CudaSlice<i32>,
5354 ex_pairs: &CudaSlice<i32>,
5355 pair_tok: &CudaSlice<i32>,
5356 aq: &CudaSlice<i8>,
5357 ad: &CudaSlice<f32>,
5358 in_f: usize,
5359 out_f: usize,
5360 n_expert: usize,
5361 n_active: usize,
5362 n_pairs: usize,
5363 qtype: i32,
5364 row_bytes: usize,
5365 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5366 let f = self.func("moe_pairs_matvec_q8_dec");
5367 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5368 const ROWS: u32 = 4;
5369 let cfg = LaunchConfig {
5370 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5371 block_dim: (32, ROWS, 1),
5372 shared_mem_bytes: 0,
5373 };
5374 let (inf, outf, ne, na, rbi) = (
5375 in_f as i32,
5376 out_f as i32,
5377 n_expert as i32,
5378 n_active as i32,
5379 row_bytes as i64,
5380 );
5381 let __s_b = self.gpu.stream();
5382 let mut b = __s_b.launch_builder(&f);
5383 b.arg(table)
5384 .arg(&proj)
5385 .arg(ex_ids)
5386 .arg(ex_off)
5387 .arg(ex_pairs)
5388 .arg(pair_tok)
5389 .arg(aq)
5390 .arg(ad)
5391 .arg(&mut y)
5392 .arg(&inf)
5393 .arg(&outf)
5394 .arg(&ne)
5395 .arg(&na)
5396 .arg(&qtype)
5397 .arg(&rbi);
5398 unsafe {
5399 b.launch(cfg)?;
5400 }
5401 Ok(y)
5402 }
5403
5404 pub fn moe_pairs_gelu_mul(
5405 &self,
5406 gate: &CudaSlice<f32>,
5407 up: &CudaSlice<f32>,
5408 n: usize,
5409 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5410 let f = self.func("moe_pairs_gelu_mul");
5411 let mut act = self.alloc_uninit::<f32>(n)?;
5412 let cfg = LaunchConfig::for_num_elems(n as u32);
5413 let nl = n as i64;
5414 let __s_b = self.gpu.stream();
5415 let mut b = __s_b.launch_builder(&f);
5416 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5417 unsafe {
5418 b.launch(cfg)?;
5419 }
5420 Ok(act)
5421 }
5422
5423 pub fn moe_pairs_silu_mul(
5424 &self,
5425 gate: &CudaSlice<f32>,
5426 up: &CudaSlice<f32>,
5427 n: usize,
5428 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5429 let f = self.func("moe_pairs_silu_mul");
5430 let mut act = self.alloc_uninit::<f32>(n)?;
5431 let cfg = LaunchConfig::for_num_elems(n as u32);
5432 let nl = n as i64;
5433 let __s_b = self.gpu.stream();
5434 let mut b = __s_b.launch_builder(&f);
5435 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5436 unsafe {
5437 b.launch(cfg)?;
5438 }
5439 Ok(act)
5440 }
5441
5442 #[allow(clippy::too_many_arguments)]
5443 pub fn moe_pairs_scatter(
5444 &self,
5445 y_down: &CudaSlice<f32>,
5446 pair_w: &CudaSlice<f32>,
5447 tok_pair_off: &CudaSlice<i32>,
5448 tok_pair_ids: &CudaSlice<i32>,
5449 moe_out: &mut CudaSlice<f32>,
5450 t: usize,
5451 n_embd: usize,
5452 ) -> Result<(), Box<dyn std::error::Error>> {
5453 let f = self.func("moe_pairs_scatter");
5454 let cfg = LaunchConfig {
5455 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5456 block_dim: (256, 1, 1),
5457 shared_mem_bytes: 0,
5458 };
5459 let ne = n_embd as i32;
5460 let __s_b = self.gpu.stream();
5461 let mut b = __s_b.launch_builder(&f);
5462 b.arg(y_down)
5463 .arg(pair_w)
5464 .arg(tok_pair_off)
5465 .arg(tok_pair_ids)
5466 .arg(moe_out)
5467 .arg(&ne);
5468 unsafe {
5469 b.launch(cfg)?;
5470 }
5471 Ok(())
5472 }
5473
5474 #[allow(clippy::too_many_arguments)]
5478 pub fn moe_gate_up_gelu8_dev_q8(
5479 &self,
5480 table: &CudaSlice<u64>,
5481 sel: &cudarc::driver::CudaView<i32>,
5482 aq: &CudaSlice<i8>,
5483 ad: &CudaSlice<f32>,
5484 in_f: usize,
5485 n_ff: usize,
5486 n_used: usize,
5487 n_expert: usize,
5488 qt_g: i32,
5489 qt_u: i32,
5490 rb_g: usize,
5491 rb_u: usize,
5492 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5493 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5494 let (inf, nff, ne, rbg, rbu) = (
5495 in_f as i32,
5496 n_ff as i32,
5497 n_expert as i32,
5498 rb_g as i64,
5499 rb_u as i64,
5500 );
5501 let f = self.func("moe_gate_up_gelu8_dev_q8");
5502 let cfg = LaunchConfig {
5503 grid_dim: (n_ff as u32, n_used as u32, 1),
5504 block_dim: (32, 1, 1),
5505 shared_mem_bytes: 0,
5506 };
5507 let __s_b = self.gpu.stream();
5508 let mut b = __s_b.launch_builder(&f);
5509 b.arg(table)
5510 .arg(sel)
5511 .arg(aq)
5512 .arg(ad)
5513 .arg(&mut act)
5514 .arg(&inf)
5515 .arg(&nff)
5516 .arg(&ne)
5517 .arg(&qt_g)
5518 .arg(&qt_u)
5519 .arg(&rbg)
5520 .arg(&rbu);
5521 unsafe {
5522 b.launch(cfg)?;
5523 }
5524 Ok(act)
5525 }
5526
5527 #[allow(clippy::too_many_arguments)]
5529 pub fn moe_gate_up_gelu8_dev_q8_rows(
5530 &self,
5531 table: &CudaSlice<u64>,
5532 sel: &CudaSlice<i32>,
5533 aq: &CudaSlice<i8>,
5534 ad: &CudaSlice<f32>,
5535 t: usize,
5536 in_f: usize,
5537 n_ff: usize,
5538 n_used: usize,
5539 n_expert: usize,
5540 qt_g: i32,
5541 qt_u: i32,
5542 rb_g: usize,
5543 rb_u: usize,
5544 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5545 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5546 let (inf, nff, ne, rbg, rbu, nu) = (
5547 in_f as i32,
5548 n_ff as i32,
5549 n_expert as i32,
5550 rb_g as i64,
5551 rb_u as i64,
5552 n_used as i32,
5553 );
5554 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5555 let cfg = LaunchConfig {
5556 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5557 block_dim: (32, 1, 1),
5558 shared_mem_bytes: 0,
5559 };
5560 let __s_b = self.gpu.stream();
5561 let mut b = __s_b.launch_builder(&f);
5562 b.arg(table)
5563 .arg(sel)
5564 .arg(aq)
5565 .arg(ad)
5566 .arg(&mut act)
5567 .arg(&inf)
5568 .arg(&nff)
5569 .arg(&ne)
5570 .arg(&qt_g)
5571 .arg(&qt_u)
5572 .arg(&rbg)
5573 .arg(&rbu)
5574 .arg(&nu);
5575 unsafe {
5576 b.launch(cfg)?;
5577 }
5578 Ok(act)
5579 }
5580
5581 #[allow(clippy::too_many_arguments)]
5583 pub fn moe_gate_up_gelu8_dev_q8_csr(
5584 &self,
5585 table: &CudaSlice<u64>,
5586 sel: &CudaSlice<i32>,
5587 aq: &CudaSlice<i8>,
5588 ad: &CudaSlice<f32>,
5589 n_pairs: usize,
5590 in_f: usize,
5591 n_ff: usize,
5592 n_used: usize,
5593 n_expert: usize,
5594 qt_g: i32,
5595 qt_u: i32,
5596 rb_g: usize,
5597 rb_u: usize,
5598 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5599 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5600 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5601 in_f as i32,
5602 n_ff as i32,
5603 n_expert as i32,
5604 rb_g as i64,
5605 rb_u as i64,
5606 n_used as i32,
5607 n_pairs as i32,
5608 );
5609 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5610 let cfg = LaunchConfig {
5611 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5612 block_dim: (32, 1, 1),
5613 shared_mem_bytes: 0,
5614 };
5615 let __s_b = self.gpu.stream();
5616 let mut b = __s_b.launch_builder(&f);
5617 b.arg(table)
5618 .arg(sel)
5619 .arg(aq)
5620 .arg(ad)
5621 .arg(&mut act)
5622 .arg(&inf)
5623 .arg(&nff)
5624 .arg(&ne)
5625 .arg(&qt_g)
5626 .arg(&qt_u)
5627 .arg(&rbg)
5628 .arg(&rbu)
5629 .arg(&nu)
5630 .arg(&npi);
5631 unsafe {
5632 b.launch(cfg)?;
5633 }
5634 Ok(act)
5635 }
5636
5637 #[allow(clippy::too_many_arguments)]
5639 pub fn moe_down8_fma_dev_q8_rows_g(
5640 &self,
5641 table: &CudaSlice<u64>,
5642 sel: &CudaSlice<i32>,
5643 w: &CudaSlice<f32>,
5644 aq2: &CudaSlice<i8>,
5645 ad2: &CudaSlice<f32>,
5646 dst: &mut CudaSlice<f32>,
5647 t: usize,
5648 in_f: usize,
5649 out_f: usize,
5650 n_used: usize,
5651 n_expert: usize,
5652 qt: i32,
5653 rb: usize,
5654 ) -> Result<(), Box<dyn std::error::Error>> {
5655 let (inf, outf, nu, ne, rbi) = (
5656 in_f as i32,
5657 out_f as i32,
5658 n_used as i32,
5659 n_expert as i32,
5660 rb as i64,
5661 );
5662 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5666 let f = self.func(if step_b1_w8 {
5667 "moe_down8_fma_dev_q8_rows_w8"
5668 } else {
5669 "moe_down8_fma_dev_q8_rows_g"
5670 });
5671 let cfg = LaunchConfig {
5672 grid_dim: (out_f as u32, 1, t as u32),
5673 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5674 shared_mem_bytes: 0,
5675 };
5676 let __s_b = self.gpu.stream();
5677 let mut b = __s_b.launch_builder(&f);
5678 b.arg(table)
5679 .arg(sel)
5680 .arg(w)
5681 .arg(aq2)
5682 .arg(ad2)
5683 .arg(dst)
5684 .arg(&inf)
5685 .arg(&outf)
5686 .arg(&nu)
5687 .arg(&ne)
5688 .arg(&qt)
5689 .arg(&rbi);
5690 unsafe {
5691 b.launch(cfg)?;
5692 }
5693 Ok(())
5694 }
5695
5696 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5700 let (out_f, in_f) = (2048usize, 2816usize);
5701 let nblk = in_f / 32;
5702 let mut seed = 0x9E3779B97F4A7C15u64;
5703 let mut rng = move || {
5704 seed = seed
5705 .wrapping_mul(6364136223846793005)
5706 .wrapping_add(1442695040888963407);
5707 (seed >> 33) as u8
5708 };
5709 let mut w = vec![0u8; out_f * nblk * 18];
5710 for b in w.iter_mut() {
5711 *b = rng();
5712 }
5713 for r in 0..out_f {
5714 for g in 0..nblk {
5715 let off = (r * nblk + g) * 18;
5716 w[off] = 0x00;
5717 w[off + 1] = 0x2C; }
5719 }
5720 let qplane = out_f * nblk * 16;
5721 let mut wrp = vec![0u8; w.len()];
5722 for r in 0..out_f {
5723 for g in 0..nblk {
5724 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5725 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5726 .copy_from_slice(&src[0..2]);
5727 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5728 }
5729 }
5730 let w_d = self.htod_bytes(&w)?;
5731 let wrp_d = self.htod_bytes(&wrp)?;
5732 let mut aq = vec![0i8; m * in_f];
5733 for v in aq.iter_mut() {
5734 *v = rng() as i8;
5735 }
5736 let aq_d = self.htod_i8(&aq)?;
5737 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5738 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5739 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5740 const RPB: u32 = 4;
5741 let cfg = LaunchConfig {
5742 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5743 block_dim: (32, RPB, 1),
5744 shared_mem_bytes: 0,
5745 };
5746 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5747 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5748 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5749 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5750 {
5751 let __s_b = self.gpu.stream();
5752 let mut b = __s_b.launch_builder(&fb);
5753 b.arg(&w_d)
5754 .arg(&aq_d)
5755 .arg(&ad_d)
5756 .arg(&mut y0)
5757 .arg(&inf)
5758 .arg(&outf)
5759 .arg(&mi)
5760 .arg(&rb);
5761 unsafe {
5762 b.launch(cfg)?;
5763 }
5764 let __s_b = self.gpu.stream();
5765 let mut b = __s_b.launch_builder(&fr);
5766 b.arg(&wrp_d)
5767 .arg(&aq_d)
5768 .arg(&ad_d)
5769 .arg(&mut y1)
5770 .arg(&inf)
5771 .arg(&outf)
5772 .arg(&mi)
5773 .arg(&qp);
5774 unsafe {
5775 b.launch(cfg)?;
5776 }
5777 }
5778 self.gpu.stream().synchronize()?;
5779 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5780 let nd = h0
5781 .iter()
5782 .zip(&h1)
5783 .filter(|(a, b)| a.to_bits() != b.to_bits())
5784 .count();
5785 if nd != 0 {
5786 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5787 }
5788 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5789 self.gpu.stream().synchronize()?;
5790 let t0 = std::time::Instant::now();
5791 for _ in 0..500 {
5792 if rp {
5793 let __s_b = self.gpu.stream();
5794 let mut b = __s_b.launch_builder(&fr);
5795 b.arg(&wrp_d)
5796 .arg(&aq_d)
5797 .arg(&ad_d)
5798 .arg(&mut y1)
5799 .arg(&inf)
5800 .arg(&outf)
5801 .arg(&mi)
5802 .arg(&qp);
5803 unsafe {
5804 b.launch(cfg)?;
5805 }
5806 } else {
5807 let __s_b = self.gpu.stream();
5808 let mut b = __s_b.launch_builder(&fb);
5809 b.arg(&w_d)
5810 .arg(&aq_d)
5811 .arg(&ad_d)
5812 .arg(&mut y0)
5813 .arg(&inf)
5814 .arg(&outf)
5815 .arg(&mi)
5816 .arg(&rb);
5817 unsafe {
5818 b.launch(cfg)?;
5819 }
5820 }
5821 }
5822 self.gpu.stream().synchronize()?;
5823 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5824 };
5825 let _ = time(false)?;
5826 let _ = time(true)?; Ok((time(false)?, time(true)?))
5828 }
5829
5830 pub fn build_q4_rp4(
5835 &self,
5836 t: &mut crate::model::GpuTensor,
5837 ) -> Result<(), Box<dyn std::error::Error>> {
5838 use crate::model::GpuTensor;
5839 let GpuTensor::Quant {
5840 bytes,
5841 qtype,
5842 row_bytes,
5843 ne,
5844 rp4,
5845 ..
5846 } = t
5847 else {
5848 return Ok(());
5849 };
5850 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5851 return Ok(());
5852 }
5853 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5854 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5855 return Ok(());
5856 }
5857 let nblk = in_f / 32;
5858 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5859 let f = self.func("q4_0_split_rp_build");
5860 let n = (out_f * nblk) as i32;
5861 let cfg = LaunchConfig {
5862 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5863 block_dim: (256, 1, 1),
5864 shared_mem_bytes: 0,
5865 };
5866 let (of, nb) = (out_f as i32, nblk as i32);
5867 let _ = n;
5868 let __s_b = self.gpu.stream();
5869 let mut b = __s_b.launch_builder(&f);
5870 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5871 unsafe {
5872 b.launch(cfg)?;
5873 }
5874 *rp4 = Some(dst);
5875 Ok(())
5876 }
5877
5878 pub fn build_q8_rp4(
5883 &self,
5884 t: &mut crate::model::GpuTensor,
5885 ) -> Result<(), Box<dyn std::error::Error>> {
5886 use crate::model::GpuTensor;
5887 let GpuTensor::Quant {
5888 bytes,
5889 qtype,
5890 row_bytes,
5891 ne,
5892 rp4,
5893 ..
5894 } = t
5895 else {
5896 return Ok(());
5897 };
5898 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5899 return Ok(());
5900 }
5901 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5902 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5903 return Ok(());
5904 }
5905 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5906 Ok(())
5907 }
5908
5909 pub fn build_q8_rp4_raw(
5912 &self,
5913 bytes: &CudaSlice<u8>,
5914 in_f: usize,
5915 out_f: usize,
5916 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5917 assert!(in_f % 32 == 0);
5918 let nblk = in_f / 32;
5919 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5920 let f = self.func("q8_0_split_rp_build");
5921 let cfg = LaunchConfig {
5922 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5923 block_dim: (256, 1, 1),
5924 shared_mem_bytes: 0,
5925 };
5926 let (of, nb) = (out_f as i32, nblk as i32);
5927 let __s_b = self.gpu.stream();
5928 let mut b = __s_b.launch_builder(&f);
5929 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5930 unsafe {
5931 b.launch(cfg)?;
5932 }
5933 Ok(dst)
5934 }
5935
5936 pub fn build_q4k_rp4(
5944 &self,
5945 t: &mut crate::model::GpuTensor,
5946 ) -> Result<(), Box<dyn std::error::Error>> {
5947 use crate::model::GpuTensor;
5948 let GpuTensor::Quant {
5949 bytes,
5950 qtype,
5951 row_bytes,
5952 ne,
5953 rp4,
5954 ..
5955 } = t
5956 else {
5957 return Ok(());
5958 };
5959 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5960 return Ok(());
5961 }
5962 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5963 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5964 return Ok(());
5965 }
5966 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5967 Ok(())
5968 }
5969
5970 pub fn build_q6k_rp4(
5971 &self,
5972 t: &mut crate::model::GpuTensor,
5973 ) -> Result<(), Box<dyn std::error::Error>> {
5974 use crate::model::GpuTensor;
5975 let GpuTensor::Quant {
5976 bytes,
5977 qtype,
5978 row_bytes,
5979 ne,
5980 rp4,
5981 ..
5982 } = t
5983 else {
5984 return Ok(());
5985 };
5986 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5987 return Ok(());
5988 }
5989 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5990 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5991 return Ok(());
5992 }
5993 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5994 Ok(())
5995 }
5996
5997 pub fn build_kq_rp4_raw(
5999 &self,
6000 bytes: &CudaSlice<u8>,
6001 in_f: usize,
6002 out_f: usize,
6003 qtype: i32,
6004 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6005 assert!(in_f % 256 == 0);
6006 let nsbk = in_f / 256;
6007 let (sb_bytes, kname) = match qtype {
6008 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6009 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6010 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6011 };
6012 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6013 let f = self.func(kname);
6014 let cfg = LaunchConfig {
6015 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6016 block_dim: (256, 1, 1),
6017 shared_mem_bytes: 0,
6018 };
6019 let (of, nb) = (out_f as i32, nsbk as i32);
6020 let __s_b = self.gpu.stream();
6021 let mut b = __s_b.launch_builder(&f);
6022 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6023 unsafe {
6024 b.launch(cfg)?;
6025 }
6026 Ok(dst)
6027 }
6028
6029 pub fn kqrp_enabled() -> bool {
6033 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6034 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6035 Ok("0") => false,
6036 Ok(_) => true,
6037 Err(_) => cfg!(memra_hopper_mma),
6038 })
6039 }
6040
6041 pub fn build_q4_rp_swap(
6047 &self,
6048 t: &mut crate::model::GpuTensor,
6049 ) -> Result<bool, Box<dyn std::error::Error>> {
6050 use crate::model::GpuTensor;
6051 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6061 return Ok(false);
6062 }
6063 self.build_q4_rp4(t)?;
6064 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6066 return Ok(false);
6067 };
6068 match rp4.take() {
6069 Some(split) => {
6070 *bytes = split; *rp = true;
6072 Ok(true)
6073 }
6074 None => Ok(false),
6075 }
6076 }
6077
6078 pub fn q4rp_enabled() -> bool {
6080 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6081 *ON.get_or_init(|| {
6082 std::env::var("MEMRA_Q4RP")
6083 .map(|v| v != "0")
6084 .unwrap_or(true)
6085 })
6086 }
6087
6088 pub fn copy_rows_strided(
6091 &self,
6092 src: &CudaSlice<f32>,
6093 dst: &mut CudaSlice<f32>,
6094 row_elems: usize,
6095 n_rows: usize,
6096 src_stride: usize,
6097 src_off: usize,
6098 ) -> Result<(), Box<dyn std::error::Error>> {
6099 let f = self.func("copy_rows_strided_f32");
6100 let cfg = LaunchConfig {
6101 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6102 block_dim: (256, 1, 1),
6103 shared_mem_bytes: 0,
6104 };
6105 let (re, nr) = (row_elems as i32, n_rows as i32);
6106 let (st, off) = (src_stride as i64, src_off as i64);
6107 let __s_b = self.gpu.stream();
6108 let mut b = __s_b.launch_builder(&f);
6109 b.arg(src)
6110 .arg(&mut *dst)
6111 .arg(&re)
6112 .arg(&nr)
6113 .arg(&st)
6114 .arg(&off);
6115 unsafe {
6116 b.launch(cfg)?;
6117 }
6118 Ok(())
6119 }
6120
6121 pub fn place_rows_strided(
6127 &self,
6128 src: &CudaSlice<f32>,
6129 dst: &mut CudaSlice<f32>,
6130 row_elems: usize,
6131 n_rows: usize,
6132 dst_stride: usize,
6133 dst_off: usize,
6134 ) -> Result<(), Box<dyn std::error::Error>> {
6135 if row_elems == 0 || n_rows == 0 {
6136 return Err("strided row placement requires nonzero rows and row width".into());
6137 }
6138 let src_len = n_rows
6139 .checked_mul(row_elems)
6140 .ok_or("strided row placement source size overflow")?;
6141 let dst_len = n_rows
6142 .checked_sub(1)
6143 .and_then(|rows| rows.checked_mul(dst_stride))
6144 .and_then(|base| base.checked_add(dst_off))
6145 .and_then(|base| base.checked_add(row_elems))
6146 .ok_or("strided row placement destination size overflow")?;
6147 let row_end = dst_off
6148 .checked_add(row_elems)
6149 .ok_or("strided row placement row size overflow")?;
6150 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6151 return Err(format!(
6152 "strided row placement geometry mismatch: src={} need_src={src_len} \
6153 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6154 dst_stride={dst_stride} dst_off={dst_off}",
6155 src.len(),
6156 dst.len(),
6157 )
6158 .into());
6159 }
6160 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6161 return Err("strided row placement exceeds CUDA kernel geometry".into());
6162 }
6163 let f = self.func("place_rows_strided_f32");
6164 let cfg = LaunchConfig {
6165 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6166 block_dim: (256, 1, 1),
6167 shared_mem_bytes: 0,
6168 };
6169 let (re, nr) = (row_elems as i32, n_rows as i32);
6170 let (st, off) = (dst_stride as i64, dst_off as i64);
6171 let __s_b = self.gpu.stream();
6172 let mut b = __s_b.launch_builder(&f);
6173 b.arg(src)
6174 .arg(&mut *dst)
6175 .arg(&re)
6176 .arg(&nr)
6177 .arg(&st)
6178 .arg(&off);
6179 unsafe {
6180 b.launch(cfg)?;
6181 }
6182 Ok(())
6183 }
6184
6185 pub fn u32_set_k(
6187 &self,
6188 dst: &mut CudaSlice<u32>,
6189 v: u32,
6190 idx: usize,
6191 ) -> Result<(), Box<dyn std::error::Error>> {
6192 let f = self.func("u32_set_k");
6193 let cfg = LaunchConfig {
6194 grid_dim: (1, 1, 1),
6195 block_dim: (1, 1, 1),
6196 shared_mem_bytes: 0,
6197 };
6198 let ii = idx as i32;
6199 let __s_b = self.gpu.stream();
6200 let mut b = __s_b.launch_builder(&f);
6201 b.arg(dst).arg(&v).arg(&ii);
6202 unsafe {
6203 b.launch(cfg)?;
6204 }
6205 Ok(())
6206 }
6207
6208 pub fn i32_add_k(
6210 &self,
6211 d: &mut CudaSlice<i32>,
6212 v: i32,
6213 ) -> Result<(), Box<dyn std::error::Error>> {
6214 let f = self.func("i32_add_k");
6215 let cfg = LaunchConfig {
6216 grid_dim: (1, 1, 1),
6217 block_dim: (32, 1, 1),
6218 shared_mem_bytes: 0,
6219 };
6220 let __s_b = self.gpu.stream();
6221 let mut b = __s_b.launch_builder(&f);
6222 b.arg(d).arg(&v);
6223 unsafe {
6224 b.launch(cfg)?;
6225 }
6226 Ok(())
6227 }
6228
6229 pub fn i32_iota_from(
6231 &self,
6232 ctr: &CudaSlice<i32>,
6233 dst: &mut CudaSlice<i32>,
6234 n: usize,
6235 ) -> Result<(), Box<dyn std::error::Error>> {
6236 let f = self.func("i32_iota_from");
6237 let cfg = LaunchConfig::for_num_elems(n as u32);
6238 let ni = n as i32;
6239 let __s_b = self.gpu.stream();
6240 let mut b = __s_b.launch_builder(&f);
6241 b.arg(ctr).arg(dst).arg(&ni);
6242 unsafe {
6243 b.launch(cfg)?;
6244 }
6245 Ok(())
6246 }
6247
6248 pub fn u32_map_k(
6250 &self,
6251 buf: &mut CudaSlice<u32>,
6252 map: &CudaSlice<u32>,
6253 idx: usize,
6254 ) -> Result<(), Box<dyn std::error::Error>> {
6255 let f = self.func("u32_map_k");
6256 let cfg = LaunchConfig {
6257 grid_dim: (1, 1, 1),
6258 block_dim: (1, 1, 1),
6259 shared_mem_bytes: 0,
6260 };
6261 let ii = idx as i32;
6262 let __s_b = self.gpu.stream();
6263 let mut b = __s_b.launch_builder(&f);
6264 b.arg(buf).arg(map).arg(&ii);
6265 unsafe {
6266 b.launch(cfg)?;
6267 }
6268 Ok(())
6269 }
6270
6271 #[allow(clippy::too_many_arguments)]
6273 pub fn u32_pack2(
6274 &self,
6275 a: &CudaSlice<u32>,
6276 off_a: usize,
6277 n1: usize,
6278 b_in: &CudaSlice<u32>,
6279 n2: usize,
6280 out: &mut CudaSlice<u32>,
6281 ) -> Result<(), Box<dyn std::error::Error>> {
6282 let f = self.func("u32_pack2");
6283 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6284 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6285 let __s_b = self.gpu.stream();
6286 let mut b = __s_b.launch_builder(&f);
6287 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6288 unsafe {
6289 b.launch(cfg)?;
6290 }
6291 Ok(())
6292 }
6293
6294 pub fn moe_w_exscale(
6296 &self,
6297 w: &mut CudaSlice<f32>,
6298 sel: &CudaSlice<i32>,
6299 s: &CudaSlice<f32>,
6300 n: usize,
6301 ) -> Result<(), Box<dyn std::error::Error>> {
6302 let f = self.func("moe_w_exscale");
6303 let cfg = LaunchConfig::for_num_elems(n as u32);
6304 let ni = n as i32;
6305 let __s_b = self.gpu.stream();
6306 let mut b = __s_b.launch_builder(&f);
6307 b.arg(w).arg(sel).arg(s).arg(&ni);
6308 unsafe {
6309 b.launch(cfg)?;
6310 }
6311 Ok(())
6312 }
6313
6314 pub fn moe_w_scale_by_expert(
6317 &self,
6318 w: &mut CudaSlice<f32>,
6319 sel: &CudaSlice<i32>,
6320 macros: &CudaSlice<f32>,
6321 n_expert: usize,
6322 n: usize,
6323 ) -> Result<(), Box<dyn std::error::Error>> {
6324 let f = self.func("moe_w_scale_by_expert");
6325 let cfg = LaunchConfig {
6326 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6327 block_dim: (64, 1, 1),
6328 shared_mem_bytes: 0,
6329 };
6330 let (ne, nn) = (n_expert as i32, n as i32);
6331 let __s_b = self.gpu.stream();
6332 let mut b = __s_b.launch_builder(&f);
6333 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6334 unsafe {
6335 b.launch(cfg)?;
6336 }
6337 Ok(())
6338 }
6339
6340 pub fn moe_gate_up_silu8_dev_q8(
6341 &self,
6342 table: &CudaSlice<u64>,
6343 sel: &cudarc::driver::CudaView<i32>,
6344 aq: &CudaSlice<i8>,
6345 ad: &CudaSlice<f32>,
6346 in_f: usize,
6347 n_ff: usize,
6348 n_used: usize,
6349 n_expert: usize,
6350 qt_g: i32,
6351 qt_u: i32,
6352 rb_g: usize,
6353 rb_u: usize,
6354 macros: &CudaSlice<f32>,
6355 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6356 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6357 let (mode, wpb) = GU.get_or_init(|| {
6358 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6359 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6360 .ok()
6361 .and_then(|v| v.parse().ok())
6362 .unwrap_or(4u32)
6363 .clamp(1, 16);
6364 (mode, wpb)
6365 });
6366 let (mode, wpb) = (mode.as_str(), *wpb);
6367 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6368 let (inf, nff, ne, rbg, rbu) = (
6369 in_f as i32,
6370 n_ff as i32,
6371 n_expert as i32,
6372 rb_g as i64,
6373 rb_u as i64,
6374 );
6375 let (f, cfg) = match mode {
6376 "1" | "2" | "4" => {
6377 let rpw: u32 = mode.parse().unwrap();
6378 let f = self.func(match rpw {
6379 1 => "moe_gate_up_silu8_dev_q8_r1",
6380 2 => "moe_gate_up_silu8_dev_q8_r2",
6381 _ => "moe_gate_up_silu8_dev_q8_r4",
6382 });
6383 let rows_per_block = (rpw * wpb) as usize;
6384 let gx = n_ff.div_ceil(rows_per_block) as u32;
6385 (
6386 f,
6387 LaunchConfig {
6388 grid_dim: (gx, n_used as u32, 1),
6389 block_dim: (32, wpb, 1),
6390 shared_mem_bytes: 0,
6391 },
6392 )
6393 }
6394 "j8" if n_used <= 32 => (
6395 self.func("moe_gate_up_silu8_dev_q8_j8"),
6396 LaunchConfig {
6397 grid_dim: (n_ff as u32, 1, 1),
6398 block_dim: (32, n_used as u32, 1),
6399 shared_mem_bytes: 0,
6400 },
6401 ),
6402 "vsm2" => {
6404 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6405 let sh = (rb_g + rb_u) as u32;
6406 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6407 f.set_attribute(
6408 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6409 sh as i32,
6410 )?;
6411 (
6412 f,
6413 LaunchConfig {
6414 grid_dim: (n_ff as u32, n_used as u32, 1),
6415 block_dim: (32, 1, 1),
6416 shared_mem_bytes: sh,
6417 },
6418 )
6419 }
6420 "vsm" => {
6421 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6422 let sh = (rb_g + rb_u) as u32;
6423 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6424 f.set_attribute(
6425 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6426 sh as i32,
6427 )?;
6428 (
6429 f,
6430 LaunchConfig {
6431 grid_dim: (n_ff as u32, n_used as u32, 1),
6432 block_dim: (32, 1, 1),
6433 shared_mem_bytes: sh,
6434 },
6435 )
6436 }
6437 "sg" => (
6438 self.func("moe_gate_up_silu8_dev_q8_sg"),
6439 LaunchConfig {
6440 grid_dim: (n_ff as u32, n_used as u32, 1),
6441 block_dim: (32, 1, 1),
6442 shared_mem_bytes: 0,
6443 },
6444 ),
6445 "j8sg" if n_used <= 32 => (
6446 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6447 LaunchConfig {
6448 grid_dim: (n_ff as u32, 1, 1),
6449 block_dim: (32, n_used as u32, 1),
6450 shared_mem_bytes: 0,
6451 },
6452 ),
6453 "u64" if in_f == 2048 => (
6454 self.func("moe_gate_up_silu8_dev_q8_u64"),
6455 LaunchConfig {
6456 grid_dim: (n_ff as u32, n_used as u32, 1),
6457 block_dim: (32, 1, 1),
6458 shared_mem_bytes: 0,
6459 },
6460 ),
6461 "gs4" if in_f == 2048 => (
6462 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6463 LaunchConfig {
6464 grid_dim: (n_ff as u32, n_used as u32, 1),
6465 block_dim: (32, 4, 1),
6466 shared_mem_bytes: 0,
6467 },
6468 ),
6469 "v" | "" => (
6471 self.func("moe_gate_up_silu8_dev_q8_v"),
6472 LaunchConfig {
6473 grid_dim: (n_ff as u32, n_used as u32, 1),
6474 block_dim: (32, 1, 1),
6475 shared_mem_bytes: 0,
6476 },
6477 ),
6478 "s2" => (
6479 self.func("moe_gate_up_silu8_dev_q8_s2"),
6480 LaunchConfig {
6481 grid_dim: (n_ff as u32, n_used as u32, 1),
6482 block_dim: (32, 2, 1),
6483 shared_mem_bytes: 0,
6484 },
6485 ),
6486 "s2z" => {
6487 let rz = wpb.min(16); (
6489 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6490 LaunchConfig {
6491 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6492 block_dim: (32, 2, rz),
6493 shared_mem_bytes: 0,
6494 },
6495 )
6496 }
6497 _ => (
6498 self.func("moe_gate_up_silu8_dev_q8"),
6499 LaunchConfig {
6500 grid_dim: (n_ff as u32, n_used as u32, 1),
6501 block_dim: (32, 1, 1),
6502 shared_mem_bytes: 0,
6503 },
6504 ),
6505 };
6506 let __s_b = self.gpu.stream();
6507 let mut b = __s_b.launch_builder(&f);
6508 b.arg(table)
6509 .arg(sel)
6510 .arg(aq)
6511 .arg(ad)
6512 .arg(&mut act)
6513 .arg(&inf)
6514 .arg(&nff)
6515 .arg(&ne)
6516 .arg(&qt_g)
6517 .arg(&qt_u)
6518 .arg(&rbg)
6519 .arg(&rbu)
6520 .arg(macros);
6521 unsafe {
6522 b.launch(cfg)?;
6523 }
6524 Ok(act)
6525 }
6526
6527 #[allow(clippy::too_many_arguments)]
6528 pub fn moe_down8_fma_dev_q8(
6529 &self,
6530 table: &CudaSlice<u64>,
6531 sel: &cudarc::driver::CudaView<i32>,
6532 w: &cudarc::driver::CudaView<f32>,
6533 aq2: &CudaSlice<i8>,
6534 ad2: &CudaSlice<f32>,
6535 dst: &mut cudarc::driver::CudaViewMut<f32>,
6536 in_f: usize,
6537 out_f: usize,
6538 n_used: usize,
6539 n_expert: usize,
6540 qt: i32,
6541 rb: usize,
6542 ) -> Result<(), Box<dyn std::error::Error>> {
6543 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6544 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6545 let (inf, outf, nu, ne, rbi) = (
6546 in_f as i32,
6547 out_f as i32,
6548 n_used as i32,
6549 n_expert as i32,
6550 rb as i64,
6551 );
6552 let (f, cfg) = match mode.as_str() {
6555 m @ ("1" | "2" | "4") if n_used <= 8 => {
6556 let rpw: usize = m.parse().unwrap();
6557 let f = self.func(match rpw {
6558 1 => "moe_down8_fma_dev_q8_w8r1",
6559 2 => "moe_down8_fma_dev_q8_w8r2",
6560 _ => "moe_down8_fma_dev_q8_w8r4",
6561 });
6562 (
6563 f,
6564 LaunchConfig {
6565 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6566 block_dim: (32, n_used as u32, 1),
6567 shared_mem_bytes: 0,
6568 },
6569 )
6570 }
6571 "h2" if in_f == 512 => (
6572 self.func("moe_down8_fma_dev_q8_h2"),
6573 LaunchConfig {
6574 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6575 block_dim: (32, 1, 1),
6576 shared_mem_bytes: 0,
6577 },
6578 ),
6579 "" if in_f == 704 && n_used <= 8 => (
6582 self.func("moe_down8_fma_dev_q8_w8r2"),
6583 LaunchConfig {
6584 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6585 block_dim: (32, n_used as u32, 1),
6586 shared_mem_bytes: 0,
6587 },
6588 ),
6589 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6593 self.func("moe_down8_fma_dev_q8_w8h2v"),
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 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6601 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6602 LaunchConfig {
6603 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6604 block_dim: (32, n_used as u32, 1),
6605 shared_mem_bytes: 0,
6606 },
6607 ),
6608 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6609 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6610 LaunchConfig {
6611 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6612 block_dim: (32, n_used as u32, 1),
6613 shared_mem_bytes: 0,
6614 },
6615 ),
6616 "w8h2" if in_f == 512 && n_used <= 8 => (
6617 self.func("moe_down8_fma_dev_q8_w8h2"),
6618 LaunchConfig {
6619 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6620 block_dim: (32, n_used as u32, 1),
6621 shared_mem_bytes: 0,
6622 },
6623 ),
6624 _ => (
6625 self.func("moe_down8_fma_dev_q8"),
6626 LaunchConfig {
6627 grid_dim: (out_f as u32, 1, 1),
6628 block_dim: (32, 1, 1),
6629 shared_mem_bytes: 0,
6630 },
6631 ),
6632 };
6633 let __s_b = self.gpu.stream();
6634 let mut b = __s_b.launch_builder(&f);
6635 b.arg(table)
6636 .arg(sel)
6637 .arg(w)
6638 .arg(aq2)
6639 .arg(ad2)
6640 .arg(dst)
6641 .arg(&inf)
6642 .arg(&outf)
6643 .arg(&nu)
6644 .arg(&ne)
6645 .arg(&qt)
6646 .arg(&rbi);
6647 unsafe {
6648 b.launch(cfg)?;
6649 }
6650 Ok(())
6651 }
6652
6653 #[allow(clippy::too_many_arguments)]
6660 pub fn moe_gate_up_silu8_dev_q8_rows(
6661 &self,
6662 table: &CudaSlice<u64>,
6663 sel: &CudaSlice<i32>,
6664 aq: &CudaSlice<i8>,
6665 ad: &CudaSlice<f32>,
6666 t: usize,
6667 in_f: usize,
6668 n_ff: usize,
6669 n_used: usize,
6670 n_expert: usize,
6671 qt_g: i32,
6672 qt_u: i32,
6673 rb_g: usize,
6674 rb_u: usize,
6675 macros: &CudaSlice<f32>,
6676 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6677 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6678 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6679 let cfg = LaunchConfig {
6680 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6681 block_dim: (32, 1, 1),
6682 shared_mem_bytes: 0,
6683 };
6684 let (inf, nff, ne, nu, rbg, rbu) = (
6685 in_f as i32,
6686 n_ff as i32,
6687 n_expert as i32,
6688 n_used as i32,
6689 rb_g as i64,
6690 rb_u as i64,
6691 );
6692 let __s_b = self.gpu.stream();
6693 let mut b = __s_b.launch_builder(&f);
6694 b.arg(table)
6695 .arg(sel)
6696 .arg(aq)
6697 .arg(ad)
6698 .arg(&mut act)
6699 .arg(&inf)
6700 .arg(&nff)
6701 .arg(&ne)
6702 .arg(&qt_g)
6703 .arg(&qt_u)
6704 .arg(&rbg)
6705 .arg(&rbu)
6706 .arg(&nu)
6707 .arg(macros);
6708 unsafe {
6709 b.launch(cfg)?;
6710 }
6711 Ok(act)
6712 }
6713
6714 #[allow(clippy::too_many_arguments)]
6719 pub fn moe_down8_fma_dev_q8_rows(
6720 &self,
6721 table: &CudaSlice<u64>,
6722 sel: &CudaSlice<i32>,
6723 w: &CudaSlice<f32>,
6724 aq2: &CudaSlice<i8>,
6725 ad2: &CudaSlice<f32>,
6726 dst: &mut CudaSlice<f32>,
6727 t: usize,
6728 in_f: usize,
6729 out_f: usize,
6730 n_used: usize,
6731 n_expert: usize,
6732 qt: i32,
6733 rb: usize,
6734 ) -> Result<(), Box<dyn std::error::Error>> {
6735 assert!(
6736 in_f == 512 && n_used <= 8,
6737 "down rows twin is w8h2v shape-gated"
6738 );
6739 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6740 let cfg = LaunchConfig {
6741 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6742 block_dim: (32, n_used as u32, 1),
6743 shared_mem_bytes: 0,
6744 };
6745 let (inf, outf, nu, ne, rbi) = (
6746 in_f as i32,
6747 out_f as i32,
6748 n_used as i32,
6749 n_expert as i32,
6750 rb as i64,
6751 );
6752 let __s_b = self.gpu.stream();
6753 let mut b = __s_b.launch_builder(&f);
6754 b.arg(table)
6755 .arg(sel)
6756 .arg(w)
6757 .arg(aq2)
6758 .arg(ad2)
6759 .arg(dst)
6760 .arg(&inf)
6761 .arg(&outf)
6762 .arg(&nu)
6763 .arg(&ne)
6764 .arg(&qt)
6765 .arg(&rbi);
6766 unsafe {
6767 b.launch(cfg)?;
6768 }
6769 Ok(())
6770 }
6771
6772 #[allow(clippy::too_many_arguments)]
6776 pub fn moe_gate_up_silu8_dev_q8_csr(
6777 &self,
6778 table: &CudaSlice<u64>,
6779 sel: &CudaSlice<i32>,
6780 aq: &CudaSlice<i8>,
6781 ad: &CudaSlice<f32>,
6782 n_pairs: usize,
6783 in_f: usize,
6784 n_ff: usize,
6785 n_used: usize,
6786 n_expert: usize,
6787 qt_g: i32,
6788 qt_u: i32,
6789 rb_g: usize,
6790 rb_u: usize,
6791 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6792 let f = if qt_g == crate::QT_NVFP4 {
6795 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6796 } else {
6797 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6798 };
6799 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6800 let cfg = LaunchConfig {
6801 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6802 block_dim: (32, 1, 1),
6803 shared_mem_bytes: 0,
6804 };
6805 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6806 in_f as i32,
6807 n_ff as i32,
6808 n_expert as i32,
6809 n_used as i32,
6810 n_pairs as i32,
6811 rb_g as i64,
6812 rb_u as i64,
6813 );
6814 let __s_b = self.gpu.stream();
6815 let mut b = __s_b.launch_builder(&f);
6816 b.arg(table)
6817 .arg(sel)
6818 .arg(aq)
6819 .arg(ad)
6820 .arg(&mut act)
6821 .arg(&inf)
6822 .arg(&nff)
6823 .arg(&ne)
6824 .arg(&qt_g)
6825 .arg(&qt_u)
6826 .arg(&rbg)
6827 .arg(&rbu)
6828 .arg(&nu)
6829 .arg(&npi);
6830 unsafe {
6831 b.launch(cfg)?;
6832 }
6833 Ok(act)
6834 }
6835
6836 #[allow(clippy::too_many_arguments)]
6840 pub fn moe_down8_fma_dev_q8_variant(
6841 &self,
6842 variant: &str,
6843 table: &CudaSlice<u64>,
6844 sel: &cudarc::driver::CudaView<i32>,
6845 w: &cudarc::driver::CudaView<f32>,
6846 aq2: &CudaSlice<i8>,
6847 ad2: &CudaSlice<f32>,
6848 dst: &mut cudarc::driver::CudaViewMut<f32>,
6849 in_f: usize,
6850 out_f: usize,
6851 n_used: usize,
6852 n_expert: usize,
6853 qt: i32,
6854 rb: usize,
6855 ) -> Result<(), Box<dyn std::error::Error>> {
6856 let (inf, outf, nu, ne, rbi) = (
6857 in_f as i32,
6858 out_f as i32,
6859 n_used as i32,
6860 n_expert as i32,
6861 rb as i64,
6862 );
6863 let (f, cfg) = match variant {
6864 "w8h2" | "w8h2v" => (
6865 self.func(if variant == "w8h2" {
6866 "moe_down8_fma_dev_q8_w8h2"
6867 } else {
6868 "moe_down8_fma_dev_q8_w8h2v"
6869 }),
6870 LaunchConfig {
6871 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6872 block_dim: (32, n_used as u32, 1),
6873 shared_mem_bytes: 0,
6874 },
6875 ),
6876 "w8h2r2" | "w8h2r2v" => (
6877 self.func(if variant == "w8h2r2" {
6878 "moe_down8_fma_dev_q8_w8h2r2"
6879 } else {
6880 "moe_down8_fma_dev_q8_w8h2r2v"
6881 }),
6882 LaunchConfig {
6883 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6884 block_dim: (32, n_used as u32, 1),
6885 shared_mem_bytes: 0,
6886 },
6887 ),
6888 _ => (
6889 self.func("moe_down8_fma_dev_q8"),
6890 LaunchConfig {
6891 grid_dim: (out_f as u32, 1, 1),
6892 block_dim: (32, 1, 1),
6893 shared_mem_bytes: 0,
6894 },
6895 ),
6896 };
6897 let __s_b = self.gpu.stream();
6898 let mut b = __s_b.launch_builder(&f);
6899 b.arg(table)
6900 .arg(sel)
6901 .arg(w)
6902 .arg(aq2)
6903 .arg(ad2)
6904 .arg(dst)
6905 .arg(&inf)
6906 .arg(&outf)
6907 .arg(&nu)
6908 .arg(&ne)
6909 .arg(&qt)
6910 .arg(&rbi);
6911 unsafe {
6912 b.launch(cfg)?;
6913 }
6914 Ok(())
6915 }
6916
6917 #[allow(clippy::too_many_arguments)]
6919 pub fn moe_gate_up_silu8_dev_q8_variant(
6920 &self,
6921 variant: &str,
6922 table: &CudaSlice<u64>,
6923 sel: &cudarc::driver::CudaView<i32>,
6924 aq: &CudaSlice<i8>,
6925 ad: &CudaSlice<f32>,
6926 in_f: usize,
6927 n_ff: usize,
6928 n_used: usize,
6929 n_expert: usize,
6930 qt_g: i32,
6931 qt_u: i32,
6932 rb_g: usize,
6933 rb_u: usize,
6934 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6935 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6936 let (inf, nff, ne, rbg, rbu) = (
6937 in_f as i32,
6938 n_ff as i32,
6939 n_expert as i32,
6940 rb_g as i64,
6941 rb_u as i64,
6942 );
6943 let f = self.func(if variant == "v" {
6944 "moe_gate_up_silu8_dev_q8_v"
6945 } else {
6946 "moe_gate_up_silu8_dev_q8"
6947 });
6948 let cfg = LaunchConfig {
6949 grid_dim: (n_ff as u32, n_used as u32, 1),
6950 block_dim: (32, 1, 1),
6951 shared_mem_bytes: 0,
6952 };
6953 let __s_b = self.gpu.stream();
6954 let mut b = __s_b.launch_builder(&f);
6955 b.arg(table)
6956 .arg(sel)
6957 .arg(aq)
6958 .arg(ad)
6959 .arg(&mut act)
6960 .arg(&inf)
6961 .arg(&nff)
6962 .arg(&ne)
6963 .arg(&qt_g)
6964 .arg(&qt_u)
6965 .arg(&rbg)
6966 .arg(&rbu);
6967 unsafe {
6968 b.launch(cfg)?;
6969 }
6970 Ok(act)
6971 }
6972
6973 pub fn moe_gate_up_silu8_dev(
6974 &self,
6975 table: &CudaSlice<u64>,
6976 sel: &cudarc::driver::CudaView<i32>,
6977 x: &cudarc::driver::CudaView<f32>,
6978 in_f: usize,
6979 n_ff: usize,
6980 n_used: usize,
6981 n_expert: usize,
6982 qt_g: i32,
6983 qt_u: i32,
6984 rb_g: usize,
6985 rb_u: usize,
6986 macros: &CudaSlice<f32>,
6987 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6988 let f = self.func("moe_gate_up_silu8_dev");
6989 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6991 grid_dim: (n_ff as u32, n_used as u32, 1),
6992 block_dim: (256, 1, 1),
6993 shared_mem_bytes: 0,
6994 };
6995 let (inf, nff, ne, rbg, rbu) = (
6996 in_f as i32,
6997 n_ff as i32,
6998 n_expert as i32,
6999 rb_g as i64,
7000 rb_u as i64,
7001 );
7002 let __s_b = self.gpu.stream();
7003 let mut b = __s_b.launch_builder(&f);
7004 b.arg(table)
7005 .arg(sel)
7006 .arg(x)
7007 .arg(&mut act)
7008 .arg(&inf)
7009 .arg(&nff)
7010 .arg(&ne)
7011 .arg(&qt_g)
7012 .arg(&qt_u)
7013 .arg(&rbg)
7014 .arg(&rbu)
7015 .arg(macros);
7016 unsafe {
7017 b.launch(cfg)?;
7018 }
7019 Ok(act)
7020 }
7021
7022 #[allow(clippy::too_many_arguments)]
7025 pub fn moe_down8_fma_dev(
7026 &self,
7027 table: &CudaSlice<u64>,
7028 sel: &cudarc::driver::CudaView<i32>,
7029 w: &cudarc::driver::CudaView<f32>,
7030 act: &CudaSlice<f32>,
7031 dst: &mut cudarc::driver::CudaViewMut<f32>,
7032 in_f: usize,
7033 out_f: usize,
7034 n_used: usize,
7035 n_expert: usize,
7036 qt: i32,
7037 rb: usize,
7038 ) -> Result<(), Box<dyn std::error::Error>> {
7039 let f = self.func("moe_down8_fma_dev");
7040 let cfg = LaunchConfig {
7041 grid_dim: (out_f as u32, 1, 1),
7042 block_dim: (256, 1, 1),
7043 shared_mem_bytes: 0,
7044 };
7045 let (inf, outf, nu, ne, rbv) = (
7046 in_f as i32,
7047 out_f as i32,
7048 n_used as i32,
7049 n_expert as i32,
7050 rb as i64,
7051 );
7052 let __s_b = self.gpu.stream();
7053 let mut b = __s_b.launch_builder(&f);
7054 b.arg(table)
7055 .arg(sel)
7056 .arg(w)
7057 .arg(act)
7058 .arg(dst)
7059 .arg(&inf)
7060 .arg(&outf)
7061 .arg(&nu)
7062 .arg(&ne)
7063 .arg(&qt)
7064 .arg(&rbv);
7065 unsafe {
7066 b.launch(cfg)?;
7067 }
7068 Ok(())
7069 }
7070
7071 pub fn axpy_into(
7073 &self,
7074 src: &CudaSlice<f32>,
7075 alpha: f32,
7076 dst: &mut cudarc::driver::CudaViewMut<f32>,
7077 n: usize,
7078 ) -> Result<(), Box<dyn std::error::Error>> {
7079 let f = self.func("axpy_f32");
7080 let cfg = LaunchConfig::for_num_elems(n as u32);
7081 let (a, ni) = (alpha, n as i32);
7082 let __s_b = self.gpu.stream();
7083 let mut b = __s_b.launch_builder(&f);
7084 b.arg(src).arg(dst).arg(&a).arg(&ni);
7085 unsafe {
7086 b.launch(cfg)?;
7087 }
7088 Ok(())
7089 }
7090
7091 pub fn axpy_host_into(
7093 &self,
7094 src: &cudarc::driver::CudaView<'_, f32>,
7095 alpha: f32,
7096 dst: &mut cudarc::driver::CudaViewMut<f32>,
7097 n: usize,
7098 ) -> Result<(), Box<dyn std::error::Error>> {
7099 let f = self.func("axpy_host_f32");
7100 let cfg = LaunchConfig::for_num_elems(n as u32);
7101 let (a, ni) = (alpha, n as i32);
7102 let __s_b = self.gpu.stream();
7103 let mut b = __s_b.launch_builder(&f);
7104 b.arg(src).arg(dst).arg(&a).arg(&ni);
7105 unsafe {
7106 b.launch(cfg)?;
7107 }
7108 Ok(())
7109 }
7110
7111 pub fn add_scaled_rows(
7113 &self,
7114 src: &CudaSlice<f32>,
7115 scale: &CudaSlice<f32>,
7116 dst: &mut CudaSlice<f32>,
7117 ncols: usize,
7118 nrows: usize,
7119 ) -> Result<(), Box<dyn std::error::Error>> {
7120 let f = self.func("add_scaled_rows_f32");
7121 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7122 let (nc, nr) = (ncols as i32, nrows as i32);
7123 let __s_b = self.gpu.stream();
7124 let mut b = __s_b.launch_builder(&f);
7125 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7126 unsafe {
7127 b.launch(cfg)?;
7128 }
7129 Ok(())
7130 }
7131
7132 pub fn gather_rows(
7136 &self,
7137 src: &CudaSlice<f32>,
7138 idx: &CudaSlice<i32>,
7139 dst: &mut CudaSlice<f32>,
7140 ncols: usize,
7141 m_e: usize,
7142 ) -> Result<(), Box<dyn std::error::Error>> {
7143 let f = self.func("gather_rows_f32");
7144 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7145 let (nc, me) = (ncols as i32, m_e as i32);
7146 let __s_b = self.gpu.stream();
7147 let mut b = __s_b.launch_builder(&f);
7148 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7149 unsafe {
7150 b.launch(cfg)?;
7151 }
7152 Ok(())
7153 }
7154
7155 pub fn scatter_slot(
7160 &self,
7161 src: &CudaSlice<f32>,
7162 tok_idx: &CudaSlice<i32>,
7163 slot_idx: &CudaSlice<i32>,
7164 weight: &CudaSlice<f32>,
7165 dst: &mut CudaSlice<f32>,
7166 wbuf: &mut CudaSlice<f32>,
7167 ncols: usize,
7168 n_used: usize,
7169 m_e: usize,
7170 ) -> Result<(), Box<dyn std::error::Error>> {
7171 let f = self.func("scatter_add_slot_f32");
7172 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7173 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7174 let __s_b = self.gpu.stream();
7175 let mut b = __s_b.launch_builder(&f);
7176 b.arg(src)
7177 .arg(tok_idx)
7178 .arg(slot_idx)
7179 .arg(weight)
7180 .arg(dst)
7181 .arg(wbuf)
7182 .arg(&nc)
7183 .arg(&nu)
7184 .arg(&me);
7185 unsafe {
7186 b.launch(cfg)?;
7187 }
7188 Ok(())
7189 }
7190
7191 pub fn reduce_slots(
7195 &self,
7196 slots: &CudaSlice<f32>,
7197 wbuf: &CudaSlice<f32>,
7198 dst: &mut CudaSlice<f32>,
7199 ncols: usize,
7200 n_used: usize,
7201 t: usize,
7202 ) -> Result<(), Box<dyn std::error::Error>> {
7203 let f = self.func("reduce_slots_f32");
7204 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7205 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7206 let __s_b = self.gpu.stream();
7207 let mut b = __s_b.launch_builder(&f);
7208 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7209 unsafe {
7210 b.launch(cfg)?;
7211 }
7212 Ok(())
7213 }
7214
7215 pub fn reduce_slots_host(
7220 &self,
7221 slots: &CudaSlice<f32>,
7222 wbuf: &CudaSlice<f32>,
7223 dst: &mut CudaSlice<f32>,
7224 ncols: usize,
7225 n_used: usize,
7226 t: usize,
7227 ) -> Result<(), Box<dyn std::error::Error>> {
7228 let f = self.func("reduce_slots_host_f32");
7229 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7230 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7231 let __s_b = self.gpu.stream();
7232 let mut b = __s_b.launch_builder(&f);
7233 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7234 unsafe {
7235 b.launch(cfg)?;
7236 }
7237 Ok(())
7238 }
7239
7240 pub fn quantize_q8_1_view(
7247 &self,
7248 x: &cudarc::driver::CudaView<f32>,
7249 m: usize,
7250 in_f: usize,
7251 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7252 let f = self.func("quantize_q8_1");
7253 let nblk = in_f / 32;
7254 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7255 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7256 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7257 let (inf, mi) = (in_f as i32, m as i32);
7258 let __s_b = self.gpu.stream();
7259 let mut b = __s_b.launch_builder(&f);
7260 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7261 unsafe {
7262 b.launch(cfg)?;
7263 }
7264 Ok((q, d))
7265 }
7266
7267 pub fn quantize_q8_1(
7268 &self,
7269 x: &CudaSlice<f32>,
7270 m: usize,
7271 in_f: usize,
7272 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7273 let nblk = in_f / 32;
7274 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);
7278 let (inf, mi) = (in_f as i32, m as i32);
7279 if Self::pdl_on() && Self::pdl_wb_on() {
7280 {
7281 use cudarc::driver::{DevicePtr, DevicePtrMut};
7282 let s = &self.gpu.stream();
7283 let (px, _g0) = x.device_ptr(s);
7284 let (pq, _g1) = q.device_ptr_mut(s);
7285 let (pd, _g2) = d.device_ptr_mut(s);
7286 let mut ps = [
7287 &px as *const _ as *mut std::ffi::c_void,
7288 &pq as *const _ as *mut _,
7289 &pd as *const _ as *mut _,
7290 &inf as *const _ as *mut _,
7291 &mi as *const _ as *mut _,
7292 ];
7293 unsafe {
7294 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7295 }
7296 }
7297 return Ok((q, d));
7298 }
7299 let f = self.func("quantize_q8_1");
7300 let __s_b = self.gpu.stream();
7301 let mut b = __s_b.launch_builder(&f);
7302 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7303 unsafe {
7304 b.launch(cfg)?;
7305 }
7306 Ok((q, d))
7307 }
7308
7309 pub fn quantize_fp4_act(
7313 &self,
7314 x: &CudaSlice<f32>,
7315 m: usize,
7316 in_f: usize,
7317 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7318 let f = self.func("quantize_fp4_act");
7319 let nb16 = in_f / 16;
7320 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);
7323 let (inf, mi) = (in_f as i32, m as i32);
7324 let __s_b = self.gpu.stream();
7325 let mut b = __s_b.launch_builder(&f);
7326 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7327 unsafe {
7328 b.launch(cfg)?;
7329 }
7330 Ok((aq4, ad4))
7331 }
7332
7333 pub fn qmatvec_gemm_nvfp4_fp4(
7338 &self,
7339 bytes: &CudaSlice<u8>,
7340 x: &CudaSlice<f32>,
7341 m: usize,
7342 in_f: usize,
7343 out_f: usize,
7344 row_bytes: usize,
7345 scale: f32,
7346 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7347 assert!(
7348 in_f % 64 == 0,
7349 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7350 );
7351 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7352 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7353 if scale != 1.0 {
7354 self.scale_inplace(&mut y, scale, m * out_f)?;
7355 }
7356 Ok(y)
7357 }
7358
7359 fn fp4_gemm_launch(
7362 &self,
7363 bytes: &CudaSlice<u8>,
7364 aq4: &CudaSlice<u32>,
7365 ad4: &CudaSlice<u8>,
7366 m: usize,
7367 in_f: usize,
7368 out_f: usize,
7369 row_bytes: usize,
7370 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7371 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7372 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7374 const BN: u32 = 256;
7375 let cfg = LaunchConfig {
7376 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7377 block_dim: (32, 4, 1),
7378 shared_mem_bytes: 0,
7379 };
7380 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7381 let __s_b = self.gpu.stream();
7382 let mut b = __s_b.launch_builder(&f);
7383 b.arg(bytes)
7384 .arg(aq4)
7385 .arg(ad4)
7386 .arg(&mut y)
7387 .arg(&inf)
7388 .arg(&outf)
7389 .arg(&mi)
7390 .arg(&rb);
7391 unsafe {
7392 b.launch(cfg)?;
7393 }
7394 Ok(y)
7395 }
7396
7397 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7399 &self,
7400 bytes: &CudaSlice<u8>,
7401 x: &CudaSlice<f32>,
7402 m: usize,
7403 in_f: usize,
7404 out_f: usize,
7405 row_bytes: usize,
7406 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7407 assert!(
7408 in_f % 64 == 0,
7409 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7410 );
7411 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7412 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7413 }
7414
7415 pub fn qmatvec_q8_0_fast(
7417 &self,
7418 w: &CudaSlice<u8>,
7419 x: &CudaSlice<f32>,
7420 m: usize,
7421 in_f: usize,
7422 out_f: usize,
7423 row_bytes: usize,
7424 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7425 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7426 let f = self.func("qmatvec_q8_0_dp4a");
7427 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7429 grid_dim: (out_f as u32, m as u32, 1),
7430 block_dim: (128, 1, 1),
7431 shared_mem_bytes: 0,
7432 };
7433 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7434 let __s_b = self.gpu.stream();
7435 let mut b = __s_b.launch_builder(&f);
7436 b.arg(w)
7437 .arg(&aq)
7438 .arg(&ad)
7439 .arg(&mut y)
7440 .arg(&inf)
7441 .arg(&outf)
7442 .arg(&mi)
7443 .arg(&rb);
7444 unsafe {
7445 b.launch(cfg)?;
7446 }
7447 Ok(y)
7448 }
7449
7450 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7453 &self,
7454 w: &CudaSlice<u8>,
7455 x: &CudaSlice<f32>,
7456 m: usize,
7457 in_f: usize,
7458 out_f: usize,
7459 row_bytes: usize,
7460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7461 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7462 let f = self.func("qmatvec_q4_K_dp4a");
7463 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7465 grid_dim: (out_f as u32, m as u32, 1),
7466 block_dim: (128, 1, 1),
7467 shared_mem_bytes: 0,
7468 };
7469 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7470 let __s_b = self.gpu.stream();
7471 let mut b = __s_b.launch_builder(&f);
7472 b.arg(w)
7473 .arg(&aq)
7474 .arg(&ad)
7475 .arg(&mut y)
7476 .arg(&inf)
7477 .arg(&outf)
7478 .arg(&mi)
7479 .arg(&rb);
7480 unsafe {
7481 b.launch(cfg)?;
7482 }
7483 Ok(y)
7484 }
7485
7486 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7489 &self,
7490 w: &CudaSlice<u8>,
7491 x: &CudaSlice<f32>,
7492 m: usize,
7493 in_f: usize,
7494 out_f: usize,
7495 row_bytes: usize,
7496 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7497 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7498 let f = self.func("qmatvec_q6_K_dp4a");
7499 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7501 grid_dim: (out_f as u32, m as u32, 1),
7502 block_dim: (128, 1, 1),
7503 shared_mem_bytes: 0,
7504 };
7505 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7506 let __s_b = self.gpu.stream();
7507 let mut b = __s_b.launch_builder(&f);
7508 b.arg(w)
7509 .arg(&aq)
7510 .arg(&ad)
7511 .arg(&mut y)
7512 .arg(&inf)
7513 .arg(&outf)
7514 .arg(&mi)
7515 .arg(&rb);
7516 unsafe {
7517 b.launch(cfg)?;
7518 }
7519 Ok(y)
7520 }
7521
7522 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7525 &self,
7526 w: &CudaSlice<u8>,
7527 x: &CudaSlice<f32>,
7528 m: usize,
7529 in_f: usize,
7530 out_f: usize,
7531 row_bytes: usize,
7532 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7533 self.qmatvec_dp4a_named(
7534 "qmatvec_q5_K_dp4a",
7535 &w.slice(0..w.len()),
7536 x,
7537 m,
7538 in_f,
7539 out_f,
7540 row_bytes,
7541 )
7542 }
7543 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7546 &self,
7547 w: &CudaSlice<u8>,
7548 x: &CudaSlice<f32>,
7549 m: usize,
7550 in_f: usize,
7551 out_f: usize,
7552 row_bytes: usize,
7553 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7554 self.qmatvec_dp4a_named(
7555 "qmatvec_q3_K_dp4a",
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_rp(
7566 &self,
7567 w: &CudaSlice<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!(
7575 in_f % 64 == 0,
7576 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7577 );
7578 self.qmatvec_dp4a_named(
7579 "qmatvec_nvfp4_dp4a_rp",
7580 &w.slice(0..w.len()),
7581 x,
7582 m,
7583 in_f,
7584 out_f,
7585 row_bytes,
7586 )
7587 }
7588 pub fn qmatvec_nvfp4_fast(
7590 &self,
7591 w: &cudarc::driver::CudaView<'_, u8>,
7592 x: &CudaSlice<f32>,
7593 m: usize,
7594 in_f: usize,
7595 out_f: usize,
7596 row_bytes: usize,
7597 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7598 assert!(
7601 in_f % 64 == 0,
7602 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7603 );
7604 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7605 }
7606 pub fn qmatvec_nvfp4_fast_v2(
7609 &self,
7610 w: &cudarc::driver::CudaView<'_, u8>,
7611 x: &CudaSlice<f32>,
7612 m: usize,
7613 in_f: usize,
7614 out_f: usize,
7615 row_bytes: usize,
7616 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7617 assert!(
7618 in_f % 64 == 0,
7619 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7620 );
7621 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7622 }
7623 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7626 &self,
7627 w: &CudaSlice<u8>,
7628 x: &CudaSlice<f32>,
7629 m: usize,
7630 in_f: usize,
7631 out_f: usize,
7632 row_bytes: usize,
7633 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7634 self.qmatvec_dp4a_named(
7635 "qmatvec_iq4_XS_dp4a",
7636 &w.slice(0..w.len()),
7637 x,
7638 m,
7639 in_f,
7640 out_f,
7641 row_bytes,
7642 )
7643 }
7644
7645 fn qmatvec_dp4a_named(
7647 &self,
7648 name: &str,
7649 w: &cudarc::driver::CudaView<'_, u8>,
7650 x: &CudaSlice<f32>,
7651 m: usize,
7652 in_f: usize,
7653 out_f: usize,
7654 row_bytes: usize,
7655 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7656 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7657 let f = self.func(name);
7658 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7660 grid_dim: (out_f as u32, m as u32, 1),
7661 block_dim: (128, 1, 1),
7662 shared_mem_bytes: 0,
7663 };
7664 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7665 let __s_b = self.gpu.stream();
7666 let mut b = __s_b.launch_builder(&f);
7667 b.arg(w)
7668 .arg(&aq)
7669 .arg(&ad)
7670 .arg(&mut y)
7671 .arg(&inf)
7672 .arg(&outf)
7673 .arg(&mi)
7674 .arg(&rb);
7675 unsafe {
7676 b.launch(cfg)?;
7677 }
7678 Ok(y)
7679 }
7680
7681 #[allow(clippy::too_many_arguments)]
7687 pub fn qmatvec_nvfp4_fast_prequant_into(
7688 &self,
7689 w: &CudaSlice<u8>,
7690 aq: &CudaSlice<i8>,
7691 ad: &CudaSlice<f32>,
7692 y: &mut CudaSlice<f32>,
7693 m: usize,
7694 in_f: usize,
7695 out_f: usize,
7696 row_bytes: usize,
7697 ) -> Result<(), Box<dyn std::error::Error>> {
7698 assert!(
7699 in_f % 64 == 0,
7700 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7701 );
7702 if y.len() < m * out_f {
7703 return Err(format!(
7704 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7705 y.len()
7706 )
7707 .into());
7708 }
7709 let f = self.func("qmatvec_nvfp4_dp4a");
7710 let cfg = LaunchConfig {
7711 grid_dim: (out_f as u32, m as u32, 1),
7712 block_dim: (128, 1, 1),
7713 shared_mem_bytes: 0,
7714 };
7715 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7716 let __s_b = self.gpu.stream();
7717 let mut b = __s_b.launch_builder(&f);
7718 b.arg(w)
7719 .arg(aq)
7720 .arg(ad)
7721 .arg(y)
7722 .arg(&inf)
7723 .arg(&outf)
7724 .arg(&mi)
7725 .arg(&rb);
7726 unsafe {
7727 b.launch(cfg)?;
7728 }
7729 Ok(())
7730 }
7731
7732 #[allow(clippy::too_many_arguments)]
7735 pub fn matvec_f32_qkv_into(
7736 &self,
7737 wq: &CudaSlice<f32>,
7738 wk: &CudaSlice<f32>,
7739 wv: &CudaSlice<f32>,
7740 wg: &CudaSlice<f32>,
7741 x: &CudaSlice<f32>,
7742 yq: &mut CudaSlice<f32>,
7743 yk: &mut CudaSlice<f32>,
7744 yv: &mut CudaSlice<f32>,
7745 yg: &mut CudaSlice<f32>,
7746 in_f: usize,
7747 out_q: usize,
7748 out_kv: usize,
7749 out_g: usize,
7750 ) -> Result<(), Box<dyn std::error::Error>> {
7751 if in_f % 4 != 0
7752 || wq.len() != out_q * in_f
7753 || wk.len() != out_kv * in_f
7754 || wv.len() != out_kv * in_f
7755 || wg.len() < out_g * in_f
7756 || x.len() < in_f
7757 || yq.len() < out_q
7758 || yk.len() < out_kv
7759 || yv.len() < out_kv
7760 || (out_g > 0 && yg.len() < out_g)
7761 {
7762 return Err(format!(
7763 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7764 wq={} wk={} wv={} wg={}",
7765 wq.len(),
7766 wk.len(),
7767 wv.len(),
7768 wg.len()
7769 )
7770 .into());
7771 }
7772 let f = self.func("matvec_f32_qkv");
7773 let cfg = LaunchConfig {
7774 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7775 block_dim: (128, 1, 1),
7776 shared_mem_bytes: 0,
7777 };
7778 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7779 let __s_b = self.gpu.stream();
7780 let mut b = __s_b.launch_builder(&f);
7781 b.arg(wq)
7782 .arg(wk)
7783 .arg(wv)
7784 .arg(wg)
7785 .arg(x)
7786 .arg(yq)
7787 .arg(yk)
7788 .arg(yv)
7789 .arg(yg)
7790 .arg(&inf)
7791 .arg(&oq)
7792 .arg(&okv)
7793 .arg(&og);
7794 unsafe {
7795 b.launch(cfg)?;
7796 }
7797 Ok(())
7798 }
7799
7800 #[allow(clippy::too_many_arguments)]
7806 #[allow(clippy::too_many_arguments)]
7809 pub fn qmatvec_nvfp4_sel_gu_into(
7810 &self,
7811 gate_bank: &CudaSlice<u8>,
7812 up_bank: &CudaSlice<u8>,
7813 sel: &CudaSlice<i32>,
7814 aq: &CudaSlice<i8>,
7815 ad: &CudaSlice<f32>,
7816 yg: &mut CudaSlice<f32>,
7817 yu: &mut CudaSlice<f32>,
7818 n_sel: usize,
7819 in_f: usize,
7820 out_f: usize,
7821 row_bytes: usize,
7822 expert_stride: usize,
7823 ) -> Result<(), Box<dyn std::error::Error>> {
7824 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7825 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7826 return Err("NVFP4 gu sel geometry".into());
7827 }
7828 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7831 let rpw = *RPW.get_or_init(|| {
7832 std::env::var("MEMRA_SEL_GU_RPW")
7833 .ok()
7834 .and_then(|v| v.parse().ok())
7835 .filter(|r| *r == 2 || *r == 4)
7836 .unwrap_or(1)
7837 });
7838 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
7839 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7842 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
7843 let f = self.func(match (wpr, rpw) {
7844 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
7845 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
7846 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
7847 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
7848 });
7849 let cfg = LaunchConfig {
7850 grid_dim: if wpr {
7851 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
7852 } else if rpw == 1 {
7853 ((2 * out_f) as u32, n_sel as u32, 1)
7854 } else {
7855 ((out_f / rpw) as u32, n_sel as u32, 1)
7856 },
7857 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
7858 shared_mem_bytes: 0,
7859 };
7860 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7861 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7862 let (ars, adrs) = (0i64, 0i64);
7863 let __s_b = self.gpu.stream();
7864 let mut b = __s_b.launch_builder(&f);
7865 b.arg(gate_bank)
7866 .arg(up_bank)
7867 .arg(sel)
7868 .arg(aq)
7869 .arg(ad)
7870 .arg(yg)
7871 .arg(yu)
7872 .arg(&inf)
7873 .arg(&outf)
7874 .arg(&ns)
7875 .arg(&rb)
7876 .arg(&es)
7877 .arg(&ars)
7878 .arg(&adrs);
7879 unsafe {
7880 b.launch(cfg)?;
7881 }
7882 Ok(())
7883 }
7884
7885 #[allow(clippy::too_many_arguments)]
7893 pub fn qmatvec_nvfp4_sel_down8_into(
7894 &self,
7895 bank: &CudaSlice<u8>,
7896 sel: &CudaSlice<i32>,
7897 aq: &CudaSlice<i8>,
7898 ad: &CudaSlice<f32>,
7899 route_w: &CudaSlice<f32>,
7900 md: &CudaSlice<f32>,
7901 dst: &mut CudaSlice<f32>,
7902 n_sel: usize,
7903 in_f: usize,
7904 out_f: usize,
7905 row_bytes: usize,
7906 expert_stride: usize,
7907 act_row_stride: usize,
7908 ad_row_stride: usize,
7909 ) -> Result<(), Box<dyn std::error::Error>> {
7910 if in_f % 64 != 0
7911 || n_sel == 0
7912 || n_sel > 8
7913 || (in_f >> 5) > 32
7914 || dst.len() < out_f
7915 || sel.len() < n_sel
7916 || route_w.len() < n_sel
7917 {
7918 return Err(format!(
7919 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
7920 dst.len()
7921 )
7922 .into());
7923 }
7924 if !crate::tp::nvfp4_bank_v2_on() {
7925 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
7926 }
7927 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
7928 let cfg = LaunchConfig {
7929 grid_dim: (out_f as u32, 1, 1),
7930 block_dim: (32, n_sel as u32, 1),
7931 shared_mem_bytes: 0,
7932 };
7933 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7934 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7935 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7936 let __s_b = self.gpu.stream();
7937 let mut b = __s_b.launch_builder(&f);
7938 b.arg(bank)
7939 .arg(sel)
7940 .arg(aq)
7941 .arg(ad)
7942 .arg(route_w)
7943 .arg(md)
7944 .arg(dst)
7945 .arg(&inf)
7946 .arg(&outf)
7947 .arg(&ns)
7948 .arg(&rb)
7949 .arg(&es)
7950 .arg(&ars)
7951 .arg(&adrs);
7952 unsafe {
7953 b.launch(cfg)?;
7954 }
7955 Ok(())
7956 }
7957
7958 #[allow(clippy::too_many_arguments)]
7962 pub fn qmatvec_nvfp4_sel_down8_rows_into(
7963 &self,
7964 bank: &CudaSlice<u8>,
7965 sel: &CudaSlice<i32>,
7966 aq: &CudaSlice<i8>,
7967 ad: &CudaSlice<f32>,
7968 route_w: &CudaSlice<f32>,
7969 md: &CudaSlice<f32>,
7970 dst: &mut CudaSlice<f32>,
7971 t: usize,
7972 n_sel_col: usize,
7973 in_f: usize,
7974 out_f: usize,
7975 row_bytes: usize,
7976 expert_stride: usize,
7977 act_row_stride: usize,
7978 ad_row_stride: usize,
7979 ) -> Result<(), Box<dyn std::error::Error>> {
7980 let n_sel = t * n_sel_col;
7981 if in_f % 64 != 0
7982 || n_sel_col == 0
7983 || n_sel_col > 8
7984 || t == 0
7985 || t > 64
7986 || (in_f >> 5) > 32
7987 || dst.len() < t * out_f
7988 || sel.len() < n_sel
7989 || route_w.len() < n_sel
7990 {
7991 return Err(format!(
7992 "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
7993 dst.len()
7994 )
7995 .into());
7996 }
7997 if !crate::tp::nvfp4_bank_v2_on() {
7998 return Err(
7999 "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8000 );
8001 }
8002 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8003 let cfg = LaunchConfig {
8004 grid_dim: (out_f as u32, t as u32, 1),
8005 block_dim: (32, n_sel_col as u32, 1),
8006 shared_mem_bytes: 0,
8007 };
8008 let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8009 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8010 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8011 let __s_b = self.gpu.stream();
8012 let mut b = __s_b.launch_builder(&f);
8013 b.arg(bank)
8014 .arg(sel)
8015 .arg(aq)
8016 .arg(ad)
8017 .arg(route_w)
8018 .arg(md)
8019 .arg(dst)
8020 .arg(&inf)
8021 .arg(&outf)
8022 .arg(&nsc)
8023 .arg(&rb)
8024 .arg(&es)
8025 .arg(&ars)
8026 .arg(&adrs);
8027 unsafe {
8028 b.launch(cfg)?;
8029 }
8030 Ok(())
8031 }
8032
8033 #[allow(clippy::too_many_arguments)]
8036 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8037 &self,
8038 gate_bank: &CudaSlice<u8>,
8039 up_bank: &CudaSlice<u8>,
8040 sel: &CudaSlice<i32>,
8041 aq: &CudaSlice<i8>,
8042 ad: &CudaSlice<f32>,
8043 yg: &mut CudaSlice<f32>,
8044 yu: &mut CudaSlice<f32>,
8045 n_sel: usize,
8046 in_f: usize,
8047 out_f: usize,
8048 row_bytes: usize,
8049 expert_stride: usize,
8050 owner: usize,
8051 ) -> Result<(), Box<dyn std::error::Error>> {
8052 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8053 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8054 return Err("NVFP4 gu ep geometry".into());
8055 }
8056 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8057 let cfg = LaunchConfig {
8058 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8059 block_dim: (128, 1, 1),
8060 shared_mem_bytes: 0,
8061 };
8062 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8063 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8064 let (ars, adrs) = (0i64, 0i64);
8065 let __s_b = self.gpu.stream();
8066 let mut b = __s_b.launch_builder(&f);
8067 b.arg(gate_bank)
8068 .arg(up_bank)
8069 .arg(sel)
8070 .arg(aq)
8071 .arg(ad)
8072 .arg(yg)
8073 .arg(yu)
8074 .arg(&inf)
8075 .arg(&outf)
8076 .arg(&ns)
8077 .arg(&rb)
8078 .arg(&es)
8079 .arg(&ars)
8080 .arg(&adrs)
8081 .arg(&own);
8082 unsafe {
8083 b.launch(cfg)?;
8084 }
8085 Ok(())
8086 }
8087
8088 #[allow(clippy::too_many_arguments)]
8090 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8091 &self,
8092 gate: &CudaSlice<f32>,
8093 up: &CudaSlice<f32>,
8094 gmac: &CudaSlice<f32>,
8095 umac: &CudaSlice<f32>,
8096 sel: &CudaSlice<i32>,
8097 limit: Option<f32>,
8098 out_q: &mut CudaSlice<i8>,
8099 out_d: &mut CudaSlice<f32>,
8100 n_per: usize,
8101 n_sel: usize,
8102 owner: usize,
8103 ) -> Result<(), Box<dyn std::error::Error>> {
8104 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8105 return Err("NVFP4 silu ep geometry".into());
8106 }
8107 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8108 let warps = n_sel * n_per / 32;
8109 let cfg = LaunchConfig {
8110 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8111 block_dim: (128, 1, 1),
8112 shared_mem_bytes: 0,
8113 };
8114 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8115 let (lim, has) = match limit {
8116 Some(l) => (l, 1i32),
8117 None => (0.0f32, 0i32),
8118 };
8119 let __s_b = self.gpu.stream();
8120 let mut b = __s_b.launch_builder(&f);
8121 b.arg(gate)
8122 .arg(up)
8123 .arg(gmac)
8124 .arg(umac)
8125 .arg(sel)
8126 .arg(&lim)
8127 .arg(&has)
8128 .arg(out_q)
8129 .arg(out_d)
8130 .arg(&np)
8131 .arg(&ns)
8132 .arg(&own);
8133 unsafe {
8134 b.launch(cfg)?;
8135 }
8136 Ok(())
8137 }
8138
8139 #[allow(clippy::too_many_arguments)]
8141 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8142 &self,
8143 bank: &CudaSlice<u8>,
8144 sel: &CudaSlice<i32>,
8145 aq: &CudaSlice<i8>,
8146 ad: &CudaSlice<f32>,
8147 route_w: &CudaSlice<f32>,
8148 md: &CudaSlice<f32>,
8149 dst: &mut CudaSlice<f32>,
8150 n_sel: usize,
8151 in_f: usize,
8152 out_f: usize,
8153 row_bytes: usize,
8154 expert_stride: usize,
8155 act_row_stride: usize,
8156 ad_row_stride: usize,
8157 owner: usize,
8158 ) -> Result<(), Box<dyn std::error::Error>> {
8159 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8160 return Err("NVFP4 down8 ep geometry".into());
8161 }
8162 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8163 let cfg = LaunchConfig {
8164 grid_dim: (out_f as u32, 1, 1),
8165 block_dim: (32, n_sel as u32, 1),
8166 shared_mem_bytes: 0,
8167 };
8168 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8169 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8170 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8171 let __s_b = self.gpu.stream();
8172 let mut b = __s_b.launch_builder(&f);
8173 b.arg(bank)
8174 .arg(sel)
8175 .arg(aq)
8176 .arg(ad)
8177 .arg(route_w)
8178 .arg(md)
8179 .arg(dst)
8180 .arg(&inf)
8181 .arg(&outf)
8182 .arg(&ns)
8183 .arg(&rb)
8184 .arg(&es)
8185 .arg(&ars)
8186 .arg(&adrs)
8187 .arg(&own);
8188 unsafe {
8189 b.launch(cfg)?;
8190 }
8191 Ok(())
8192 }
8193
8194 pub fn qmatvec_nvfp4_sel_into(
8195 &self,
8196 bank: &CudaSlice<u8>,
8197 sel: &CudaSlice<i32>,
8198 aq: &CudaSlice<i8>,
8199 ad: &CudaSlice<f32>,
8200 y: &mut CudaSlice<f32>,
8201 n_sel: usize,
8202 in_f: usize,
8203 out_f: usize,
8204 row_bytes: usize,
8205 expert_stride: usize,
8206 act_row_stride: usize,
8207 ad_row_stride: usize,
8208 ) -> Result<(), Box<dyn std::error::Error>> {
8209 assert!(
8210 in_f % 64 == 0,
8211 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8212 );
8213 if y.len() < n_sel * out_f || sel.len() < n_sel {
8214 return Err(format!(
8215 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8216 y.len(),
8217 sel.len()
8218 )
8219 .into());
8220 }
8221 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8228 let mode = *MR.get_or_init(|| {
8229 if crate::tp::nvfp4_bank_v2_on() {
8230 3
8231 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8232 2
8233 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8234 1
8235 } else {
8236 0
8237 }
8238 });
8239 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8240 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8244 let v2s = mode == 3
8245 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8246 && row_bytes % 16 == 0
8247 && in_f <= 4096;
8248 let f = match (mode, v2s) {
8249 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8250 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8251 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8252 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8253 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8254 };
8255 let nsb = in_f >> 5;
8260 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8261 32
8262 } else if mode == 1 {
8263 512
8264 } else {
8265 128
8266 };
8267 let cfg = LaunchConfig {
8268 grid_dim: (
8269 if v2s {
8270 (out_f as u32).div_ceil(8)
8271 } else {
8272 match mode {
8273 2 => (out_f as u32).div_ceil(16),
8274 1 => (out_f as u32).div_ceil(4),
8275 _ => out_f as u32,
8276 }
8277 },
8278 n_sel as u32,
8279 1,
8280 ),
8281 block_dim: (fit_block, 1, 1),
8282 shared_mem_bytes: 0,
8283 };
8284 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8285 let (rb, es, ars, adrs) = (
8286 row_bytes as i64,
8287 expert_stride as i64,
8288 act_row_stride as i64,
8289 ad_row_stride as i64,
8290 );
8291 let __s_b = self.gpu.stream();
8292 let mut b = __s_b.launch_builder(&f);
8293 b.arg(bank)
8294 .arg(sel)
8295 .arg(aq)
8296 .arg(ad)
8297 .arg(y)
8298 .arg(&inf)
8299 .arg(&outf)
8300 .arg(&ns)
8301 .arg(&rb)
8302 .arg(&es)
8303 .arg(&ars)
8304 .arg(&adrs);
8305 unsafe {
8306 b.launch(cfg)?;
8307 }
8308 Ok(())
8309 }
8310
8311 #[allow(clippy::too_many_arguments)]
8316 pub fn silu_mul_scaled_q8_1_sel_into(
8317 &self,
8318 gate: &CudaSlice<f32>,
8319 up: &CudaSlice<f32>,
8320 gmac: &CudaSlice<f32>,
8321 umac: &CudaSlice<f32>,
8322 sel: &CudaSlice<i32>,
8323 limit: Option<f32>,
8324 out_q: &mut CudaSlice<i8>,
8325 out_d: &mut CudaSlice<f32>,
8326 n_per: usize,
8327 n_sel: usize,
8328 ) -> Result<(), Box<dyn std::error::Error>> {
8329 let n = n_per * n_sel;
8330 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8331 return Err(format!(
8332 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8333 out_q.len(),
8334 out_d.len()
8335 )
8336 .into());
8337 }
8338 if let Some(limit) = limit {
8339 if limit <= 1e-6 {
8340 return Err(format!(
8341 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8342 )
8343 .into());
8344 }
8345 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8346 let cfg = LaunchConfig::for_num_elems(n as u32);
8347 let (np, ns) = (n_per as i32, n_sel as i32);
8348 let __s_b = self.gpu.stream();
8349 let mut b = __s_b.launch_builder(&f);
8350 b.arg(gate)
8351 .arg(up)
8352 .arg(gmac)
8353 .arg(umac)
8354 .arg(sel)
8355 .arg(&limit)
8356 .arg(out_q)
8357 .arg(out_d)
8358 .arg(&np)
8359 .arg(&ns);
8360 unsafe {
8361 b.launch(cfg)?;
8362 }
8363 return Ok(());
8364 }
8365 let f = self.func("silu_mul_scaled_q8_1_sel");
8366 let cfg = LaunchConfig::for_num_elems(n as u32);
8367 let (np, ns) = (n_per as i32, n_sel as i32);
8368 let __s_b = self.gpu.stream();
8369 let mut b = __s_b.launch_builder(&f);
8370 b.arg(gate)
8371 .arg(up)
8372 .arg(gmac)
8373 .arg(umac)
8374 .arg(sel)
8375 .arg(out_q)
8376 .arg(out_d)
8377 .arg(&np)
8378 .arg(&ns);
8379 unsafe {
8380 b.launch(cfg)?;
8381 }
8382 Ok(())
8383 }
8384
8385 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8386 Ok(self.gpu.stream().clone_htod(v)?)
8387 }
8388 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8389 Ok(self.gpu.stream().clone_htod(v)?)
8390 }
8391 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8393 Ok(self.gpu.stream().clone_htod(v)?)
8394 }
8395 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8396 Ok(self.gpu.stream().clone_htod(v)?)
8397 }
8398 pub fn dtoh_view(
8400 &self,
8401 d: &cudarc::driver::CudaView<f32>,
8402 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8403 let v = self.gpu.stream().clone_dtoh(d)?;
8404 self.gpu.stream().synchronize()?;
8405 Ok(v)
8406 }
8407 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8408 let v = self.gpu.stream().clone_dtoh(d)?;
8409 self.gpu.stream().synchronize()?;
8410 Ok(v)
8411 }
8412 pub fn dtoh_pair(
8416 &self,
8417 a: &CudaSlice<f32>,
8418 b: &CudaSlice<f32>,
8419 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8420 let av = self.gpu.stream().clone_dtoh(a)?;
8421 let bv = self.gpu.stream().clone_dtoh(b)?;
8422 self.gpu.stream().synchronize()?;
8423 Ok((av, bv))
8424 }
8425 pub fn dtoh_pair_views(
8428 &self,
8429 a: &cudarc::driver::CudaView<f32>,
8430 b: &cudarc::driver::CudaView<f32>,
8431 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8432 let av = self.gpu.stream().clone_dtoh(a)?;
8433 let bv = self.gpu.stream().clone_dtoh(b)?;
8434 self.gpu.stream().synchronize()?;
8435 Ok((av, bv))
8436 }
8437 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8439 let v = self.gpu.stream().clone_dtoh(d)?;
8440 self.gpu.stream().synchronize()?;
8441 Ok(v)
8442 }
8443 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8445 let v = self.gpu.stream().clone_dtoh(d)?;
8446 self.gpu.stream().synchronize()?;
8447 Ok(v)
8448 }
8449 pub fn dtoh_u8_view(
8450 &self,
8451 d: &cudarc::driver::CudaView<u8>,
8452 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8453 let v = self.gpu.stream().clone_dtoh(d)?;
8454 self.gpu.stream().synchronize()?;
8455 Ok(v)
8456 }
8457 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8458 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8459 self.keep_if_capturing(&s);
8460 Ok(s)
8461 }
8462
8463 pub fn prob_of_token_device(
8472 &self,
8473 logits: &CudaSlice<f32>,
8474 tok: &CudaSlice<u32>,
8475 n_vocab: usize,
8476 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8477 let nb = ARGMAX_NB;
8478 let mut part = self.alloc_uninit::<f32>(nb)?;
8479 let mut p = self.alloc_uninit::<f32>(1)?;
8480 let f1 = self.func("prob_of_token_partial_f32");
8481 let cfg1 = LaunchConfig {
8482 grid_dim: (nb as u32, 1, 1),
8483 block_dim: (256, 1, 1),
8484 shared_mem_bytes: 0,
8485 };
8486 let nv = n_vocab as i32;
8487 let __s_b1 = self.gpu.stream();
8488 let mut b1 = __s_b1.launch_builder(&f1);
8489 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8490 unsafe {
8491 b1.launch(cfg1)?;
8492 }
8493 let f2 = self.func("prob_of_token_final_f32");
8494 let cfg2 = LaunchConfig {
8495 grid_dim: (1, 1, 1),
8496 block_dim: (256, 1, 1),
8497 shared_mem_bytes: 0,
8498 };
8499 let nbi = nb as i32;
8500 let __s_b2 = self.gpu.stream();
8501 let mut b2 = __s_b2.launch_builder(&f2);
8502 b2.arg(&part).arg(&mut p).arg(&nbi);
8503 unsafe {
8504 b2.launch(cfg2)?;
8505 }
8506 Ok(p)
8507 }
8508
8509 pub fn prob_of_token_device_col(
8516 &self,
8517 logits: &CudaSlice<f32>,
8518 tok_all: &CudaSlice<u32>,
8519 tok_idx: usize,
8520 p_out: &mut CudaSlice<f32>,
8521 p_idx: usize,
8522 n_vocab: usize,
8523 ) -> Result<(), Box<dyn std::error::Error>> {
8524 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8525 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8526 let nb = ARGMAX_NB;
8527 let mut part = self.alloc_uninit::<f32>(nb)?;
8528 let f1 = self.func("prob_of_token_partial_f32");
8529 let cfg1 = LaunchConfig {
8530 grid_dim: (nb as u32, 1, 1),
8531 block_dim: (256, 1, 1),
8532 shared_mem_bytes: 0,
8533 };
8534 let nv = n_vocab as i32;
8535 let __s_b1 = self.gpu.stream();
8536 let mut b1 = __s_b1.launch_builder(&f1);
8537 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8538 unsafe {
8539 b1.launch(cfg1)?;
8540 }
8541 let f2 = self.func("prob_of_token_final_f32");
8542 let cfg2 = LaunchConfig {
8543 grid_dim: (1, 1, 1),
8544 block_dim: (256, 1, 1),
8545 shared_mem_bytes: 0,
8546 };
8547 let nbi = nb as i32;
8548 let __s_b2 = self.gpu.stream();
8549 let mut b2 = __s_b2.launch_builder(&f2);
8550 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8551 unsafe {
8552 b2.launch(cfg2)?;
8553 }
8554 Ok(())
8555 }
8556
8557 pub fn prob_of_token_device_into(
8558 &self,
8559 logits: &CudaSlice<f32>,
8560 tok: &CudaSlice<u32>,
8561 p_out: &mut CudaSlice<f32>,
8562 n_vocab: usize,
8563 ) -> Result<(), Box<dyn std::error::Error>> {
8564 let nb = ARGMAX_NB;
8565 let mut part = self.alloc_uninit::<f32>(nb)?;
8566 let f1 = self.func("prob_of_token_partial_f32");
8567 let cfg1 = LaunchConfig {
8568 grid_dim: (nb as u32, 1, 1),
8569 block_dim: (256, 1, 1),
8570 shared_mem_bytes: 0,
8571 };
8572 let nv = n_vocab as i32;
8573 let __s_b1 = self.gpu.stream();
8574 let mut b1 = __s_b1.launch_builder(&f1);
8575 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8576 unsafe {
8577 b1.launch(cfg1)?;
8578 }
8579 let f2 = self.func("prob_of_token_final_f32");
8580 let cfg2 = LaunchConfig {
8581 grid_dim: (1, 1, 1),
8582 block_dim: (256, 1, 1),
8583 shared_mem_bytes: 0,
8584 };
8585 let nbi = nb as i32;
8586 let __s_b2 = self.gpu.stream();
8587 let mut b2 = __s_b2.launch_builder(&f2);
8588 b2.arg(&part).arg(p_out).arg(&nbi);
8589 unsafe {
8590 b2.launch(cfg2)?;
8591 }
8592 Ok(())
8593 }
8594
8595 pub fn u32_hist_append(
8598 &self,
8599 tok: &CudaSlice<u32>,
8600 hist: &mut CudaSlice<u32>,
8601 idx: &mut CudaSlice<i32>,
8602 ) -> Result<(), Box<dyn std::error::Error>> {
8603 let f = self.func("u32_hist_append");
8604 let cfg = LaunchConfig {
8605 grid_dim: (1, 1, 1),
8606 block_dim: (32, 1, 1),
8607 shared_mem_bytes: 0,
8608 };
8609 let __s_b = self.gpu.stream();
8610 let mut b = __s_b.launch_builder(&f);
8611 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8612 unsafe {
8613 b.launch(cfg)?;
8614 }
8615 Ok(())
8616 }
8617
8618 pub fn argmax_token_device(
8619 &self,
8620 logits: &CudaSlice<f32>,
8621 n_vocab: usize,
8622 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8623 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8624 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8625 Ok(tok)
8626 }
8627 pub fn argmax_token_device_into(
8634 &self,
8635 logits: &CudaSlice<f32>,
8636 tok: &mut CudaSlice<u32>,
8637 n_vocab: usize,
8638 ) -> Result<(), Box<dyn std::error::Error>> {
8639 let nb = ARGMAX_NB;
8640 let f1 = self.func("argmax_partial_f32");
8641 let f2 = self.func("argmax_final_f32");
8642 let mut guard = self.argmax_partials.lock().unwrap();
8643 if guard.is_none() {
8644 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8647 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8648 *guard = Some((pv, pi));
8649 }
8650 let (part_v, part_i) = guard.as_mut().unwrap();
8651 let nv = n_vocab as i32;
8652 let nbi = nb as i32;
8653 let cfg1 = LaunchConfig {
8655 grid_dim: (nb as u32, 1, 1),
8656 block_dim: (256, 1, 1),
8657 shared_mem_bytes: 0,
8658 };
8659 let __s_b1 = self.gpu.stream();
8660 let mut b1 = __s_b1.launch_builder(&f1);
8661 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8662 unsafe {
8663 b1.launch(cfg1)?;
8664 }
8665 let cfg2 = LaunchConfig {
8667 grid_dim: (1, 1, 1),
8668 block_dim: (256, 1, 1),
8669 shared_mem_bytes: 0,
8670 };
8671 let __s_b2 = self.gpu.stream();
8672 let mut b2 = __s_b2.launch_builder(&f2);
8673 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8674 unsafe {
8675 b2.launch(cfg2)?;
8676 }
8677 Ok(())
8678 }
8679 pub fn argmax_token_device_col(
8685 &self,
8686 logits: &CudaSlice<f32>,
8687 col: usize,
8688 n_vocab: usize,
8689 toks: &mut CudaSlice<u32>,
8690 out_idx: usize,
8691 ) -> Result<(), Box<dyn std::error::Error>> {
8692 let nb = ARGMAX_NB;
8693 let f1 = self.func("argmax_partial_f32");
8694 let f2 = self.func("argmax_final_f32");
8695 let mut guard = self.argmax_partials.lock().unwrap();
8696 if guard.is_none() {
8697 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8698 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8699 *guard = Some((pv, pi));
8700 }
8701 let (part_v, part_i) = guard.as_mut().unwrap();
8702 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8703 let nv = n_vocab as i32;
8704 let nbi = nb as i32;
8705 let cfg1 = LaunchConfig {
8706 grid_dim: (nb as u32, 1, 1),
8707 block_dim: (256, 1, 1),
8708 shared_mem_bytes: 0,
8709 };
8710 let __s_b1 = self.gpu.stream();
8711 let mut b1 = __s_b1.launch_builder(&f1);
8712 b1.arg(&col_view)
8713 .arg(&mut *part_v)
8714 .arg(&mut *part_i)
8715 .arg(&nv);
8716 unsafe {
8717 b1.launch(cfg1)?;
8718 }
8719 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8720 let cfg2 = LaunchConfig {
8721 grid_dim: (1, 1, 1),
8722 block_dim: (256, 1, 1),
8723 shared_mem_bytes: 0,
8724 };
8725 let __s_b2 = self.gpu.stream();
8726 let mut b2 = __s_b2.launch_builder(&f2);
8727 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8728 unsafe {
8729 b2.launch(cfg2)?;
8730 }
8731 Ok(())
8732 }
8733 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8735 Ok(self.gpu.stream().clone_htod(v)?)
8736 }
8737 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8738 let v = self.gpu.stream().clone_dtoh(d)?;
8739 self.gpu.stream().synchronize()?;
8740 Ok(v)
8741 }
8742
8743 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8744 let v = self.gpu.stream().clone_dtoh(d)?;
8745 self.gpu.stream().synchronize()?;
8746 Ok(v)
8747 }
8748 pub fn htod_u32_into(
8752 &self,
8753 dst: &mut CudaSlice<u32>,
8754 src: &[u32],
8755 ) -> Result<(), Box<dyn std::error::Error>> {
8756 let mut view = dst.slice_mut(0..src.len());
8757 self.gpu.stream().memcpy_htod(src, &mut view)?;
8758 Ok(())
8759 }
8760
8761 pub fn htod_i32_into(
8764 &self,
8765 dst: &mut CudaSlice<i32>,
8766 src: &[i32],
8767 ) -> Result<(), Box<dyn std::error::Error>> {
8768 let mut view = dst.slice_mut(0..src.len());
8769 self.gpu.stream().memcpy_htod(src, &mut view)?;
8770 Ok(())
8771 }
8772
8773 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8774 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8775 self.keep_if_capturing(&s);
8776 Ok(s)
8777 }
8778 pub fn embed_gather_device_into(
8781 &self,
8782 embd: &CudaSlice<u8>,
8783 token_d: &CudaSlice<u32>,
8784 x_out: &mut CudaSlice<f32>,
8785 n_embd: usize,
8786 qtype: i32,
8787 row_bytes: usize,
8788 ) -> Result<(), Box<dyn std::error::Error>> {
8789 let f = self.func("embed_gather_u32");
8790 let cfg = LaunchConfig {
8791 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8792 block_dim: (256, 1, 1),
8793 shared_mem_bytes: 0,
8794 };
8795 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8796 let __s_b = self.gpu.stream();
8797 let mut b = __s_b.launch_builder(&f);
8798 b.arg(embd)
8799 .arg(token_d)
8800 .arg(x_out)
8801 .arg(&ne)
8802 .arg(&qt)
8803 .arg(&rb);
8804 unsafe {
8805 b.launch(cfg)?;
8806 }
8807 Ok(())
8808 }
8809 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8811 let v = self.gpu.stream().clone_dtoh(d)?;
8812 self.gpu.stream().synchronize()?;
8813 Ok(v[0])
8814 }
8815 pub fn i32_set_k(
8822 &self,
8823 dst: &mut CudaSlice<i32>,
8824 v: i32,
8825 ) -> Result<(), Box<dyn std::error::Error>> {
8826 let f = self.func("i32_set_k");
8827 let cfg = LaunchConfig {
8828 grid_dim: (1, 1, 1),
8829 block_dim: (1, 1, 1),
8830 shared_mem_bytes: 0,
8831 };
8832 let idx = 0i32;
8833 let __s_b = self.gpu.stream();
8834 let mut b = __s_b.launch_builder(&f);
8835 b.arg(dst).arg(&v).arg(&idx);
8836 unsafe {
8837 b.launch(cfg)?;
8838 }
8839 Ok(())
8840 }
8841
8842 pub fn set_i32_one(
8843 &self,
8844 d: &mut CudaSlice<i32>,
8845 v: i32,
8846 ) -> Result<(), Box<dyn std::error::Error>> {
8847 self.gpu.stream().memcpy_htod(&[v], d)?;
8848 Ok(())
8849 }
8850 pub fn set_u32_one(
8853 &self,
8854 d: &mut CudaSlice<u32>,
8855 v: u32,
8856 ) -> Result<(), Box<dyn std::error::Error>> {
8857 self.gpu.stream().memcpy_htod(&[v], d)?;
8858 Ok(())
8859 }
8860 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8862 let v = self.gpu.stream().clone_dtoh(d)?;
8863 self.gpu.stream().synchronize()?;
8864 Ok(v[0])
8865 }
8866 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8868 Ok(self.gpu.stream().clone_htod(bytes)?)
8869 }
8870 pub fn embed_gather_device(
8874 &self,
8875 embd: &CudaSlice<u8>,
8876 token_d: &CudaSlice<u32>,
8877 n_embd: usize,
8878 qtype: i32,
8879 row_bytes: usize,
8880 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8881 let f = self.func("embed_gather_u32");
8882 let mut x = self.alloc_uninit::<f32>(n_embd)?;
8883 let cfg = LaunchConfig {
8884 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8885 block_dim: (256, 1, 1),
8886 shared_mem_bytes: 0,
8887 };
8888 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8889 let __s_b = self.gpu.stream();
8890 let mut b = __s_b.launch_builder(&f);
8891 b.arg(embd)
8892 .arg(token_d)
8893 .arg(&mut x)
8894 .arg(&ne)
8895 .arg(&qt)
8896 .arg(&rb);
8897 unsafe {
8898 b.launch(cfg)?;
8899 }
8900 Ok(x)
8901 }
8902
8903 pub fn embed_gather_device_t(
8907 &self,
8908 embd: &CudaSlice<u8>,
8909 tokens: &[u32],
8910 n_embd: usize,
8911 qtype: i32,
8912 row_bytes: usize,
8913 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8914 let t = tokens.len();
8915 let tok_d = self.gpu.stream().clone_htod(tokens)?;
8916 let f = self.func("embed_gather_u32_t");
8917 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8918 let cfg = LaunchConfig {
8919 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8920 block_dim: (256, 1, 1),
8921 shared_mem_bytes: 0,
8922 };
8923 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8924 let __s_b = self.gpu.stream();
8925 let mut b = __s_b.launch_builder(&f);
8926 b.arg(embd)
8927 .arg(&tok_d)
8928 .arg(&mut x)
8929 .arg(&ne)
8930 .arg(&qt)
8931 .arg(&rb)
8932 .arg(&ti);
8933 unsafe {
8934 b.launch(cfg)?;
8935 }
8936 Ok(x)
8937 }
8938
8939 pub fn embed_gather_device_tv(
8944 &self,
8945 embd: &CudaSlice<u8>,
8946 tok_v: &cudarc::driver::CudaView<u32>,
8947 t: usize,
8948 n_embd: usize,
8949 qtype: i32,
8950 row_bytes: usize,
8951 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8952 let f = self.func("embed_gather_u32_t");
8953 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8954 let cfg = LaunchConfig {
8955 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8956 block_dim: (256, 1, 1),
8957 shared_mem_bytes: 0,
8958 };
8959 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8960 let __s_b = self.gpu.stream();
8961 let mut b = __s_b.launch_builder(&f);
8962 b.arg(embd)
8963 .arg(tok_v)
8964 .arg(&mut x)
8965 .arg(&ne)
8966 .arg(&qt)
8967 .arg(&rb)
8968 .arg(&ti);
8969 unsafe {
8970 b.launch(cfg)?;
8971 }
8972 Ok(x)
8973 }
8974
8975 pub fn embed_gather_device_td(
8976 &self,
8977 embd: &CudaSlice<u8>,
8978 tok_d: &CudaSlice<u32>,
8979 t: usize,
8980 n_embd: usize,
8981 qtype: i32,
8982 row_bytes: usize,
8983 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8984 let f = self.func("embed_gather_u32_t");
8985 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8986 let cfg = LaunchConfig {
8987 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8988 block_dim: (256, 1, 1),
8989 shared_mem_bytes: 0,
8990 };
8991 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8992 let __s_b = self.gpu.stream();
8993 let mut b = __s_b.launch_builder(&f);
8994 b.arg(embd)
8995 .arg(tok_d)
8996 .arg(&mut x)
8997 .arg(&ne)
8998 .arg(&qt)
8999 .arg(&rb)
9000 .arg(&ti);
9001 unsafe {
9002 b.launch(cfg)?;
9003 }
9004 Ok(x)
9005 }
9006
9007 #[inline]
9013 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9015 if self
9016 .capture_keep_on
9017 .load(std::sync::atomic::Ordering::Relaxed)
9018 {
9019 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9020 }
9021 }
9022
9023 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9024 &self,
9025 n: usize,
9026 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9027 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9028 {
9032 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9033 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9034 use cudarc::driver::DevicePtrMut;
9036 let n_bytes = s.len() * std::mem::size_of::<T>();
9037 let stream = self.gpu.stream();
9038 let (p_, _g) = s.device_ptr_mut(&stream);
9039 unsafe {
9040 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9041 .result()?;
9042 }
9043 }
9044 }
9045 self.keep_if_capturing(&s);
9046 Ok(s)
9047 }
9048
9049 pub fn uninit_q8_pair(
9054 &self,
9055 n: usize,
9056 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9057 Ok((
9058 self.alloc_uninit::<i8>(n)?,
9059 self.alloc_uninit::<f32>(n / 32)?,
9060 ))
9061 }
9062
9063 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9064 self.alloc_uninit::<f32>(n)
9065 }
9066
9067 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9069 self.alloc_uninit::<i8>(n)
9070 }
9071
9072 #[allow(clippy::too_many_arguments)]
9076 pub fn rms_norm3(
9077 &self,
9078 x: &CudaSlice<f32>,
9079 w0: &CudaSlice<f32>,
9080 w1: &CudaSlice<f32>,
9081 w2: &CudaSlice<f32>,
9082 d0: &mut CudaSlice<f32>,
9083 d1: &mut CudaSlice<f32>,
9084 d2: &mut CudaSlice<f32>,
9085 ncols: usize,
9086 nrows: usize,
9087 eps: f32,
9088 ) -> Result<(), Box<dyn std::error::Error>> {
9089 let f = self.func("rms_norm3_f32");
9090 let cfg = LaunchConfig {
9091 grid_dim: (nrows as u32, 1, 1),
9092 block_dim: (rms_block(), 1, 1),
9093 shared_mem_bytes: 0,
9094 };
9095 let (nc, e) = (ncols as i32, eps);
9096 let __s_b = self.gpu.stream();
9097 let mut b = __s_b.launch_builder(&f);
9098 b.arg(x)
9099 .arg(w0)
9100 .arg(w1)
9101 .arg(w2)
9102 .arg(d0)
9103 .arg(d1)
9104 .arg(d2)
9105 .arg(&nc)
9106 .arg(&e);
9107 unsafe {
9108 b.launch(cfg)?;
9109 }
9110 Ok(())
9111 }
9112
9113 #[allow(clippy::too_many_arguments)]
9115 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9118 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9119 *WARP_ON.get_or_init(|| {
9120 std::env::var("MEMRA_QKVNORM_W")
9121 .map(|v| v != "0")
9122 .unwrap_or(true)
9123 }) && ncols % 4 == 0
9124 && rows >= 64
9125 }
9126
9127 #[allow(clippy::too_many_arguments)]
9130 pub fn rms_norm_qkv_w4b(
9131 &self,
9132 q: &CudaSlice<f32>,
9133 k: &CudaSlice<f32>,
9134 v: &CudaSlice<f32>,
9135 wq: &CudaSlice<f32>,
9136 wk: &CudaSlice<f32>,
9137 wv: &CudaSlice<f32>,
9138 dq: &mut CudaSlice<f32>,
9139 dk: &mut CudaSlice<f32>,
9140 dv: &mut CudaSlice<f32>,
9141 dvb: &mut CudaSlice<u8>,
9142 ncols: usize,
9143 rq: usize,
9144 rk: usize,
9145 eps: f32,
9146 vf16: bool,
9147 ) -> Result<(), Box<dyn std::error::Error>> {
9148 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9149 let f = self.func("rms_norm_qkv_w4b_f32");
9150 let rows = (rq + 2 * rk) as u32;
9151 let cfg = LaunchConfig {
9152 grid_dim: (rows.div_ceil(8), 1, 1),
9153 block_dim: (256, 1, 1),
9154 shared_mem_bytes: 0,
9155 };
9156 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9157 let vf = vf16 as i32;
9158 let __s_b = self.gpu.stream();
9159 let mut b = __s_b.launch_builder(&f);
9160 b.arg(q)
9161 .arg(k)
9162 .arg(v)
9163 .arg(wq)
9164 .arg(wk)
9165 .arg(wv)
9166 .arg(dq)
9167 .arg(dk)
9168 .arg(dv)
9169 .arg(&mut *dvb)
9170 .arg(&nc)
9171 .arg(&rqi)
9172 .arg(&rki)
9173 .arg(&rvi)
9174 .arg(&e)
9175 .arg(&vf);
9176 unsafe {
9177 b.launch(cfg)?;
9178 }
9179 Ok(())
9180 }
9181
9182 pub fn rms_norm_qkv(
9183 &self,
9184 q: &CudaSlice<f32>,
9185 k: &CudaSlice<f32>,
9186 v: &CudaSlice<f32>,
9187 wq: &CudaSlice<f32>,
9188 wk: &CudaSlice<f32>,
9189 wv: &CudaSlice<f32>,
9190 dq: &mut CudaSlice<f32>,
9191 dk: &mut CudaSlice<f32>,
9192 dv: &mut CudaSlice<f32>,
9193 ncols: usize,
9194 rq: usize,
9195 rk: usize,
9196 eps: f32,
9197 ) -> Result<(), Box<dyn std::error::Error>> {
9198 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9202 let warp_on = *WARP_ON.get_or_init(|| {
9203 std::env::var("MEMRA_QKVNORM_W")
9204 .map(|v| v != "0")
9205 .unwrap_or(true)
9206 });
9207 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9210 let f = self.func("rms_norm_qkv_w4_f32");
9211 let rows = (rq + 2 * rk) as u32;
9212 let cfg = LaunchConfig {
9213 grid_dim: (rows.div_ceil(8), 1, 1),
9214 block_dim: (256, 1, 1),
9215 shared_mem_bytes: 0,
9216 };
9217 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9218 let __s_b = self.gpu.stream();
9219 let mut b = __s_b.launch_builder(&f);
9220 b.arg(q)
9221 .arg(k)
9222 .arg(v)
9223 .arg(wq)
9224 .arg(wk)
9225 .arg(wv)
9226 .arg(dq)
9227 .arg(dk)
9228 .arg(dv)
9229 .arg(&nc)
9230 .arg(&rqi)
9231 .arg(&rki)
9232 .arg(&rvi)
9233 .arg(&e);
9234 unsafe {
9235 b.launch(cfg)?;
9236 }
9237 return Ok(());
9238 }
9239 let f = self.func("rms_norm_qkv_f32");
9240 let grid = (rq + 2 * rk) as u32;
9241 let cfg = LaunchConfig {
9242 grid_dim: (grid, 1, 1),
9243 block_dim: (rms_block(), 1, 1),
9244 shared_mem_bytes: 0,
9245 };
9246 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9247 let __s_b = self.gpu.stream();
9248 let mut b = __s_b.launch_builder(&f);
9249 b.arg(q)
9250 .arg(k)
9251 .arg(v)
9252 .arg(wq)
9253 .arg(wk)
9254 .arg(wv)
9255 .arg(dq)
9256 .arg(dk)
9257 .arg(dv)
9258 .arg(&nc)
9259 .arg(&rqi)
9260 .arg(&rki)
9261 .arg(&e);
9262 unsafe {
9263 b.launch(cfg)?;
9264 }
9265 Ok(())
9266 }
9267
9268 #[allow(clippy::too_many_arguments)]
9270 pub fn rms_norm2x(
9271 &self,
9272 a: &CudaSlice<f32>,
9273 bb: &CudaSlice<f32>,
9274 wa: &CudaSlice<f32>,
9275 wb: &CudaSlice<f32>,
9276 da: &mut CudaSlice<f32>,
9277 db: &mut CudaSlice<f32>,
9278 ncols: usize,
9279 nrows: usize,
9280 eps: f32,
9281 ) -> Result<(), Box<dyn std::error::Error>> {
9282 let f = self.func("rms_norm2x_f32");
9283 let cfg = LaunchConfig {
9284 grid_dim: (2 * nrows as u32, 1, 1),
9285 block_dim: (rms_block(), 1, 1),
9286 shared_mem_bytes: 0,
9287 };
9288 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9289 let __s_b = self.gpu.stream();
9290 let mut b = __s_b.launch_builder(&f);
9291 b.arg(a)
9292 .arg(bb)
9293 .arg(wa)
9294 .arg(wb)
9295 .arg(da)
9296 .arg(db)
9297 .arg(&nc)
9298 .arg(&nr)
9299 .arg(&e);
9300 unsafe {
9301 b.launch(cfg)?;
9302 }
9303 Ok(())
9304 }
9305
9306 pub fn softcap(
9308 &self,
9309 y: &mut CudaSlice<f32>,
9310 cap: f32,
9311 n: usize,
9312 ) -> Result<(), Box<dyn std::error::Error>> {
9313 let f = self.func("softcap_f32");
9314 let cfg = LaunchConfig::for_num_elems(n as u32);
9315 let ni = n as i32;
9316 let __s_b = self.gpu.stream();
9317 let mut b = __s_b.launch_builder(&f);
9318 b.arg(y).arg(&cap).arg(&ni);
9319 unsafe {
9320 b.launch(cfg)?;
9321 }
9322 Ok(())
9323 }
9324
9325 pub fn mask_ids_rows(
9328 &self,
9329 y: &mut CudaSlice<f32>,
9330 ids: &CudaSlice<i32>,
9331 n_ids: usize,
9332 n_vocab: usize,
9333 t: usize,
9334 ) -> Result<(), Box<dyn std::error::Error>> {
9335 let f = self.func("mask_ids_rows_f32");
9336 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9337 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9338 let __s_b = self.gpu.stream();
9339 let mut b = __s_b.launch_builder(&f);
9340 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9341 unsafe {
9342 b.launch(cfg)?;
9343 }
9344 Ok(())
9345 }
9346
9347 #[allow(clippy::too_many_arguments)]
9349 pub fn add_scale_rms_norm(
9350 &self,
9351 a: &CudaSlice<f32>,
9352 b_in: &CudaSlice<f32>,
9353 c: f32,
9354 w: &CudaSlice<f32>,
9355 res: &mut CudaSlice<f32>,
9356 dst: &mut CudaSlice<f32>,
9357 ncols: usize,
9358 nrows: usize,
9359 eps: f32,
9360 ) -> Result<(), Box<dyn std::error::Error>> {
9361 let f = self.func("add_scale_rms_norm_f32");
9362 let cfg = LaunchConfig {
9363 grid_dim: (nrows as u32, 1, 1),
9364 block_dim: (rms_block(), 1, 1),
9365 shared_mem_bytes: 0,
9366 };
9367 let (nc, e2) = (ncols as i32, eps);
9368 let __s_b = self.gpu.stream();
9369 let mut b = __s_b.launch_builder(&f);
9370 b.arg(a)
9371 .arg(b_in)
9372 .arg(&c)
9373 .arg(w)
9374 .arg(res)
9375 .arg(dst)
9376 .arg(&nc)
9377 .arg(&e2);
9378 unsafe {
9379 b.launch(cfg)?;
9380 }
9381 Ok(())
9382 }
9383
9384 #[allow(clippy::too_many_arguments)]
9387 pub fn add_scale_rms_norm_q8_1(
9388 &self,
9389 a: &CudaSlice<f32>,
9390 b_in: &CudaSlice<f32>,
9391 c: f32,
9392 w: &CudaSlice<f32>,
9393 res: &mut CudaSlice<f32>,
9394 ncols: usize,
9395 nrows: usize,
9396 eps: f32,
9397 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9398 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9399 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9400 let (nc, e2) = (ncols as i32, eps);
9401 if Self::pdl_on() && Self::pdl_wb_on() {
9402 {
9403 use cudarc::driver::{DevicePtr, DevicePtrMut};
9404 let s = &self.gpu.stream();
9405 let (pa, _g0) = a.device_ptr(s);
9406 let (pb, _g1) = b_in.device_ptr(s);
9407 let (pw, _g2) = w.device_ptr(s);
9408 let (pr, _g3) = res.device_ptr_mut(s);
9409 let (pq, _g4) = out_q.device_ptr_mut(s);
9410 let (pd, _g5) = out_d.device_ptr_mut(s);
9411 let mut ps = [
9412 &pa as *const _ as *mut std::ffi::c_void,
9413 &pb as *const _ as *mut _,
9414 &c as *const _ as *mut _,
9415 &pw as *const _ as *mut _,
9416 &pr as *const _ as *mut _,
9417 &pq as *const _ as *mut _,
9418 &pd as *const _ as *mut _,
9419 &nc as *const _ as *mut _,
9420 &e2 as *const _ as *mut _,
9421 ];
9422 unsafe {
9423 self.launch_pdl(
9424 "add_scale_rms_norm_q8_1",
9425 (nrows as u32, 1, 1),
9426 (rms_block(), 1, 1),
9427 &mut ps,
9428 )?;
9429 }
9430 }
9431 return Ok((out_q, out_d));
9432 }
9433 let f = self.func("add_scale_rms_norm_q8_1");
9434 let cfg = LaunchConfig {
9435 grid_dim: (nrows as u32, 1, 1),
9436 block_dim: (rms_block(), 1, 1),
9437 shared_mem_bytes: 0,
9438 };
9439 let __s_b = self.gpu.stream();
9440 let mut b = __s_b.launch_builder(&f);
9441 b.arg(a)
9442 .arg(b_in)
9443 .arg(&c)
9444 .arg(w)
9445 .arg(res)
9446 .arg(&mut out_q)
9447 .arg(&mut out_d)
9448 .arg(&nc)
9449 .arg(&e2);
9450 unsafe {
9451 b.launch(cfg)?;
9452 }
9453 Ok((out_q, out_d))
9454 }
9455
9456 #[allow(clippy::too_many_arguments)]
9458 pub fn add_scale_rms_norm_q8_1_into(
9459 &self,
9460 a: &CudaSlice<f32>,
9461 b_in: &CudaSlice<f32>,
9462 c: f32,
9463 w: &CudaSlice<f32>,
9464 res: &mut CudaSlice<f32>,
9465 ncols: usize,
9466 nrows: usize,
9467 eps: f32,
9468 out_q: &mut CudaSlice<i8>,
9469 out_d: &mut CudaSlice<f32>,
9470 ) -> Result<(), Box<dyn std::error::Error>> {
9471 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9472 let (nc, e2) = (ncols as i32, eps);
9473 if Self::pdl_on() && Self::pdl_wb_on() {
9474 use cudarc::driver::{DevicePtr, DevicePtrMut};
9475 let s = &self.gpu.stream();
9476 let (pa, _g0) = a.device_ptr(s);
9477 let (pb, _g1) = b_in.device_ptr(s);
9478 let (pw, _g2) = w.device_ptr(s);
9479 let (pr, _g3) = res.device_ptr_mut(s);
9480 let (pq, _g4) = out_q.device_ptr_mut(s);
9481 let (pd, _g5) = out_d.device_ptr_mut(s);
9482 let mut ps = [
9483 &pa as *const _ as *mut std::ffi::c_void,
9484 &pb as *const _ as *mut _,
9485 &c as *const _ as *mut _,
9486 &pw as *const _ as *mut _,
9487 &pr as *const _ as *mut _,
9488 &pq as *const _ as *mut _,
9489 &pd as *const _ as *mut _,
9490 &nc as *const _ as *mut _,
9491 &e2 as *const _ as *mut _,
9492 ];
9493 unsafe {
9494 self.launch_pdl(
9495 "add_scale_rms_norm_q8_1",
9496 (nrows as u32, 1, 1),
9497 (rms_block(), 1, 1),
9498 &mut ps,
9499 )?;
9500 }
9501 return Ok(());
9502 }
9503 let f = self.func("add_scale_rms_norm_q8_1");
9504 let cfg = LaunchConfig {
9505 grid_dim: (nrows as u32, 1, 1),
9506 block_dim: (rms_block(), 1, 1),
9507 shared_mem_bytes: 0,
9508 };
9509 let __s_b = self.gpu.stream();
9510 let mut b = __s_b.launch_builder(&f);
9511 b.arg(a)
9512 .arg(b_in)
9513 .arg(&c)
9514 .arg(w)
9515 .arg(res)
9516 .arg(&mut *out_q)
9517 .arg(&mut *out_d)
9518 .arg(&nc)
9519 .arg(&e2);
9520 unsafe {
9521 b.launch(cfg)?;
9522 }
9523 Ok(())
9524 }
9525
9526 #[allow(clippy::too_many_arguments)]
9529 pub fn rms_pre_add_scale_rms_norm_q8_1(
9530 &self,
9531 a: &CudaSlice<f32>,
9532 wa: &CudaSlice<f32>,
9533 b_in: &CudaSlice<f32>,
9534 c: f32,
9535 w: &CudaSlice<f32>,
9536 res: &mut CudaSlice<f32>,
9537 ncols: usize,
9538 nrows: usize,
9539 eps: f32,
9540 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9541 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9542 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9543 let (nc, e2) = (ncols as i32, eps);
9544 if Self::pdl_on() {
9545 {
9546 use cudarc::driver::{DevicePtr, DevicePtrMut};
9547 let s = &self.gpu.stream();
9548 let (pa, _g0) = a.device_ptr(s);
9549 let (pwa, _g1) = wa.device_ptr(s);
9550 let (pb, _g2) = b_in.device_ptr(s);
9551 let (pw, _g3) = w.device_ptr(s);
9552 let (pr, _g4) = res.device_ptr_mut(s);
9553 let (pq, _g5) = out_q.device_ptr_mut(s);
9554 let (pd, _g6) = out_d.device_ptr_mut(s);
9555 let mut ps = [
9556 &pa as *const _ as *mut std::ffi::c_void,
9557 &pwa as *const _ as *mut _,
9558 &pb as *const _ as *mut _,
9559 &c as *const _ as *mut _,
9560 &pw as *const _ as *mut _,
9561 &pr as *const _ as *mut _,
9562 &pq as *const _ as *mut _,
9563 &pd as *const _ as *mut _,
9564 &nc as *const _ as *mut _,
9565 &e2 as *const _ as *mut _,
9566 ];
9567 unsafe {
9568 self.launch_pdl(
9569 "rms_pre_add_scale_rms_norm_q8_1",
9570 (nrows as u32, 1, 1),
9571 (rms_block(), 1, 1),
9572 &mut ps,
9573 )?;
9574 }
9575 }
9576 return Ok((out_q, out_d));
9577 }
9578 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9579 let cfg = LaunchConfig {
9580 grid_dim: (nrows as u32, 1, 1),
9581 block_dim: (rms_block(), 1, 1),
9582 shared_mem_bytes: 0,
9583 };
9584 let __s_b = self.gpu.stream();
9585 let mut b = __s_b.launch_builder(&f);
9586 b.arg(a)
9587 .arg(wa)
9588 .arg(b_in)
9589 .arg(&c)
9590 .arg(w)
9591 .arg(res)
9592 .arg(&mut out_q)
9593 .arg(&mut out_d)
9594 .arg(&nc)
9595 .arg(&e2);
9596 unsafe {
9597 b.launch(cfg)?;
9598 }
9599 Ok((out_q, out_d))
9600 }
9601
9602 pub fn gelu_tanh_mul_q8_1(
9605 &self,
9606 gate: &CudaSlice<f32>,
9607 up: &cudarc::driver::CudaView<f32>,
9608 act: &mut CudaSlice<f32>,
9609 ncols: usize,
9610 nrows: usize,
9611 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9612 debug_assert!(ncols % 128 == 0);
9613 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9614 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9615 let nc = ncols as i32;
9616 if Self::pdl_on() {
9617 {
9618 use cudarc::driver::{DevicePtr, DevicePtrMut};
9619 let s = &self.gpu.stream();
9620 let (pg, _g0) = gate.device_ptr(s);
9621 let (pu, _g1) = up.device_ptr(s);
9622 let (pact, _g2) = act.device_ptr_mut(s);
9623 let (pq, _g3) = out_q.device_ptr_mut(s);
9624 let (pd, _g4) = out_d.device_ptr_mut(s);
9625 let mut ps = [
9626 &pg as *const _ as *mut std::ffi::c_void,
9627 &pu as *const _ as *mut _,
9628 &pact as *const _ as *mut _,
9629 &pq as *const _ as *mut _,
9630 &pd as *const _ as *mut _,
9631 &nc as *const _ as *mut _,
9632 ];
9633 unsafe {
9634 self.launch_pdl(
9635 "gelu_tanh_mul_q8_1",
9636 (nrows as u32, 1, 1),
9637 (rms_block(), 1, 1),
9638 &mut ps,
9639 )?;
9640 }
9641 }
9642 return Ok((out_q, out_d));
9643 }
9644 let f = self.func("gelu_tanh_mul_q8_1");
9645 let cfg = LaunchConfig {
9646 grid_dim: (nrows as u32, 1, 1),
9647 block_dim: (rms_block(), 1, 1),
9648 shared_mem_bytes: 0,
9649 };
9650 let __s_b = self.gpu.stream();
9651 let mut b = __s_b.launch_builder(&f);
9652 b.arg(gate)
9653 .arg(up)
9654 .arg(act)
9655 .arg(&mut out_q)
9656 .arg(&mut out_d)
9657 .arg(&nc);
9658 unsafe {
9659 b.launch(cfg)?;
9660 }
9661 Ok((out_q, out_d))
9662 }
9663
9664 #[allow(clippy::too_many_arguments)]
9666 pub fn gelu_tanh_mul_q8_1_into(
9667 &self,
9668 gate: &CudaSlice<f32>,
9669 up: &cudarc::driver::CudaView<f32>,
9670 act: &mut CudaSlice<f32>,
9671 ncols: usize,
9672 nrows: usize,
9673 out_q: &mut CudaSlice<i8>,
9674 out_d: &mut CudaSlice<f32>,
9675 ) -> Result<(), Box<dyn std::error::Error>> {
9676 debug_assert!(ncols % 128 == 0);
9677 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9678 let nc = ncols as i32;
9679 if Self::pdl_on() {
9680 use cudarc::driver::{DevicePtr, DevicePtrMut};
9681 let s = &self.gpu.stream();
9682 let (pg, _g0) = gate.device_ptr(s);
9683 let (pu, _g1) = up.device_ptr(s);
9684 let (pact, _g2) = act.device_ptr_mut(s);
9685 let (pq, _g3) = out_q.device_ptr_mut(s);
9686 let (pd, _g4) = out_d.device_ptr_mut(s);
9687 let mut ps = [
9688 &pg as *const _ as *mut std::ffi::c_void,
9689 &pu as *const _ as *mut _,
9690 &pact as *const _ as *mut _,
9691 &pq as *const _ as *mut _,
9692 &pd as *const _ as *mut _,
9693 &nc as *const _ as *mut _,
9694 ];
9695 unsafe {
9696 self.launch_pdl(
9697 "gelu_tanh_mul_q8_1",
9698 (nrows as u32, 1, 1),
9699 (rms_block(), 1, 1),
9700 &mut ps,
9701 )?;
9702 }
9703 return Ok(());
9704 }
9705 let f = self.func("gelu_tanh_mul_q8_1");
9706 let cfg = LaunchConfig {
9707 grid_dim: (nrows as u32, 1, 1),
9708 block_dim: (rms_block(), 1, 1),
9709 shared_mem_bytes: 0,
9710 };
9711 let __s_b = self.gpu.stream();
9712 let mut b = __s_b.launch_builder(&f);
9713 b.arg(gate)
9714 .arg(up)
9715 .arg(&mut *act)
9716 .arg(&mut *out_q)
9717 .arg(&mut *out_d)
9718 .arg(&nc);
9719 unsafe {
9720 b.launch(cfg)?;
9721 }
9722 Ok(())
9723 }
9724
9725 #[allow(clippy::too_many_arguments)]
9727 pub fn add_rms_norm3_q8z(
9728 &self,
9729 a: &CudaSlice<f32>,
9730 b_in: &CudaSlice<f32>,
9731 w0: &CudaSlice<f32>,
9732 w1: &CudaSlice<f32>,
9733 w2: &CudaSlice<f32>,
9734 res: &mut CudaSlice<f32>,
9735 out1: &mut CudaSlice<f32>,
9736 ncols: usize,
9737 nrows: usize,
9738 eps: f32,
9739 ) -> Result<
9740 (
9741 (CudaSlice<i8>, CudaSlice<f32>),
9742 (CudaSlice<i8>, CudaSlice<f32>),
9743 ),
9744 Box<dyn std::error::Error>,
9745 > {
9746 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9747 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9748 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9749 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9750 let f = self.func("add_rms_norm3_q8z_f32");
9751 let cfg = LaunchConfig {
9752 grid_dim: (nrows as u32, 1, 1),
9753 block_dim: (rms_block(), 1, 1),
9754 shared_mem_bytes: 0,
9755 };
9756 let (nc, e2) = (ncols as i32, eps);
9757 let __s_b = self.gpu.stream();
9758 let mut b = __s_b.launch_builder(&f);
9759 b.arg(a)
9760 .arg(b_in)
9761 .arg(w0)
9762 .arg(w1)
9763 .arg(w2)
9764 .arg(res)
9765 .arg(&mut q0)
9766 .arg(&mut d0)
9767 .arg(out1)
9768 .arg(&mut q2)
9769 .arg(&mut d2)
9770 .arg(&nc)
9771 .arg(&e2);
9772 unsafe {
9773 b.launch(cfg)?;
9774 }
9775 Ok(((q0, d0), (q2, d2)))
9776 }
9777
9778 #[allow(clippy::too_many_arguments)]
9780 pub fn add_rms_norm3(
9781 &self,
9782 a: &CudaSlice<f32>,
9783 b_in: &CudaSlice<f32>,
9784 w0: &CudaSlice<f32>,
9785 w1: &CudaSlice<f32>,
9786 w2: &CudaSlice<f32>,
9787 res: &mut CudaSlice<f32>,
9788 d0: &mut CudaSlice<f32>,
9789 d1: &mut CudaSlice<f32>,
9790 d2: &mut CudaSlice<f32>,
9791 ncols: usize,
9792 nrows: usize,
9793 eps: f32,
9794 ) -> Result<(), Box<dyn std::error::Error>> {
9795 let f = self.func("add_rms_norm3_f32");
9796 let cfg = LaunchConfig {
9797 grid_dim: (nrows as u32, 1, 1),
9798 block_dim: (rms_block(), 1, 1),
9799 shared_mem_bytes: 0,
9800 };
9801 let (nc, e2) = (ncols as i32, eps);
9802 let __s_b = self.gpu.stream();
9803 let mut b = __s_b.launch_builder(&f);
9804 b.arg(a)
9805 .arg(b_in)
9806 .arg(w0)
9807 .arg(w1)
9808 .arg(w2)
9809 .arg(res)
9810 .arg(d0)
9811 .arg(d1)
9812 .arg(d2)
9813 .arg(&nc)
9814 .arg(&e2);
9815 unsafe {
9816 b.launch(cfg)?;
9817 }
9818 Ok(())
9819 }
9820
9821 pub fn add_scale(
9823 &self,
9824 a: &CudaSlice<f32>,
9825 b_in: &CudaSlice<f32>,
9826 c: f32,
9827 dst: &mut CudaSlice<f32>,
9828 n: usize,
9829 ) -> Result<(), Box<dyn std::error::Error>> {
9830 let f = self.func("add_scale_f32");
9831 let cfg = LaunchConfig::for_num_elems(n as u32);
9832 let ni = n as i32;
9833 let __s_b = self.gpu.stream();
9834 let mut b = __s_b.launch_builder(&f);
9835 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9836 unsafe {
9837 b.launch(cfg)?;
9838 }
9839 Ok(())
9840 }
9841
9842 pub fn layer_norm_bias(
9844 &self,
9845 x: &CudaSlice<f32>,
9846 w: &CudaSlice<f32>,
9847 b: &CudaSlice<f32>,
9848 dst: &mut CudaSlice<f32>,
9849 ncols: usize,
9850 nrows: usize,
9851 eps: f32,
9852 ) -> Result<(), Box<dyn std::error::Error>> {
9853 let f = self.func("layer_norm_bias_f32");
9854 let (nc, e) = (ncols as i32, eps);
9855 let cfg = LaunchConfig {
9856 grid_dim: (nrows as u32, 1, 1),
9857 block_dim: (256, 1, 1),
9858 shared_mem_bytes: 0,
9859 };
9860 let __s_b = self.gpu.stream();
9861 let mut lb = __s_b.launch_builder(&f);
9862 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9863 unsafe {
9864 lb.launch(cfg)?;
9865 }
9866 Ok(())
9867 }
9868
9869 pub fn gelu_tanh(
9871 &self,
9872 x: &CudaSlice<f32>,
9873 dst: &mut CudaSlice<f32>,
9874 n: usize,
9875 ) -> Result<(), Box<dyn std::error::Error>> {
9876 let f = self.func("gelu_tanh_f32");
9877 let ni = n as i64;
9878 let cfg = LaunchConfig {
9879 grid_dim: (n.div_ceil(256) as u32, 1, 1),
9880 block_dim: (256, 1, 1),
9881 shared_mem_bytes: 0,
9882 };
9883 let __s_b = self.gpu.stream();
9884 let mut lb = __s_b.launch_builder(&f);
9885 lb.arg(x).arg(&mut *dst).arg(&ni);
9886 unsafe {
9887 lb.launch(cfg)?;
9888 }
9889 Ok(())
9890 }
9891
9892 pub fn row_softmax(
9894 &self,
9895 x: &mut CudaSlice<f32>,
9896 ncols: usize,
9897 nrows: usize,
9898 ) -> Result<(), Box<dyn std::error::Error>> {
9899 let f = self.func("row_softmax_f32");
9900 let nc = ncols as i32;
9901 let cfg = LaunchConfig {
9902 grid_dim: (nrows as u32, 1, 1),
9903 block_dim: (256, 1, 1),
9904 shared_mem_bytes: 0,
9905 };
9906 let __s_b = self.gpu.stream();
9907 let mut lb = __s_b.launch_builder(&f);
9908 lb.arg(&mut *x).arg(&nc);
9909 unsafe {
9910 lb.launch(cfg)?;
9911 }
9912 Ok(())
9913 }
9914
9915 pub fn rms_norm(
9916 &self,
9917 x: &CudaSlice<f32>,
9918 w: &CudaSlice<f32>,
9919 dst: &mut CudaSlice<f32>,
9920 ncols: usize,
9921 nrows: usize,
9922 eps: f32,
9923 ) -> Result<(), Box<dyn std::error::Error>> {
9924 let (nc, e) = (ncols as i32, eps);
9925 let kname = if Self::norm_ilp_on() {
9926 "rms_norm_f32_v2"
9927 } else {
9928 "rms_norm_f32"
9929 };
9930 if Self::pdl_on() && Self::pdl_wb_on() {
9931 use cudarc::driver::{DevicePtr, DevicePtrMut};
9932 let s = &self.gpu.stream();
9933 let (px, _g0) = x.device_ptr(s);
9934 let (pw, _g1) = w.device_ptr(s);
9935 let (pd, _g2) = dst.device_ptr_mut(s);
9936 let mut ps = [
9937 &px as *const _ as *mut std::ffi::c_void,
9938 &pw as *const _ as *mut _,
9939 &pd as *const _ as *mut _,
9940 &nc as *const _ as *mut _,
9941 &e as *const _ as *mut _,
9942 ];
9943 unsafe {
9944 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9945 }
9946 return Ok(());
9947 }
9948 let f = self.func(kname);
9949 let cfg = LaunchConfig {
9950 grid_dim: (nrows as u32, 1, 1),
9951 block_dim: (rms_block(), 1, 1),
9952 shared_mem_bytes: 0,
9953 };
9954 let __s_b = self.gpu.stream();
9955 let mut b = __s_b.launch_builder(&f);
9956 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9957 unsafe {
9958 b.launch(cfg)?;
9959 }
9960 Ok(())
9961 }
9962
9963 pub fn rms_norm_decode(
9971 &self,
9972 x: &CudaSlice<f32>,
9973 w: &CudaSlice<f32>,
9974 dst: &mut CudaSlice<f32>,
9975 ncols: usize,
9976 nrows: usize,
9977 eps: f32,
9978 ) -> Result<(), Box<dyn std::error::Error>> {
9979 let f = self.func(if Self::norm_ilp_on() {
9980 "rms_norm_f32_v2"
9981 } else {
9982 "rms_norm_f32"
9983 });
9984 let cfg = LaunchConfig {
9985 grid_dim: (nrows as u32, 1, 1),
9986 block_dim: (1024, 1, 1),
9987 shared_mem_bytes: 0,
9988 };
9989 let (nc, e) = (ncols as i32, eps);
9990 let __s_b = self.gpu.stream();
9991 let mut b = __s_b.launch_builder(&f);
9992 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9993 unsafe {
9994 b.launch(cfg)?;
9995 }
9996 Ok(())
9997 }
9998
9999 pub fn rms_norm_q8_1(
10003 &self,
10004 x: &CudaSlice<f32>,
10005 w: &CudaSlice<f32>,
10006 ncols: usize,
10007 nrows: usize,
10008 eps: f32,
10009 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10010 let nblk = ncols / 32;
10011 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10012 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10013 let (nc, e) = (ncols as i32, eps);
10014 if Self::pdl_on() {
10015 {
10016 use cudarc::driver::{DevicePtr, DevicePtrMut};
10017 let s = &self.gpu.stream();
10018 let (px, _g0) = x.device_ptr(s);
10019 let (pw, _g1) = w.device_ptr(s);
10020 let (pq, _g2) = q.device_ptr_mut(s);
10021 let (pd, _g3) = d.device_ptr_mut(s);
10022 let mut ps = [
10023 &px as *const _ as *mut std::ffi::c_void,
10024 &pw as *const _ as *mut _,
10025 &pq as *const _ as *mut _,
10026 &pd as *const _ as *mut _,
10027 &nc as *const _ as *mut _,
10028 &e as *const _ as *mut _,
10029 ];
10030 unsafe {
10031 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10032 }
10033 }
10034 return Ok((q, d));
10035 }
10036 let f = self.func("rms_norm_q8_1");
10037 let cfg = LaunchConfig {
10040 grid_dim: (nrows as u32, 1, 1),
10041 block_dim: (1024, 1, 1),
10042 shared_mem_bytes: 0,
10043 };
10044 let __s_b = self.gpu.stream();
10045 let mut b = __s_b.launch_builder(&f);
10046 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10047 unsafe {
10048 b.launch(cfg)?;
10049 }
10050 Ok((q, d))
10051 }
10052
10053 pub fn rms_norm_q8_1_into(
10056 &self,
10057 x: &CudaSlice<f32>,
10058 w: &CudaSlice<f32>,
10059 ncols: usize,
10060 nrows: usize,
10061 eps: f32,
10062 q: &mut CudaSlice<i8>,
10063 d: &mut CudaSlice<f32>,
10064 ) -> Result<(), Box<dyn std::error::Error>> {
10065 let nblk = ncols / 32;
10066 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10067 let (nc, e) = (ncols as i32, eps);
10068 if Self::pdl_on() {
10069 use cudarc::driver::{DevicePtr, DevicePtrMut};
10070 let s = &self.gpu.stream();
10071 let (px, _g0) = x.device_ptr(s);
10072 let (pw, _g1) = w.device_ptr(s);
10073 let (pq, _g2) = q.device_ptr_mut(s);
10074 let (pd, _g3) = d.device_ptr_mut(s);
10075 let mut ps = [
10076 &px as *const _ as *mut std::ffi::c_void,
10077 &pw as *const _ as *mut _,
10078 &pq as *const _ as *mut _,
10079 &pd as *const _ as *mut _,
10080 &nc as *const _ as *mut _,
10081 &e as *const _ as *mut _,
10082 ];
10083 unsafe {
10084 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10085 }
10086 return Ok(());
10087 }
10088 let f = self.func("rms_norm_q8_1");
10089 let cfg = LaunchConfig {
10090 grid_dim: (nrows as u32, 1, 1),
10091 block_dim: (1024, 1, 1),
10092 shared_mem_bytes: 0,
10093 };
10094 let __s_b = self.gpu.stream();
10095 let mut b = __s_b.launch_builder(&f);
10096 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10097 unsafe {
10098 b.launch(cfg)?;
10099 }
10100 Ok(())
10101 }
10102
10103 pub fn quantize_q8_1_into(
10105 &self,
10106 x: &CudaSlice<f32>,
10107 m: usize,
10108 in_f: usize,
10109 q: &mut CudaSlice<i8>,
10110 d: &mut CudaSlice<f32>,
10111 ) -> Result<(), Box<dyn std::error::Error>> {
10112 let nblk = in_f / 32;
10113 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10114 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10115 let (inf, mi) = (in_f as i32, m as i32);
10116 if Self::pdl_on() && Self::pdl_wb_on() {
10117 use cudarc::driver::{DevicePtr, DevicePtrMut};
10118 let s = &self.gpu.stream();
10119 let (px, _g0) = x.device_ptr(s);
10120 let (pq, _g1) = q.device_ptr_mut(s);
10121 let (pd, _g2) = d.device_ptr_mut(s);
10122 let mut ps = [
10123 &px as *const _ as *mut std::ffi::c_void,
10124 &pq as *const _ as *mut _,
10125 &pd as *const _ as *mut _,
10126 &inf as *const _ as *mut _,
10127 &mi as *const _ as *mut _,
10128 ];
10129 unsafe {
10130 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10131 }
10132 return Ok(());
10133 }
10134 let f = self.func("quantize_q8_1");
10135 let __s_b = self.gpu.stream();
10136 let mut b = __s_b.launch_builder(&f);
10137 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10138 unsafe {
10139 b.launch(cfg)?;
10140 }
10141 Ok(())
10142 }
10143
10144 pub fn add_rms_norm_q8_1(
10148 &self,
10149 a: &CudaSlice<f32>,
10150 b_in: &CudaSlice<f32>,
10151 w: &CudaSlice<f32>,
10152 res: &mut CudaSlice<f32>,
10153 ncols: usize,
10154 nrows: usize,
10155 eps: f32,
10156 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10157 let nblk = ncols / 32;
10158 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10159 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10160 let f = self.func("add_rms_norm_q8_1");
10161 let cfg = LaunchConfig {
10163 grid_dim: (nrows as u32, 1, 1),
10164 block_dim: (1024, 1, 1),
10165 shared_mem_bytes: 0,
10166 };
10167 let (nc, e) = (ncols as i32, eps);
10168 let __s_bld = self.gpu.stream();
10169 let mut bld = __s_bld.launch_builder(&f);
10170 bld.arg(a)
10171 .arg(b_in)
10172 .arg(w)
10173 .arg(res)
10174 .arg(&mut q)
10175 .arg(&mut d)
10176 .arg(&nc)
10177 .arg(&e);
10178 unsafe {
10179 bld.launch(cfg)?;
10180 }
10181 Ok((q, d))
10182 }
10183
10184 #[allow(clippy::too_many_arguments)]
10190 pub fn join_add_rms_norm_raw(
10191 &self,
10192 a0_raw: u64,
10193 a1_raw: u64,
10194 x: &CudaSlice<f32>,
10195 w: &CudaSlice<f32>,
10196 res: &mut CudaSlice<f32>,
10197 dst: &mut CudaSlice<f32>,
10198 ncols: usize,
10199 eps: f32,
10200 ) -> Result<(), Box<dyn std::error::Error>> {
10201 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10202 return Err("join_add_rms_norm geometry".into());
10203 }
10204 let f = self.func("join_add_rms_norm_f32");
10205 let cfg = LaunchConfig {
10206 grid_dim: (1, 1, 1),
10207 block_dim: (rms_block(), 1, 1),
10208 shared_mem_bytes: 0,
10209 };
10210 let (nc, e) = (ncols as i32, eps);
10211 let __s_b = self.gpu.stream();
10212 let mut b = __s_b.launch_builder(&f);
10213 b.arg(&a0_raw)
10214 .arg(&a1_raw)
10215 .arg(x)
10216 .arg(w)
10217 .arg(&mut *res)
10218 .arg(&mut *dst)
10219 .arg(&nc)
10220 .arg(&e);
10221 unsafe {
10222 b.launch(cfg)?;
10223 }
10224 Ok(())
10225 }
10226
10227 pub fn add_rms_norm(
10228 &self,
10229 a: &CudaSlice<f32>,
10230 b: &CudaSlice<f32>,
10231 w: &CudaSlice<f32>,
10232 res: &mut CudaSlice<f32>,
10233 dst: &mut CudaSlice<f32>,
10234 ncols: usize,
10235 nrows: usize,
10236 eps: f32,
10237 ) -> Result<(), Box<dyn std::error::Error>> {
10238 let (nc, e) = (ncols as i32, eps);
10239 let kname = if Self::norm_ilp_on() {
10240 "add_rms_norm_f32_v2"
10241 } else {
10242 "add_rms_norm_f32"
10243 };
10244 if Self::pdl_on() && Self::pdl_wb_on() {
10245 use cudarc::driver::{DevicePtr, DevicePtrMut};
10246 let s = &self.gpu.stream();
10247 let (pa, _g0) = a.device_ptr(s);
10248 let (pb, _g1) = b.device_ptr(s);
10249 let (pw, _g2) = w.device_ptr(s);
10250 let (pr, _g3) = res.device_ptr_mut(s);
10251 let (pd, _g4) = dst.device_ptr_mut(s);
10252 let mut ps = [
10253 &pa as *const _ as *mut std::ffi::c_void,
10254 &pb as *const _ as *mut _,
10255 &pw as *const _ as *mut _,
10256 &pr as *const _ as *mut _,
10257 &pd as *const _ as *mut _,
10258 &nc as *const _ as *mut _,
10259 &e as *const _ as *mut _,
10260 ];
10261 unsafe {
10262 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10263 }
10264 return Ok(());
10265 }
10266 let f = self.func(kname);
10267 let cfg = LaunchConfig {
10268 grid_dim: (nrows as u32, 1, 1),
10269 block_dim: (rms_block(), 1, 1),
10270 shared_mem_bytes: 0,
10271 };
10272 let __s_b2 = self.gpu.stream();
10273 let mut b2 = __s_b2.launch_builder(&f);
10274 b2.arg(a)
10275 .arg(b)
10276 .arg(w)
10277 .arg(&mut *res)
10278 .arg(&mut *dst)
10279 .arg(&nc)
10280 .arg(&e);
10281 unsafe {
10282 b2.launch(cfg)?;
10283 }
10284 Ok(())
10285 }
10286
10287 #[allow(clippy::too_many_arguments)]
10290 pub fn rms_pre_add_rms_norm(
10291 &self,
10292 a: &CudaSlice<f32>,
10293 wa: &CudaSlice<f32>,
10294 b: &CudaSlice<f32>,
10295 w: &CudaSlice<f32>,
10296 res: &mut CudaSlice<f32>,
10297 dst: &mut CudaSlice<f32>,
10298 ncols: usize,
10299 nrows: usize,
10300 eps: f32,
10301 ) -> Result<(), Box<dyn std::error::Error>> {
10302 let f = self.func("rms_pre_add_rms_norm_f32");
10303 let cfg = LaunchConfig {
10304 grid_dim: (nrows as u32, 1, 1),
10305 block_dim: (rms_block(), 1, 1),
10306 shared_mem_bytes: 0,
10307 };
10308 let (nc, e) = (ncols as i32, eps);
10309 let __s_b2 = self.gpu.stream();
10310 let mut b2 = __s_b2.launch_builder(&f);
10311 b2.arg(a)
10312 .arg(wa)
10313 .arg(b)
10314 .arg(w)
10315 .arg(&mut *res)
10316 .arg(&mut *dst)
10317 .arg(&nc)
10318 .arg(&e);
10319 unsafe {
10320 b2.launch(cfg)?;
10321 }
10322 Ok(())
10323 }
10324
10325 #[allow(clippy::too_many_arguments)]
10327 pub fn rms_pre_add_rms_norm_q8z(
10328 &self,
10329 a: &CudaSlice<f32>,
10330 wa: &CudaSlice<f32>,
10331 b: &CudaSlice<f32>,
10332 w: &CudaSlice<f32>,
10333 res: &mut CudaSlice<f32>,
10334 dst: &mut CudaSlice<f32>,
10335 ncols: usize,
10336 nrows: usize,
10337 eps: f32,
10338 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10339 debug_assert!(ncols % 128 == 0);
10340 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10341 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10342 let (nc, e) = (ncols as i32, eps);
10343 if Self::pdl_on() {
10344 {
10345 use cudarc::driver::{DevicePtr, DevicePtrMut};
10346 let s = &self.gpu.stream();
10347 let (pa, _g0) = a.device_ptr(s);
10348 let (pwa, _g1) = wa.device_ptr(s);
10349 let (pb, _g2) = b.device_ptr(s);
10350 let (pw, _g3) = w.device_ptr(s);
10351 let (pr, _g4) = res.device_ptr_mut(s);
10352 let (pdst, _g5) = dst.device_ptr_mut(s);
10353 let (pq, _g6) = out_q.device_ptr_mut(s);
10354 let (pd, _g7) = out_d.device_ptr_mut(s);
10355 let mut ps = [
10356 &pa as *const _ as *mut std::ffi::c_void,
10357 &pwa as *const _ as *mut _,
10358 &pb as *const _ as *mut _,
10359 &pw as *const _ as *mut _,
10360 &pr as *const _ as *mut _,
10361 &pdst as *const _ as *mut _,
10362 &pq as *const _ as *mut _,
10363 &pd as *const _ as *mut _,
10364 &nc as *const _ as *mut _,
10365 &e as *const _ as *mut _,
10366 ];
10367 unsafe {
10368 self.launch_pdl(
10369 "rms_pre_add_rms_norm_q8z_f32",
10370 (nrows as u32, 1, 1),
10371 (rms_block(), 1, 1),
10372 &mut ps,
10373 )?;
10374 }
10375 }
10376 return Ok((out_q, out_d));
10377 }
10378 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10379 let cfg = LaunchConfig {
10380 grid_dim: (nrows as u32, 1, 1),
10381 block_dim: (rms_block(), 1, 1),
10382 shared_mem_bytes: 0,
10383 };
10384 let __s_b2 = self.gpu.stream();
10385 let mut b2 = __s_b2.launch_builder(&f);
10386 b2.arg(a)
10387 .arg(wa)
10388 .arg(b)
10389 .arg(w)
10390 .arg(&mut *res)
10391 .arg(&mut *dst)
10392 .arg(&mut out_q)
10393 .arg(&mut out_d)
10394 .arg(&nc)
10395 .arg(&e);
10396 unsafe {
10397 b2.launch(cfg)?;
10398 }
10399 Ok((out_q, out_d))
10400 }
10401
10402 #[allow(clippy::too_many_arguments)]
10406 pub fn rms_pre_add_rms_norm_q8z_into(
10407 &self,
10408 a: &CudaSlice<f32>,
10409 wa: &CudaSlice<f32>,
10410 b: &CudaSlice<f32>,
10411 w: &CudaSlice<f32>,
10412 res: &mut CudaSlice<f32>,
10413 dst: &mut CudaSlice<f32>,
10414 ncols: usize,
10415 nrows: usize,
10416 eps: f32,
10417 out_q: &mut CudaSlice<i8>,
10418 out_d: &mut CudaSlice<f32>,
10419 ) -> Result<(), Box<dyn std::error::Error>> {
10420 debug_assert!(ncols % 128 == 0);
10421 let (nc, e) = (ncols as i32, eps);
10422 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10423 let cfg = LaunchConfig {
10424 grid_dim: (nrows as u32, 1, 1),
10425 block_dim: (rms_block(), 1, 1),
10426 shared_mem_bytes: 0,
10427 };
10428 let __s_b = self.gpu.stream();
10429 let mut b2 = __s_b.launch_builder(&f);
10430 b2.arg(a)
10431 .arg(wa)
10432 .arg(b)
10433 .arg(w)
10434 .arg(&mut *res)
10435 .arg(&mut *dst)
10436 .arg(&mut *out_q)
10437 .arg(&mut *out_d)
10438 .arg(&nc)
10439 .arg(&e);
10440 unsafe {
10441 b2.launch(cfg)?;
10442 }
10443 Ok(())
10444 }
10445
10446 #[allow(clippy::too_many_arguments)]
10449 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10450 &self,
10451 a: &CudaSlice<f32>,
10452 wa: &CudaSlice<f32>,
10453 b_in: &CudaSlice<f32>,
10454 c: f32,
10455 w: &CudaSlice<f32>,
10456 res: &mut CudaSlice<f32>,
10457 ncols: usize,
10458 nrows: usize,
10459 eps: f32,
10460 out_q: &mut CudaSlice<i8>,
10461 out_d: &mut CudaSlice<f32>,
10462 ) -> Result<(), Box<dyn std::error::Error>> {
10463 debug_assert!(ncols % 128 == 0);
10464 let (nc, e2) = (ncols as i32, eps);
10465 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10466 let cfg = LaunchConfig {
10467 grid_dim: (nrows as u32, 1, 1),
10468 block_dim: (rms_block(), 1, 1),
10469 shared_mem_bytes: 0,
10470 };
10471 let __s_b = self.gpu.stream();
10472 let mut b2 = __s_b.launch_builder(&f);
10473 b2.arg(a)
10474 .arg(wa)
10475 .arg(b_in)
10476 .arg(&c)
10477 .arg(w)
10478 .arg(&mut *res)
10479 .arg(&mut *out_q)
10480 .arg(&mut *out_d)
10481 .arg(&nc)
10482 .arg(&e2);
10483 unsafe {
10484 b2.launch(cfg)?;
10485 }
10486 Ok(())
10487 }
10488
10489 pub fn g4_pnfold_on() -> bool {
10497 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10498 *ON.get_or_init(|| {
10499 std::env::var("MEMRA_G4_PNFOLD")
10500 .map(|v| v != "0")
10501 .unwrap_or(true)
10502 })
10503 }
10504
10505 pub fn build_q4_out_concat3(
10509 &self,
10510 w0: &crate::model::GpuTensor,
10511 w1: &crate::model::GpuTensor,
10512 w2: &crate::model::GpuTensor,
10513 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10514 use crate::model::GpuTensor;
10515 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10516 match w {
10517 GpuTensor::Quant {
10518 qtype,
10519 row_bytes,
10520 rp,
10521 ..
10522 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10523 _ => None,
10524 }
10525 };
10526 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10527 else {
10528 return Ok(None);
10529 };
10530 if rb0 != rb1
10531 || rb0 != rb2
10532 || w0.in_features() != w1.in_features()
10533 || w0.in_features() != w2.in_features()
10534 {
10535 return Ok(None);
10536 }
10537 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10538 match w {
10539 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10540 _ => unreachable!(),
10541 }
10542 }
10543 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10544 let total = rb0 * (o0 + o1 + o2);
10545 let mut cat = self.alloc_u8(total)?;
10546 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10547 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10548 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10549 Ok(Some(GpuTensor::Quant {
10550 bytes: cat,
10551 qtype: QT_Q4_0,
10552 row_bytes: rb0,
10553 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10554 scale: 1.0,
10555 rp: false,
10556 #[cfg(memra_cutlass)]
10557 cutlass: None,
10558 fp8: None,
10559 blk: None,
10560 rp4: None,
10561 f16: None,
10562 }))
10563 }
10564
10565 fn full_width_rope_only(
10583 kernel: &str,
10584 n_rot: usize,
10585 head_dim: usize,
10586 ) -> Result<(), Box<dyn std::error::Error>> {
10587 if n_rot == head_dim {
10588 return Ok(());
10589 }
10590 Err(format!(
10591 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10592 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10593 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10594 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10595 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10596 )
10597 .into())
10598 }
10599
10600 #[allow(clippy::too_many_arguments)]
10604 pub fn rms_norm_qkv_rope_cat(
10605 &self,
10606 qkv: &CudaSlice<f32>,
10607 wq: &CudaSlice<f32>,
10608 wk: &CudaSlice<f32>,
10609 wv: &CudaSlice<f32>,
10610 q: &mut CudaSlice<f32>,
10611 k: &mut CudaSlice<f32>,
10612 v: &mut CudaSlice<f32>,
10613 head_dim: usize,
10614 n_rot: usize,
10615 rq: usize,
10616 rk: usize,
10617 pos: &CudaSlice<i32>,
10618 nh_q: usize,
10619 nh_k: usize,
10620 base: f32,
10621 freq_scale: f32,
10622 ff: Option<&CudaSlice<f32>>,
10623 eps: f32,
10624 ) -> Result<(), Box<dyn std::error::Error>> {
10625 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10626 let rows = rq + rk + rk;
10627 let theta_scale = base.powf(-2.0 / head_dim as f32);
10628 let (nc, rqi, rki, nhq, nhk) = (
10629 head_dim as i32,
10630 rq as i32,
10631 rk as i32,
10632 nh_q as i32,
10633 nh_k as i32,
10634 );
10635 if Self::pdl_on() {
10636 use cudarc::driver::{DevicePtr, DevicePtrMut};
10637 let s = &self.gpu.stream();
10638 let (pqkv, _g0) = qkv.device_ptr(s);
10639 let (pwq, _g1) = wq.device_ptr(s);
10640 let (pwk, _g2) = wk.device_ptr(s);
10641 let (pwv, _g3) = wv.device_ptr(s);
10642 let (pq, _g4) = q.device_ptr_mut(s);
10643 let (pk, _g5) = k.device_ptr_mut(s);
10644 let (pv, _g6) = v.device_ptr_mut(s);
10645 let (ppos, _g7) = pos.device_ptr(s);
10646 let (pff, _g8) = match ff {
10647 Some(t) => {
10648 let (p, g) = t.device_ptr(s);
10649 (p, Some(g))
10650 }
10651 None => (0, None),
10652 };
10653 let mut ps = [
10654 &pqkv as *const _ as *mut std::ffi::c_void,
10655 &pwq as *const _ as *mut _,
10656 &pwk as *const _ as *mut _,
10657 &pwv as *const _ as *mut _,
10658 &pq as *const _ as *mut _,
10659 &pk as *const _ as *mut _,
10660 &pv as *const _ as *mut _,
10661 &nc as *const _ as *mut _,
10662 &rqi as *const _ as *mut _,
10663 &rki as *const _ as *mut _,
10664 &ppos as *const _ as *mut _,
10665 &nhq as *const _ as *mut _,
10666 &nhk as *const _ as *mut _,
10667 &theta_scale as *const _ as *mut _,
10668 &freq_scale as *const _ as *mut _,
10669 &pff as *const _ as *mut _,
10670 &eps as *const _ as *mut _,
10671 ];
10672 unsafe {
10673 self.launch_pdl(
10674 "rms_norm_qkv_rope_cat_f32",
10675 (rows as u32, 1, 1),
10676 (rms_block(), 1, 1),
10677 &mut ps,
10678 )?;
10679 }
10680 return Ok(());
10681 }
10682 let f = self.func("rms_norm_qkv_rope_cat_f32");
10683 let cfg = LaunchConfig {
10684 grid_dim: (rows as u32, 1, 1),
10685 block_dim: (rms_block(), 1, 1),
10686 shared_mem_bytes: 0,
10687 };
10688 let __s_b = self.gpu.stream();
10689 let mut b = __s_b.launch_builder(&f);
10690 match ff {
10691 Some(t) => {
10692 b.arg(qkv)
10693 .arg(wq)
10694 .arg(wk)
10695 .arg(wv)
10696 .arg(&mut *q)
10697 .arg(&mut *k)
10698 .arg(&mut *v)
10699 .arg(&nc)
10700 .arg(&rqi)
10701 .arg(&rki)
10702 .arg(pos)
10703 .arg(&nhq)
10704 .arg(&nhk)
10705 .arg(&theta_scale)
10706 .arg(&freq_scale)
10707 .arg(t)
10708 .arg(&eps);
10709 unsafe {
10710 b.launch(cfg)?;
10711 }
10712 }
10713 None => {
10714 let null: u64 = 0;
10715 b.arg(qkv)
10716 .arg(wq)
10717 .arg(wk)
10718 .arg(wv)
10719 .arg(&mut *q)
10720 .arg(&mut *k)
10721 .arg(&mut *v)
10722 .arg(&nc)
10723 .arg(&rqi)
10724 .arg(&rki)
10725 .arg(pos)
10726 .arg(&nhq)
10727 .arg(&nhk)
10728 .arg(&theta_scale)
10729 .arg(&freq_scale)
10730 .arg(&null)
10731 .arg(&eps);
10732 unsafe {
10733 b.launch(cfg)?;
10734 }
10735 }
10736 }
10737 Ok(())
10738 }
10739
10740 #[allow(clippy::too_many_arguments)]
10744 pub fn rms_norm_qkv_rope(
10745 &self,
10746 q0: &CudaSlice<f32>,
10747 k0: &CudaSlice<f32>,
10748 v0: &CudaSlice<f32>,
10749 wq: &CudaSlice<f32>,
10750 wk: &CudaSlice<f32>,
10751 wv: &CudaSlice<f32>,
10752 q: &mut CudaSlice<f32>,
10753 k: &mut CudaSlice<f32>,
10754 v: &mut CudaSlice<f32>,
10755 head_dim: usize,
10756 n_rot: usize,
10757 rq: usize,
10758 rk: usize,
10759 pos: &CudaSlice<i32>,
10760 nh_q: usize,
10761 nh_k: usize,
10762 base: f32,
10763 freq_scale: f32,
10764 ff: Option<&CudaSlice<f32>>,
10765 eps: f32,
10766 ) -> Result<(), Box<dyn std::error::Error>> {
10767 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10768 let f = self.func("rms_norm_qkv_rope_f32");
10769 let rows = rq + rk + rk; let cfg = LaunchConfig {
10771 grid_dim: (rows as u32, 1, 1),
10772 block_dim: (rms_block(), 1, 1),
10773 shared_mem_bytes: 0,
10774 };
10775 let theta_scale = base.powf(-2.0 / head_dim as f32);
10776 let (nc, rqi, rki, nhq, nhk) = (
10777 head_dim as i32,
10778 rq as i32,
10779 rk as i32,
10780 nh_q as i32,
10781 nh_k as i32,
10782 );
10783 let __s_b = self.gpu.stream();
10784 let mut b = __s_b.launch_builder(&f);
10785 match ff {
10786 Some(t) => {
10787 b.arg(q0)
10788 .arg(k0)
10789 .arg(v0)
10790 .arg(wq)
10791 .arg(wk)
10792 .arg(wv)
10793 .arg(&mut *q)
10794 .arg(&mut *k)
10795 .arg(&mut *v)
10796 .arg(&nc)
10797 .arg(&rqi)
10798 .arg(&rki)
10799 .arg(pos)
10800 .arg(&nhq)
10801 .arg(&nhk)
10802 .arg(&theta_scale)
10803 .arg(&freq_scale)
10804 .arg(t)
10805 .arg(&eps);
10806 unsafe {
10807 b.launch(cfg)?;
10808 }
10809 }
10810 None => {
10811 let null: u64 = 0;
10812 b.arg(q0)
10813 .arg(k0)
10814 .arg(v0)
10815 .arg(wq)
10816 .arg(wk)
10817 .arg(wv)
10818 .arg(&mut *q)
10819 .arg(&mut *k)
10820 .arg(&mut *v)
10821 .arg(&nc)
10822 .arg(&rqi)
10823 .arg(&rki)
10824 .arg(pos)
10825 .arg(&nhq)
10826 .arg(&nhk)
10827 .arg(&theta_scale)
10828 .arg(&freq_scale)
10829 .arg(&null)
10830 .arg(&eps);
10831 unsafe {
10832 b.launch(cfg)?;
10833 }
10834 }
10835 }
10836 Ok(())
10837 }
10838
10839 #[allow(clippy::too_many_arguments)]
10845 pub fn rms_norm_qkv_rope_append_dc(
10846 &self,
10847 q0: &CudaSlice<f32>,
10848 k0: &CudaSlice<f32>,
10849 v0: &CudaSlice<f32>,
10850 wq: &CudaSlice<f32>,
10851 wk: &CudaSlice<f32>,
10852 wv: &CudaSlice<f32>,
10853 q: &mut CudaSlice<f32>,
10854 k: &mut CudaSlice<f32>,
10855 v: &mut CudaSlice<f32>,
10856 head_dim: usize,
10857 n_rot: usize,
10858 rq: usize,
10859 rk: usize,
10860 pos: &CudaSlice<i32>,
10861 nh_q: usize,
10862 nh_k: usize,
10863 base: f32,
10864 freq_scale: f32,
10865 ff: Option<&CudaSlice<f32>>,
10866 eps: f32,
10867 kc: &mut CudaSlice<u8>,
10868 vc: &mut CudaSlice<u8>,
10869 t_dev: &CudaSlice<i32>,
10870 k_tok_bytes: usize,
10871 v_tok_bytes: usize,
10872 g: bool,
10873 ) -> Result<(), Box<dyn std::error::Error>> {
10874 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10875 let rows = rq + rk + rk;
10876 let theta_scale = base.powf(-2.0 / head_dim as f32);
10877 let (nc, rqi, rki, nhq, nhk) = (
10878 head_dim as i32,
10879 rq as i32,
10880 rk as i32,
10881 nh_q as i32,
10882 nh_k as i32,
10883 );
10884 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10885 if Self::pdl_on() && Self::pdl_wb_on() {
10886 use cudarc::driver::{DevicePtr, DevicePtrMut};
10887 let s = &self.gpu.stream();
10888 let (p0, _a0) = q0.device_ptr(s);
10889 let (p1, _a1) = k0.device_ptr(s);
10890 let (p2, _a2) = v0.device_ptr(s);
10891 let (pwq, _a3) = wq.device_ptr(s);
10892 let (pwk, _a4) = wk.device_ptr(s);
10893 let (pwv, _a5) = wv.device_ptr(s);
10894 let (pq, _a6) = q.device_ptr_mut(s);
10895 let (pk, _a7) = k.device_ptr_mut(s);
10896 let (pv, _a8) = v.device_ptr_mut(s);
10897 let (pp, _a9) = pos.device_ptr(s);
10898 let pff: u64 = match ff {
10899 Some(t) => {
10900 let (p, _gg) = t.device_ptr(s);
10901 p as u64
10902 }
10903 None => 0,
10904 };
10905 let (pkc, _a10) = kc.device_ptr_mut(s);
10906 let (pvc, _a11) = vc.device_ptr_mut(s);
10907 let (pt, _a12) = t_dev.device_ptr(s);
10908 let mut ps = [
10909 &p0 as *const _ as *mut std::ffi::c_void,
10910 &p1 as *const _ as *mut _,
10911 &p2 as *const _ as *mut _,
10912 &pwq as *const _ as *mut _,
10913 &pwk as *const _ as *mut _,
10914 &pwv as *const _ as *mut _,
10915 &pq as *const _ as *mut _,
10916 &pk as *const _ as *mut _,
10917 &pv as *const _ as *mut _,
10918 &nc as *const _ as *mut _,
10919 &rqi as *const _ as *mut _,
10920 &rki as *const _ as *mut _,
10921 &pp as *const _ as *mut _,
10922 &nhq as *const _ as *mut _,
10923 &nhk as *const _ as *mut _,
10924 &theta_scale as *const _ as *mut _,
10925 &freq_scale as *const _ as *mut _,
10926 &pff as *const _ as *mut _,
10927 &eps as *const _ as *mut _,
10928 &pkc as *const _ as *mut _,
10929 &pvc as *const _ as *mut _,
10930 &pt as *const _ as *mut _,
10931 &ktb as *const _ as *mut _,
10932 &vtb as *const _ as *mut _,
10933 ];
10934 unsafe {
10935 self.launch_pdl_flash(
10936 g,
10937 "rms_norm_qkv_rope_append_dc_f32",
10938 (rows as u32, 1, 1),
10939 (rms_block(), 1, 1),
10940 0,
10941 &mut ps,
10942 )?;
10943 }
10944 return Ok(());
10945 }
10946 let f = if g {
10947 self.func_g("rms_norm_qkv_rope_append_dc_f32")
10948 } else {
10949 self.func("rms_norm_qkv_rope_append_dc_f32")
10950 };
10951 let cfg = LaunchConfig {
10952 grid_dim: (rows as u32, 1, 1),
10953 block_dim: (rms_block(), 1, 1),
10954 shared_mem_bytes: 0,
10955 };
10956 let __s_b = self.gpu.stream();
10957 let mut b = __s_b.launch_builder(&f);
10958 match ff {
10959 Some(t) => {
10960 b.arg(q0)
10961 .arg(k0)
10962 .arg(v0)
10963 .arg(wq)
10964 .arg(wk)
10965 .arg(wv)
10966 .arg(&mut *q)
10967 .arg(&mut *k)
10968 .arg(&mut *v)
10969 .arg(&nc)
10970 .arg(&rqi)
10971 .arg(&rki)
10972 .arg(pos)
10973 .arg(&nhq)
10974 .arg(&nhk)
10975 .arg(&theta_scale)
10976 .arg(&freq_scale)
10977 .arg(t)
10978 .arg(&eps)
10979 .arg(&mut *kc)
10980 .arg(&mut *vc)
10981 .arg(t_dev)
10982 .arg(&ktb)
10983 .arg(&vtb);
10984 unsafe {
10985 b.launch(cfg)?;
10986 }
10987 }
10988 None => {
10989 let null: u64 = 0;
10990 b.arg(q0)
10991 .arg(k0)
10992 .arg(v0)
10993 .arg(wq)
10994 .arg(wk)
10995 .arg(wv)
10996 .arg(&mut *q)
10997 .arg(&mut *k)
10998 .arg(&mut *v)
10999 .arg(&nc)
11000 .arg(&rqi)
11001 .arg(&rki)
11002 .arg(pos)
11003 .arg(&nhq)
11004 .arg(&nhk)
11005 .arg(&theta_scale)
11006 .arg(&freq_scale)
11007 .arg(&null)
11008 .arg(&eps)
11009 .arg(&mut *kc)
11010 .arg(&mut *vc)
11011 .arg(t_dev)
11012 .arg(&ktb)
11013 .arg(&vtb);
11014 unsafe {
11015 b.launch(cfg)?;
11016 }
11017 }
11018 }
11019 Ok(())
11020 }
11021
11022 #[allow(clippy::too_many_arguments)]
11030 pub fn rms_norm_qkv_rope_append(
11031 &self,
11032 q0: &CudaSlice<f32>,
11033 k0: &CudaSlice<f32>,
11034 v0: &CudaSlice<f32>,
11035 wq: &CudaSlice<f32>,
11036 wk: &CudaSlice<f32>,
11037 wv: &CudaSlice<f32>,
11038 q: &mut CudaSlice<f32>,
11039 k: &mut CudaSlice<f32>,
11040 v: &mut CudaSlice<f32>,
11041 head_dim: usize,
11042 n_rot: usize,
11043 rq: usize,
11044 rk: usize,
11045 pos: &CudaSlice<i32>,
11046 nh_q: usize,
11047 nh_k: usize,
11048 base: f32,
11049 freq_scale: f32,
11050 ff: Option<&CudaSlice<f32>>,
11051 eps: f32,
11052 kc: &mut CudaSlice<u8>,
11053 vc: &mut CudaSlice<u8>,
11054 t: usize,
11055 k_tok_bytes: usize,
11056 v_tok_bytes: usize,
11057 g: bool,
11058 ) -> Result<(), Box<dyn std::error::Error>> {
11059 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11060 let rows = rq + rk + rk;
11061 let theta_scale = base.powf(-2.0 / head_dim as f32);
11062 let (nc, rqi, rki, nhq, nhk) = (
11063 head_dim as i32,
11064 rq as i32,
11065 rk as i32,
11066 nh_q as i32,
11067 nh_k as i32,
11068 );
11069 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11070 let ti = t as i32;
11071 if Self::pdl_on() && Self::pdl_wb_on() {
11072 use cudarc::driver::{DevicePtr, DevicePtrMut};
11073 let s = &self.gpu.stream();
11074 let (p0, _a0) = q0.device_ptr(s);
11075 let (p1, _a1) = k0.device_ptr(s);
11076 let (p2, _a2) = v0.device_ptr(s);
11077 let (pwq, _a3) = wq.device_ptr(s);
11078 let (pwk, _a4) = wk.device_ptr(s);
11079 let (pwv, _a5) = wv.device_ptr(s);
11080 let (pq, _a6) = q.device_ptr_mut(s);
11081 let (pk, _a7) = k.device_ptr_mut(s);
11082 let (pv, _a8) = v.device_ptr_mut(s);
11083 let (pp, _a9) = pos.device_ptr(s);
11084 let pff: u64 = match ff {
11085 Some(t) => {
11086 let (p, _gg) = t.device_ptr(s);
11087 p as u64
11088 }
11089 None => 0,
11090 };
11091 let (pkc, _a10) = kc.device_ptr_mut(s);
11092 let (pvc, _a11) = vc.device_ptr_mut(s);
11093 let mut ps = [
11094 &p0 as *const _ as *mut std::ffi::c_void,
11095 &p1 as *const _ as *mut _,
11096 &p2 as *const _ as *mut _,
11097 &pwq as *const _ as *mut _,
11098 &pwk as *const _ as *mut _,
11099 &pwv as *const _ as *mut _,
11100 &pq as *const _ as *mut _,
11101 &pk as *const _ as *mut _,
11102 &pv as *const _ as *mut _,
11103 &nc as *const _ as *mut _,
11104 &rqi as *const _ as *mut _,
11105 &rki as *const _ as *mut _,
11106 &pp as *const _ as *mut _,
11107 &nhq as *const _ as *mut _,
11108 &nhk as *const _ as *mut _,
11109 &theta_scale as *const _ as *mut _,
11110 &freq_scale as *const _ as *mut _,
11111 &pff as *const _ as *mut _,
11112 &eps as *const _ as *mut _,
11113 &pkc as *const _ as *mut _,
11114 &pvc as *const _ as *mut _,
11115 &ti as *const _ as *mut _,
11116 &ktb as *const _ as *mut _,
11117 &vtb as *const _ as *mut _,
11118 ];
11119 unsafe {
11120 self.launch_pdl_flash(
11121 g,
11122 "rms_norm_qkv_rope_append_f32",
11123 (rows as u32, 1, 1),
11124 (rms_block(), 1, 1),
11125 0,
11126 &mut ps,
11127 )?;
11128 }
11129 return Ok(());
11130 }
11131 let f = if g {
11132 self.func_g("rms_norm_qkv_rope_append_f32")
11133 } else {
11134 self.func("rms_norm_qkv_rope_append_f32")
11135 };
11136 let cfg = LaunchConfig {
11137 grid_dim: (rows as u32, 1, 1),
11138 block_dim: (rms_block(), 1, 1),
11139 shared_mem_bytes: 0,
11140 };
11141 let __s_b = self.gpu.stream();
11142 let mut b = __s_b.launch_builder(&f);
11143 let null: u64 = 0;
11144 b.arg(q0)
11145 .arg(k0)
11146 .arg(v0)
11147 .arg(wq)
11148 .arg(wk)
11149 .arg(wv)
11150 .arg(&mut *q)
11151 .arg(&mut *k)
11152 .arg(&mut *v)
11153 .arg(&nc)
11154 .arg(&rqi)
11155 .arg(&rki)
11156 .arg(pos)
11157 .arg(&nhq)
11158 .arg(&nhk)
11159 .arg(&theta_scale)
11160 .arg(&freq_scale);
11161 match ff {
11162 Some(t) => {
11163 b.arg(t);
11164 }
11165 None => {
11166 b.arg(&null);
11167 }
11168 }
11169 b.arg(&eps)
11170 .arg(&mut *kc)
11171 .arg(&mut *vc)
11172 .arg(&ti)
11173 .arg(&ktb)
11174 .arg(&vtb);
11175 unsafe {
11176 b.launch(cfg)?;
11177 }
11178 Ok(())
11179 }
11180
11181 pub fn add_q8_1(
11182 &self,
11183 a: &CudaSlice<f32>,
11184 b: &CudaSlice<f32>,
11185 res: &mut CudaSlice<f32>,
11186 ncols: usize,
11187 nrows: usize,
11188 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11189 debug_assert!(ncols % 128 == 0);
11190 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11191 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11192 let f = self.func("add_q8_1_f32");
11193 let cfg = LaunchConfig {
11194 grid_dim: (nrows as u32, 1, 1),
11195 block_dim: (rms_block(), 1, 1),
11196 shared_mem_bytes: 0,
11197 };
11198 let nc = ncols as i32;
11199 let __s_b2 = self.gpu.stream();
11200 let mut b2 = __s_b2.launch_builder(&f);
11201 b2.arg(a)
11202 .arg(b)
11203 .arg(&mut *res)
11204 .arg(&mut out_q)
11205 .arg(&mut out_d)
11206 .arg(&nc);
11207 unsafe {
11208 b2.launch(cfg)?;
11209 }
11210 Ok((out_q, out_d))
11211 }
11212
11213 pub fn rms_pre_add_q8_1(
11217 &self,
11218 a: &CudaSlice<f32>,
11219 wa: &CudaSlice<f32>,
11220 b: &CudaSlice<f32>,
11221 res: &mut CudaSlice<f32>,
11222 ncols: usize,
11223 nrows: usize,
11224 eps: f32,
11225 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11226 debug_assert!(ncols % 128 == 0);
11227 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11228 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11229 let f = self.func("rms_pre_add_q8_1_f32");
11230 let cfg = LaunchConfig {
11231 grid_dim: (nrows as u32, 1, 1),
11232 block_dim: (rms_block(), 1, 1),
11233 shared_mem_bytes: 0,
11234 };
11235 let (nc, ep) = (ncols as i32, eps);
11236 let __s_b2 = self.gpu.stream();
11237 let mut b2 = __s_b2.launch_builder(&f);
11238 b2.arg(a)
11239 .arg(wa)
11240 .arg(b)
11241 .arg(&mut *res)
11242 .arg(&mut out_q)
11243 .arg(&mut out_d)
11244 .arg(&nc)
11245 .arg(&ep);
11246 unsafe {
11247 b2.launch(cfg)?;
11248 }
11249 Ok((out_q, out_d))
11250 }
11251
11252 pub fn l2_v2_on(ncols: usize) -> bool {
11256 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11257 }
11258
11259 pub fn l2_norm_pp(
11260 &self,
11261 x: &CudaSlice<f32>,
11262 dst: &mut CudaSlice<f32>,
11263 dst16: Option<&mut CudaSlice<u8>>,
11264 ncols: usize,
11265 nrows: usize,
11266 eps: f32,
11267 ) -> Result<(), Box<dyn std::error::Error>> {
11268 if Self::l2_v2_on(ncols) {
11269 let f = self.func("l2_norm_pp_v2_f32");
11270 let rows_per_block = 8u32; let cfg = LaunchConfig {
11272 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11273 block_dim: (256, 1, 1),
11274 shared_mem_bytes: 0,
11275 };
11276 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11277 let d16: u64 = match dst16 {
11279 Some(d) => self.addr_u8(d),
11280 None => 0,
11281 };
11282 let __s_b = self.gpu.stream();
11283 let mut b = __s_b.launch_builder(&f);
11284 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11285 unsafe {
11286 b.launch(cfg)?;
11287 }
11288 return Ok(());
11289 }
11290 self.l2_norm(x, dst, ncols, nrows, eps)
11291 }
11292
11293 pub fn l2_norm(
11294 &self,
11295 x: &CudaSlice<f32>,
11296 dst: &mut CudaSlice<f32>,
11297 ncols: usize,
11298 nrows: usize,
11299 eps: f32,
11300 ) -> Result<(), Box<dyn std::error::Error>> {
11301 let f = self.func("l2_norm_f32");
11302 let cfg = LaunchConfig {
11303 grid_dim: (nrows as u32, 1, 1),
11304 block_dim: (256, 1, 1),
11305 shared_mem_bytes: 0,
11306 };
11307 let (nc, e) = (ncols as i32, eps);
11308 let __s_b = self.gpu.stream();
11309 let mut b = __s_b.launch_builder(&f);
11310 b.arg(x).arg(dst).arg(&nc).arg(&e);
11311 unsafe {
11312 b.launch(cfg)?;
11313 }
11314 Ok(())
11315 }
11316
11317 pub fn l2_norm_decode(
11323 &self,
11324 x: &CudaSlice<f32>,
11325 dst: &mut CudaSlice<f32>,
11326 ncols: usize,
11327 nrows: usize,
11328 eps: f32,
11329 ) -> Result<(), Box<dyn std::error::Error>> {
11330 let f = self.func("l2_norm_f32");
11331 let cfg = LaunchConfig {
11332 grid_dim: (nrows as u32, 1, 1),
11333 block_dim: (32, 1, 1),
11334 shared_mem_bytes: 0,
11335 };
11336 let (nc, e) = (ncols as i32, eps);
11337 let __s_b = self.gpu.stream();
11338 let mut b = __s_b.launch_builder(&f);
11339 b.arg(x).arg(dst).arg(&nc).arg(&e);
11340 unsafe {
11341 b.launch(cfg)?;
11342 }
11343 Ok(())
11344 }
11345
11346 pub fn rope_neox(
11348 &self,
11349 x: &mut CudaSlice<f32>,
11350 pos: &CudaSlice<i32>,
11351 head_dim: usize,
11352 n_dims: usize,
11353 n_heads: usize,
11354 n_tokens: usize,
11355 freq_base: f32,
11356 freq_scale: f32,
11357 ) -> Result<(), Box<dyn std::error::Error>> {
11358 let f = self.func("rope_neox_f32");
11359 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11360 let grid = (n_heads * n_tokens) as u32;
11361 let cfg = LaunchConfig {
11362 grid_dim: (grid, 1, 1),
11363 block_dim: ((head_dim / 2) as u32, 1, 1),
11364 shared_mem_bytes: 0,
11365 };
11366 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11367 let __s_b = self.gpu.stream();
11368 let mut b = __s_b.launch_builder(&f);
11369 b.arg(x)
11370 .arg(pos)
11371 .arg(&hd)
11372 .arg(&nd)
11373 .arg(&nh)
11374 .arg(&theta_scale)
11375 .arg(&freq_scale);
11376 unsafe {
11377 b.launch(cfg)?;
11378 }
11379 Ok(())
11380 }
11381
11382 pub fn rope_neox_ff(
11384 &self,
11385 x: &mut CudaSlice<f32>,
11386 pos: &CudaSlice<i32>,
11387 head_dim: usize,
11388 n_dims: usize,
11389 n_heads: usize,
11390 n_tokens: usize,
11391 freq_base: f32,
11392 freq_scale: f32,
11393 ff: &CudaSlice<f32>,
11394 ) -> Result<(), Box<dyn std::error::Error>> {
11395 let f = self.func("rope_neox_ff_f32");
11396 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11397 let grid = (n_heads * n_tokens) as u32;
11398 let cfg = LaunchConfig {
11399 grid_dim: (grid, 1, 1),
11400 block_dim: ((head_dim / 2) as u32, 1, 1),
11401 shared_mem_bytes: 0,
11402 };
11403 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11404 let __s_b = self.gpu.stream();
11405 let mut b = __s_b.launch_builder(&f);
11406 b.arg(x)
11407 .arg(pos)
11408 .arg(&hd)
11409 .arg(&nd)
11410 .arg(&nh)
11411 .arg(&theta_scale)
11412 .arg(&freq_scale)
11413 .arg(ff);
11414 unsafe {
11415 b.launch(cfg)?;
11416 }
11417 Ok(())
11418 }
11419
11420 #[allow(clippy::too_many_arguments)]
11422 pub fn rope_neox2(
11423 &self,
11424 q: &mut CudaSlice<f32>,
11425 k: &mut CudaSlice<f32>,
11426 pos: &CudaSlice<i32>,
11427 head_dim: usize,
11428 n_dims: usize,
11429 nh_q: usize,
11430 nh_k: usize,
11431 n_tokens: usize,
11432 freq_base: f32,
11433 freq_scale: f32,
11434 ff: Option<&CudaSlice<f32>>,
11435 ) -> Result<(), Box<dyn std::error::Error>> {
11436 let f = self.func("rope_neox2_f32");
11437 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11438 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11439 let cfg = LaunchConfig {
11440 grid_dim: (grid, 1, 1),
11441 block_dim: ((head_dim / 2) as u32, 1, 1),
11442 shared_mem_bytes: 0,
11443 };
11444 let (hd, nd, nq, nk, nt) = (
11445 head_dim as i32,
11446 n_dims as i32,
11447 nh_q as i32,
11448 nh_k as i32,
11449 n_tokens as i32,
11450 );
11451 let __s_b = self.gpu.stream();
11452 let mut b = __s_b.launch_builder(&f);
11453 b.arg(q)
11454 .arg(k)
11455 .arg(pos)
11456 .arg(&hd)
11457 .arg(&nd)
11458 .arg(&nq)
11459 .arg(&nk)
11460 .arg(&nt)
11461 .arg(&theta_scale)
11462 .arg(&freq_scale);
11463 match ff {
11464 Some(ffv) => {
11465 b.arg(ffv);
11466 unsafe {
11467 b.launch(cfg)?;
11468 }
11469 }
11470 None => {
11471 let null: u64 = 0;
11472 b.arg(&null);
11473 unsafe {
11474 b.launch(cfg)?;
11475 }
11476 }
11477 }
11478 Ok(())
11479 }
11480
11481 pub fn gelu_tanh_mul(
11483 &self,
11484 gate: &CudaSlice<f32>,
11485 up: &CudaSlice<f32>,
11486 dst: &mut CudaSlice<f32>,
11487 n: usize,
11488 ) -> Result<(), Box<dyn std::error::Error>> {
11489 let f = self.func("gelu_tanh_mul_f32");
11490 let cfg = LaunchConfig::for_num_elems(n as u32);
11491 let ni = n as i32;
11492 let __s_b = self.gpu.stream();
11493 let mut b = __s_b.launch_builder(&f);
11494 b.arg(gate).arg(up).arg(dst).arg(&ni);
11495 unsafe {
11496 b.launch(cfg)?;
11497 }
11498 Ok(())
11499 }
11500
11501 pub fn silu_mul(
11502 &self,
11503 gate: &CudaSlice<f32>,
11504 up: &CudaSlice<f32>,
11505 dst: &mut CudaSlice<f32>,
11506 n: usize,
11507 ) -> Result<(), Box<dyn std::error::Error>> {
11508 let f = self.func("silu_mul_f32");
11509 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11511 let ni = n as i32;
11512 let __s_b = self.gpu.stream();
11513 let mut b = __s_b.launch_builder(&f);
11514 b.arg(gate).arg(up).arg(dst).arg(&ni);
11515 unsafe {
11516 b.launch(cfg)?;
11517 }
11518 Ok(())
11519 }
11520
11521 pub fn silu_mul_host_expf(
11523 &self,
11524 gate: &CudaSlice<f32>,
11525 up: &CudaSlice<f32>,
11526 dst: &mut CudaSlice<f32>,
11527 n: usize,
11528 ) -> Result<(), Box<dyn std::error::Error>> {
11529 let f = self.func("silu_mul_host_expf_f32");
11530 let cfg = LaunchConfig::for_num_elems(n as u32);
11531 let ni = n as i32;
11532 let __s_b = self.gpu.stream();
11533 let mut b = __s_b.launch_builder(&f);
11534 b.arg(gate).arg(up).arg(dst).arg(&ni);
11535 unsafe {
11536 b.launch(cfg)?;
11537 }
11538 Ok(())
11539 }
11540
11541 pub fn silu_clamped_mul_host_expf(
11543 &self,
11544 gate: &CudaSlice<f32>,
11545 up: &CudaSlice<f32>,
11546 limit: f32,
11547 dst: &mut CudaSlice<f32>,
11548 n: usize,
11549 ) -> Result<(), Box<dyn std::error::Error>> {
11550 if !limit.is_finite() || limit <= 0.0 {
11551 return Err(
11552 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11553 );
11554 }
11555 let f = self.func("silu_clamped_mul_host_expf_f32");
11556 let cfg = LaunchConfig::for_num_elems(n as u32);
11557 let ni = n as i32;
11558 let __s_b = self.gpu.stream();
11559 let mut b = __s_b.launch_builder(&f);
11560 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11561 unsafe {
11562 b.launch(cfg)?;
11563 }
11564 Ok(())
11565 }
11566
11567 pub fn silu_mul_f16out(
11570 &self,
11571 gate: &CudaSlice<f32>,
11572 up: &CudaSlice<f32>,
11573 dst: &mut CudaSlice<f32>,
11574 dst16: &mut CudaSlice<u8>,
11575 n: usize,
11576 ) -> Result<(), Box<dyn std::error::Error>> {
11577 let f = self.func("silu_mul_f16out_f32");
11578 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11579 let ni = n as i32;
11580 let __s_b = self.gpu.stream();
11581 let mut b = __s_b.launch_builder(&f);
11582 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11583 unsafe {
11584 b.launch(cfg)?;
11585 }
11586 Ok(())
11587 }
11588
11589 pub fn silu_mul_scaled(
11596 &self,
11597 gate: &CudaSlice<f32>,
11598 up: &CudaSlice<f32>,
11599 gs: f32,
11600 us: f32,
11601 dst: &mut CudaSlice<f32>,
11602 n: usize,
11603 ) -> Result<(), Box<dyn std::error::Error>> {
11604 let f = self.func("silu_mul_scaled_f32");
11605 let cfg = LaunchConfig::for_num_elems(n as u32);
11606 let ni = n as i32;
11607 let (gsf, usf) = (gs, us);
11608 let __s_b = self.gpu.stream();
11609 let mut b = __s_b.launch_builder(&f);
11610 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11611 unsafe {
11612 b.launch(cfg)?;
11613 }
11614 Ok(())
11615 }
11616
11617 #[allow(clippy::too_many_arguments)]
11621 pub fn swigluoai_mul_scaled(
11622 &self,
11623 gate: &CudaSlice<f32>,
11624 up: &CudaSlice<f32>,
11625 gs: f32,
11626 us: f32,
11627 alpha: f32,
11628 limit: f32,
11629 dst: &mut CudaSlice<f32>,
11630 n: usize,
11631 ) -> Result<(), Box<dyn std::error::Error>> {
11632 let f = self.func("swigluoai_mul_scaled_f32");
11633 let cfg = LaunchConfig::for_num_elems(n as u32);
11634 let ni = n as i32;
11635 let __s_b = self.gpu.stream();
11636 let mut b = __s_b.launch_builder(&f);
11637 b.arg(gate)
11638 .arg(up)
11639 .arg(&gs)
11640 .arg(&us)
11641 .arg(&alpha)
11642 .arg(&limit)
11643 .arg(dst)
11644 .arg(&ni);
11645 unsafe {
11646 b.launch(cfg)?;
11647 }
11648 Ok(())
11649 }
11650
11651 pub fn silu_mul_scaled_q8_1(
11659 &self,
11660 gate: &CudaSlice<f32>,
11661 up: &CudaSlice<f32>,
11662 gs: f32,
11663 us: f32,
11664 n: usize,
11665 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11666 let f = self.func("silu_mul_scaled_q8_1");
11667 let nblk = n / 32;
11668 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);
11672 let (gsf, usf, ni) = (gs, us, n as i32);
11673 let __s_b = self.gpu.stream();
11674 let mut b = __s_b.launch_builder(&f);
11675 b.arg(gate)
11676 .arg(up)
11677 .arg(&gsf)
11678 .arg(&usf)
11679 .arg(&mut aq)
11680 .arg(&mut ad)
11681 .arg(&ni);
11682 unsafe {
11683 b.launch(cfg)?;
11684 }
11685 Ok((aq, ad))
11686 }
11687
11688 pub fn add(
11689 &self,
11690 a: &CudaSlice<f32>,
11691 b_in: &CudaSlice<f32>,
11692 dst: &mut CudaSlice<f32>,
11693 n: usize,
11694 ) -> Result<(), Box<dyn std::error::Error>> {
11695 let f = self.func("add_f32");
11696 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11698 let ni = n as i32;
11699 let __s_bld = self.gpu.stream();
11700 let mut bld = __s_bld.launch_builder(&f);
11701 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11702 unsafe {
11703 bld.launch(cfg)?;
11704 }
11705 Ok(())
11706 }
11707
11708 pub fn mul(
11709 &self,
11710 a: &CudaSlice<f32>,
11711 b_in: &CudaSlice<f32>,
11712 dst: &mut CudaSlice<f32>,
11713 n: usize,
11714 ) -> Result<(), Box<dyn std::error::Error>> {
11715 let f = self.func("mul_f32");
11716 let cfg = LaunchConfig::for_num_elems(n as u32);
11717 let ni = n as i32;
11718 let __s_bld = self.gpu.stream();
11719 let mut bld = __s_bld.launch_builder(&f);
11720 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11721 unsafe {
11722 bld.launch(cfg)?;
11723 }
11724 Ok(())
11725 }
11726
11727 pub fn matmul(
11730 &self,
11731 w: &crate::model::GpuTensor,
11732 x: &CudaSlice<f32>,
11733 m: usize,
11734 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11735 use crate::model::GpuTensor;
11736 let in_f = w.in_features();
11737 let out_f = w.out_features();
11738 #[allow(non_snake_case)]
11746 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11749 usize::MAX
11750 } else {
11751 16usize
11752 };
11753
11754 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11779 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11780 return Ok(y);
11781 }
11782 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11789 return Ok(y);
11790 }
11791 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11794 return Ok(y);
11795 }
11796 }
11797 if let GpuTensor::Quant { qtype, .. } = w {
11812 if *qtype == QT_F8_E4M3_BLK {
11813 if m >= GEMM_M_THRESHOLD {
11814 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11815 return Ok(y);
11816 }
11817 }
11818 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11819 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11820 return Ok(y);
11821 }
11822 }
11823 }
11824 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11825 return self.qmatvec_mmq(w, x, m);
11826 }
11827 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11828 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11829 return self.qmatvec_gemm(w, &aq, &ad, m);
11830 }
11831 if m >= GEMM_M_THRESHOLD {
11834 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11835 return Ok(y);
11836 }
11837 }
11838 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11842 if m == 1 && fast {
11847 if let GpuTensor::Quant {
11848 bytes,
11849 qtype,
11850 row_bytes,
11851 rp,
11852 rp4,
11853 scale,
11854 ..
11855 } = w
11856 {
11857 if self.mmvq_supports(*qtype) {
11858 let (bytes, rp) = match rp4 {
11862 Some(m4) => (m4, true),
11863 None => (bytes, *rp),
11864 };
11865 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11866 return self.qmatvec_mmvq(
11867 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11868 );
11869 }
11870 }
11871 }
11872 if (2..=16).contains(&m)
11888 && fast
11889 && std::env::var("MEMRA_NO_BATCHED").is_err()
11890 && (m <= 4 || Self::b8_enabled())
11891 {
11892 let m_ok = m <= 8
11902 || matches!(w, GpuTensor::Quant { qtype, .. }
11903 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11904 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11905 if m_ok {
11906 if let GpuTensor::Quant {
11907 bytes,
11908 qtype,
11909 row_bytes,
11910 rp,
11911 rp4,
11912 ..
11913 } = w
11914 {
11915 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11916 let (bytes, rp) = match rp4 {
11917 Some(m4) => (m4, true),
11918 None => (bytes, *rp),
11919 };
11920 let mcols = Self::batched_mcols(m);
11921 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11922 let mut y = self.qmatvec_mmvq_batched(
11923 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11924 )?;
11925 if let GpuTensor::Quant { scale, .. } = w {
11926 if *scale != 1.0 {
11927 self.scale_inplace(&mut y, *scale, m * out_f)?;
11928 }
11929 }
11930 return Ok(y);
11931 }
11932 }
11933 }
11934 }
11935 if fast {
11941 if let GpuTensor::Quant {
11942 bytes,
11943 qtype,
11944 row_bytes,
11945 scale,
11946 ..
11947 } = w
11948 {
11949 if *qtype == QT_F8_E4M3 {
11950 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11951 return self.qmatvec_mmvq(
11952 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11953 );
11954 }
11955 }
11956 }
11957 let mut y = match w {
11958 GpuTensor::Quant {
11959 bytes,
11960 qtype,
11961 row_bytes,
11962 ..
11963 } if fast && *qtype == QT_Q8_0 => {
11964 self.qmatvec_q8_0_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_Q4_K => {
11972 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11973 }
11974 GpuTensor::Quant {
11975 bytes,
11976 qtype,
11977 row_bytes,
11978 ..
11979 } if fast && *qtype == QT_Q6_K => {
11980 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11981 }
11982 GpuTensor::Quant {
11983 bytes,
11984 qtype,
11985 row_bytes,
11986 ..
11987 } if fast && *qtype == QT_Q5_K => {
11988 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11989 }
11990 GpuTensor::Quant {
11991 bytes,
11992 qtype,
11993 row_bytes,
11994 ..
11995 } if fast && *qtype == QT_Q3_K => {
11996 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11997 }
11998 GpuTensor::Quant {
11999 bytes,
12000 qtype,
12001 row_bytes,
12002 rp,
12003 ..
12004 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12005 if *rp {
12006 "qmatvec_nvfp4_dp4a_rp"
12007 } else {
12008 "qmatvec_nvfp4_dp4a"
12009 },
12010 &bytes.slice(0..bytes.len()),
12011 x,
12012 m,
12013 in_f,
12014 out_f,
12015 *row_bytes,
12016 )?,
12017 GpuTensor::Quant {
12021 bytes,
12022 qtype,
12023 row_bytes,
12024 ..
12025 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12026 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12027 }
12028 GpuTensor::Quant {
12033 bytes,
12034 qtype,
12035 row_bytes,
12036 rp,
12037 ..
12038 } =>
12039 {
12042 self.qmatvec(
12043 bytes,
12044 x,
12045 m,
12046 in_f,
12047 out_f,
12048 if *rp && *qtype == QT_NVFP4 {
12049 QT_NVFP4_RP
12050 } else {
12051 *qtype
12052 },
12053 *row_bytes,
12054 )?
12055 }
12056 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12057 GpuTensor::FloatBf16 { data, .. } => {
12060 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12067 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12068 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12069 y
12070 } else {
12071 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12072 }
12073 }
12074 };
12075 if let GpuTensor::Quant { scale, .. } = w {
12077 if *scale != 1.0 {
12078 self.scale_inplace(&mut y, *scale, m * out_f)?;
12079 }
12080 }
12081 Ok(y)
12082 }
12083
12084 pub fn stage_a_raw_needed() -> bool {
12094 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12095 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12096 }
12097
12098 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12101 use crate::model::GpuTensor;
12102 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12103 return false;
12104 }
12105 match w {
12106 GpuTensor::Quant { qtype, .. } => {
12113 matches!(
12114 *qtype,
12115 QT_Q8_0
12116 | QT_Q4_K
12117 | QT_Q6_K
12118 | QT_Q5_K
12119 | QT_Q3_K
12120 | QT_NVFP4
12121 | QT_F8_E4M3
12122 | QT_F8_E4M3_BLK
12123 | QT_Q4_0
12124 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12125 }
12126 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12127 }
12128 }
12129
12130 pub fn matmul_pre(
12135 &self,
12136 w: &crate::model::GpuTensor,
12137 aq: &CudaSlice<i8>,
12138 ad: &CudaSlice<f32>,
12139 x_fallback: &CudaSlice<f32>,
12140 m: usize,
12141 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12142 use crate::model::GpuTensor;
12143 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12149 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12152 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12153 return Ok(y);
12154 }
12155 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12158 return Ok(y);
12159 }
12160 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12162 return Ok(y);
12163 }
12164 }
12165 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12171 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12172 return Ok(y);
12173 }
12174 }
12175 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12176 return Ok(y);
12177 }
12178 if m >= 16
12183 && w.out_features() >= 128
12184 && self.mmq_supports(w)
12185 && !self.verify_exact_on()
12186 && x_raw_ok
12187 {
12188 return self.qmatvec_mmq(w, x_fallback, m);
12189 }
12190 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12193 if let Some(y) =
12194 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12195 {
12196 return Ok(y);
12197 }
12198 }
12199 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12202 return self.qmatvec_gemm(w, aq, ad, m);
12203 }
12204 if !self.uses_q8_1_fast(w) {
12223 if !x_raw_ok {
12224 return Err(format!(
12225 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12226 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12227 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12228 activation (see Engine::rms_norm_decode, which is bit-identical to \
12229 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12230 x_fallback.len(),
12231 m,
12232 w.in_features(),
12233 m * w.in_features()
12234 )
12235 .into());
12236 }
12237 return self.matmul(w, x_fallback, m);
12238 }
12239 let in_f = w.in_features();
12240 let out_f = w.out_features();
12241 let (bytes, qtype, row_bytes, scale, rp) = match w {
12242 GpuTensor::Quant {
12243 bytes,
12244 qtype,
12245 row_bytes,
12246 scale,
12247 rp,
12248 ..
12249 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12250 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12251 };
12252 let (mbytes, mrp) = match w {
12255 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12256 _ => (bytes, rp),
12257 };
12258 if m == 1 && self.mmvq_supports(qtype) {
12262 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12263 }
12264 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12277 && std::env::var("MEMRA_NO_BATCHED").is_err()
12278 && (m <= 4 || Self::b8_enabled())
12279 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12283 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12284 {
12285 let mcols = Self::batched_mcols(m);
12286 return self.qmatvec_mmvq_batched(
12287 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12288 );
12289 }
12290 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12296 let (b2, r2) = if qtype == QT_Q4_0 {
12297 (mbytes, mrp)
12298 } else {
12299 (bytes, rp)
12300 };
12301 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12302 }
12303 let name = match qtype {
12304 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12305 QT_Q4_K => "qmatvec_q4_K_dp4a",
12306 QT_Q6_K => "qmatvec_q6_K_dp4a",
12307 QT_Q5_K => "qmatvec_q5_K_dp4a",
12308 QT_Q3_K => "qmatvec_q3_K_dp4a",
12309 QT_NVFP4 => {
12310 if rp {
12311 "qmatvec_nvfp4_dp4a_rp"
12312 } else {
12313 "qmatvec_nvfp4_dp4a"
12314 }
12315 }
12316 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12317 _ => unreachable!(),
12318 };
12319 let f = self.func(name);
12320 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12322 grid_dim: (out_f as u32, m as u32, 1),
12323 block_dim: (128, 1, 1),
12324 shared_mem_bytes: 0,
12325 };
12326 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12327 let __s_b = self.gpu.stream();
12328 let mut b = __s_b.launch_builder(&f);
12329 b.arg(bytes)
12330 .arg(aq)
12331 .arg(ad)
12332 .arg(&mut y)
12333 .arg(&inf)
12334 .arg(&outf)
12335 .arg(&mi)
12336 .arg(&rb);
12337 unsafe {
12338 b.launch(cfg)?;
12339 }
12340 if scale != 1.0 {
12341 self.scale_inplace(&mut y, scale, m * out_f)?;
12342 }
12343 Ok(y)
12344 }
12345
12346 pub fn matmul_decode_exact(
12354 &self,
12355 w: &crate::model::GpuTensor,
12356 x: &CudaSlice<f32>,
12357 m: usize,
12358 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12359 use crate::model::GpuTensor;
12360 if let GpuTensor::Float { data, .. } = w {
12368 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12369 }
12370 if let GpuTensor::FloatBf16 { data, .. } = w {
12373 let (in_f, out_f) = (w.in_features(), w.out_features());
12374 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12377 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12378 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12379 return Ok(y);
12380 }
12381 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12382 }
12383 if !self.uses_q8_1_fast(w) {
12384 return self.matmul(w, x, m);
12385 }
12386 let in_f = w.in_features();
12387 let out_f = w.out_features();
12388 let (bytes, qtype, row_bytes, scale, rp) = match w {
12389 GpuTensor::Quant {
12390 bytes,
12391 qtype,
12392 row_bytes,
12393 scale,
12394 rp,
12395 ..
12396 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12397 _ => return self.matmul(w, x, m),
12398 };
12399 let (bytes, rp) = match w {
12402 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12403 _ => (bytes, rp),
12404 };
12405 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12406 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12410 return Ok(y);
12411 }
12412 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12421 && std::env::var("MEMRA_NO_BATCHED").is_err()
12422 && (m <= 4 || Self::b8_enabled())
12423 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12426 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12427 {
12428 let mcols = Self::batched_mcols(m);
12429 return self.qmatvec_mmvq_batched(
12430 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12431 );
12432 }
12433 if self.mmvq_supports(qtype) {
12434 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12437 }
12438 self.matmul_pre(w, &aq, &ad, x, m)
12441 }
12442
12443 pub fn matmul_decode_exact_pre(
12453 &self,
12454 w: &crate::model::GpuTensor,
12455 aq: &CudaSlice<i8>,
12456 ad: &CudaSlice<f32>,
12457 m: usize,
12458 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12459 use crate::model::GpuTensor;
12460 debug_assert!(
12461 self.uses_q8_1_fast(w),
12462 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12463 );
12464 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12466 return Ok(y);
12467 }
12468 let in_f = w.in_features();
12469 let out_f = w.out_features();
12470 let (bytes, qtype, row_bytes, scale, rp) = match w {
12471 GpuTensor::Quant {
12472 bytes,
12473 qtype,
12474 row_bytes,
12475 scale,
12476 rp,
12477 ..
12478 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12479 _ => {
12480 return Err(
12481 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12482 );
12483 }
12484 };
12485 let (bytes, rp) = match w {
12487 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12488 _ => (bytes, rp),
12489 };
12490 if (2..=16).contains(&m)
12492 && self.batched_supports(qtype)
12493 && self.mmvq_supports(qtype)
12494 && std::env::var("MEMRA_NO_BATCHED").is_err()
12495 && (m <= 4 || Self::b8_enabled())
12496 && (m <= 8
12497 || qtype == QT_Q4_0
12498 || qtype == QT_Q6_K
12499 || qtype == QT_F8_E4M3
12500 || qtype == QT_NVFP4
12501 || qtype == QT_Q4_K
12502 || qtype == QT_Q5_K
12503 || qtype == QT_Q8_0)
12504 {
12505 let mcols = Self::batched_mcols(m);
12506 return self.qmatvec_mmvq_batched(
12507 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12508 );
12509 }
12510 if self.mmvq_supports(qtype) {
12511 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12512 }
12513 let x0 = self.zeros(0)?;
12516 self.matmul_pre(w, aq, ad, &x0, m)
12517 }
12518
12519 pub fn matmul_decode_exact_dual_pre(
12528 &self,
12529 w0: &crate::model::GpuTensor,
12530 w1: &crate::model::GpuTensor,
12531 aq: &CudaSlice<i8>,
12532 ad: &CudaSlice<f32>,
12533 m: usize,
12534 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12535 {
12536 use crate::model::GpuTensor;
12537 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12538 let on = *ON.get_or_init(|| {
12539 std::env::var("MEMRA_SPEC_DUAL_T")
12540 .map(|v| v != "0")
12541 .unwrap_or(true)
12542 });
12543 if !on
12544 || !(2..=7).contains(&m)
12545 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12546 || !self.uses_q8_1_fast(w0)
12547 || !self.uses_q8_1_fast(w1)
12548 {
12549 return Ok(None);
12550 }
12551 if !self.mmvq_supports(QT_NVFP4) {
12556 return Ok(None);
12557 }
12558 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12559 if w1.in_features() != in_f || w1.out_features() != out_f {
12560 return Ok(None);
12561 }
12562 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12563 (
12564 GpuTensor::Quant {
12565 bytes: b0,
12566 qtype: q0,
12567 row_bytes: rb0,
12568 scale: s0,
12569 rp: rp0,
12570 rp4: None,
12571 ..
12572 },
12573 GpuTensor::Quant {
12574 bytes: b1,
12575 qtype: q1,
12576 row_bytes: rb1,
12577 scale: s1,
12578 rp: rp1,
12579 rp4: None,
12580 ..
12581 },
12582 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12583 (b0, b1, *rb0, *s0, *s1, *rp0)
12584 }
12585 _ => return Ok(None),
12586 };
12587 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12590 {
12591 return Ok(None);
12592 }
12593 let (y0, y1) =
12594 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12595 Ok(Some(((y0, s0), (y1, s1))))
12596 }
12597
12598 pub fn matmul_decode_exact_group4_pre(
12610 &self,
12611 ws: [&crate::model::GpuTensor; 4],
12612 aq: &CudaSlice<i8>,
12613 ad: &CudaSlice<f32>,
12614 m: usize,
12615 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12616 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12617 let on = *ON.get_or_init(|| {
12618 std::env::var("MEMRA_TK_GDN_GROUP")
12619 .map(|v| v != "0")
12620 .unwrap_or(true)
12621 });
12622 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12623 }
12624
12625 pub fn matmul_decode_exact_group3_pre(
12630 &self,
12631 ws: [&crate::model::GpuTensor; 3],
12632 aq: &CudaSlice<i8>,
12633 ad: &CudaSlice<f32>,
12634 m: usize,
12635 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12636 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12637 let on = *ON.get_or_init(|| {
12638 std::env::var("MEMRA_TK_FA_GROUP")
12639 .map(|v| v != "0")
12640 .unwrap_or(true)
12641 });
12642 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12643 }
12644
12645 fn matmul_decode_exact_group_pre(
12649 &self,
12650 ws: &[&crate::model::GpuTensor],
12651 aq: &CudaSlice<i8>,
12652 ad: &CudaSlice<f32>,
12653 m: usize,
12654 on: bool,
12655 tag: &'static str,
12656 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12657 use crate::model::GpuTensor;
12658 if !on
12659 || !(2..=16).contains(&m)
12660 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12661 || (m > 4 && !Self::b8_enabled())
12662 || !self.mmvq_supports(QT_NVFP4)
12663 || !self.batched_supports(QT_NVFP4)
12664 {
12665 return Ok(None);
12666 }
12667 let in_f = ws[0].in_features();
12668 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12669 for w in ws {
12670 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12671 return Ok(None);
12672 }
12673 match w {
12674 GpuTensor::Quant {
12675 bytes,
12676 qtype,
12677 scale,
12678 rp: true,
12679 rp4: None,
12680 ..
12681 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12682 parts.push((bytes, w.out_features(), *scale));
12683 }
12684 _ => return Ok(None),
12685 }
12686 }
12687 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12689 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12690 let mcols = if (5..=7).contains(&m) && b567 {
12691 m
12692 } else {
12693 Self::batched_mcols(m)
12694 };
12695 let kname: &'static str = match mcols {
12696 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12697 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12698 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12699 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12700 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12701 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12702 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12703 _ => return Ok(None),
12704 };
12705 if std::env::var("MEMRA_DEBUG").is_ok() {
12708 use std::sync::Mutex;
12709 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12710 let mut seen = SEEN.lock().unwrap();
12711 if !seen.contains(&tag) {
12712 seen.push(tag);
12713 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12714 }
12715 }
12716 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12719 let three = parts.len() == 3;
12720 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12721 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12722 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12723 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12726 let cfg = LaunchConfig {
12727 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12728 block_dim: (32, ROWS_PER_BLOCK, 1),
12729 shared_mem_bytes: 0,
12730 };
12731 let (inf, mi) = (in_f as i32, m as i32);
12732 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12733 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12734 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12735 let s3 = if three { 1.0f32 } else { parts[3].2 };
12736 let w3 = if three { parts[0].0 } else { parts[3].0 };
12737 let f = self.func(kname);
12738 let __s_b = self.gpu.stream();
12739 let mut b = __s_b.launch_builder(&f);
12740 b.arg(parts[0].0)
12741 .arg(parts[1].0)
12742 .arg(parts[2].0)
12743 .arg(w3)
12744 .arg(aq)
12745 .arg(ad)
12746 .arg(&mut y0)
12747 .arg(&mut y1)
12748 .arg(&mut y2)
12749 .arg(&mut y3)
12750 .arg(&inf)
12751 .arg(&n0)
12752 .arg(&n1)
12753 .arg(&n2)
12754 .arg(&n3)
12755 .arg(&mi)
12756 .arg(&s0)
12757 .arg(&s1)
12758 .arg(&s2)
12759 .arg(&s3);
12760 unsafe {
12761 b.launch(cfg)?;
12762 }
12763 Ok(Some(if three {
12764 vec![y0, y1, y2]
12765 } else {
12766 vec![y0, y1, y2, y3]
12767 }))
12768 }
12769
12770 pub fn matmul_decode_exact_dual(
12786 &self,
12787 w0: &crate::model::GpuTensor,
12788 w1: &crate::model::GpuTensor,
12789 x: &CudaSlice<f32>,
12790 m: usize,
12791 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12792 use crate::model::GpuTensor;
12793 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12794 let on = *ON.get_or_init(|| {
12795 std::env::var("MEMRA_SPEC_DUAL_T")
12796 .map(|v| v != "0")
12797 .unwrap_or(true)
12798 });
12799 if !on
12800 || !(2..=4).contains(&m)
12801 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12802 || !self.uses_q8_1_fast(w0)
12803 || !self.uses_q8_1_fast(w1)
12804 {
12805 return Ok(None);
12806 }
12807 if !self.mmvq_supports(QT_NVFP4) {
12812 return Ok(None);
12813 }
12814 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12815 if w1.in_features() != in_f || w1.out_features() != out_f {
12816 return Ok(None);
12817 }
12818 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12819 (
12820 GpuTensor::Quant {
12821 bytes: b0,
12822 qtype: q0,
12823 row_bytes: rb0,
12824 scale: s0,
12825 rp: rp0,
12826 rp4: None,
12827 ..
12828 },
12829 GpuTensor::Quant {
12830 bytes: b1,
12831 qtype: q1,
12832 row_bytes: rb1,
12833 scale: s1,
12834 rp: rp1,
12835 rp4: None,
12836 ..
12837 },
12838 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12839 (b0, b1, *rb0, *s0, *s1, *rp0)
12840 }
12841 _ => return Ok(None),
12842 };
12843 if std::env::var("MEMRA_DEBUG").is_ok() {
12846 static ONCE: std::sync::Once = std::sync::Once::new();
12847 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12848 }
12849 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12850 let (y0, y1) =
12851 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12852 let mut y0 = y0;
12853 let mut y1 = y1;
12854 if s0 != 1.0 {
12855 self.scale_inplace(&mut y0, s0, m * out_f)?;
12856 }
12857 if s1 != 1.0 {
12858 self.scale_inplace(&mut y1, s1, m * out_f)?;
12859 }
12860 Ok(Some((y0, y1)))
12861 }
12862
12863 #[allow(clippy::too_many_arguments)]
12868 pub fn qmatvec_batched_dual_raw(
12869 &self,
12870 b0: &CudaSlice<u8>,
12871 b1: &CudaSlice<u8>,
12872 aq: &CudaSlice<i8>,
12873 ad: &CudaSlice<f32>,
12874 m: usize,
12875 in_f: usize,
12876 out_f: usize,
12877 row_bytes: usize,
12878 rp: bool,
12879 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12880 const ROWS_PER_BLOCK: u32 = 4;
12881 let mcols = Self::batched_mcols(m);
12882 let tiny_rp1 = rp
12885 && mcols == 4
12886 && out_f <= 128
12887 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12888 let (name, rows_per_block) = if tiny_rp1 {
12889 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12890 } else {
12891 match (mcols, rp, m) {
12892 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12893 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12894 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12895 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12896 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12897 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12898 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12899 _ => {
12900 return Err(
12901 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12902 );
12903 }
12904 }
12905 };
12906 let f = self.func(name);
12907 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12908 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12909 let cfg = LaunchConfig {
12910 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12911 block_dim: (32, ROWS_PER_BLOCK, 1),
12912 shared_mem_bytes: 0,
12913 };
12914 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12915 let __s_b = self.gpu.stream();
12916 let mut b = __s_b.launch_builder(&f);
12917 b.arg(b0)
12918 .arg(b1)
12919 .arg(aq)
12920 .arg(ad)
12921 .arg(&mut y0)
12922 .arg(&mut y1)
12923 .arg(&inf)
12924 .arg(&outf)
12925 .arg(&mi)
12926 .arg(&rb);
12927 unsafe {
12928 b.launch(cfg)?;
12929 }
12930 Ok((y0, y1))
12931 }
12932
12933 pub fn matmul_pre_dual_noscale(
12945 &self,
12946 w0: &crate::model::GpuTensor,
12947 w1: &crate::model::GpuTensor,
12948 aq: &CudaSlice<i8>,
12949 ad: &CudaSlice<f32>,
12950 m: usize,
12951 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12952 {
12953 use crate::model::GpuTensor;
12954 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12955 return Ok(None);
12956 }
12957 if !self.mmvq_supports(QT_NVFP4) {
12967 return Ok(None);
12968 }
12969 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12970 if w1.in_features() != in_f || w1.out_features() != out_f {
12971 return Ok(None);
12972 }
12973 let no_mirror =
12986 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12987 if self.q8_ffn_fuse2_on()
12988 && no_mirror(w0)
12989 && no_mirror(w1)
12990 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12991 {
12992 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12993 return Ok(Some(((y0, 1.0), (y1, 1.0))));
12994 }
12995 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13005 let (y0, y1) =
13006 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13007 return Ok(Some(((y0, p0.3), (y1, p1.3))));
13008 }
13009 let (b0, q0, rb0, s0, rp0) = match w0 {
13010 GpuTensor::Quant {
13011 bytes,
13012 qtype,
13013 row_bytes,
13014 scale,
13015 rp,
13016 ..
13017 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13018 _ => return Ok(None),
13019 };
13020 let (b1, q1, rb1, s1, rp1) = match w1 {
13021 GpuTensor::Quant {
13022 bytes,
13023 qtype,
13024 row_bytes,
13025 scale,
13026 rp,
13027 ..
13028 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13029 _ => return Ok(None),
13030 };
13031 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13032 return Ok(None);
13033 }
13034 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13036 let rows_per_block = ROWS_PER_BLOCK * RPW;
13037 let f = self.func(if rp0 {
13038 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13039 } else {
13040 "qmatvec_nvfp4_mmvq_dual_mr2"
13041 });
13042 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13043 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13044 let cfg = LaunchConfig {
13045 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13046 block_dim: (32, ROWS_PER_BLOCK, 1),
13047 shared_mem_bytes: 0,
13048 };
13049 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13050 let one = 1.0f32;
13053 let __s_b = self.gpu.stream();
13054 let mut b = __s_b.launch_builder(&f);
13055 b.arg(b0)
13056 .arg(b1)
13057 .arg(aq)
13058 .arg(ad)
13059 .arg(&mut y0)
13060 .arg(&mut y1)
13061 .arg(&inf)
13062 .arg(&outf)
13063 .arg(&mi)
13064 .arg(&rb)
13065 .arg(&one)
13066 .arg(&one);
13067 unsafe {
13068 b.launch(cfg)?;
13069 }
13070 Ok(Some(((y0, s0), (y1, s1))))
13071 }
13072
13073 #[allow(clippy::too_many_arguments)]
13081 pub fn matmul_nvfp4_fused3(
13082 &self,
13083 w0: &crate::model::GpuTensor,
13084 w1: &crate::model::GpuTensor,
13085 w2: &crate::model::GpuTensor,
13086 aq: &CudaSlice<i8>,
13087 ad: &CudaSlice<f32>,
13088 m: usize,
13089 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13090 {
13091 use crate::model::GpuTensor;
13092 if !self.mmvq_supports(QT_NVFP4)
13099 || !self.uses_q8_1_fast(w0)
13100 || !self.uses_q8_1_fast(w1)
13101 || !self.uses_q8_1_fast(w2)
13102 {
13103 return Ok(None);
13104 }
13105 if (9..=16).contains(&m) {
13108 return Ok(
13109 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13110 Some(mut ys) => {
13111 let y2 = ys.pop().unwrap();
13112 let y1 = ys.pop().unwrap();
13113 let y0 = ys.pop().unwrap();
13114 Some((y0, y1, y2))
13115 }
13116 None => None,
13117 },
13118 );
13119 }
13120 if !(1..=8).contains(&m) {
13121 return Ok(None);
13122 }
13123 if m > 1 {
13124 let in_f = w0.in_features();
13125 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13126 || !self.batched_supports(QT_NVFP4)
13127 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13128 || (m > 4 && !Self::b8_enabled())
13129 || in_f % 512 != 0
13130 || in_f / 64 > 272
13131 {
13132 return Ok(None);
13133 }
13134 }
13135 let unpack = |w: &crate::model::GpuTensor| match w {
13136 GpuTensor::Quant {
13137 bytes,
13138 qtype,
13139 scale,
13140 rp,
13141 ..
13142 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13143 _ => None,
13144 };
13145 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13146 return Ok(None);
13147 };
13148 let in_f = w0.in_features();
13149 if w1.in_features() != in_f || w2.in_features() != in_f {
13150 return Ok(None);
13151 }
13152 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13153 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13155 let rows_pb = ROWS_PER_BLOCK * RPW;
13156 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13157 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13158 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13159 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13160 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13161 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13164 if m > 1 {
13165 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13167 return Ok(None);
13168 }
13169 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13170 let cfg = LaunchConfig {
13171 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13172 block_dim: (32, ROWS_PER_BLOCK, 1),
13173 shared_mem_bytes: 0,
13174 };
13175 let __s_b = self.gpu.stream();
13176 let mut b = __s_b.launch_builder(&f);
13177 b.arg(b0)
13178 .arg(b1)
13179 .arg(b2)
13180 .arg(aq)
13181 .arg(ad)
13182 .arg(&mut y0)
13183 .arg(&mut y1)
13184 .arg(&mut y2)
13185 .arg(&inf)
13186 .arg(&oi0)
13187 .arg(&oi1)
13188 .arg(&oi2)
13189 .arg(&mi);
13190 unsafe {
13191 b.launch(cfg)?;
13192 }
13193 return Ok(Some((y0, y1, y2)));
13194 }
13195 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13196 let cfg = LaunchConfig {
13197 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13198 block_dim: (32, ROWS_PER_BLOCK, 1),
13199 shared_mem_bytes: 0,
13200 };
13201 let __s_b = self.gpu.stream();
13202 let mut b = __s_b.launch_builder(&f);
13203 b.arg(b0)
13204 .arg(b1)
13205 .arg(b2)
13206 .arg(aq)
13207 .arg(ad)
13208 .arg(&mut y0)
13209 .arg(&mut y1)
13210 .arg(&mut y2)
13211 .arg(&inf)
13212 .arg(&oi0)
13213 .arg(&oi1)
13214 .arg(&oi2)
13215 .arg(&mi)
13216 .arg(&p0.1)
13217 .arg(&p1.1)
13218 .arg(&p2.1);
13219 unsafe {
13220 b.launch(cfg)?;
13221 }
13222 Ok(Some((y0, y1, y2)))
13223 }
13224
13225 pub fn matmul_nvfp4_fused2(
13234 &self,
13235 w0: &crate::model::GpuTensor,
13236 w1: &crate::model::GpuTensor,
13237 aq: &CudaSlice<i8>,
13238 ad: &CudaSlice<f32>,
13239 m: usize,
13240 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13241 use crate::model::GpuTensor;
13242 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13243 let off =
13244 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13245 if off
13248 || m != 1
13249 || !self.mmvq_supports(QT_NVFP4)
13250 || !self.uses_q8_1_fast(w0)
13251 || !self.uses_q8_1_fast(w1)
13252 {
13253 return Ok(None);
13254 }
13255 let unpack = |w: &crate::model::GpuTensor| match w {
13256 GpuTensor::Quant {
13257 bytes,
13258 qtype,
13259 scale,
13260 rp,
13261 ..
13262 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13263 _ => None,
13264 };
13265 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13266 return Ok(None);
13267 };
13268 let in_f = w0.in_features();
13269 if w1.in_features() != in_f {
13270 return Ok(None);
13271 }
13272 let (o0, o1) = (w0.out_features(), w1.out_features());
13273 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13275 let rows_pb = ROWS_PER_BLOCK * RPW;
13276 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13277 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13278 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13279 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13280 let cfg = LaunchConfig {
13281 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13282 block_dim: (32, ROWS_PER_BLOCK, 1),
13283 shared_mem_bytes: 0,
13284 };
13285 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13286 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13289 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13292 {
13293 use cudarc::driver::{DevicePtr, DevicePtrMut};
13294 let s = &self.gpu.stream();
13295 let (pw0, _g0) = b0.device_ptr(s);
13296 let (pw1, _g1) = b1.device_ptr(s);
13297 let (paq, _g2) = aq.device_ptr(s);
13298 let (pad, _g3) = ad.device_ptr(s);
13299 let (py0, _g4) = y0.device_ptr_mut(s);
13300 let (py1, _g5) = y1.device_ptr_mut(s);
13301 let (s0, s1) = (p0.1, p1.1);
13302 let mut ps = [
13303 &pw0 as *const _ as *mut std::ffi::c_void,
13304 &pw1 as *const _ as *mut _,
13305 &paq as *const _ as *mut _,
13306 &pad as *const _ as *mut _,
13307 &py0 as *const _ as *mut _,
13308 &py1 as *const _ as *mut _,
13309 &inf as *const _ as *mut _,
13310 &oi0 as *const _ as *mut _,
13311 &oi1 as *const _ as *mut _,
13312 &mi as *const _ as *mut _,
13313 &s0 as *const _ as *mut _,
13314 &s1 as *const _ as *mut _,
13315 ];
13316 unsafe {
13317 self.launch_pdl(
13318 "qmatvec_nvfp4_mmvq_fused2_rp",
13319 cfg.grid_dim,
13320 cfg.block_dim,
13321 &mut ps,
13322 )?;
13323 }
13324 }
13325 return Ok(Some((y0, y1)));
13326 }
13327 let __s_b = self.gpu.stream();
13328 let mut b = __s_b.launch_builder(&f);
13329 b.arg(b0)
13330 .arg(b1)
13331 .arg(aq)
13332 .arg(ad)
13333 .arg(&mut y0)
13334 .arg(&mut y1)
13335 .arg(&inf)
13336 .arg(&oi0)
13337 .arg(&oi1)
13338 .arg(&mi)
13339 .arg(&p0.1)
13340 .arg(&p1.1);
13341 unsafe {
13342 b.launch(cfg)?;
13343 }
13344 Ok(Some((y0, y1)))
13345 }
13346
13347 pub fn matmul_nvfp4_fused2_into(
13352 &self,
13353 w0: &crate::model::GpuTensor,
13354 w1: &crate::model::GpuTensor,
13355 aq: &CudaSlice<i8>,
13356 ad: &CudaSlice<f32>,
13357 y0: &mut CudaSlice<f32>,
13358 y1: &mut CudaSlice<f32>,
13359 ) -> Result<bool, Box<dyn std::error::Error>> {
13360 use crate::model::GpuTensor;
13361 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13362 let off =
13363 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13364 if off
13365 || !self.mmvq_supports(QT_NVFP4)
13366 || !self.uses_q8_1_fast(w0)
13367 || !self.uses_q8_1_fast(w1)
13368 {
13369 return Ok(false);
13370 }
13371 let unpack = |w: &crate::model::GpuTensor| match w {
13372 GpuTensor::Quant {
13373 bytes,
13374 qtype,
13375 scale,
13376 rp,
13377 ..
13378 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13379 _ => None,
13380 };
13381 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13382 return Ok(false);
13383 };
13384 let in_f = w0.in_features();
13385 if w1.in_features() != in_f {
13386 return Ok(false);
13387 }
13388 let (o0, o1) = (w0.out_features(), w1.out_features());
13389 if y0.len() < o0 || y1.len() < o1 {
13390 return Ok(false);
13391 }
13392 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13394 let rows_pb = ROWS_PER_BLOCK * RPW;
13395 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13396 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13397 let cfg = LaunchConfig {
13398 grid_dim: (nb(o0) + nb(o1), 1, 1),
13399 block_dim: (32, ROWS_PER_BLOCK, 1),
13400 shared_mem_bytes: 0,
13401 };
13402 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13403 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13406 let __s_b = self.gpu.stream();
13407 let mut b = __s_b.launch_builder(&f);
13408 b.arg(b0)
13409 .arg(b1)
13410 .arg(aq)
13411 .arg(ad)
13412 .arg(&mut *y0)
13413 .arg(&mut *y1)
13414 .arg(&inf)
13415 .arg(&oi0)
13416 .arg(&oi1)
13417 .arg(&mi)
13418 .arg(&p0.1)
13419 .arg(&p1.1);
13420 unsafe {
13421 b.launch(cfg)?;
13422 }
13423 Ok(true)
13424 }
13425
13426 #[allow(clippy::type_complexity)]
13431 pub fn matmul_nvfp4_fused4(
13432 &self,
13433 w0: &crate::model::GpuTensor,
13434 w1: &crate::model::GpuTensor,
13435 w2: &crate::model::GpuTensor,
13436 w3: &crate::model::GpuTensor,
13437 aq: &CudaSlice<i8>,
13438 ad: &CudaSlice<f32>,
13439 m: usize,
13440 ) -> Result<
13441 Option<(
13442 CudaSlice<f32>,
13443 CudaSlice<f32>,
13444 CudaSlice<f32>,
13445 CudaSlice<f32>,
13446 )>,
13447 Box<dyn std::error::Error>,
13448 > {
13449 use crate::model::GpuTensor;
13450 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13457 || !self.mmvq_supports(QT_NVFP4)
13458 || !self.uses_q8_1_fast(w0)
13459 || !self.uses_q8_1_fast(w1)
13460 || !self.uses_q8_1_fast(w2)
13461 || !self.uses_q8_1_fast(w3)
13462 {
13463 return Ok(None);
13464 }
13465 if (9..=16).contains(&m) {
13470 return Ok(
13471 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13472 Some(mut ys) => {
13473 let y3 = ys.pop().unwrap();
13474 let y2 = ys.pop().unwrap();
13475 let y1 = ys.pop().unwrap();
13476 let y0 = ys.pop().unwrap();
13477 Some((y0, y1, y2, y3))
13478 }
13479 None => None,
13480 },
13481 );
13482 }
13483 if !(1..=8).contains(&m) {
13484 return Ok(None);
13485 }
13486 if m > 1 {
13487 let in_f = w0.in_features();
13490 if !self.batched_supports(QT_NVFP4)
13491 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13492 || (m > 4 && !Self::b8_enabled())
13493 || in_f % 512 != 0
13494 || in_f / 64 > 272
13495 {
13496 return Ok(None);
13497 }
13498 }
13499 let unpack = |w: &crate::model::GpuTensor| match w {
13500 GpuTensor::Quant {
13501 bytes,
13502 qtype,
13503 scale,
13504 rp,
13505 ..
13506 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13507 _ => None,
13508 };
13509 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13510 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13511 else {
13512 return Ok(None);
13513 };
13514 let in_f = w0.in_features();
13515 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13516 return Ok(None);
13517 }
13518 let (o0, o1, o2, o3) = (
13519 w0.out_features(),
13520 w1.out_features(),
13521 w2.out_features(),
13522 w3.out_features(),
13523 );
13524 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13526 let rows_pb = ROWS_PER_BLOCK * RPW;
13527 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13528 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13529 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13530 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13531 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13532 let (inf, oi0, oi1, oi2, oi3, mi) = (
13533 in_f as i32,
13534 o0 as i32,
13535 o1 as i32,
13536 o2 as i32,
13537 o3 as i32,
13538 m as i32,
13539 );
13540 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13543 if m > 1 {
13544 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13547 return Ok(None);
13548 }
13549 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13550 let cfg = LaunchConfig {
13551 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13552 block_dim: (32, ROWS_PER_BLOCK, 1),
13553 shared_mem_bytes: 0,
13554 };
13555 let __s_b = self.gpu.stream();
13556 let mut b = __s_b.launch_builder(&f);
13557 b.arg(b0)
13558 .arg(b1)
13559 .arg(b2)
13560 .arg(b3)
13561 .arg(aq)
13562 .arg(ad)
13563 .arg(&mut y0)
13564 .arg(&mut y1)
13565 .arg(&mut y2)
13566 .arg(&mut y3)
13567 .arg(&inf)
13568 .arg(&oi0)
13569 .arg(&oi1)
13570 .arg(&oi2)
13571 .arg(&oi3)
13572 .arg(&mi);
13573 unsafe {
13574 b.launch(cfg)?;
13575 }
13576 return Ok(Some((y0, y1, y2, y3)));
13577 }
13578 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13579 let cfg = LaunchConfig {
13580 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13581 block_dim: (32, ROWS_PER_BLOCK, 1),
13582 shared_mem_bytes: 0,
13583 };
13584 let __s_b = self.gpu.stream();
13585 let mut b = __s_b.launch_builder(&f);
13586 b.arg(b0)
13587 .arg(b1)
13588 .arg(b2)
13589 .arg(b3)
13590 .arg(aq)
13591 .arg(ad)
13592 .arg(&mut y0)
13593 .arg(&mut y1)
13594 .arg(&mut y2)
13595 .arg(&mut y3)
13596 .arg(&inf)
13597 .arg(&oi0)
13598 .arg(&oi1)
13599 .arg(&oi2)
13600 .arg(&oi3)
13601 .arg(&mi)
13602 .arg(&p0.1)
13603 .arg(&p1.1)
13604 .arg(&p2.1)
13605 .arg(&p3.1);
13606 unsafe {
13607 b.launch(cfg)?;
13608 }
13609 Ok(Some((y0, y1, y2, y3)))
13610 }
13611
13612 pub fn matmul_q8_fused2(
13620 &self,
13621 w0: &crate::model::GpuTensor,
13622 w1: &crate::model::GpuTensor,
13623 aq: &CudaSlice<i8>,
13624 ad: &CudaSlice<f32>,
13625 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13626 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13632 return Ok(Some(self.e4m3_fused2_core(
13633 p0.0,
13634 p1.0,
13635 aq,
13636 ad,
13637 w0.in_features(),
13638 p0.1,
13639 p1.1,
13640 p0.2,
13641 p0.3,
13642 p1.3,
13643 )?));
13644 }
13645 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13646 return Ok(None);
13647 };
13648 Ok(Some(self.q8_fused2_core(
13649 p0.0,
13650 p1.0,
13651 aq,
13652 ad,
13653 w0.in_features(),
13654 p0.1,
13655 p1.1,
13656 p0.2,
13657 )?))
13658 }
13659
13660 #[allow(clippy::too_many_arguments)]
13661 fn q8_fused2_core(
13662 &self,
13663 b0: &CudaSlice<u8>,
13664 b1: &CudaSlice<u8>,
13665 aq: &CudaSlice<i8>,
13666 ad: &CudaSlice<f32>,
13667 in_f: usize,
13668 out0: usize,
13669 out1: usize,
13670 row_bytes: usize,
13671 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13672 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13674 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13675 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13676 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13677 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13678 let cfg = LaunchConfig {
13679 grid_dim: (nb0 + nb1, 1, 1),
13680 block_dim: (32, ROWS_PER_BLOCK, 1),
13681 shared_mem_bytes: 0,
13682 };
13683 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13684 let __s_b = self.gpu.stream();
13685 let mut b = __s_b.launch_builder(&f);
13686 b.arg(b0)
13687 .arg(b1)
13688 .arg(aq)
13689 .arg(ad)
13690 .arg(&mut y0)
13691 .arg(&mut y1)
13692 .arg(&inf)
13693 .arg(&o0)
13694 .arg(&o1)
13695 .arg(&rbl);
13696 unsafe {
13697 b.launch(cfg)?;
13698 }
13699 Ok((y0, y1))
13700 }
13701
13702 pub fn matmul_q8_fused2_x(
13708 &self,
13709 w0: &crate::model::GpuTensor,
13710 w1: &crate::model::GpuTensor,
13711 x: &CudaSlice<f32>,
13712 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13713 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13714 return Ok(None);
13715 }
13716 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13717 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13718 return Ok(Some(self.e4m3_fused2_core(
13719 p0.0,
13720 p1.0,
13721 &aq,
13722 &ad,
13723 w0.in_features(),
13724 p0.1,
13725 p1.1,
13726 p0.2,
13727 p0.3,
13728 p1.3,
13729 )?));
13730 }
13731 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13732 return Ok(None);
13733 };
13734 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13735 Ok(Some(self.q8_fused2_core(
13736 p0.0,
13737 p1.0,
13738 &aq,
13739 &ad,
13740 w0.in_features(),
13741 p0.1,
13742 p1.1,
13743 p0.2,
13744 )?))
13745 }
13746
13747 #[allow(clippy::too_many_arguments)]
13750 pub fn qmatvec_q8_fused2_raw(
13751 &self,
13752 b0: &CudaSlice<u8>,
13753 b1: &CudaSlice<u8>,
13754 x: &CudaSlice<f32>,
13755 in_f: usize,
13756 out0: usize,
13757 out1: usize,
13758 row_bytes: usize,
13759 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13760 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13761 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13762 }
13763
13764 pub fn matmul_q4_fused3(
13770 &self,
13771 w0: &crate::model::GpuTensor,
13772 w1: &crate::model::GpuTensor,
13773 w2: &crate::model::GpuTensor,
13774 aq: &CudaSlice<i8>,
13775 ad: &CudaSlice<f32>,
13776 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13777 {
13778 use crate::model::GpuTensor;
13779 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13780 match w {
13781 GpuTensor::Quant {
13782 qtype, row_bytes, ..
13783 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13784 _ => None,
13785 }
13786 };
13787 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13788 return Ok(None);
13789 };
13790 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13791 return Ok(None);
13792 }
13793 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13797 match w {
13798 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13799 Some(m) => (m, true),
13800 None => (bytes, *rp),
13801 },
13802 _ => unreachable!(),
13803 }
13804 }
13805 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13806 if rp0 != rp1 || rp1 != rp2 {
13807 return Ok(None);
13808 }
13809 let rp = rp0;
13810 let rpb: u32 = 4;
13811 let mr1 = rp && Self::q40_mr1_on();
13815 let nb = |o: usize| {
13816 if mr1 {
13817 (o as u32).div_ceil(rpb)
13818 } else {
13819 (o as u32).div_ceil(2).div_ceil(rpb)
13820 }
13821 };
13822 let grid = nb(o0) + nb(o1) + nb(o2);
13823 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13824 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13825 let mut y2 = self.alloc_uninit::<f32>(o2)?;
13826 let f = self.func(if mr1 {
13827 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13828 } else if rp {
13829 "qmatvec_q4_0_mmvq_fused3_rp"
13830 } else {
13831 "qmatvec_q4_0_mmvq_fused3"
13832 });
13833 let cfg = LaunchConfig {
13834 grid_dim: (grid, 1, 1),
13835 block_dim: (32, rpb, 1),
13836 shared_mem_bytes: 0,
13837 };
13838 let inf = w0.in_features() as i32;
13839 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13840 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13841 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13844 {
13845 use cudarc::driver::{DevicePtr, DevicePtrMut};
13846 let s = &self.gpu.stream();
13847 let (p0, _g0) = b0.device_ptr(s);
13848 let (p1, _g1) = b1.device_ptr(s);
13849 let (p2, _g2) = b2.device_ptr(s);
13850 let (paq, _g3) = aq.device_ptr(s);
13851 let (pad, _g4) = ad.device_ptr(s);
13852 let (py0, _g5) = y0.device_ptr_mut(s);
13853 let (py1, _g6) = y1.device_ptr_mut(s);
13854 let (py2, _g7) = y2.device_ptr_mut(s);
13855 let mut ps = [
13856 &p0 as *const _ as *mut std::ffi::c_void,
13857 &p1 as *const _ as *mut _,
13858 &p2 as *const _ as *mut _,
13859 &paq as *const _ as *mut _,
13860 &pad as *const _ as *mut _,
13861 &py0 as *const _ as *mut _,
13862 &py1 as *const _ as *mut _,
13863 &py2 as *const _ as *mut _,
13864 &inf as *const _ as *mut _,
13865 &oo0 as *const _ as *mut _,
13866 &oo1 as *const _ as *mut _,
13867 &oo2 as *const _ as *mut _,
13868 &r0 as *const _ as *mut _,
13869 &r1 as *const _ as *mut _,
13870 &r2 as *const _ as *mut _,
13871 ];
13872 unsafe {
13873 self.launch_pdl(
13874 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13875 (grid, 1, 1),
13876 (32, rpb, 1),
13877 &mut ps,
13878 )?;
13879 }
13880 }
13881 return Ok(Some((y0, y1, y2)));
13882 }
13883 let __s_b = self.gpu.stream();
13884 let mut b = __s_b.launch_builder(&f);
13885 b.arg(b0)
13886 .arg(b1)
13887 .arg(b2)
13888 .arg(aq)
13889 .arg(ad)
13890 .arg(&mut y0)
13891 .arg(&mut y1)
13892 .arg(&mut y2)
13893 .arg(&inf)
13894 .arg(&oo0)
13895 .arg(&oo1)
13896 .arg(&oo2)
13897 .arg(&r0)
13898 .arg(&r1)
13899 .arg(&r2);
13900 unsafe {
13901 b.launch(cfg)?;
13902 }
13903 Ok(Some((y0, y1, y2)))
13904 }
13905
13906 #[allow(clippy::too_many_arguments)]
13909 pub fn matmul_q4_fused3_into(
13910 &self,
13911 w0: &crate::model::GpuTensor,
13912 w1: &crate::model::GpuTensor,
13913 w2: &crate::model::GpuTensor,
13914 aq: &CudaSlice<i8>,
13915 ad: &CudaSlice<f32>,
13916 y0: &mut CudaSlice<f32>,
13917 y1: &mut CudaSlice<f32>,
13918 y2: &mut CudaSlice<f32>,
13919 ) -> Result<bool, Box<dyn std::error::Error>> {
13920 use crate::model::GpuTensor;
13921 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13922 match w {
13923 GpuTensor::Quant {
13924 qtype, row_bytes, ..
13925 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13926 _ => None,
13927 }
13928 };
13929 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13930 return Ok(false);
13931 };
13932 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13933 return Ok(false);
13934 }
13935 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13936 match w {
13937 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13938 Some(m) => (m, true),
13939 None => (bytes, *rp),
13940 },
13941 _ => unreachable!(),
13942 }
13943 }
13944 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13945 if rp0 != rp1 || rp1 != rp2 {
13946 return Ok(false);
13947 }
13948 let rp = rp0;
13949 let rpb: u32 = 4;
13950 let mr1 = rp && Self::q40_mr1_on();
13951 let nb = |o: usize| {
13952 if mr1 {
13953 (o as u32).div_ceil(rpb)
13954 } else {
13955 (o as u32).div_ceil(2).div_ceil(rpb)
13956 }
13957 };
13958 let grid = nb(o0) + nb(o1) + nb(o2);
13959 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13960 let f = self.func(if mr1 {
13961 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13962 } else if rp {
13963 "qmatvec_q4_0_mmvq_fused3_rp"
13964 } else {
13965 "qmatvec_q4_0_mmvq_fused3"
13966 });
13967 let cfg = LaunchConfig {
13968 grid_dim: (grid, 1, 1),
13969 block_dim: (32, rpb, 1),
13970 shared_mem_bytes: 0,
13971 };
13972 let inf = w0.in_features() as i32;
13973 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13974 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13975 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13977 use cudarc::driver::{DevicePtr, DevicePtrMut};
13978 let s = &self.gpu.stream();
13979 let (p0, _g0) = b0.device_ptr(s);
13980 let (p1, _g1) = b1.device_ptr(s);
13981 let (p2, _g2) = b2.device_ptr(s);
13982 let (paq, _g3) = aq.device_ptr(s);
13983 let (pad, _g4) = ad.device_ptr(s);
13984 let (py0, _g5) = y0.device_ptr_mut(s);
13985 let (py1, _g6) = y1.device_ptr_mut(s);
13986 let (py2, _g7) = y2.device_ptr_mut(s);
13987 let mut ps = [
13988 &p0 as *const _ as *mut std::ffi::c_void,
13989 &p1 as *const _ as *mut _,
13990 &p2 as *const _ as *mut _,
13991 &paq as *const _ as *mut _,
13992 &pad as *const _ as *mut _,
13993 &py0 as *const _ as *mut _,
13994 &py1 as *const _ as *mut _,
13995 &py2 as *const _ as *mut _,
13996 &inf as *const _ as *mut _,
13997 &oo0 as *const _ as *mut _,
13998 &oo1 as *const _ as *mut _,
13999 &oo2 as *const _ as *mut _,
14000 &r0 as *const _ as *mut _,
14001 &r1 as *const _ as *mut _,
14002 &r2 as *const _ as *mut _,
14003 ];
14004 unsafe {
14005 self.launch_pdl(
14006 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14007 (grid, 1, 1),
14008 (32, rpb, 1),
14009 &mut ps,
14010 )?;
14011 }
14012 return Ok(true);
14013 }
14014 let __s_b = self.gpu.stream();
14015 let mut b = __s_b.launch_builder(&f);
14016 b.arg(b0)
14017 .arg(b1)
14018 .arg(b2)
14019 .arg(aq)
14020 .arg(ad)
14021 .arg(&mut *y0)
14022 .arg(&mut *y1)
14023 .arg(&mut *y2)
14024 .arg(&inf)
14025 .arg(&oo0)
14026 .arg(&oo1)
14027 .arg(&oo2)
14028 .arg(&r0)
14029 .arg(&r1)
14030 .arg(&r2);
14031 unsafe {
14032 b.launch(cfg)?;
14033 }
14034 Ok(true)
14035 }
14036
14037 pub fn matmul_q4_fused2(
14039 &self,
14040 w0: &crate::model::GpuTensor,
14041 w1: &crate::model::GpuTensor,
14042 aq: &CudaSlice<i8>,
14043 ad: &CudaSlice<f32>,
14044 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14045 use crate::model::GpuTensor;
14046 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14047 match w {
14048 GpuTensor::Quant {
14049 qtype, row_bytes, ..
14050 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14051 _ => None,
14052 }
14053 };
14054 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14055 return Ok(None);
14056 };
14057 if w0.in_features() != w1.in_features() {
14058 return Ok(None);
14059 }
14060 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14062 match w {
14063 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14064 Some(m) => (m, true),
14065 None => (bytes, *rp),
14066 },
14067 _ => unreachable!(),
14068 }
14069 }
14070 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14071 if rp0 != rp1 {
14072 return Ok(None);
14073 }
14074 let rp = rp0;
14075 let rpb: u32 = 4;
14076 let mr1 = rp && Self::q40_mr1_on();
14078 let nb = |o: usize| {
14079 if mr1 {
14080 (o as u32).div_ceil(rpb)
14081 } else {
14082 (o as u32).div_ceil(2).div_ceil(rpb)
14083 }
14084 };
14085 let grid = nb(o0) + nb(o1);
14086 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14087 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14088 let f = self.func(if mr1 {
14089 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14090 } else if rp {
14091 "qmatvec_q4_0_mmvq_fused2_rp"
14092 } else {
14093 "qmatvec_q4_0_mmvq_fused2"
14094 });
14095 let cfg = LaunchConfig {
14096 grid_dim: (grid, 1, 1),
14097 block_dim: (32, rpb, 1),
14098 shared_mem_bytes: 0,
14099 };
14100 let inf = w0.in_features() as i32;
14101 let (oo0, oo1) = (o0 as i32, o1 as i32);
14102 let (r0, r1) = (rb0 as i64, rb1 as i64);
14103 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14105 {
14106 use cudarc::driver::{DevicePtr, DevicePtrMut};
14107 let s = &self.gpu.stream();
14108 let (p0, _g0) = b0.device_ptr(s);
14109 let (p1, _g1) = b1.device_ptr(s);
14110 let (paq, _g2) = aq.device_ptr(s);
14111 let (pad, _g3) = ad.device_ptr(s);
14112 let (py0, _g4) = y0.device_ptr_mut(s);
14113 let (py1, _g5) = y1.device_ptr_mut(s);
14114 let mut ps = [
14115 &p0 as *const _ as *mut std::ffi::c_void,
14116 &p1 as *const _ as *mut _,
14117 &paq as *const _ as *mut _,
14118 &pad as *const _ as *mut _,
14119 &py0 as *const _ as *mut _,
14120 &py1 as *const _ as *mut _,
14121 &inf as *const _ as *mut _,
14122 &oo0 as *const _ as *mut _,
14123 &oo1 as *const _ as *mut _,
14124 &r0 as *const _ as *mut _,
14125 &r1 as *const _ as *mut _,
14126 ];
14127 unsafe {
14128 self.launch_pdl(
14129 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14130 (grid, 1, 1),
14131 (32, rpb, 1),
14132 &mut ps,
14133 )?;
14134 }
14135 }
14136 return Ok(Some((y0, y1)));
14137 }
14138 let __s_b = self.gpu.stream();
14139 let mut b = __s_b.launch_builder(&f);
14140 b.arg(b0)
14141 .arg(b1)
14142 .arg(aq)
14143 .arg(ad)
14144 .arg(&mut y0)
14145 .arg(&mut y1)
14146 .arg(&inf)
14147 .arg(&oo0)
14148 .arg(&oo1)
14149 .arg(&r0)
14150 .arg(&r1);
14151 unsafe {
14152 b.launch(cfg)?;
14153 }
14154 Ok(Some((y0, y1)))
14155 }
14156
14157 pub fn matmul_q4_fused2_into(
14159 &self,
14160 w0: &crate::model::GpuTensor,
14161 w1: &crate::model::GpuTensor,
14162 aq: &CudaSlice<i8>,
14163 ad: &CudaSlice<f32>,
14164 y0: &mut CudaSlice<f32>,
14165 y1: &mut CudaSlice<f32>,
14166 ) -> Result<bool, Box<dyn std::error::Error>> {
14167 use crate::model::GpuTensor;
14168 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14169 match w {
14170 GpuTensor::Quant {
14171 qtype, row_bytes, ..
14172 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14173 _ => None,
14174 }
14175 };
14176 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14177 return Ok(false);
14178 };
14179 if w0.in_features() != w1.in_features() {
14180 return Ok(false);
14181 }
14182 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14183 match w {
14184 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14185 Some(m) => (m, true),
14186 None => (bytes, *rp),
14187 },
14188 _ => unreachable!(),
14189 }
14190 }
14191 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14192 if rp0 != rp1 {
14193 return Ok(false);
14194 }
14195 let rp = rp0;
14196 let rpb: u32 = 4;
14197 let mr1 = rp && Self::q40_mr1_on();
14198 let nb = |o: usize| {
14199 if mr1 {
14200 (o as u32).div_ceil(rpb)
14201 } else {
14202 (o as u32).div_ceil(2).div_ceil(rpb)
14203 }
14204 };
14205 let grid = nb(o0) + nb(o1);
14206 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14207 let f = self.func(if mr1 {
14208 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14209 } else if rp {
14210 "qmatvec_q4_0_mmvq_fused2_rp"
14211 } else {
14212 "qmatvec_q4_0_mmvq_fused2"
14213 });
14214 let cfg = LaunchConfig {
14215 grid_dim: (grid, 1, 1),
14216 block_dim: (32, rpb, 1),
14217 shared_mem_bytes: 0,
14218 };
14219 let inf = w0.in_features() as i32;
14220 let (oo0, oo1) = (o0 as i32, o1 as i32);
14221 let (r0, r1) = (rb0 as i64, rb1 as i64);
14222 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14224 use cudarc::driver::{DevicePtr, DevicePtrMut};
14225 let s = &self.gpu.stream();
14226 let (p0, _g0) = b0.device_ptr(s);
14227 let (p1, _g1) = b1.device_ptr(s);
14228 let (paq, _g2) = aq.device_ptr(s);
14229 let (pad, _g3) = ad.device_ptr(s);
14230 let (py0, _g4) = y0.device_ptr_mut(s);
14231 let (py1, _g5) = y1.device_ptr_mut(s);
14232 let mut ps = [
14233 &p0 as *const _ as *mut std::ffi::c_void,
14234 &p1 as *const _ as *mut _,
14235 &paq as *const _ as *mut _,
14236 &pad as *const _ as *mut _,
14237 &py0 as *const _ as *mut _,
14238 &py1 as *const _ as *mut _,
14239 &inf as *const _ as *mut _,
14240 &oo0 as *const _ as *mut _,
14241 &oo1 as *const _ as *mut _,
14242 &r0 as *const _ as *mut _,
14243 &r1 as *const _ as *mut _,
14244 ];
14245 unsafe {
14246 self.launch_pdl(
14247 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14248 (grid, 1, 1),
14249 (32, rpb, 1),
14250 &mut ps,
14251 )?;
14252 }
14253 return Ok(true);
14254 }
14255 let __s_b = self.gpu.stream();
14256 let mut b = __s_b.launch_builder(&f);
14257 b.arg(b0)
14258 .arg(b1)
14259 .arg(aq)
14260 .arg(ad)
14261 .arg(&mut *y0)
14262 .arg(&mut *y1)
14263 .arg(&inf)
14264 .arg(&oo0)
14265 .arg(&oo1)
14266 .arg(&r0)
14267 .arg(&r1);
14268 unsafe {
14269 b.launch(cfg)?;
14270 }
14271 Ok(true)
14272 }
14273
14274 pub fn matmul_q4_fused2_batched(
14279 &self,
14280 w0: &crate::model::GpuTensor,
14281 w1: &crate::model::GpuTensor,
14282 aq: &CudaSlice<i8>,
14283 ad: &CudaSlice<f32>,
14284 m: usize,
14285 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14286 use crate::model::GpuTensor;
14287 if m < 2 || m > 8 {
14288 return Ok(None);
14289 }
14290 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14291 match w {
14292 GpuTensor::Quant {
14293 qtype, row_bytes, ..
14294 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14295 _ => None,
14296 }
14297 };
14298 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14299 return Ok(None);
14300 };
14301 if w0.in_features() != w1.in_features() {
14302 return Ok(None);
14303 }
14304 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14305 match w {
14306 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14307 Some(mr) => (mr, true),
14308 None => (bytes, *rp),
14309 },
14310 _ => unreachable!(),
14311 }
14312 }
14313 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14314 if !rp0 || !rp1 {
14315 return Ok(None);
14316 }
14317 let mcols = Self::batched_mcols(m);
14318 let rpb: u32 = 4;
14319 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14320 let grid = nb(o0) + nb(o1);
14321 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14322 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14323 let f = self.func(match mcols {
14324 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14325 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14326 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14327 });
14328 let cfg = LaunchConfig {
14329 grid_dim: (grid, 1, 1),
14330 block_dim: (32, rpb, 1),
14331 shared_mem_bytes: 0,
14332 };
14333 let inf = w0.in_features() as i32;
14334 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14335 let rb = rb0 as i64;
14336 let __s_b = self.gpu.stream();
14337 let mut b = __s_b.launch_builder(&f);
14338 b.arg(b0)
14339 .arg(b1)
14340 .arg(aq)
14341 .arg(ad)
14342 .arg(&mut y0)
14343 .arg(&mut y1)
14344 .arg(&inf)
14345 .arg(&oo0)
14346 .arg(&oo1)
14347 .arg(&mi)
14348 .arg(&rb);
14349 unsafe {
14350 b.launch(cfg)?;
14351 }
14352 Ok(Some((y0, y1)))
14353 }
14354
14355 #[allow(clippy::too_many_arguments)]
14358 pub fn matmul_q4_fused3_batched(
14359 &self,
14360 w0: &crate::model::GpuTensor,
14361 w1: &crate::model::GpuTensor,
14362 w2: &crate::model::GpuTensor,
14363 aq: &CudaSlice<i8>,
14364 ad: &CudaSlice<f32>,
14365 m: usize,
14366 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14367 {
14368 use crate::model::GpuTensor;
14369 if m < 2 || m > 8 {
14370 return Ok(None);
14371 }
14372 let q4 = |w: &GpuTensor| -> Option<usize> {
14373 match w {
14374 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14375 _ => None,
14376 }
14377 };
14378 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14379 return Ok(None);
14380 };
14381 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14382 return Ok(None);
14383 }
14384 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14385 match w {
14386 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14387 Some(mr) => (mr, true),
14388 None => (bytes, *rp),
14389 },
14390 _ => unreachable!(),
14391 }
14392 }
14393 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14394 if !rp0 || !rp1 || !rp2 {
14395 return Ok(None);
14396 }
14397 let mcols = Self::batched_mcols(m);
14398 let rpb: u32 = 4;
14399 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14400 let grid = nb(o0) + nb(o1) + nb(o2);
14401 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14402 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14403 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14404 let f = self.func(match mcols {
14405 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14406 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14407 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14408 });
14409 let cfg = LaunchConfig {
14410 grid_dim: (grid, 1, 1),
14411 block_dim: (32, rpb, 1),
14412 shared_mem_bytes: 0,
14413 };
14414 let inf = w0.in_features() as i32;
14415 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14416 let rb = 0i64;
14417 let __s_b = self.gpu.stream();
14418 let mut b = __s_b.launch_builder(&f);
14419 b.arg(b0)
14420 .arg(b1)
14421 .arg(b2)
14422 .arg(aq)
14423 .arg(ad)
14424 .arg(&mut y0)
14425 .arg(&mut y1)
14426 .arg(&mut y2)
14427 .arg(&inf)
14428 .arg(&oo0)
14429 .arg(&oo1)
14430 .arg(&oo2)
14431 .arg(&mi)
14432 .arg(&rb);
14433 unsafe {
14434 b.launch(cfg)?;
14435 }
14436 Ok(Some((y0, y1, y2)))
14437 }
14438
14439 pub fn matmul_q8_fused3(
14440 &self,
14441 w0: &crate::model::GpuTensor,
14442 w1: &crate::model::GpuTensor,
14443 w2: &crate::model::GpuTensor,
14444 aq: &CudaSlice<i8>,
14445 ad: &CudaSlice<f32>,
14446 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14447 {
14448 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14451 return Ok(Some(self.e4m3_fused3_core(
14452 p0.0,
14453 p1.0,
14454 p2.0,
14455 aq,
14456 ad,
14457 w0.in_features(),
14458 p0.1,
14459 p1.1,
14460 p2.1,
14461 p0.2,
14462 p0.3,
14463 p1.3,
14464 p2.3,
14465 )?));
14466 }
14467 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14468 return Ok(None);
14469 };
14470 Ok(Some(self.q8_fused3_core(
14471 p0.0,
14472 p1.0,
14473 p2.0,
14474 aq,
14475 ad,
14476 w0.in_features(),
14477 p0.1,
14478 p1.1,
14479 p2.1,
14480 p0.2,
14481 )?))
14482 }
14483
14484 #[allow(clippy::too_many_arguments)]
14485 fn q8_fused3_core(
14486 &self,
14487 b0: &CudaSlice<u8>,
14488 b1: &CudaSlice<u8>,
14489 b2: &CudaSlice<u8>,
14490 aq: &CudaSlice<i8>,
14491 ad: &CudaSlice<f32>,
14492 in_f: usize,
14493 out0: usize,
14494 out1: usize,
14495 out2: usize,
14496 row_bytes: usize,
14497 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14498 const ROWS_PER_BLOCK: u32 = 4;
14499 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14500 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14501 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14502 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14503 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14504 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14505 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14506 let cfg = LaunchConfig {
14507 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14508 block_dim: (32, ROWS_PER_BLOCK, 1),
14509 shared_mem_bytes: 0,
14510 };
14511 let (inf, o0, o1, o2, rbl) = (
14512 in_f as i32,
14513 out0 as i32,
14514 out1 as i32,
14515 out2 as i32,
14516 row_bytes as i64,
14517 );
14518 let __s_b = self.gpu.stream();
14519 let mut b = __s_b.launch_builder(&f);
14520 b.arg(b0)
14521 .arg(b1)
14522 .arg(b2)
14523 .arg(aq)
14524 .arg(ad)
14525 .arg(&mut y0)
14526 .arg(&mut y1)
14527 .arg(&mut y2)
14528 .arg(&inf)
14529 .arg(&o0)
14530 .arg(&o1)
14531 .arg(&o2)
14532 .arg(&rbl);
14533 unsafe {
14534 b.launch(cfg)?;
14535 }
14536 Ok((y0, y1, y2))
14537 }
14538
14539 #[allow(clippy::too_many_arguments)]
14541 pub fn qmatvec_q8_fused3_raw(
14542 &self,
14543 b0: &CudaSlice<u8>,
14544 b1: &CudaSlice<u8>,
14545 b2: &CudaSlice<u8>,
14546 x: &CudaSlice<f32>,
14547 in_f: usize,
14548 out0: usize,
14549 out1: usize,
14550 out2: usize,
14551 row_bytes: usize,
14552 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14553 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14554 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14555 }
14556
14557 pub fn matmul_q8_fused2_t(
14568 &self,
14569 w0: &crate::model::GpuTensor,
14570 w1: &crate::model::GpuTensor,
14571 aq: &CudaSlice<i8>,
14572 ad: &CudaSlice<f32>,
14573 m: usize,
14574 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14575 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14579 return Ok(None);
14580 }
14581 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14584 if m > 4 && !Self::b8_enabled() {
14585 return Ok(None);
14586 }
14587 return Ok(Some(self.e4m3_fused2_t_core(
14588 p0.0,
14589 p1.0,
14590 aq,
14591 ad,
14592 m,
14593 w0.in_features(),
14594 p0.1,
14595 p1.1,
14596 p0.2,
14597 p0.3,
14598 p1.3,
14599 )?));
14600 }
14601 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14602 return Ok(None);
14603 };
14604 Ok(Some(self.q8_fused2_t_core(
14605 p0.0,
14606 p1.0,
14607 aq,
14608 ad,
14609 m,
14610 w0.in_features(),
14611 p0.1,
14612 p1.1,
14613 p0.2,
14614 )?))
14615 }
14616
14617 #[allow(clippy::too_many_arguments)]
14618 fn q8_fused2_t_core(
14619 &self,
14620 b0: &CudaSlice<u8>,
14621 b1: &CudaSlice<u8>,
14622 aq: &CudaSlice<i8>,
14623 ad: &CudaSlice<f32>,
14624 m: usize,
14625 in_f: usize,
14626 out0: usize,
14627 out1: usize,
14628 row_bytes: usize,
14629 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14630 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14632 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14633 let f = self.func(match Self::batched_mcols(m) {
14634 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14635 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14636 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14638 });
14639 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14640 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14641 let cfg = LaunchConfig {
14642 grid_dim: (nb0 + nb1, 1, 1),
14643 block_dim: (32, ROWS_PER_BLOCK, 1),
14644 shared_mem_bytes: 0,
14645 };
14646 let (inf, o0, o1, mi, rbl) = (
14647 in_f as i32,
14648 out0 as i32,
14649 out1 as i32,
14650 m as i32,
14651 row_bytes as i64,
14652 );
14653 let __s_b = self.gpu.stream();
14654 let mut b = __s_b.launch_builder(&f);
14655 b.arg(b0)
14656 .arg(b1)
14657 .arg(aq)
14658 .arg(ad)
14659 .arg(&mut y0)
14660 .arg(&mut y1)
14661 .arg(&inf)
14662 .arg(&o0)
14663 .arg(&o1)
14664 .arg(&mi)
14665 .arg(&rbl);
14666 unsafe {
14667 b.launch(cfg)?;
14668 }
14669 Ok((y0, y1))
14670 }
14671
14672 #[allow(clippy::too_many_arguments)]
14675 pub fn qmatvec_q8_fused2_t_raw(
14676 &self,
14677 b0: &CudaSlice<u8>,
14678 b1: &CudaSlice<u8>,
14679 x: &CudaSlice<f32>,
14680 m: usize,
14681 in_f: usize,
14682 out0: usize,
14683 out1: usize,
14684 row_bytes: usize,
14685 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14686 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14687 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14688 }
14689
14690 #[allow(clippy::too_many_arguments)]
14693 pub fn matmul_q8_fused3_t(
14694 &self,
14695 w0: &crate::model::GpuTensor,
14696 w1: &crate::model::GpuTensor,
14697 w2: &crate::model::GpuTensor,
14698 aq: &CudaSlice<i8>,
14699 ad: &CudaSlice<f32>,
14700 m: usize,
14701 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14702 {
14703 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14704 return Ok(None);
14705 }
14706 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14707 return Ok(Some(self.e4m3_fused3_t_core(
14708 p0.0,
14709 p1.0,
14710 p2.0,
14711 aq,
14712 ad,
14713 m,
14714 w0.in_features(),
14715 p0.1,
14716 p1.1,
14717 p2.1,
14718 p0.2,
14719 p0.3,
14720 p1.3,
14721 p2.3,
14722 )?));
14723 }
14724 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14725 return Ok(None);
14726 };
14727 Ok(Some(self.q8_fused3_t_core(
14728 p0.0,
14729 p1.0,
14730 p2.0,
14731 aq,
14732 ad,
14733 m,
14734 w0.in_features(),
14735 p0.1,
14736 p1.1,
14737 p2.1,
14738 p0.2,
14739 )?))
14740 }
14741
14742 #[allow(clippy::too_many_arguments)]
14743 fn q8_fused3_t_core(
14744 &self,
14745 b0: &CudaSlice<u8>,
14746 b1: &CudaSlice<u8>,
14747 b2: &CudaSlice<u8>,
14748 aq: &CudaSlice<i8>,
14749 ad: &CudaSlice<f32>,
14750 m: usize,
14751 in_f: usize,
14752 out0: usize,
14753 out1: usize,
14754 out2: usize,
14755 row_bytes: usize,
14756 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14757 const ROWS_PER_BLOCK: u32 = 4;
14758 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14759 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14760 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14761 let f = self.func(if Self::batched_mcols(m) == 2 {
14762 "qmatvec_q8_0_mmvq_fused3_b2"
14763 } else {
14764 "qmatvec_q8_0_mmvq_fused3_b4"
14765 });
14766 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14767 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14768 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14769 let cfg = LaunchConfig {
14770 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14771 block_dim: (32, ROWS_PER_BLOCK, 1),
14772 shared_mem_bytes: 0,
14773 };
14774 let (inf, o0, o1, o2, mi, rbl) = (
14775 in_f as i32,
14776 out0 as i32,
14777 out1 as i32,
14778 out2 as i32,
14779 m as i32,
14780 row_bytes as i64,
14781 );
14782 let __s_b = self.gpu.stream();
14783 let mut b = __s_b.launch_builder(&f);
14784 b.arg(b0)
14785 .arg(b1)
14786 .arg(b2)
14787 .arg(aq)
14788 .arg(ad)
14789 .arg(&mut y0)
14790 .arg(&mut y1)
14791 .arg(&mut y2)
14792 .arg(&inf)
14793 .arg(&o0)
14794 .arg(&o1)
14795 .arg(&o2)
14796 .arg(&mi)
14797 .arg(&rbl);
14798 unsafe {
14799 b.launch(cfg)?;
14800 }
14801 Ok((y0, y1, y2))
14802 }
14803
14804 #[allow(clippy::too_many_arguments)]
14806 pub fn qmatvec_q8_fused3_t_raw(
14807 &self,
14808 b0: &CudaSlice<u8>,
14809 b1: &CudaSlice<u8>,
14810 b2: &CudaSlice<u8>,
14811 x: &CudaSlice<f32>,
14812 m: usize,
14813 in_f: usize,
14814 out0: usize,
14815 out1: usize,
14816 out2: usize,
14817 row_bytes: usize,
14818 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14819 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14820 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14821 }
14822
14823 pub fn q8_ffn_fuse2_on(&self) -> bool {
14827 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14828 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14829 }
14830
14831 #[allow(clippy::type_complexity)]
14837 fn q8_fused_params<'w, const N: usize>(
14838 &self,
14839 ws: &[&'w crate::model::GpuTensor; N],
14840 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14841 use crate::model::GpuTensor;
14842 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14843 return None;
14844 }
14845 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14846 return None;
14847 }
14848 let in_f = ws[0].in_features();
14849 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14850 for (i, w) in ws.iter().enumerate() {
14851 match w {
14852 GpuTensor::Quant {
14853 bytes,
14854 qtype,
14855 row_bytes,
14856 scale,
14857 ..
14858 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14859 out[i] = Some((bytes, w.out_features(), *row_bytes))
14860 }
14861 _ => return None,
14862 }
14863 }
14864 Some(out.map(|o| o.unwrap()))
14865 }
14866
14867 pub fn e4m3_dual_on(&self) -> bool {
14870 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14871 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14872 }
14873
14874 #[allow(clippy::type_complexity)]
14886 fn e4m3_fused_params<'w, const N: usize>(
14887 &self,
14888 ws: &[&'w crate::model::GpuTensor; N],
14889 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14890 use crate::model::GpuTensor;
14891 if !self.e4m3_dual_on() {
14892 return None;
14893 }
14894 let in_f = ws[0].in_features();
14895 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14896 for (i, w) in ws.iter().enumerate() {
14897 match w {
14898 GpuTensor::Quant {
14899 bytes,
14900 qtype,
14901 row_bytes,
14902 scale,
14903 rp,
14904 rp4,
14905 ..
14906 } if *qtype == QT_F8_E4M3
14907 && w.in_features() == in_f
14908 && *row_bytes == in_f
14909 && !*rp
14910 && rp4.is_none() =>
14911 {
14912 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14913 }
14914 _ => return None,
14915 }
14916 }
14917 Some(out.map(|o| o.unwrap()))
14918 }
14919
14920 #[allow(clippy::too_many_arguments)]
14924 fn e4m3_fused2_core(
14925 &self,
14926 b0: &CudaSlice<u8>,
14927 b1: &CudaSlice<u8>,
14928 aq: &CudaSlice<i8>,
14929 ad: &CudaSlice<f32>,
14930 in_f: usize,
14931 out0: usize,
14932 out1: usize,
14933 row_bytes: usize,
14934 ws0: f32,
14935 ws1: f32,
14936 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14937 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14939 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14940 let f = self.func("qmatvec_e4m3_mmvq_fused2");
14941 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14942 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14943 let cfg = LaunchConfig {
14944 grid_dim: (nb0 + nb1, 1, 1),
14945 block_dim: (32, ROWS_PER_BLOCK, 1),
14946 shared_mem_bytes: 0,
14947 };
14948 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14949 let __s_b = self.gpu.stream();
14950 let mut b = __s_b.launch_builder(&f);
14951 b.arg(b0)
14952 .arg(b1)
14953 .arg(aq)
14954 .arg(ad)
14955 .arg(&mut y0)
14956 .arg(&mut y1)
14957 .arg(&inf)
14958 .arg(&o0)
14959 .arg(&o1)
14960 .arg(&rbl)
14961 .arg(&ws0)
14962 .arg(&ws1);
14963 unsafe {
14964 b.launch(cfg)?;
14965 }
14966 Ok((y0, y1))
14967 }
14968
14969 #[allow(clippy::too_many_arguments)]
14971 fn e4m3_fused3_core(
14972 &self,
14973 b0: &CudaSlice<u8>,
14974 b1: &CudaSlice<u8>,
14975 b2: &CudaSlice<u8>,
14976 aq: &CudaSlice<i8>,
14977 ad: &CudaSlice<f32>,
14978 in_f: usize,
14979 out0: usize,
14980 out1: usize,
14981 out2: usize,
14982 row_bytes: usize,
14983 ws0: f32,
14984 ws1: f32,
14985 ws2: f32,
14986 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14987 const ROWS_PER_BLOCK: u32 = 4;
14988 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14989 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14990 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14991 let f = self.func("qmatvec_e4m3_mmvq_fused3");
14992 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14993 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14994 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14995 let cfg = LaunchConfig {
14996 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14997 block_dim: (32, ROWS_PER_BLOCK, 1),
14998 shared_mem_bytes: 0,
14999 };
15000 let (inf, o0, o1, o2, rbl) = (
15001 in_f as i32,
15002 out0 as i32,
15003 out1 as i32,
15004 out2 as i32,
15005 row_bytes as i64,
15006 );
15007 let __s_b = self.gpu.stream();
15008 let mut b = __s_b.launch_builder(&f);
15009 b.arg(b0)
15010 .arg(b1)
15011 .arg(b2)
15012 .arg(aq)
15013 .arg(ad)
15014 .arg(&mut y0)
15015 .arg(&mut y1)
15016 .arg(&mut y2)
15017 .arg(&inf)
15018 .arg(&o0)
15019 .arg(&o1)
15020 .arg(&o2)
15021 .arg(&rbl)
15022 .arg(&ws0)
15023 .arg(&ws1)
15024 .arg(&ws2);
15025 unsafe {
15026 b.launch(cfg)?;
15027 }
15028 Ok((y0, y1, y2))
15029 }
15030
15031 #[allow(clippy::too_many_arguments)]
15035 fn e4m3_fused2_t_core(
15036 &self,
15037 b0: &CudaSlice<u8>,
15038 b1: &CudaSlice<u8>,
15039 aq: &CudaSlice<i8>,
15040 ad: &CudaSlice<f32>,
15041 m: usize,
15042 in_f: usize,
15043 out0: usize,
15044 out1: usize,
15045 row_bytes: usize,
15046 ws0: f32,
15047 ws1: f32,
15048 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15049 const ROWS_PER_BLOCK: u32 = 4;
15050 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15051 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15052 let f = self.func(match Self::batched_mcols(m) {
15053 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15054 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15055 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15056 });
15057 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15058 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15059 let cfg = LaunchConfig {
15060 grid_dim: (nb0 + nb1, 1, 1),
15061 block_dim: (32, ROWS_PER_BLOCK, 1),
15062 shared_mem_bytes: 0,
15063 };
15064 let (inf, o0, o1, mi, rbl) = (
15065 in_f as i32,
15066 out0 as i32,
15067 out1 as i32,
15068 m as i32,
15069 row_bytes as i64,
15070 );
15071 let __s_b = self.gpu.stream();
15072 let mut b = __s_b.launch_builder(&f);
15073 b.arg(b0)
15074 .arg(b1)
15075 .arg(aq)
15076 .arg(ad)
15077 .arg(&mut y0)
15078 .arg(&mut y1)
15079 .arg(&inf)
15080 .arg(&o0)
15081 .arg(&o1)
15082 .arg(&mi)
15083 .arg(&rbl);
15084 unsafe {
15085 b.launch(cfg)?;
15086 }
15087 if ws0 != 1.0 {
15088 self.scale_inplace(&mut y0, ws0, m * out0)?;
15089 }
15090 if ws1 != 1.0 {
15091 self.scale_inplace(&mut y1, ws1, m * out1)?;
15092 }
15093 Ok((y0, y1))
15094 }
15095
15096 #[allow(clippy::too_many_arguments)]
15098 fn e4m3_fused3_t_core(
15099 &self,
15100 b0: &CudaSlice<u8>,
15101 b1: &CudaSlice<u8>,
15102 b2: &CudaSlice<u8>,
15103 aq: &CudaSlice<i8>,
15104 ad: &CudaSlice<f32>,
15105 m: usize,
15106 in_f: usize,
15107 out0: usize,
15108 out1: usize,
15109 out2: usize,
15110 row_bytes: usize,
15111 ws0: f32,
15112 ws1: f32,
15113 ws2: f32,
15114 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15115 const ROWS_PER_BLOCK: u32 = 4;
15116 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15117 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15118 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15119 let f = self.func(if Self::batched_mcols(m) == 2 {
15120 "qmatvec_e4m3_mmvq_fused3_b2"
15121 } else {
15122 "qmatvec_e4m3_mmvq_fused3_b4"
15123 });
15124 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15125 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15126 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15127 let cfg = LaunchConfig {
15128 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15129 block_dim: (32, ROWS_PER_BLOCK, 1),
15130 shared_mem_bytes: 0,
15131 };
15132 let (inf, o0, o1, o2, mi, rbl) = (
15133 in_f as i32,
15134 out0 as i32,
15135 out1 as i32,
15136 out2 as i32,
15137 m as i32,
15138 row_bytes as i64,
15139 );
15140 let __s_b = self.gpu.stream();
15141 let mut b = __s_b.launch_builder(&f);
15142 b.arg(b0)
15143 .arg(b1)
15144 .arg(b2)
15145 .arg(aq)
15146 .arg(ad)
15147 .arg(&mut y0)
15148 .arg(&mut y1)
15149 .arg(&mut y2)
15150 .arg(&inf)
15151 .arg(&o0)
15152 .arg(&o1)
15153 .arg(&o2)
15154 .arg(&mi)
15155 .arg(&rbl);
15156 unsafe {
15157 b.launch(cfg)?;
15158 }
15159 if ws0 != 1.0 {
15160 self.scale_inplace(&mut y0, ws0, m * out0)?;
15161 }
15162 if ws1 != 1.0 {
15163 self.scale_inplace(&mut y1, ws1, m * out1)?;
15164 }
15165 if ws2 != 1.0 {
15166 self.scale_inplace(&mut y2, ws2, m * out2)?;
15167 }
15168 Ok((y0, y1, y2))
15169 }
15170
15171 pub fn qmatvec_e4m3_blk_mmvq(
15181 &self,
15182 bytes: &CudaSlice<u8>,
15183 aq: &CudaSlice<i8>,
15184 ad: &CudaSlice<f32>,
15185 scales: &CudaSlice<f32>,
15186 m: usize,
15187 in_f: usize,
15188 out_f: usize,
15189 row_bytes: usize,
15190 scale_cols: usize,
15191 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15192 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15194 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15195 )?;
15196 Ok(y)
15197 }
15198
15199 #[allow(clippy::too_many_arguments)]
15201 pub fn qmatvec_e4m3_blk_mmvq_into(
15202 &self,
15203 bytes: &CudaSlice<u8>,
15204 aq: &CudaSlice<i8>,
15205 ad: &CudaSlice<f32>,
15206 scales: &CudaSlice<f32>,
15207 m: usize,
15208 in_f: usize,
15209 out_f: usize,
15210 row_bytes: usize,
15211 scale_cols: usize,
15212 y: &mut CudaSlice<f32>,
15213 ) -> Result<(), Box<dyn std::error::Error>> {
15214 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15216 let cfg = LaunchConfig {
15217 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15218 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15221 let (inf, outf, mi, rb, sc) = (
15222 in_f as i32,
15223 out_f as i32,
15224 m as i32,
15225 row_bytes as i64,
15226 scale_cols as i32,
15227 );
15228 let __s_b = self.gpu.stream();
15229 let mut b = __s_b.launch_builder(&f);
15230 b.arg(bytes)
15231 .arg(aq)
15232 .arg(ad)
15233 .arg(scales)
15234 .arg(&mut *y)
15235 .arg(&inf)
15236 .arg(&outf)
15237 .arg(&mi)
15238 .arg(&rb)
15239 .arg(&sc);
15240 unsafe {
15241 b.launch(cfg)?;
15242 }
15243 Ok(())
15244 }
15245
15246 #[allow(clippy::too_many_arguments)]
15252 pub fn qmatvec_e4m3_blk_mmvq_batched(
15253 &self,
15254 bytes: &CudaSlice<u8>,
15255 aq: &CudaSlice<i8>,
15256 ad: &CudaSlice<f32>,
15257 scales: &CudaSlice<f32>,
15258 m: usize,
15259 in_f: usize,
15260 out_f: usize,
15261 row_bytes: usize,
15262 scale_cols: usize,
15263 mcols: usize,
15264 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15265 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15267 let name = match mcols {
15268 2 => "qmatvec_e4m3_blk_mmvq_b2",
15269 4 => "qmatvec_e4m3_blk_mmvq_b4",
15270 8 => "qmatvec_e4m3_blk_mmvq_b8",
15271 16 => "qmatvec_e4m3_blk_mmvq_b16",
15272 _ => {
15273 return Err(
15274 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15275 );
15276 }
15277 };
15278 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15279 let f = self.func(name);
15280 let cfg = LaunchConfig {
15281 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15282 block_dim: (32, ROWS_PER_BLOCK, 1),
15283 shared_mem_bytes: 0,
15284 };
15285 let (inf, outf, mi, rb, sc) = (
15286 in_f as i32,
15287 out_f as i32,
15288 m as i32,
15289 row_bytes as i64,
15290 scale_cols as i32,
15291 );
15292 let __s_b = self.gpu.stream();
15293 let mut b = __s_b.launch_builder(&f);
15294 b.arg(bytes)
15295 .arg(aq)
15296 .arg(ad)
15297 .arg(scales)
15298 .arg(&mut y)
15299 .arg(&inf)
15300 .arg(&outf)
15301 .arg(&mi)
15302 .arg(&rb)
15303 .arg(&sc);
15304 unsafe {
15305 b.launch(cfg)?;
15306 }
15307 Ok(y)
15308 }
15309
15310 #[allow(clippy::too_many_arguments)]
15313 pub fn qmatvec_e4m3_blk_batched_raw(
15314 &self,
15315 bytes: &CudaSlice<u8>,
15316 x: &CudaSlice<f32>,
15317 scales: &CudaSlice<f32>,
15318 m: usize,
15319 in_f: usize,
15320 out_f: usize,
15321 row_bytes: usize,
15322 scale_cols: usize,
15323 mcols: usize,
15324 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15325 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15326 self.qmatvec_e4m3_blk_mmvq_batched(
15327 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15328 )
15329 }
15330
15331 #[allow(clippy::too_many_arguments)]
15334 pub fn qmatvec_e4m3_blk_mmvq_raw(
15335 &self,
15336 bytes: &CudaSlice<u8>,
15337 x: &CudaSlice<f32>,
15338 scales: &CudaSlice<f32>,
15339 m: usize,
15340 in_f: usize,
15341 out_f: usize,
15342 row_bytes: usize,
15343 scale_cols: usize,
15344 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15345 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15346 self.qmatvec_e4m3_blk_mmvq(
15347 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15348 )
15349 }
15350
15351 #[allow(clippy::too_many_arguments)]
15354 pub fn qmatvec_e4m3_fused2_raw(
15355 &self,
15356 b0: &CudaSlice<u8>,
15357 b1: &CudaSlice<u8>,
15358 x: &CudaSlice<f32>,
15359 in_f: usize,
15360 out0: usize,
15361 out1: usize,
15362 row_bytes: usize,
15363 ws0: f32,
15364 ws1: f32,
15365 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15366 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15367 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15368 }
15369
15370 #[allow(clippy::too_many_arguments)]
15371 pub fn qmatvec_e4m3_fused3_raw(
15372 &self,
15373 b0: &CudaSlice<u8>,
15374 b1: &CudaSlice<u8>,
15375 b2: &CudaSlice<u8>,
15376 x: &CudaSlice<f32>,
15377 in_f: usize,
15378 out0: usize,
15379 out1: usize,
15380 out2: usize,
15381 row_bytes: usize,
15382 ws0: f32,
15383 ws1: f32,
15384 ws2: f32,
15385 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15386 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15387 self.e4m3_fused3_core(
15388 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15389 )
15390 }
15391
15392 #[allow(clippy::too_many_arguments)]
15393 pub fn qmatvec_e4m3_fused2_t_raw(
15394 &self,
15395 b0: &CudaSlice<u8>,
15396 b1: &CudaSlice<u8>,
15397 x: &CudaSlice<f32>,
15398 m: usize,
15399 in_f: usize,
15400 out0: usize,
15401 out1: usize,
15402 row_bytes: usize,
15403 ws0: f32,
15404 ws1: f32,
15405 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15406 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15407 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15408 }
15409
15410 #[allow(clippy::too_many_arguments)]
15411 pub fn qmatvec_e4m3_fused3_t_raw(
15412 &self,
15413 b0: &CudaSlice<u8>,
15414 b1: &CudaSlice<u8>,
15415 b2: &CudaSlice<u8>,
15416 x: &CudaSlice<f32>,
15417 m: usize,
15418 in_f: usize,
15419 out0: usize,
15420 out1: usize,
15421 out2: usize,
15422 row_bytes: usize,
15423 ws0: f32,
15424 ws1: f32,
15425 ws2: f32,
15426 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15427 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15428 self.e4m3_fused3_t_core(
15429 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15430 )
15431 }
15432
15433 fn try_e4m3_blk_pre(
15444 &self,
15445 w: &crate::model::GpuTensor,
15446 aq: &CudaSlice<i8>,
15447 ad: &CudaSlice<f32>,
15448 m: usize,
15449 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15450 use crate::model::GpuTensor;
15451 if let GpuTensor::Quant {
15452 bytes,
15453 qtype,
15454 row_bytes,
15455 blk: Some(g),
15456 ..
15457 } = w
15458 {
15459 if *qtype == QT_F8_E4M3_BLK {
15460 if (2..=16).contains(&m)
15466 && std::env::var("MEMRA_NO_BATCHED").is_err()
15467 && (m <= 4 || Self::b8_enabled())
15468 {
15469 let mcols = Self::batched_mcols(m);
15470 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15471 bytes,
15472 aq,
15473 ad,
15474 &g.scales,
15475 m,
15476 w.in_features(),
15477 w.out_features(),
15478 *row_bytes,
15479 g.cols,
15480 mcols,
15481 )?));
15482 }
15483 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15484 bytes,
15485 aq,
15486 ad,
15487 &g.scales,
15488 m,
15489 w.in_features(),
15490 w.out_features(),
15491 *row_bytes,
15492 g.cols,
15493 )?));
15494 }
15495 }
15496 Ok(None)
15497 }
15498
15499 fn try_e4m3_blk_prefill(
15546 &self,
15547 w: &crate::model::GpuTensor,
15548 x: &CudaSlice<f32>,
15549 m: usize,
15550 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15551 use crate::model::GpuTensor;
15552 let GpuTensor::Quant {
15553 bytes,
15554 qtype,
15555 blk: Some(g),
15556 ..
15557 } = w
15558 else {
15559 return Ok(None);
15560 };
15561 if *qtype != QT_F8_E4M3_BLK {
15562 return Ok(None);
15563 }
15564 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15569 return Ok(Some(y));
15570 }
15571 let (in_f, out_f) = (w.in_features(), w.out_features());
15572 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15573 let tmp = GpuTensor::Quant {
15574 bytes: slab,
15575 qtype: QT_Q8_0,
15576 row_bytes: in_f / 32 * 34,
15577 ne: vec![in_f as u64, out_f as u64],
15578 scale: 1.0,
15579 rp: false,
15580 #[cfg(memra_cutlass)]
15581 cutlass: None,
15582 fp8: None,
15583 blk: None,
15584 f16: None,
15585 rp4: None,
15586 };
15587 Ok(Some(self.matmul(&tmp, x, m)?))
15589 }
15590
15591 pub fn matmul_pre_noscale(
15592 &self,
15593 w: &crate::model::GpuTensor,
15594 aq: &CudaSlice<i8>,
15595 ad: &CudaSlice<f32>,
15596 m: usize,
15597 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15598 use crate::model::GpuTensor;
15599 if m == 1 {
15603 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15604 return Ok(Some((y, 1.0)));
15605 }
15606 }
15607 if m != 1 || !self.uses_q8_1_fast(w) {
15609 return Ok(None);
15610 }
15611 let in_f = w.in_features();
15612 let out_f = w.out_features();
15613 let (bytes, qtype, row_bytes, scale, rp) = match w {
15614 GpuTensor::Quant {
15615 bytes,
15616 qtype,
15617 row_bytes,
15618 scale,
15619 rp,
15620 ..
15621 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15622 _ => return Ok(None),
15623 };
15624 if self.mmvq_supports(qtype) {
15626 let (mbytes, mrp) = match w {
15628 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15629 _ => (bytes, rp),
15630 };
15631 let y = self.qmatvec_mmvq(
15632 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15633 )?;
15634 return Ok(Some((y, scale)));
15635 }
15636 let name = match qtype {
15638 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15639 QT_Q4_K => "qmatvec_q4_K_dp4a",
15640 QT_Q6_K => "qmatvec_q6_K_dp4a",
15641 QT_Q5_K => "qmatvec_q5_K_dp4a",
15642 QT_Q3_K => "qmatvec_q3_K_dp4a",
15643 QT_NVFP4 => {
15644 if rp {
15645 "qmatvec_nvfp4_dp4a_rp"
15646 } else {
15647 "qmatvec_nvfp4_dp4a"
15648 }
15649 }
15650 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15651 _ => return Ok(None),
15652 };
15653 let f = self.func(name);
15654 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15655 let cfg = LaunchConfig {
15656 grid_dim: (out_f as u32, m as u32, 1),
15657 block_dim: (128, 1, 1),
15658 shared_mem_bytes: 0,
15659 };
15660 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15661 let __s_b = self.gpu.stream();
15662 let mut b = __s_b.launch_builder(&f);
15663 b.arg(bytes)
15664 .arg(aq)
15665 .arg(ad)
15666 .arg(&mut y)
15667 .arg(&inf)
15668 .arg(&outf)
15669 .arg(&mi)
15670 .arg(&rb);
15671 unsafe {
15672 b.launch(cfg)?;
15673 }
15674 Ok(Some((y, scale)))
15675 }
15676
15677 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15680 if qtype == QT_F8_E4M3 {
15685 return true;
15686 }
15687 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15688 return false;
15689 }
15690 matches!(
15691 qtype,
15692 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15693 )
15694 }
15695
15696 pub fn qmatvec_mmvq(
15701 &self,
15702 bytes: &CudaSlice<u8>,
15703 aq: &CudaSlice<i8>,
15704 ad: &CudaSlice<f32>,
15705 m: usize,
15706 in_f: usize,
15707 out_f: usize,
15708 qtype: i32,
15709 row_bytes: usize,
15710 scale: f32,
15711 rp: bool,
15712 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15713 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15715 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15716 )?;
15717 Ok(y)
15718 }
15719
15720 #[allow(clippy::too_many_arguments)]
15722 pub fn qmatvec_mmvq_into(
15723 &self,
15724 bytes: &CudaSlice<u8>,
15725 aq: &CudaSlice<i8>,
15726 ad: &CudaSlice<f32>,
15727 m: usize,
15728 in_f: usize,
15729 out_f: usize,
15730 qtype: i32,
15731 row_bytes: usize,
15732 scale: f32,
15733 rp: bool,
15734 y: &mut CudaSlice<f32>,
15735 ) -> Result<(), Box<dyn std::error::Error>> {
15736 debug_assert!(y.len() >= m * out_f);
15737 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15743 && rp
15744 && m == 1
15745 && out_f >= 64
15746 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15747 && {
15748 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15749 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15750 }
15751 {
15752 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15753 let cfg = LaunchConfig {
15754 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15755 block_dim: (32, 2, 1),
15756 shared_mem_bytes: 0,
15757 };
15758 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15759 let __s_b = self.gpu.stream();
15760 let mut b = __s_b.launch_builder(&f);
15761 b.arg(bytes)
15762 .arg(aq)
15763 .arg(ad)
15764 .arg(&mut *y)
15765 .arg(&inf)
15766 .arg(&outf)
15767 .arg(&mi)
15768 .arg(&rb);
15769 unsafe {
15770 b.launch(cfg)?;
15771 }
15772 if scale != 1.0 {
15773 self.scale_inplace(y, scale, out_f)?;
15774 }
15775 return Ok(());
15776 }
15777 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15786 2
15787 } else {
15788 1
15789 };
15790 if m == 1 && qtype == QT_Q4_0 {
15795 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15796 mr = *Q40MR.get_or_init(|| {
15799 std::env::var("MEMRA_Q40_MR")
15800 .ok()
15801 .and_then(|v| v.parse().ok())
15802 .unwrap_or(1)
15803 });
15804 }
15805 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15816 let q5_force = q5_mode.as_deref() == Some("2");
15817 let q5_il = qtype == QT_Q5_K
15820 && m == 1
15821 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15822 if q5_il && !q5_force && out_f > 65536 {
15823 mr = 1;
15824 }
15825 if qtype == QT_Q4_0 && rp && mr != 1 {
15828 mr = 2;
15829 }
15830 if qtype == QT_Q8_0 && rp {
15834 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15835 mr = *Q80MR.get_or_init(|| {
15836 std::env::var("MEMRA_Q80_MR")
15837 .ok()
15838 .and_then(|v| v.parse().ok())
15839 .unwrap_or(1)
15840 });
15841 }
15842 let name = match (qtype, mr, rp) {
15843 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15844 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15845 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15846 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15847 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15848 (QT_Q5_K, 2, _) => {
15849 if q5_il {
15850 "qmatvec_q5_K_mmvq_mr2_il"
15851 } else {
15852 "qmatvec_q5_K_mmvq_mr2"
15853 }
15854 }
15855 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15856 (QT_Q8_0, _, true)
15861 if in_f % 1024 == 0 && {
15862 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15863 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15864 } =>
15865 {
15866 "qmatvec_q8_0_mmvq_rpca"
15867 }
15868 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15869 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15870 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15874 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15875 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15876 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15877 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15878 (QT_Q5_K, _, _) => {
15879 if q5_il {
15880 "qmatvec_q5_K_mmvq_il"
15881 } else {
15882 "qmatvec_q5_K_mmvq"
15883 }
15884 }
15885 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15886 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15887 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15888 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15889 };
15890 let f = self.func(name);
15891 let rows_per_block = ROWS_PER_BLOCK * mr;
15893 let cfg = LaunchConfig {
15894 grid_dim: (
15895 (out_f as u32 + rows_per_block - 1) / rows_per_block,
15896 m as u32,
15897 1,
15898 ),
15899 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15902 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15903 let __s_b = self.gpu.stream();
15904 let mut b = __s_b.launch_builder(&f);
15905 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15910 if Self::pdl_on()
15913 && Self::pdl_mmvq_on()
15914 && Self::pdl_nvfp4q8_on()
15915 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15916 {
15917 use cudarc::driver::{DevicePtr, DevicePtrMut};
15918 let s = &self.gpu.stream();
15919 let (pw, _g0) = bytes.device_ptr(s);
15920 let (paq, _g1) = aq.device_ptr(s);
15921 let (pad, _g2) = ad.device_ptr(s);
15922 let (py, _g3) = y.device_ptr_mut(s);
15923 let mut ps = [
15924 &pw as *const _ as *mut std::ffi::c_void,
15925 &paq as *const _ as *mut _,
15926 &pad as *const _ as *mut _,
15927 &py as *const _ as *mut _,
15928 &inf as *const _ as *mut _,
15929 &outf as *const _ as *mut _,
15930 &mi as *const _ as *mut _,
15931 &rb as *const _ as *mut _,
15932 &scale as *const _ as *mut _,
15933 ];
15934 unsafe {
15935 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15936 }
15937 return Ok(());
15938 }
15939 b.arg(bytes)
15940 .arg(aq)
15941 .arg(ad)
15942 .arg(&mut *y)
15943 .arg(&inf)
15944 .arg(&outf)
15945 .arg(&mi)
15946 .arg(&rb)
15947 .arg(&scale);
15948 unsafe {
15949 b.launch(cfg)?;
15950 }
15951 } else if Self::pdl_on()
15952 && Self::pdl_mmvq_on()
15953 && (matches!(
15954 name,
15955 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15956 ) || (Self::pdl_nvfp4q8_on()
15957 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15958 {
15959 {
15963 use cudarc::driver::{DevicePtr, DevicePtrMut};
15964 let s = &self.gpu.stream();
15965 let (pw, _g0) = bytes.device_ptr(s);
15966 let (paq, _g1) = aq.device_ptr(s);
15967 let (pad, _g2) = ad.device_ptr(s);
15968 let (py, _g3) = y.device_ptr_mut(s);
15969 let mut ps = [
15970 &pw as *const _ as *mut std::ffi::c_void,
15971 &paq as *const _ as *mut _,
15972 &pad as *const _ as *mut _,
15973 &py as *const _ as *mut _,
15974 &inf as *const _ as *mut _,
15975 &outf as *const _ as *mut _,
15976 &mi as *const _ as *mut _,
15977 &rb as *const _ as *mut _,
15978 ];
15979 unsafe {
15980 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15981 }
15982 }
15983 if scale != 1.0 {
15984 self.scale_inplace(y, scale, m * out_f)?;
15985 }
15986 } else {
15987 b.arg(bytes)
15988 .arg(aq)
15989 .arg(ad)
15990 .arg(&mut *y)
15991 .arg(&inf)
15992 .arg(&outf)
15993 .arg(&mi)
15994 .arg(&rb);
15995 unsafe {
15996 b.launch(cfg)?;
15997 }
15998 if scale != 1.0 {
15999 self.scale_inplace(y, scale, m * out_f)?;
16000 }
16001 }
16002 Ok(())
16003 }
16004
16005 pub fn qmatvec_mmvq_raw(
16009 &self,
16010 bytes: &CudaSlice<u8>,
16011 x: &CudaSlice<f32>,
16012 m: usize,
16013 in_f: usize,
16014 out_f: usize,
16015 qtype: i32,
16016 row_bytes: usize,
16017 rp: bool,
16018 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16019 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16020 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16021 }
16022
16023 pub fn batched_supports(&self, qtype: i32) -> bool {
16027 matches!(
16028 qtype,
16029 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16030 )
16031 }
16032
16033 pub fn iq_fast_enabled() -> bool {
16041 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16042 *ON.get_or_init(|| {
16043 std::env::var("MEMRA_IQ_FAST")
16044 .map(|v| v != "0")
16045 .unwrap_or(true)
16046 })
16047 }
16048
16049 pub fn b8_enabled() -> bool {
16052 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16053 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16054 }
16055
16056 pub fn batched_mcols(m: usize) -> usize {
16058 if m == 2 {
16059 2
16060 } else if m <= 4 {
16061 4
16062 } else if m <= 8 {
16063 8
16064 } else {
16065 16
16066 }
16067 }
16068
16069 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16074 Some(match (qtype, mcols) {
16075 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16076 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16077 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16078 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16084 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16085 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16086 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16087 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16090 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16091 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16092 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16093 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16096 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16097 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16098 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16099 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16100 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16101 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16102 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16103 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16107 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16108 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16109 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16110 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16114 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16115 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16116 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16117 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16118 _ => return None,
16119 })
16120 }
16121
16122 pub fn sm_count(&self) -> i32 {
16157 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16158 *SMS.get_or_init(|| {
16159 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16160 self.gpu
16161 .ctx
16162 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16163 .unwrap_or(82)
16164 })
16165 }
16166
16167 pub fn batched_variant(
16168 &self,
16169 _m: usize,
16170 in_f: usize,
16171 out_f: usize,
16172 qtype: i32,
16173 row_bytes: usize,
16174 mcols: usize,
16175 rp: bool,
16176 ) -> &'static str {
16177 if qtype == QT_Q8_0 {
16182 return if rp { "rp" } else { "base" };
16183 }
16184 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16185 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16186 Ok("base") => "base",
16187 Ok("pf") => "pf",
16188 Ok("r2") => "r2",
16189 Ok("r2w8") => "r2w8",
16190 Ok("pfr2") => "pfr2",
16191 Ok("ca") => "ca",
16192 Ok("car2") => "car2",
16193 Ok("rp") => "rp",
16196 Ok("rpr2") => "rpr2",
16197 Ok("rpr2w8") => "rpr2w8",
16198 Ok("rpca") => "rpca",
16201 Ok("rpcar2") => "rpcar2",
16202 Ok("rpsc") => "rpsc",
16209 Ok("rpms") => "rpms",
16210 Ok("rpmsc") => "rpmsc",
16211 Ok("rpks") => "rpks",
16212 Ok("rpksc") => "rpksc",
16213 _ => "auto",
16214 });
16215 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16219 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16224 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16225 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16226 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16227 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16228 let sms = *SMS.get_or_init(|| {
16229 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16230 self.gpu
16231 .ctx
16232 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16233 .unwrap_or(82)
16234 });
16235 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16255 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16258 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16259 Ok("base") => "base",
16260 Ok("r2") => "r2",
16261 Ok("r2w8") => "r2w8",
16262 _ => "auto",
16263 });
16264 let variant: &'static str = if qtype == QT_Q4_0 {
16265 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16269 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16270 Ok("base") => "base",
16276 Ok("r2") => "r2",
16277 Ok("ms") => "ms",
16278 Ok("sm") => "sm",
16279 Ok("la") => "la",
16280 _ => "auto",
16281 });
16282 let v = if q40 != "auto" {
16283 q40
16284 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16285 "r2"
16286 } else {
16287 "base"
16288 };
16289 if rp {
16294 match v {
16295 "ms" => "r2ms_rp",
16296 "sm" => "r2sm_rp",
16297 "la" => "r2la_rp",
16298 "r2" => "r2_rp",
16299 _ => "rp",
16300 }
16301 } else if matches!(v, "ms" | "sm" | "la") {
16302 "r2"
16303 } else {
16304 v
16305 }
16306 } else if qtype != QT_NVFP4 && !kq_r2 {
16307 "base"
16308 } else if kq_r2 && rp {
16309 "rp"
16313 } else if kq_r2 {
16314 if kq_bv != "auto" {
16317 if kq_bv == "r2w8" && mcols != 4 {
16318 "r2"
16319 } else {
16320 kq_bv
16321 }
16322 } else if bv != "auto" {
16323 match bv {
16324 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16325 "r2w8" | "rpr2w8" => {
16326 if mcols != 4 {
16327 "r2"
16328 } else {
16329 "r2w8"
16330 }
16331 }
16332 _ => "base", }
16334 } else {
16335 let blocks = (out_f + 7) / 8;
16336 let waves = blocks as f64 / (7 * sms as usize) as f64;
16337 let filled = blocks >= 4 * sms as usize;
16338 let use_r2 = if qtype == QT_Q4_K {
16339 filled
16340 } else {
16341 waves >= 2.0
16342 };
16343 if use_r2 { "r2" } else { "base" }
16344 }
16345 } else if bv != "auto" {
16346 let v = if bv == "r2w8" && mcols == 2 {
16351 "r2"
16352 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16353 "pf"
16354 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16355 "r2"
16356 } else if bv == "pfr2" && mcols == 8 {
16357 "r2"
16358 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16359 "rpr2"
16360 }
16361 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16363 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16364 } else if bv == "rpcar2" && mcols == 2 {
16365 "rpca"
16366 }
16367 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16370 "rpr2"
16371 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16372 "rpr2"
16373 } else {
16374 bv
16375 };
16376 if rp {
16377 match v {
16378 "base" | "pf" | "ca" | "rp" => "rp",
16379 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16380 "r2w8" | "rpr2w8" => {
16381 if mcols == 2 {
16382 "rpr2"
16383 } else {
16384 "rpr2w8"
16385 }
16386 }
16387 other => other, }
16389 } else {
16390 v
16391 }
16392 } else if mcols == 8 {
16393 if rp {
16404 if sc_ok { "rpsc" } else { "rpr2w8" }
16405 } else {
16406 "r2w8"
16407 }
16408 } else if mcols >= 4 {
16409 let blocks = (out_f + 7) / 8;
16413 let r7 = 7 * sms as usize;
16414 let r8 = 8 * sms as usize;
16415 let waves = blocks as f64 / r7 as f64;
16416 let filled = blocks >= 4 * sms as usize;
16417 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16421 if rp { "rpr2w8" } else { "r2w8" }
16425 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16426 if rp { "rpr2" } else { "r2" }
16429 } else {
16430 if rp { "rp" } else { "pf" }
16434 }
16435 } else if in_f >= 6144 {
16436 if rp { "rpr2" } else { "r2" }
16440 } else if rp {
16441 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16446 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16447 "rpsc"
16448 } else {
16449 "rp"
16450 }
16451 } else {
16452 "base"
16453 };
16454 variant
16455 }
16456
16457 pub fn qmatvec_mmvq_batched(
16458 &self,
16459 bytes: &CudaSlice<u8>,
16460 aq: &CudaSlice<i8>,
16461 ad: &CudaSlice<f32>,
16462 m: usize,
16463 in_f: usize,
16464 out_f: usize,
16465 qtype: i32,
16466 row_bytes: usize,
16467 mcols: usize,
16468 scale: f32,
16469 rp: bool,
16470 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16471 const ROWS_PER_BLOCK: u32 = 4;
16472 let forced: Option<&'static str> = {
16477 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16478 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16479 .as_deref()
16480 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16481 };
16482 let variant = match forced {
16483 Some(v) if !rp || v.contains("rp") => v,
16484 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16485 };
16486 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16487 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16488 })?;
16489 let variant = if mcols == 16 {
16493 if rp { "rp" } else { "base" }
16494 } else {
16495 variant
16496 };
16497 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16504 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16505 if b567
16506 && qtype == QT_NVFP4
16507 && rp
16508 && mcols == 8
16509 && (5..=7).contains(&m)
16510 && matches!(variant, "rpsc" | "rpr2w8")
16511 {
16512 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16513 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16515 let cfg = LaunchConfig {
16516 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16517 block_dim: (32, ROWS_PER_BLOCK, 1),
16518 shared_mem_bytes: 0,
16519 };
16520 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16521 let __s_b = self.gpu.stream();
16522 let mut b = __s_b.launch_builder(&f);
16523 b.arg(bytes)
16524 .arg(aq)
16525 .arg(ad)
16526 .arg(&mut y)
16527 .arg(&inf)
16528 .arg(&outf)
16529 .arg(&mi)
16530 .arg(&rb);
16531 unsafe {
16532 b.launch(cfg)?;
16533 }
16534 if scale != 1.0 {
16535 self.scale_inplace(&mut y, scale, m * out_f)?;
16536 }
16537 return Ok(y);
16538 }
16539 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16540 "base" => (base_name.into(), ROWS_PER_BLOCK),
16541 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16542 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16543 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16544 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16548 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16549 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16550 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16551 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16552 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16553 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16554 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16556 debug_assert!(
16557 !rp || name.contains("_rp"),
16558 "rp weight dispatched to a GGUF-layout kernel"
16559 );
16560 let f = self.func(&name);
16561 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16562 let smem = if name.contains("_r2sm_rp") {
16564 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16565 } else {
16566 0
16567 };
16568 let cfg = LaunchConfig {
16569 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16570 block_dim: (32, ROWS_PER_BLOCK, 1),
16571 shared_mem_bytes: smem,
16572 };
16573 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16574 let __s_b = self.gpu.stream();
16575 let mut b = __s_b.launch_builder(&f);
16576 b.arg(bytes)
16577 .arg(aq)
16578 .arg(ad)
16579 .arg(&mut y)
16580 .arg(&inf)
16581 .arg(&outf)
16582 .arg(&mi)
16583 .arg(&rb);
16584 unsafe {
16585 b.launch(cfg)?;
16586 }
16587 if scale != 1.0 {
16588 self.scale_inplace(&mut y, scale, m * out_f)?;
16589 }
16590 Ok(y)
16591 }
16592
16593 pub fn qmatvec_batched_raw(
16597 &self,
16598 bytes: &CudaSlice<u8>,
16599 x: &CudaSlice<f32>,
16600 m: usize,
16601 in_f: usize,
16602 out_f: usize,
16603 qtype: i32,
16604 row_bytes: usize,
16605 mcols: usize,
16606 rp: bool,
16607 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16608 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16609 self.qmatvec_mmvq_batched(
16610 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16611 )
16612 }
16613
16614 pub fn qmatvec_nvfp4_batched_raw(
16616 &self,
16617 bytes: &CudaSlice<u8>,
16618 x: &CudaSlice<f32>,
16619 m: usize,
16620 in_f: usize,
16621 out_f: usize,
16622 row_bytes: usize,
16623 mcols: usize,
16624 rp: bool,
16625 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16626 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16627 }
16628
16629 fn try_fp4_gemm(
16633 &self,
16634 w: &crate::model::GpuTensor,
16635 x: &CudaSlice<f32>,
16636 m: usize,
16637 in_f: usize,
16638 out_f: usize,
16639 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16640 use crate::model::GpuTensor;
16641 if cfg!(memra_portable_cuda) {
16642 return Ok(None);
16643 }
16644 if std::env::var("MEMRA_FP4").is_ok() {
16647 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16648 }
16649 if std::env::var("MEMRA_FP4").is_err() {
16650 return Ok(None);
16651 }
16652 #[cfg(memra_cutlass)]
16661 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16662 if let GpuTensor::Quant {
16663 bytes,
16664 qtype,
16665 scale,
16666 row_bytes,
16667 cutlass,
16668 ..
16669 } = w
16670 {
16671 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16672 if let Some(cw) = cutlass {
16673 let y = self.cutlass_fp4_gemm(
16675 &cw.b_packed,
16676 &cw.sfb_swizzled,
16677 x,
16678 *scale,
16679 m,
16680 out_f,
16681 in_f,
16682 )?;
16683 return Ok(Some(y));
16684 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16685 let (b_packed, sfb_sw) =
16690 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16691 let y =
16692 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16693 return Ok(Some(y));
16694 }
16695 }
16696 }
16697 }
16698 if let GpuTensor::Quant {
16699 bytes,
16700 qtype,
16701 row_bytes,
16702 scale,
16703 rp,
16704 ..
16705 } = w
16706 {
16707 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16710 let y =
16711 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16712 return Ok(Some(y));
16713 }
16714 }
16715 Ok(None)
16716 }
16717
16718 pub fn rms_norm_f16out(
16722 &self,
16723 x: &CudaSlice<f32>,
16724 w: &CudaSlice<f32>,
16725 dst: &mut CudaSlice<f32>,
16726 dst16: &mut CudaSlice<u8>,
16727 ncols: usize,
16728 nrows: usize,
16729 eps: f32,
16730 ) -> Result<(), Box<dyn std::error::Error>> {
16731 let f = self.func("rms_norm_f16out_f32");
16732 let cfg = LaunchConfig {
16733 grid_dim: (nrows as u32, 1, 1),
16734 block_dim: (rms_block(), 1, 1),
16735 shared_mem_bytes: 0,
16736 };
16737 let (nc, e) = (ncols as i32, eps);
16738 let __s_b = self.gpu.stream();
16739 let mut b = __s_b.launch_builder(&f);
16740 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16741 unsafe {
16742 b.launch(cfg)?;
16743 }
16744 Ok(())
16745 }
16746
16747 #[allow(clippy::too_many_arguments)]
16750 pub fn add_rms_norm_f16out(
16751 &self,
16752 a: &CudaSlice<f32>,
16753 b: &CudaSlice<f32>,
16754 w: &CudaSlice<f32>,
16755 res: &mut CudaSlice<f32>,
16756 dst: &mut CudaSlice<f32>,
16757 dst16: &mut CudaSlice<u8>,
16758 ncols: usize,
16759 nrows: usize,
16760 eps: f32,
16761 ) -> Result<(), Box<dyn std::error::Error>> {
16762 let f = self.func("add_rms_norm_f16out_f32");
16763 let cfg = LaunchConfig {
16764 grid_dim: (nrows as u32, 1, 1),
16765 block_dim: (rms_block(), 1, 1),
16766 shared_mem_bytes: 0,
16767 };
16768 let (nc, e) = (ncols as i32, eps);
16769 let __s_lb = self.gpu.stream();
16770 let mut lb = __s_lb.launch_builder(&f);
16771 lb.arg(a)
16772 .arg(b)
16773 .arg(w)
16774 .arg(res)
16775 .arg(dst)
16776 .arg(dst16)
16777 .arg(&nc)
16778 .arg(&e);
16779 unsafe {
16780 lb.launch(cfg)?;
16781 }
16782 Ok(())
16783 }
16784
16785 pub fn matmul_group_xh(
16788 &self,
16789 ws: &[&crate::model::GpuTensor],
16790 x: &CudaSlice<f32>,
16791 xh: &CudaSlice<u8>,
16792 m: usize,
16793 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16794 let mut out = Vec::with_capacity(ws.len());
16795 let in_f = ws[0].in_features();
16796 for w in ws {
16797 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16798 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16799 out.push(y);
16800 continue;
16801 }
16802 }
16803 out.push(self.matmul(w, x, m)?);
16804 }
16805 Ok(out)
16806 }
16807
16808 pub fn gdn_pad_mask(
16811 &self,
16812 beta: &mut CudaSlice<f32>,
16813 g_log: &mut CudaSlice<f32>,
16814 len_d: &CudaSlice<i32>,
16815 h: usize,
16816 t: usize,
16817 ) -> Result<(), Box<dyn std::error::Error>> {
16818 let f = self.func("gdn_pad_mask_f32");
16819 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16820 let (hi, ti) = (h as i32, t as i32);
16821 let __s_b = self.gpu.stream();
16822 let mut b = __s_b.launch_builder(&f);
16823 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16824 unsafe {
16825 b.launch(cfg)?;
16826 }
16827 Ok(())
16828 }
16829
16830 pub fn row_gather_dev(
16833 &self,
16834 src: &CudaSlice<f32>,
16835 dst: &mut CudaSlice<f32>,
16836 len_d: &CudaSlice<i32>,
16837 ncols: usize,
16838 ) -> Result<(), Box<dyn std::error::Error>> {
16839 let f = self.func("row_gather_dev_f32");
16840 let cfg = LaunchConfig::for_num_elems(ncols as u32);
16841 let nc = ncols as i32;
16842 let __s_b = self.gpu.stream();
16843 let mut b = __s_b.launch_builder(&f);
16844 b.arg(src).arg(dst).arg(len_d).arg(&nc);
16845 unsafe {
16846 b.launch(cfg)?;
16847 }
16848 Ok(())
16849 }
16850
16851 pub fn matmul_group(
16858 &self,
16859 ws: &[&crate::model::GpuTensor],
16860 x: &CudaSlice<f32>,
16861 m: usize,
16862 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16863 use crate::model::GpuTensor;
16864 let mut out = Vec::with_capacity(ws.len());
16865 let any_mirror = ws
16866 .iter()
16867 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16868 if m >= 16 && any_mirror && !self.verify_exact_on() {
16869 let in_f = ws[0].in_features();
16870 let xh = self.f16_act(x, m * in_f, in_f)?;
16871 for w in ws {
16872 if w.in_features() == in_f {
16873 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16874 out.push(y);
16875 continue;
16876 }
16877 }
16878 out.push(self.matmul(w, x, m)?);
16879 }
16880 return Ok(out);
16881 }
16882 for w in ws {
16883 out.push(self.matmul(w, x, m)?);
16884 }
16885 Ok(out)
16886 }
16887
16888 pub fn matmul_group_multi(
16895 &self,
16896 ws: &[&crate::model::GpuTensor],
16897 xs: &[&CudaSlice<f32>],
16898 ms: &[usize],
16899 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16900 assert_eq!(xs.len(), ms.len());
16901 let in_f = ws[0].in_features();
16902 let total: usize = ms.iter().sum();
16903 let mut xcat = self.uninit(total * in_f)?;
16904 let mut off = 0usize;
16905 for (x, &m) in xs.iter().zip(ms) {
16906 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16907 off += m;
16908 }
16909 let ys = self.matmul_group(ws, &xcat, total)?;
16910 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16911 for (w, y) in ws.iter().zip(ys) {
16912 let out_f = w.out_features();
16913 let mut off = 0usize;
16914 for (s, &m) in ms.iter().enumerate() {
16915 let mut ys_s = self.uninit(m * out_f)?;
16916 let src = y.slice(off * out_f..(off + m) * out_f);
16917 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16918 out[s].push(ys_s);
16919 off += m;
16920 }
16921 }
16922 Ok(out)
16923 }
16924
16925 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16935 use crate::model::GpuTensor;
16936 if !legacy_quant_gemm_allowed(
16937 cfg!(memra_portable_cuda),
16938 cfg!(memra_hopper_mma),
16939 std::env::var_os("MEMRA_NO_GEMM").is_some(),
16940 ) {
16941 return false;
16942 }
16943 match w {
16944 GpuTensor::Quant { qtype, .. } => {
16945 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16946 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16947 }
16948 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16949 }
16950 }
16951
16952 pub fn qmatvec_gemm(
16959 &self,
16960 w: &crate::model::GpuTensor,
16961 aq: &CudaSlice<i8>,
16962 ad: &CudaSlice<f32>,
16963 m: usize,
16964 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16965 use crate::model::GpuTensor;
16966 let in_f = w.in_features();
16967 let out_f = w.out_features();
16968 let (bytes, qtype, row_bytes, scale, rp) = match w {
16969 GpuTensor::Quant {
16970 bytes,
16971 qtype,
16972 row_bytes,
16973 scale,
16974 rp,
16975 ..
16976 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16977 _ => unreachable!("gemm_supports guaranteed Quant"),
16978 };
16979 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16985 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16986 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16987 if scale != 1.0 {
16988 self.scale_inplace(&mut y, scale, m * out_f)?;
16989 }
16990 return Ok(y);
16991 }
16992 }
16993 let name = match qtype {
16994 QT_Q8_0 => "qmatvec_gemm_q8_0",
16995 QT_Q4_K => "qmatvec_gemm_q4_K",
16996 QT_Q4_0 => {
16997 if rp {
16998 "qmatvec_gemm_q4_0_rp"
16999 } else {
17000 "qmatvec_gemm_q4_0"
17001 }
17002 }
17003 QT_Q5_K => "qmatvec_gemm_q5_K",
17004 QT_Q6_K => "qmatvec_gemm_q6_K",
17005 QT_NVFP4 => {
17006 if rp {
17007 "qmatvec_gemm_nvfp4_rp"
17008 } else {
17009 "qmatvec_gemm_nvfp4"
17010 }
17011 }
17012 _ => unreachable!(),
17013 };
17014 let f = self.func(name);
17015 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);
17020 let k1_tile = if is_k1 {
17022 k1_launch_override().unwrap_or((128, 128, 8))
17023 } else {
17024 (128, 128, 8)
17025 };
17026 let (bm, bn): (u32, u32) = if is_k1 {
17027 (k1_tile.0, k1_tile.1)
17028 } else {
17029 (64, 256)
17030 };
17031 let warps: u32 = if is_k1 {
17032 k1_tile.2
17033 } else {
17034 match qtype {
17035 QT_NVFP4 => 8,
17036 _ => 4,
17037 }
17038 };
17039 let cfg = LaunchConfig {
17040 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17041 block_dim: (32, warps, 1),
17042 shared_mem_bytes: 0,
17043 };
17044 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17045 let __s_b = self.gpu.stream();
17046 let mut b = __s_b.launch_builder(&f);
17047 b.arg(bytes)
17048 .arg(aq)
17049 .arg(ad)
17050 .arg(&mut y)
17051 .arg(&inf)
17052 .arg(&outf)
17053 .arg(&mi)
17054 .arg(&rb);
17055 unsafe {
17056 b.launch(cfg)?;
17057 }
17058 if scale != 1.0 {
17059 self.scale_inplace(&mut y, scale, m * out_f)?;
17060 }
17061 Ok(y)
17062 }
17063
17064 pub fn qmatvec_gemm_raw(
17069 &self,
17070 bytes: &CudaSlice<u8>,
17071 x: &CudaSlice<f32>,
17072 m: usize,
17073 in_f: usize,
17074 out_f: usize,
17075 qtype: i32,
17076 row_bytes: usize,
17077 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17078 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17079 let name = match qtype {
17080 QT_Q8_0 => "qmatvec_gemm_q8_0",
17081 QT_Q4_K => "qmatvec_gemm_q4_K",
17082 QT_Q4_0 => "qmatvec_gemm_q4_0",
17083 QT_Q5_K => "qmatvec_gemm_q5_K",
17084 QT_Q6_K => "qmatvec_gemm_q6_K",
17085 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17086 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17087 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17088 };
17089 let f = self.func(name);
17090 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);
17094 let k1_tile = if is_k1 {
17096 k1_launch_override().unwrap_or((128, 128, 8))
17097 } else {
17098 (128, 128, 8)
17099 };
17100 let (bm, bn): (u32, u32) = if is_k1 {
17101 (k1_tile.0, k1_tile.1)
17102 } else {
17103 (64, 256)
17104 };
17105 let warps: u32 = if is_k1 {
17106 k1_tile.2
17107 } else {
17108 match qtype {
17109 QT_NVFP4 | QT_NVFP4_RP => 8,
17110 _ => 4,
17111 }
17112 };
17113 let cfg = LaunchConfig {
17114 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17115 block_dim: (32, warps, 1),
17116 shared_mem_bytes: 0,
17117 };
17118 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17119 let __s_b = self.gpu.stream();
17120 let mut b = __s_b.launch_builder(&f);
17121 b.arg(bytes)
17122 .arg(&aq)
17123 .arg(&ad)
17124 .arg(&mut y)
17125 .arg(&inf)
17126 .arg(&outf)
17127 .arg(&mi)
17128 .arg(&rb);
17129 unsafe {
17130 b.launch(cfg)?;
17131 }
17132 Ok(y)
17133 }
17134
17135 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17142 &self,
17143 rp4: &CudaSlice<u8>,
17144 aq: &CudaSlice<i8>,
17145 ad: &CudaSlice<f32>,
17146 m: usize,
17147 in_f: usize,
17148 out_f: usize,
17149 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17150 assert!(
17151 out_f % 64 == 0 && in_f % 32 == 0,
17152 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17153 );
17154 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17155 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17157 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17158 block_dim: (128, 1, 1),
17159 shared_mem_bytes: 0,
17160 };
17161 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17162 let __s_b = self.gpu.stream();
17163 let mut b = __s_b.launch_builder(&f);
17164 b.arg(rp4)
17165 .arg(aq)
17166 .arg(ad)
17167 .arg(&mut y)
17168 .arg(&inf)
17169 .arg(&outf)
17170 .arg(&mi);
17171 unsafe {
17172 b.launch(cfg)?;
17173 }
17174 Ok(y)
17175 }
17176
17177 pub fn scale_inplace(
17179 &self,
17180 y: &mut CudaSlice<f32>,
17181 s: f32,
17182 n: usize,
17183 ) -> Result<(), Box<dyn std::error::Error>> {
17184 let f = self.func("scale_f32");
17185 let cfg = LaunchConfig::for_num_elems(n as u32);
17186 let (sf, ni) = (s, n as i32);
17187 let __s_b = self.gpu.stream();
17188 let mut b = __s_b.launch_builder(&f);
17189 b.arg(y).arg(&sf).arg(&ni);
17190 unsafe {
17191 b.launch(cfg)?;
17192 }
17193 Ok(())
17194 }
17195
17196 pub fn bf16_to_f32(
17201 &self,
17202 data: &cudarc::driver::CudaView<'_, u8>,
17203 n: usize,
17204 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17205 let mut out = self.alloc_uninit::<f32>(n)?;
17206 let f = self.func("bf16_to_f32");
17207 let cfg = LaunchConfig::for_num_elems(n as u32);
17208 let ni = n as i32;
17209 let __s_b = self.gpu.stream();
17210 let mut b = __s_b.launch_builder(&f);
17211 b.arg(data).arg(&mut out).arg(&ni);
17212 unsafe {
17213 b.launch(cfg)?;
17214 }
17215 Ok(out)
17216 }
17217
17218 fn linear_bf16_chunked(
17225 &self,
17226 x: &CudaSlice<f32>,
17227 data: &CudaSlice<u8>,
17228 m: usize,
17229 in_f: usize,
17230 out_f: usize,
17231 exact: bool,
17232 canonical_chunk_rows: Option<usize>,
17233 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17234 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17238 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17239 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17240 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17241 let started = timing.then(std::time::Instant::now);
17242 let result =
17243 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17244 if let Some(started) = started {
17245 use std::sync::atomic::Ordering;
17246 self.stream().synchronize()?;
17247 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17248 + started.elapsed().as_nanos() as u64;
17249 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17250 + (in_f * out_f * 2) as u64;
17251 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17252 if calls % 1024 == 0 {
17253 eprintln!(
17254 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17255 weight_gb={:.2}",
17256 ns as f64 / 1.0e6,
17257 ns as f64 / calls as f64 / 1.0e3,
17258 wb as f64 / 1.0e9,
17259 );
17260 }
17261 }
17262 result
17263 }
17264
17265 pub(crate) fn bf16_mmv_on() -> bool {
17270 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17271 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17272 }
17273
17274 fn matvec_bf16(
17277 &self,
17278 data: &CudaSlice<u8>,
17279 x: &CudaSlice<f32>,
17280 in_f: usize,
17281 out_f: usize,
17282 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17283 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17284 return Err(format!(
17285 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17286 data.len(),
17287 x.len()
17288 )
17289 .into());
17290 }
17291 let mut y = self.alloc_uninit::<f32>(out_f)?;
17292 let f = self.func("matvec_bf16_f32acc");
17293 let cfg = LaunchConfig {
17294 grid_dim: (out_f as u32, 1, 1),
17295 block_dim: (mmv_block(), 1, 1),
17296 shared_mem_bytes: 0,
17297 };
17298 let ini = in_f as i32;
17299 let __s_bld = self.gpu.stream();
17300 let mut bld = __s_bld.launch_builder(&f);
17301 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17302 unsafe {
17303 bld.launch(cfg)?;
17304 }
17305 Ok(y)
17306 }
17307
17308 #[allow(clippy::too_many_arguments)]
17312 #[allow(clippy::too_many_arguments)]
17317 #[allow(clippy::too_many_arguments)]
17322 pub fn qk_norm_rope_append_inc_dcw_rows(
17323 &self,
17324 q_raw_t: &CudaSlice<f32>,
17325 k_raw_t: &CudaSlice<f32>,
17326 v_raw_t: &CudaSlice<f32>,
17327 qw: &CudaSlice<f32>,
17328 kw: &CudaSlice<f32>,
17329 q_out_t: &mut CudaSlice<f32>,
17330 k_out_t: &mut CudaSlice<f32>,
17331 tab: &CudaSlice<u64>,
17332 pos_t: &CudaSlice<i32>,
17333 same_session: bool,
17334 t: usize,
17335 kv_dim_k: usize,
17336 kv_dim_v: usize,
17337 k_tok_bytes: usize,
17338 v_tok_bytes: usize,
17339 head_dim: usize,
17340 n_dims: usize,
17341 nh_q: usize,
17342 nh_k: usize,
17343 eps: f32,
17344 freq_base: f32,
17345 freq_scale: f32,
17346 ff: Option<&CudaSlice<f32>>,
17347 ) -> Result<(), Box<dyn std::error::Error>> {
17348 if head_dim != 128
17349 || kv_dim_v != kv_dim_k
17350 || kv_dim_k != nh_k * head_dim
17351 || t == 0
17352 || t > 32
17353 || tab.len() < t * 6
17354 || pos_t.len() < t
17355 || q_raw_t.len() < t * nh_q * head_dim
17356 || k_raw_t.len() < t * nh_k * head_dim
17357 || v_raw_t.len() < t * kv_dim_v
17358 || q_out_t.len() < t * nh_q * head_dim
17359 || k_out_t.len() < t * nh_k * head_dim
17360 {
17361 return Err(format!(
17362 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17363 nh_q={nh_q} nh_k={nh_k}"
17364 )
17365 .into());
17366 }
17367 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17368 let same_t: i32 = if same_session { t as i32 } else { 0 };
17369 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17370 let cfg = LaunchConfig {
17371 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17372 block_dim: (128, 1, 1),
17373 shared_mem_bytes: 0,
17374 };
17375 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17376 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17377 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17378 let null: u64 = 0;
17379 let __s_b = self.gpu.stream();
17380 let mut b = __s_b.launch_builder(&f);
17381 b.arg(q_raw_t)
17382 .arg(k_raw_t)
17383 .arg(v_raw_t)
17384 .arg(qw)
17385 .arg(kw)
17386 .arg(q_out_t)
17387 .arg(k_out_t)
17388 .arg(tab)
17389 .arg(pos_t)
17390 .arg(&same_t)
17391 .arg(&kvk)
17392 .arg(&kvv)
17393 .arg(&ktb)
17394 .arg(&vtb)
17395 .arg(&hd)
17396 .arg(&nd)
17397 .arg(&nq)
17398 .arg(&nk)
17399 .arg(&eps)
17400 .arg(&theta_scale)
17401 .arg(&freq_scale);
17402 match ff {
17403 Some(freqs) => {
17404 b.arg(freqs);
17405 }
17406 None => {
17407 b.arg(&null);
17408 }
17409 }
17410 unsafe {
17411 b.launch(cfg)?;
17412 }
17413 Ok(())
17414 }
17415
17416 pub fn qk_norm_rope_append_inc_dcw(
17417 &self,
17418 q_raw: &CudaSlice<f32>,
17419 k_raw: &CudaSlice<f32>,
17420 v_raw: &CudaSlice<f32>,
17421 qw: &CudaSlice<f32>,
17422 kw: &CudaSlice<f32>,
17423 q_out: &mut CudaSlice<f32>,
17424 k_out: &mut CudaSlice<f32>,
17425 pos: &CudaSlice<i32>,
17426 k_plane: &mut CudaSlice<u8>,
17427 v_plane: &mut CudaSlice<u8>,
17428 len_dev: &CudaSlice<i32>,
17431 base_dev: Option<&CudaSlice<i32>>,
17432 done_ctr: &mut CudaSlice<u32>,
17433 kv_dim_k: usize,
17434 kv_dim_v: usize,
17435 k_tok_bytes: usize,
17436 v_tok_bytes: usize,
17437 head_dim: usize,
17438 n_dims: usize,
17439 nh_q: usize,
17440 nh_k: usize,
17441 eps: f32,
17442 freq_base: f32,
17443 freq_scale: f32,
17444 ff: Option<&CudaSlice<f32>>,
17445 ) -> Result<(), Box<dyn std::error::Error>> {
17446 if head_dim != 128
17447 || kv_dim_v != kv_dim_k
17448 || kv_dim_k != nh_k * head_dim
17449 || q_raw.len() < nh_q * head_dim
17450 || k_raw.len() < nh_k * head_dim
17451 || v_raw.len() < kv_dim_v
17452 || q_out.len() < nh_q * head_dim
17453 || k_out.len() < nh_k * head_dim
17454 || pos.is_empty()
17455 || done_ctr.is_empty()
17456 {
17457 return Err(format!(
17458 "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}"
17459 )
17460 .into());
17461 }
17462 let f = self.func("qk_norm_rope_append_inc_dcw");
17463 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17464 let cfg = LaunchConfig {
17465 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17466 block_dim: (128, 1, 1),
17467 shared_mem_bytes: 0,
17468 };
17469 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17470 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17471 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17472 let null: u64 = 0;
17473 let __s_b = self.gpu.stream();
17474 let mut b = __s_b.launch_builder(&f);
17475 b.arg(q_raw)
17476 .arg(k_raw)
17477 .arg(v_raw)
17478 .arg(qw)
17479 .arg(kw)
17480 .arg(q_out)
17481 .arg(k_out)
17482 .arg(pos)
17483 .arg(&mut *k_plane)
17484 .arg(&mut *v_plane)
17485 .arg(len_dev);
17486 match base_dev {
17487 Some(base) => {
17488 b.arg(base);
17489 }
17490 None => {
17491 b.arg(&null);
17492 }
17493 }
17494 b.arg(&mut *done_ctr)
17495 .arg(&kvk)
17496 .arg(&kvv)
17497 .arg(&ktb)
17498 .arg(&vtb)
17499 .arg(&hd)
17500 .arg(&nd)
17501 .arg(&nq)
17502 .arg(&eps)
17503 .arg(&theta_scale)
17504 .arg(&freq_scale);
17505 match ff {
17506 Some(freqs) => {
17507 b.arg(freqs);
17508 }
17509 None => {
17510 b.arg(&null);
17511 }
17512 }
17513 unsafe {
17514 b.launch(cfg)?;
17515 }
17516 Ok(())
17517 }
17518
17519 pub fn qk_norm_rope_into(
17520 &self,
17521 q_raw: &CudaSlice<f32>,
17522 k_raw: &CudaSlice<f32>,
17523 qw: &CudaSlice<f32>,
17524 kw: &CudaSlice<f32>,
17525 q_out: &mut CudaSlice<f32>,
17526 k_out: &mut CudaSlice<f32>,
17527 pos: &CudaSlice<i32>,
17528 head_dim: usize,
17529 n_dims: usize,
17530 nh_q: usize,
17531 nh_k: usize,
17532 eps: f32,
17533 freq_base: f32,
17534 freq_scale: f32,
17535 ff: Option<&CudaSlice<f32>>,
17536 ) -> Result<(), Box<dyn std::error::Error>> {
17537 if head_dim > 512
17538 || q_raw.len() < nh_q * head_dim
17539 || k_raw.len() < nh_k * head_dim
17540 || q_out.len() < nh_q * head_dim
17541 || k_out.len() < nh_k * head_dim
17542 || qw.len() < head_dim
17543 || kw.len() < head_dim
17544 || pos.is_empty()
17545 {
17546 return Err(format!(
17547 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17548 )
17549 .into());
17550 }
17551 let f = self.func("qk_norm_rope_f32");
17552 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17553 let cfg = LaunchConfig {
17554 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17555 block_dim: (128, 1, 1),
17556 shared_mem_bytes: 0,
17557 };
17558 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17559 let __s_b = self.gpu.stream();
17560 let mut b = __s_b.launch_builder(&f);
17561 b.arg(q_raw)
17562 .arg(k_raw)
17563 .arg(qw)
17564 .arg(kw)
17565 .arg(q_out)
17566 .arg(k_out)
17567 .arg(pos)
17568 .arg(&hd)
17569 .arg(&nd)
17570 .arg(&nq)
17571 .arg(&eps)
17572 .arg(&theta_scale)
17573 .arg(&freq_scale);
17574 match ff {
17575 Some(ffv) => {
17576 b.arg(ffv);
17577 unsafe {
17578 b.launch(cfg)?;
17579 }
17580 }
17581 None => {
17582 let null: u64 = 0;
17583 b.arg(&null);
17584 unsafe {
17585 b.launch(cfg)?;
17586 }
17587 }
17588 }
17589 Ok(())
17590 }
17591
17592 #[allow(clippy::too_many_arguments)]
17595 pub fn matvec_f32_b4_into(
17596 &self,
17597 w: [&CudaSlice<f32>; 4],
17598 x: &CudaSlice<f32>,
17599 y: &mut CudaSlice<f32>,
17600 block_cols: usize,
17601 out_f: usize,
17602 ) -> Result<(), Box<dyn std::error::Error>> {
17603 if block_cols % 4 != 0
17604 || x.len() < 4 * block_cols
17605 || y.len() < out_f
17606 || w.iter().any(|w| w.len() != out_f * block_cols)
17607 {
17608 return Err(format!(
17609 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17610 x.len()
17611 )
17612 .into());
17613 }
17614 let f = self.func("matvec_f32_b4");
17615 let cfg = LaunchConfig {
17616 grid_dim: (out_f as u32, 1, 1),
17617 block_dim: (128, 1, 1),
17618 shared_mem_bytes: 0,
17619 };
17620 let (bc, of) = (block_cols as i32, out_f as i32);
17621 let __s_b = self.gpu.stream();
17622 let mut b = __s_b.launch_builder(&f);
17623 b.arg(w[0])
17624 .arg(w[1])
17625 .arg(w[2])
17626 .arg(w[3])
17627 .arg(x)
17628 .arg(y)
17629 .arg(&bc)
17630 .arg(&of);
17631 unsafe {
17632 b.launch(cfg)?;
17633 }
17634 Ok(())
17635 }
17636
17637 pub fn axpy_rows_seq_into(
17640 &self,
17641 x: &CudaSlice<f32>,
17642 w: &CudaSlice<f32>,
17643 y: &mut CudaSlice<f32>,
17644 width: usize,
17645 n_rows: usize,
17646 ) -> Result<(), Box<dyn std::error::Error>> {
17647 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17648 return Err(format!(
17649 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17650 x.len(),
17651 w.len(),
17652 y.len()
17653 )
17654 .into());
17655 }
17656 let f = self.func("axpy_rows_seq_f32");
17657 let cfg = LaunchConfig::for_num_elems(width as u32);
17658 let (wi, nr) = (width as i32, n_rows as i32);
17659 let __s_b = self.gpu.stream();
17660 let mut b = __s_b.launch_builder(&f);
17661 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17662 unsafe {
17663 b.launch(cfg)?;
17664 }
17665 Ok(())
17666 }
17667
17668 #[allow(clippy::too_many_arguments)]
17672 pub fn axpy_rows_seq_md_off_into(
17673 &self,
17674 x: &CudaSlice<f32>,
17675 w_route: &CudaSlice<f32>,
17676 md: &CudaSlice<f32>,
17677 sel: &CudaSlice<i32>,
17678 y: &mut CudaSlice<f32>,
17679 width: usize,
17680 n_rows: usize,
17681 row0: usize,
17682 ) -> Result<(), Box<dyn std::error::Error>> {
17683 if x.len() < (row0 + n_rows) * width
17684 || w_route.len() < row0 + n_rows
17685 || sel.len() < row0 + n_rows
17686 || y.len() < width
17687 {
17688 return Err(format!(
17689 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17690 rows={n_rows} row0={row0}",
17691 x.len(),
17692 w_route.len(),
17693 sel.len(),
17694 y.len()
17695 )
17696 .into());
17697 }
17698 let f = self.func("axpy_rows_seq_md_off_f32");
17699 let cfg = LaunchConfig::for_num_elems(width as u32);
17700 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17701 let __s_b = self.gpu.stream();
17702 let mut b = __s_b.launch_builder(&f);
17703 b.arg(x)
17704 .arg(w_route)
17705 .arg(md)
17706 .arg(sel)
17707 .arg(y)
17708 .arg(&wi)
17709 .arg(&nr)
17710 .arg(&r0);
17711 unsafe {
17712 b.launch(cfg)?;
17713 }
17714 Ok(())
17715 }
17716
17717 #[allow(clippy::too_many_arguments)]
17722 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17723 &self,
17724 gate_bank: &CudaSlice<u8>,
17725 up_bank: &CudaSlice<u8>,
17726 sel: &CudaSlice<i32>,
17727 aq: &CudaSlice<i8>,
17728 ad: &CudaSlice<f32>,
17729 yg: &mut CudaSlice<f32>,
17730 yu: &mut CudaSlice<f32>,
17731 n_sel: usize,
17732 n_sel_col: usize,
17733 in_f: usize,
17734 out_f: usize,
17735 row_bytes: usize,
17736 expert_stride: usize,
17737 act_row_stride: usize,
17738 ad_row_stride: usize,
17739 ) -> Result<(), Box<dyn std::error::Error>> {
17740 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17741 if yg.len() < n_sel * out_f
17742 || yu.len() < n_sel * out_f
17743 || sel.len() < n_sel
17744 || n_sel_col == 0
17745 || n_sel % n_sel_col != 0
17746 {
17747 return Err("NVFP4 gu tcol geometry".into());
17748 }
17749 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17750 let cfg = LaunchConfig {
17751 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17752 block_dim: (128, 1, 1),
17753 shared_mem_bytes: 0,
17754 };
17755 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17756 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17757 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17758 let __s_b = self.gpu.stream();
17759 let mut b = __s_b.launch_builder(&f);
17760 b.arg(gate_bank)
17761 .arg(up_bank)
17762 .arg(sel)
17763 .arg(aq)
17764 .arg(ad)
17765 .arg(yg)
17766 .arg(yu)
17767 .arg(&inf)
17768 .arg(&outf)
17769 .arg(&ns)
17770 .arg(&rb)
17771 .arg(&es)
17772 .arg(&ars)
17773 .arg(&adrs)
17774 .arg(&nsc);
17775 unsafe {
17776 b.launch(cfg)?;
17777 }
17778 Ok(())
17779 }
17780
17781 #[allow(clippy::too_many_arguments)]
17784 pub fn axpy_rows_seq_md_into(
17785 &self,
17786 x: &CudaSlice<f32>,
17787 w_route: &CudaSlice<f32>,
17788 md: &CudaSlice<f32>,
17789 sel: &CudaSlice<i32>,
17790 y: &mut CudaSlice<f32>,
17791 width: usize,
17792 n_rows: usize,
17793 ) -> Result<(), Box<dyn std::error::Error>> {
17794 if x.len() < n_rows * width
17795 || w_route.len() < n_rows
17796 || sel.len() < n_rows
17797 || y.len() < width
17798 {
17799 return Err(format!(
17800 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17801 x.len(),
17802 w_route.len(),
17803 sel.len(),
17804 y.len()
17805 )
17806 .into());
17807 }
17808 let f = self.func("axpy_rows_seq_md_f32");
17809 let cfg = LaunchConfig::for_num_elems(width as u32);
17810 let (wi, nr) = (width as i32, n_rows as i32);
17811 let __s_b = self.gpu.stream();
17812 let mut b = __s_b.launch_builder(&f);
17813 b.arg(x)
17814 .arg(w_route)
17815 .arg(md)
17816 .arg(sel)
17817 .arg(y)
17818 .arg(&wi)
17819 .arg(&nr);
17820 unsafe {
17821 b.launch(cfg)?;
17822 }
17823 Ok(())
17824 }
17825
17826 #[allow(clippy::too_many_arguments)]
17828 #[allow(clippy::too_many_arguments)]
17832 pub fn matvec_bf16_qkvg_tcol_into(
17833 &self,
17834 wq: &CudaSlice<u8>,
17835 wk: &CudaSlice<u8>,
17836 wv: &CudaSlice<u8>,
17837 wg: &CudaSlice<u8>,
17838 x_t: &CudaSlice<f32>,
17839 yq: &mut CudaSlice<f32>,
17840 yk: &mut CudaSlice<f32>,
17841 yv: &mut CudaSlice<f32>,
17842 yg: &mut CudaSlice<f32>,
17843 in_f: usize,
17844 out_q: usize,
17845 out_kv: usize,
17846 out_g: usize,
17847 t: usize,
17848 ) -> Result<(), Box<dyn std::error::Error>> {
17849 if t == 0
17850 || t > 8
17851 || in_f % 8 != 0
17852 || x_t.len() < t * in_f
17853 || yq.len() < t * out_q
17854 || yk.len() < t * out_kv
17855 || yv.len() < t * out_kv
17856 || (out_g > 0 && yg.len() < t * out_g)
17857 {
17858 return Err("matvec_bf16_qkvg_tcol geometry".into());
17859 }
17860 let grid = out_q + 2 * out_kv + out_g;
17861 let cfg = LaunchConfig {
17862 grid_dim: (grid as u32, 1, 1),
17863 block_dim: (mmv_block(), 1, 1),
17864 shared_mem_bytes: 0,
17865 };
17866 let (ini, oq, okv, og, ti) = (
17867 in_f as i32,
17868 out_q as i32,
17869 out_kv as i32,
17870 out_g as i32,
17871 t as i32,
17872 );
17873 let __s_b = self.gpu.stream();
17874 if let Some(name) = match t {
17876 2 => Some("matvec_bf16_qkvg_tcol_t2"),
17877 4 => Some("matvec_bf16_qkvg_tcol_t4"),
17878 8 => Some("matvec_bf16_qkvg_tcol_t8"),
17879 _ => None,
17880 } {
17881 let f = self.func(name);
17882 let mut b = __s_b.launch_builder(&f);
17883 b.arg(wq)
17884 .arg(wk)
17885 .arg(wv)
17886 .arg(wg)
17887 .arg(x_t)
17888 .arg(yq)
17889 .arg(yk)
17890 .arg(yv)
17891 .arg(yg)
17892 .arg(&ini)
17893 .arg(&oq)
17894 .arg(&okv)
17895 .arg(&og);
17896 unsafe {
17897 b.launch(cfg)?;
17898 }
17899 return Ok(());
17900 }
17901 let f = self.func("matvec_bf16_qkvg_tcol");
17902 let mut b = __s_b.launch_builder(&f);
17903 b.arg(wq)
17904 .arg(wk)
17905 .arg(wv)
17906 .arg(wg)
17907 .arg(x_t)
17908 .arg(yq)
17909 .arg(yk)
17910 .arg(yv)
17911 .arg(yg)
17912 .arg(&ini)
17913 .arg(&oq)
17914 .arg(&okv)
17915 .arg(&og)
17916 .arg(&ti);
17917 unsafe {
17918 b.launch(cfg)?;
17919 }
17920 Ok(())
17921 }
17922
17923 pub fn matvec_bf16_qkvg_into(
17924 &self,
17925 wq: &CudaSlice<u8>,
17926 wk: &CudaSlice<u8>,
17927 wv: &CudaSlice<u8>,
17928 wg: &CudaSlice<u8>,
17929 x: &CudaSlice<f32>,
17930 yq: &mut CudaSlice<f32>,
17931 yk: &mut CudaSlice<f32>,
17932 yv: &mut CudaSlice<f32>,
17933 yg: &mut CudaSlice<f32>,
17934 in_f: usize,
17935 out_q: usize,
17936 out_kv: usize,
17937 out_g: usize,
17938 ) -> Result<(), Box<dyn std::error::Error>> {
17939 if in_f % 8 != 0
17940 || wq.len() != out_q * in_f * 2
17941 || wk.len() != out_kv * in_f * 2
17942 || wv.len() != out_kv * in_f * 2
17943 || wg.len() < out_g * in_f * 2
17944 || x.len() < in_f
17945 || yq.len() < out_q
17946 || yk.len() < out_kv
17947 || yv.len() < out_kv
17948 || (out_g > 0 && yg.len() < out_g)
17949 {
17950 return Err(format!(
17951 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17952 )
17953 .into());
17954 }
17955 let f = self.func("matvec_bf16_qkvg");
17956 let cfg = LaunchConfig {
17957 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17958 block_dim: (mmv_block(), 1, 1),
17959 shared_mem_bytes: 0,
17960 };
17961 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17962 let __s_b = self.gpu.stream();
17963 let mut b = __s_b.launch_builder(&f);
17964 b.arg(wq)
17965 .arg(wk)
17966 .arg(wv)
17967 .arg(wg)
17968 .arg(x)
17969 .arg(yq)
17970 .arg(yk)
17971 .arg(yv)
17972 .arg(yg)
17973 .arg(&inf)
17974 .arg(&oq)
17975 .arg(&okv)
17976 .arg(&og);
17977 unsafe {
17978 b.launch(cfg)?;
17979 }
17980 Ok(())
17981 }
17982
17983 pub fn matvec_bf16_b4_into(
17985 &self,
17986 w: [&CudaSlice<u8>; 4],
17987 x: &CudaSlice<f32>,
17988 y: &mut CudaSlice<f32>,
17989 block_cols: usize,
17990 out_f: usize,
17991 ) -> Result<(), Box<dyn std::error::Error>> {
17992 if block_cols % 8 != 0
17993 || x.len() < 4 * block_cols
17994 || y.len() < out_f
17995 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17996 {
17997 return Err(format!(
17998 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17999 x.len()
18000 )
18001 .into());
18002 }
18003 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18006 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18007 let f = self.func(if x2 {
18008 "matvec_bf16_b4_x2"
18009 } else {
18010 "matvec_bf16_b4"
18011 });
18012 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18013 let cfg = LaunchConfig {
18014 grid_dim: (grid as u32, 1, 1),
18015 block_dim: (mmv_block(), 1, 1),
18016 shared_mem_bytes: 0,
18017 };
18018 let (bc, of) = (block_cols as i32, out_f as i32);
18019 let __s_b = self.gpu.stream();
18020 let mut b = __s_b.launch_builder(&f);
18021 b.arg(w[0])
18022 .arg(w[1])
18023 .arg(w[2])
18024 .arg(w[3])
18025 .arg(x)
18026 .arg(y)
18027 .arg(&bc)
18028 .arg(&of);
18029 unsafe {
18030 b.launch(cfg)?;
18031 }
18032 Ok(())
18033 }
18034
18035 pub fn matvec_bf16_b4_tcol_into(
18041 &self,
18042 w: [&CudaSlice<u8>; 4],
18043 x_t: &CudaSlice<f32>,
18044 y_t: &mut CudaSlice<f32>,
18045 block_cols: usize,
18046 out_f: usize,
18047 t: usize,
18048 ) -> Result<(), Box<dyn std::error::Error>> {
18049 if block_cols % 8 != 0
18050 || t == 0
18051 || t > 8
18052 || x_t.len() < t * 4 * block_cols
18053 || y_t.len() < t * out_f
18054 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18055 {
18056 return Err(format!(
18057 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18058 x_t.len()
18059 )
18060 .into());
18061 }
18062 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18063 return Err(
18064 "b4 tcol verify is qualified against the plain b4 kernel only \
18065 (MEMRA_B4_X2=1 is a different t=1 program)"
18066 .into(),
18067 );
18068 }
18069 let cfg = LaunchConfig {
18073 grid_dim: (out_f as u32, 1, 1),
18074 block_dim: (mmv_block(), 1, 1),
18075 shared_mem_bytes: 0,
18076 };
18077 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18078 let __s_b = self.gpu.stream();
18079 if let Some(name) = match t {
18080 2 => Some("matvec_bf16_b4_tcol_t2"),
18081 4 => Some("matvec_bf16_b4_tcol_t4"),
18082 8 => Some("matvec_bf16_b4_tcol_t8"),
18083 _ => None,
18084 } {
18085 let f = self.func(name);
18086 let mut b = __s_b.launch_builder(&f);
18087 b.arg(w[0])
18088 .arg(w[1])
18089 .arg(w[2])
18090 .arg(w[3])
18091 .arg(x_t)
18092 .arg(y_t)
18093 .arg(&bc)
18094 .arg(&of);
18095 unsafe {
18096 b.launch(cfg)?;
18097 }
18098 return Ok(());
18099 }
18100 let f = self.func("matvec_bf16_b4_tcol");
18101 let mut b = __s_b.launch_builder(&f);
18102 b.arg(w[0])
18103 .arg(w[1])
18104 .arg(w[2])
18105 .arg(w[3])
18106 .arg(x_t)
18107 .arg(y_t)
18108 .arg(&bc)
18109 .arg(&of)
18110 .arg(&ti);
18111 unsafe {
18112 b.launch(cfg)?;
18113 }
18114 Ok(())
18115 }
18116
18117 pub fn q8_0_row_bytes(in_f: usize) -> usize {
18120 in_f / 32 * 34
18121 }
18122
18123 pub fn encode_q8_0_from_bf16(
18127 &self,
18128 w_bf16: &CudaSlice<u8>,
18129 out: &mut CudaSlice<u8>,
18130 in_f: usize,
18131 out_f: usize,
18132 ) -> Result<(), Box<dyn std::error::Error>> {
18133 if in_f % 32 != 0
18134 || w_bf16.len() < in_f * out_f * 2
18135 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18136 {
18137 return Err(format!(
18138 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18139 w_bf16.len(),
18140 out.len()
18141 )
18142 .into());
18143 }
18144 let f = self.func("encode_q8_0_rows_from_bf16");
18145 const PAIRS_PER_BLOCK: u32 = 4;
18148 let pairs = (out_f * (in_f / 32)) as u64;
18149 let cfg = LaunchConfig {
18150 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18151 block_dim: (32, PAIRS_PER_BLOCK, 1),
18152 shared_mem_bytes: 0,
18153 };
18154 let (ini, outi) = (in_f as i32, out_f as i32);
18155 let __s_b = self.gpu.stream();
18156 let mut b = __s_b.launch_builder(&f);
18157 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18158 unsafe {
18159 b.launch(cfg)?;
18160 }
18161 Ok(())
18162 }
18163
18164 #[allow(clippy::too_many_arguments)]
18169 pub fn qmatvec_q8_0_qkv_rp_into(
18170 &self,
18171 wq: &CudaSlice<u8>,
18172 wk: &CudaSlice<u8>,
18173 wv: &CudaSlice<u8>,
18174 aq: &CudaSlice<i8>,
18175 ad: &CudaSlice<f32>,
18176 yq: &mut CudaSlice<f32>,
18177 yk: &mut CudaSlice<f32>,
18178 yv: &mut CudaSlice<f32>,
18179 in_f: usize,
18180 out_q: usize,
18181 out_kv: usize,
18182 ) -> Result<(), Box<dyn std::error::Error>> {
18183 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
18185 let nblk = in_f / 32;
18186 if in_f % 32 != 0
18187 || aq.len() < in_f
18188 || ad.len() < nblk
18189 || yq.len() < out_q
18190 || yk.len() < out_kv
18191 || yv.len() < out_kv
18192 || wq.len() < out_q * nblk * 34
18193 || wk.len() < out_kv * nblk * 34
18194 || wv.len() < out_kv * nblk * 34
18195 {
18196 return Err(
18197 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18198 );
18199 }
18200 let f = self.func("qmatvec_q8_0_qkv_rp");
18201 let cfg = LaunchConfig {
18202 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18203 block_dim: (32, ROWS_PER_BLOCK, 1),
18204 shared_mem_bytes: 0,
18205 };
18206 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18207 let __s_b = self.gpu.stream();
18208 let mut b = __s_b.launch_builder(&f);
18209 b.arg(wq)
18210 .arg(wk)
18211 .arg(wv)
18212 .arg(aq)
18213 .arg(ad)
18214 .arg(yq)
18215 .arg(yk)
18216 .arg(yv)
18217 .arg(&ini)
18218 .arg(&oq)
18219 .arg(&okv);
18220 unsafe {
18221 b.launch(cfg)?;
18222 }
18223 Ok(())
18224 }
18225
18226 #[allow(clippy::too_many_arguments)]
18230 pub fn qmatvec_q8_0_b4_rp_into(
18231 &self,
18232 w: [&CudaSlice<u8>; 4],
18233 aq: &CudaSlice<i8>,
18234 ad: &CudaSlice<f32>,
18235 y: &mut CudaSlice<f32>,
18236 block_cols: usize,
18237 out_f: usize,
18238 ) -> Result<(), Box<dyn std::error::Error>> {
18239 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
18241 if block_cols % 32 != 0
18242 || aq.len() < 4 * block_cols
18243 || ad.len() < 4 * nblk
18244 || y.len() < out_f
18245 || w.iter().any(|p| p.len() < out_f * nblk * 34)
18246 {
18247 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18248 }
18249 let f = self.func("qmatvec_q8_0_b4_rp");
18250 let cfg = LaunchConfig {
18251 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18252 block_dim: (32, ROWS_PER_BLOCK, 1),
18253 shared_mem_bytes: 0,
18254 };
18255 let (bc, of) = (block_cols as i32, out_f as i32);
18256 let __s_b = self.gpu.stream();
18257 let mut b = __s_b.launch_builder(&f);
18258 b.arg(w[0])
18259 .arg(w[1])
18260 .arg(w[2])
18261 .arg(w[3])
18262 .arg(aq)
18263 .arg(ad)
18264 .arg(y)
18265 .arg(&bc)
18266 .arg(&of);
18267 unsafe {
18268 b.launch(cfg)?;
18269 }
18270 Ok(())
18271 }
18272
18273 fn matvec_bf16_via_q8_mirror(
18276 &self,
18277 data: &CudaSlice<u8>,
18278 x: &CudaSlice<f32>,
18279 y: &mut CudaSlice<f32>,
18280 in_f: usize,
18281 out_f: usize,
18282 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18283 use cudarc::driver::DevicePtr;
18284 let key = {
18285 let s = self.gpu.stream();
18286 let (p, _g) = data.device_ptr(&s);
18287 p as u64
18288 };
18289 {
18290 let mut mirrors = self
18291 .w8_mirrors
18292 .lock()
18293 .map_err(|_| "w8 mirror map is poisoned")?;
18294 if !mirrors.contains_key(&key) {
18295 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18296 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18297 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18298 mirrors.insert(key, planar);
18299 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18305 eprintln!(
18306 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18307 mirrors.len()
18308 );
18309 }
18310 }
18311 }
18312 let nblk = in_f / 32;
18313 {
18314 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18315 if !act.contains_key(&in_f) {
18316 let aq = self.alloc_uninit::<i8>(in_f)?;
18317 let ad = self.alloc_uninit::<f32>(nblk)?;
18318 act.insert(in_f, (aq, ad));
18319 }
18320 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18321 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18322 }
18323 let mirrors = self
18324 .w8_mirrors
18325 .lock()
18326 .map_err(|_| "w8 mirror map is poisoned")?;
18327 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18328 let mirror = mirrors.get(&key).expect("built above");
18329 let (aq, ad) = act.get(&in_f).expect("built above");
18330 self.qmatvec_mmvq_into(
18331 mirror,
18332 aq,
18333 ad,
18334 1,
18335 in_f,
18336 out_f,
18337 QT_Q8_0,
18338 Self::q8_0_row_bytes(in_f),
18339 1.0,
18340 true,
18341 y,
18342 )?;
18343 Ok(Some(()))
18344 }
18345
18346 pub fn matvec_bf16_into(
18347 &self,
18348 data: &CudaSlice<u8>,
18349 x: &CudaSlice<f32>,
18350 y: &mut CudaSlice<f32>,
18351 in_f: usize,
18352 out_f: usize,
18353 ) -> Result<(), Box<dyn std::error::Error>> {
18354 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18355 return Err(format!(
18356 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18357 data.len(),
18358 x.len(),
18359 y.len()
18360 )
18361 .into());
18362 }
18363 if step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18370 if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
18371 return Ok(());
18372 }
18373 }
18374 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18378 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18379 && in_f <= 2048;
18380 if x4 {
18381 let f = self.func("matvec_bf16_f32acc_x4");
18382 let cfg = LaunchConfig {
18383 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18384 block_dim: (mmv_block(), 1, 1),
18385 shared_mem_bytes: 0,
18386 };
18387 let (ini, outi) = (in_f as i32, out_f as i32);
18388 let __s_b = self.gpu.stream();
18389 let mut b = __s_b.launch_builder(&f);
18390 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18391 unsafe {
18392 b.launch(cfg)?;
18393 }
18394 return Ok(());
18395 }
18396 let f = self.func("matvec_bf16_f32acc");
18397 let cfg = LaunchConfig {
18398 grid_dim: (out_f as u32, 1, 1),
18399 block_dim: (mmv_block(), 1, 1),
18400 shared_mem_bytes: 0,
18401 };
18402 let ini = in_f as i32;
18403 let __s_b = self.gpu.stream();
18404 let mut b = __s_b.launch_builder(&f);
18405 b.arg(data).arg(x).arg(y).arg(&ini);
18406 unsafe {
18407 b.launch(cfg)?;
18408 }
18409 Ok(())
18410 }
18411
18412 pub fn matvec_bf16_view_into(
18415 &self,
18416 data: &cudarc::driver::CudaView<'_, u8>,
18417 x: &CudaSlice<f32>,
18418 y: &mut CudaSlice<f32>,
18419 in_f: usize,
18420 out_f: usize,
18421 ) -> Result<(), Box<dyn std::error::Error>> {
18422 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18423 return Err(format!(
18424 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18425 data.len(),
18426 x.len(),
18427 y.len()
18428 )
18429 .into());
18430 }
18431 let f = self.func("matvec_bf16_f32acc");
18432 let cfg = LaunchConfig {
18433 grid_dim: (out_f as u32, 1, 1),
18434 block_dim: (mmv_block(), 1, 1),
18435 shared_mem_bytes: 0,
18436 };
18437 let ini = in_f as i32;
18438 let __s_b = self.gpu.stream();
18439 let mut b = __s_b.launch_builder(&f);
18440 b.arg(data).arg(x).arg(y).arg(&ini);
18441 unsafe {
18442 b.launch(cfg)?;
18443 }
18444 Ok(())
18445 }
18446
18447 pub fn matvec_bf16_raw_out(
18450 &self,
18451 w: &CudaSlice<u8>,
18452 x: &CudaSlice<f32>,
18453 y_raw: u64,
18454 in_f: usize,
18455 out_f: usize,
18456 ) -> Result<(), Box<dyn std::error::Error>> {
18457 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18458 return Err("matvec_bf16_raw_out geometry".into());
18459 }
18460 let f = self.func("matvec_bf16_f32acc");
18461 let cfg = LaunchConfig {
18462 grid_dim: (out_f as u32, 1, 1),
18463 block_dim: (mmv_block(), 1, 1),
18464 shared_mem_bytes: 0,
18465 };
18466 let ini = in_f as i32;
18467 let __s_b = self.gpu.stream();
18468 let mut b = __s_b.launch_builder(&f);
18469 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18470 unsafe {
18471 b.launch(cfg)?;
18472 }
18473 Ok(())
18474 }
18475
18476 pub fn add3_raw(
18480 &self,
18481 a: &CudaSlice<f32>,
18482 b: &CudaSlice<f32>,
18483 sh_raw: u64,
18484 scale_raw: u64,
18485 dst: &mut CudaSlice<f32>,
18486 n: usize,
18487 ) -> Result<(), Box<dyn std::error::Error>> {
18488 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18489 return Err("add3_raw geometry".into());
18490 }
18491 let f = self.func("add3_f32");
18492 let cfg = LaunchConfig {
18493 grid_dim: ((n as u32).div_ceil(256), 1, 1),
18494 block_dim: (256, 1, 1),
18495 shared_mem_bytes: 0,
18496 };
18497 let ni = n as i32;
18498 let __s_b = self.gpu.stream();
18499 let mut bld = __s_b.launch_builder(&f);
18500 bld.arg(a)
18501 .arg(b)
18502 .arg(&sh_raw)
18503 .arg(&scale_raw)
18504 .arg(dst)
18505 .arg(&ni);
18506 unsafe {
18507 bld.launch(cfg)?;
18508 }
18509 Ok(())
18510 }
18511
18512 pub fn matvec_bf16_down_addscale_into(
18515 &self,
18516 w: &CudaSlice<u8>,
18517 x: &CudaSlice<f32>,
18518 scale: &CudaSlice<f32>,
18519 dst: &mut CudaSlice<f32>,
18520 in_f: usize,
18521 out_f: usize,
18522 ) -> Result<(), Box<dyn std::error::Error>> {
18523 if w.len() != in_f * out_f * 2
18524 || x.len() < in_f
18525 || in_f % 8 != 0
18526 || dst.len() < out_f
18527 || scale.is_empty()
18528 {
18529 return Err("matvec_bf16_down_addscale geometry".into());
18530 }
18531 let f = self.func("matvec_bf16_down_addscale");
18532 let cfg = LaunchConfig {
18533 grid_dim: (out_f as u32, 1, 1),
18534 block_dim: (mmv_block(), 1, 1),
18535 shared_mem_bytes: 0,
18536 };
18537 let ini = in_f as i32;
18538 let __s_b = self.gpu.stream();
18539 let mut b = __s_b.launch_builder(&f);
18540 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18541 unsafe {
18542 b.launch(cfg)?;
18543 }
18544 Ok(())
18545 }
18546
18547 #[allow(clippy::too_many_arguments)]
18551 pub fn matvec_bf16_dual_silu_rows_into(
18552 &self,
18553 wg: &CudaSlice<u8>,
18554 wu: &CudaSlice<u8>,
18555 x: &CudaSlice<f32>,
18556 act: &mut CudaSlice<f32>,
18557 in_f: usize,
18558 out_f: usize,
18559 limit: Option<f32>,
18560 t: usize,
18561 ) -> Result<(), Box<dyn std::error::Error>> {
18562 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18563 return Err("matvec_bf16_dual_silu_rows geometry".into());
18564 }
18565 let f = self.func("matvec_bf16_dual_silu_rows");
18566 let cfg = LaunchConfig {
18567 grid_dim: (out_f as u32, t as u32, 1),
18568 block_dim: (mmv_block(), 1, 1),
18569 shared_mem_bytes: 0,
18570 };
18571 let (ini, outi) = (in_f as i32, out_f as i32);
18572 let lim = limit.unwrap_or(0.0);
18573 let __s_b = self.gpu.stream();
18574 let mut b = __s_b.launch_builder(&f);
18575 b.arg(wg)
18576 .arg(wu)
18577 .arg(x)
18578 .arg(&mut *act)
18579 .arg(&ini)
18580 .arg(&outi)
18581 .arg(&lim);
18582 unsafe {
18583 b.launch(cfg)?;
18584 }
18585 Ok(())
18586 }
18587
18588 pub fn matvec_bf16_rows_into(
18590 &self,
18591 w: &CudaSlice<u8>,
18592 x: &CudaSlice<f32>,
18593 y: &mut CudaSlice<f32>,
18594 in_f: usize,
18595 out_f: usize,
18596 t: usize,
18597 ) -> Result<(), Box<dyn std::error::Error>> {
18598 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18599 return Err("matvec_bf16_rows geometry".into());
18600 }
18601 if t == 1 && step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18607 if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
18608 return Ok(());
18609 }
18610 }
18611 let f = self.func("matvec_bf16_f32acc_x4_rows");
18612 let cfg = LaunchConfig {
18613 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18614 block_dim: (mmv_block(), 1, 1),
18615 shared_mem_bytes: 0,
18616 };
18617 let (ini, outi) = (in_f as i32, out_f as i32);
18618 let __s_b = self.gpu.stream();
18619 let mut b = __s_b.launch_builder(&f);
18620 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18621 unsafe {
18622 b.launch(cfg)?;
18623 }
18624 Ok(())
18625 }
18626
18627 pub fn matvec_bf16_dual_silu_into(
18628 &self,
18629 wg: &CudaSlice<u8>,
18630 wu: &CudaSlice<u8>,
18631 x: &CudaSlice<f32>,
18632 act: &mut CudaSlice<f32>,
18633 in_f: usize,
18634 out_f: usize,
18635 limit: Option<f32>,
18636 ) -> Result<(), Box<dyn std::error::Error>> {
18637 if wg.len() != in_f * out_f * 2
18638 || wu.len() != in_f * out_f * 2
18639 || x.len() < in_f
18640 || in_f % 8 != 0
18641 || act.len() < out_f
18642 {
18643 return Err("matvec_bf16_dual_silu geometry".into());
18644 }
18645 let f = self.func("matvec_bf16_dual_silu");
18646 let cfg = LaunchConfig {
18647 grid_dim: (out_f as u32, 1, 1),
18648 block_dim: (mmv_block(), 1, 1),
18649 shared_mem_bytes: 0,
18650 };
18651 let (ini, outi) = (in_f as i32, out_f as i32);
18652 let lim = limit.unwrap_or(0.0);
18653 let __s_b = self.gpu.stream();
18654 let mut b = __s_b.launch_builder(&f);
18655 b.arg(wg)
18656 .arg(wu)
18657 .arg(x)
18658 .arg(act)
18659 .arg(&ini)
18660 .arg(&outi)
18661 .arg(&lim);
18662 unsafe {
18663 b.launch(cfg)?;
18664 }
18665 Ok(())
18666 }
18667
18668 #[allow(clippy::too_many_arguments)]
18671 pub fn matvec_bf16_dual_view_into(
18672 &self,
18673 wg: &cudarc::driver::CudaView<'_, u8>,
18674 wu: &cudarc::driver::CudaView<'_, u8>,
18675 x: &CudaSlice<f32>,
18676 yg: &mut CudaSlice<f32>,
18677 yu: &mut CudaSlice<f32>,
18678 in_f: usize,
18679 out_f: usize,
18680 ) -> Result<(), Box<dyn std::error::Error>> {
18681 if wg.len() != in_f * out_f * 2
18682 || wu.len() != in_f * out_f * 2
18683 || x.len() < in_f
18684 || in_f % 8 != 0
18685 || yg.len() < out_f
18686 || yu.len() < out_f
18687 {
18688 return Err(format!(
18689 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18690 wg.len(),
18691 wu.len(),
18692 x.len()
18693 )
18694 .into());
18695 }
18696 let f = self.func("matvec_bf16_dual");
18697 let cfg = LaunchConfig {
18698 grid_dim: ((2 * out_f) as u32, 1, 1),
18699 block_dim: (mmv_block(), 1, 1),
18700 shared_mem_bytes: 0,
18701 };
18702 let (ini, outi) = (in_f as i32, out_f as i32);
18703 let __s_b = self.gpu.stream();
18704 let mut b = __s_b.launch_builder(&f);
18705 b.arg(wg)
18706 .arg(wu)
18707 .arg(x)
18708 .arg(yg)
18709 .arg(yu)
18710 .arg(&ini)
18711 .arg(&outi);
18712 unsafe {
18713 b.launch(cfg)?;
18714 }
18715 Ok(())
18716 }
18717
18718 #[allow(clippy::too_many_arguments)]
18720 pub fn matvec_bf16_dual_into(
18721 &self,
18722 wg: &CudaSlice<u8>,
18723 wu: &CudaSlice<u8>,
18724 x: &CudaSlice<f32>,
18725 yg: &mut CudaSlice<f32>,
18726 yu: &mut CudaSlice<f32>,
18727 in_f: usize,
18728 out_f: usize,
18729 ) -> Result<(), Box<dyn std::error::Error>> {
18730 if wg.len() != in_f * out_f * 2
18731 || wu.len() != in_f * out_f * 2
18732 || x.len() < in_f
18733 || in_f % 8 != 0
18734 || yg.len() < out_f
18735 || yu.len() < out_f
18736 {
18737 return Err(format!(
18738 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18739 wg.len(),
18740 wu.len(),
18741 x.len()
18742 )
18743 .into());
18744 }
18745 let f = self.func("matvec_bf16_dual");
18746 let cfg = LaunchConfig {
18747 grid_dim: ((2 * out_f) as u32, 1, 1),
18748 block_dim: (mmv_block(), 1, 1),
18749 shared_mem_bytes: 0,
18750 };
18751 let (ini, outi) = (in_f as i32, out_f as i32);
18752 let __s_b = self.gpu.stream();
18753 let mut b = __s_b.launch_builder(&f);
18754 b.arg(wg)
18755 .arg(wu)
18756 .arg(x)
18757 .arg(yg)
18758 .arg(yu)
18759 .arg(&ini)
18760 .arg(&outi);
18761 unsafe {
18762 b.launch(cfg)?;
18763 }
18764 Ok(())
18765 }
18766
18767 pub(crate) fn matvec_bf16_dual(
18770 &self,
18771 wg: &CudaSlice<u8>,
18772 wu: &CudaSlice<u8>,
18773 x: &CudaSlice<f32>,
18774 in_f: usize,
18775 out_f: usize,
18776 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18777 if wg.len() != in_f * out_f * 2
18778 || wu.len() != in_f * out_f * 2
18779 || x.len() < in_f
18780 || in_f % 8 != 0
18781 {
18782 return Err(format!(
18783 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18784 wg.len(),
18785 wu.len(),
18786 x.len()
18787 )
18788 .into());
18789 }
18790 let mut yg = self.alloc_uninit::<f32>(out_f)?;
18791 let mut yu = self.alloc_uninit::<f32>(out_f)?;
18792 let f = self.func("matvec_bf16_dual");
18793 let cfg = LaunchConfig {
18794 grid_dim: ((2 * out_f) as u32, 1, 1),
18795 block_dim: (mmv_block(), 1, 1),
18796 shared_mem_bytes: 0,
18797 };
18798 let (ini, outi) = (in_f as i32, out_f as i32);
18799 let __s_b = self.gpu.stream();
18800 let mut b = __s_b.launch_builder(&f);
18801 b.arg(wg)
18802 .arg(wu)
18803 .arg(x)
18804 .arg(&mut yg)
18805 .arg(&mut yu)
18806 .arg(&ini)
18807 .arg(&outi);
18808 unsafe {
18809 b.launch(cfg)?;
18810 }
18811 Ok((yg, yu))
18812 }
18813
18814 #[allow(clippy::too_many_arguments)]
18815 fn linear_bf16_chunked_inner(
18816 &self,
18817 x: &CudaSlice<f32>,
18818 data: &CudaSlice<u8>,
18819 m: usize,
18820 in_f: usize,
18821 out_f: usize,
18822 exact: bool,
18823 canonical_chunk_rows: Option<usize>,
18824 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18825 const CHUNK_BYTES: usize = 256 << 20;
18826 if m == 1
18829 && !exact
18830 && canonical_chunk_rows.is_none()
18831 && in_f % 8 == 0
18832 && Self::bf16_mmv_on()
18833 {
18834 return self.matvec_bf16(data, x, in_f, out_f);
18835 }
18836 let row_bytes = in_f
18837 .checked_mul(std::mem::size_of::<f32>())
18838 .ok_or("BF16 chunk row byte count overflow")?;
18839 if row_bytes == 0 || out_f == 0 {
18840 return Err("BF16 chunk dimensions must be nonzero".into());
18841 }
18842 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18843 let chunk_rows = match canonical_chunk_rows {
18844 Some(rows) if rows == 0 => {
18845 return Err("canonical BF16 chunk rows must be nonzero".into());
18846 }
18847 Some(rows) if rows > max_chunk_rows => {
18848 return Err(format!(
18849 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18850 )
18851 .into());
18852 }
18853 Some(rows) if out_f % rows != 0 => {
18854 return Err(format!(
18855 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18856 )
18857 .into());
18858 }
18859 Some(rows) => rows,
18860 None => max_chunk_rows,
18861 };
18862 if chunk_rows >= out_f {
18863 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18864 return if exact {
18865 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
18866 } else {
18867 self.linear(x, &wf32, m, in_f, out_f)
18868 };
18869 }
18870 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18871 let mut r0 = 0usize;
18872 while r0 < out_f {
18873 let rows = chunk_rows.min(out_f - r0);
18874 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
18875 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
18876 let yc = if exact {
18877 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
18878 } else {
18879 self.linear(x, &wf32, m, in_f, rows)?
18880 };
18881 for mi in 0..m {
18883 let src = yc.slice(mi * rows..(mi + 1) * rows);
18884 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
18885 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
18886 }
18887 r0 += rows;
18888 }
18889 Ok(y)
18890 }
18891
18892 pub fn linear_bf16_resident(
18896 &self,
18897 x: &CudaSlice<f32>,
18898 data: &CudaSlice<u8>,
18899 m: usize,
18900 in_f: usize,
18901 out_f: usize,
18902 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18903 if data.len() != in_f * out_f * 2 {
18904 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18905 }
18906 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
18907 }
18908
18909 pub fn linear_bf16_resident_canonical_rows(
18915 &self,
18916 x: &CudaSlice<f32>,
18917 data: &CudaSlice<u8>,
18918 m: usize,
18919 in_f: usize,
18920 out_f: usize,
18921 canonical_chunk_rows: usize,
18922 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18923 if data.len() != in_f * out_f * 2 {
18924 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18925 }
18926 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18927 }
18928
18929 pub fn linear_f32_resident_canonical_rows(
18934 &self,
18935 x: &CudaSlice<f32>,
18936 data: &CudaSlice<f32>,
18937 m: usize,
18938 in_f: usize,
18939 out_f: usize,
18940 canonical_chunk_rows: usize,
18941 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18942 self.linear_f32_resident_canonical_rows_inner(
18943 x,
18944 data,
18945 m,
18946 in_f,
18947 out_f,
18948 canonical_chunk_rows,
18949 false,
18950 )
18951 }
18952
18953 pub fn linear_f32_resident_canonical_rows_strided(
18959 &self,
18960 x: &CudaSlice<f32>,
18961 data: &CudaSlice<f32>,
18962 m: usize,
18963 in_f: usize,
18964 out_f: usize,
18965 canonical_chunk_rows: usize,
18966 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18967 self.linear_f32_resident_canonical_rows_inner(
18968 x,
18969 data,
18970 m,
18971 in_f,
18972 out_f,
18973 canonical_chunk_rows,
18974 true,
18975 )
18976 }
18977
18978 fn linear_f32_resident_canonical_rows_inner(
18979 &self,
18980 x: &CudaSlice<f32>,
18981 data: &CudaSlice<f32>,
18982 m: usize,
18983 in_f: usize,
18984 out_f: usize,
18985 canonical_chunk_rows: usize,
18986 strided_output: bool,
18987 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18988 if data.len() != in_f * out_f {
18989 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18990 }
18991 if canonical_chunk_rows == 0
18992 || canonical_chunk_rows > out_f
18993 || out_f % canonical_chunk_rows != 0
18994 {
18995 return Err(format!(
18996 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18997 )
18998 .into());
18999 }
19000 if canonical_chunk_rows == out_f {
19001 return self.linear(x, data, m, in_f, out_f);
19002 }
19003
19004 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19005 let input = x.slice(0..x.len());
19006 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19007 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19008 if m == 1 {
19009 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19010 self.linear_device_into(
19011 &input,
19012 &weights,
19013 &mut destination,
19014 1,
19015 in_f,
19016 canonical_chunk_rows,
19017 )?;
19018 continue;
19019 }
19020 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19021 if strided_output {
19022 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19023 } else {
19024 for token in 0..m {
19025 let source = chunk
19026 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19027 let mut destination =
19028 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19029 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19030 }
19031 }
19032 }
19033 Ok(y)
19034 }
19035
19036 pub fn linear_f32_resident_canonical_rows_t1_into(
19041 &self,
19042 x: &CudaSlice<f32>,
19043 data: &CudaSlice<f32>,
19044 y: &mut CudaSlice<f32>,
19045 in_f: usize,
19046 out_f: usize,
19047 canonical_chunk_rows: usize,
19048 ) -> Result<(), Box<dyn std::error::Error>> {
19049 if data.len() != in_f * out_f {
19050 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19051 }
19052 if y.len() != out_f || x.len() != in_f {
19053 return Err(format!(
19054 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19055 x.len(),
19056 y.len()
19057 )
19058 .into());
19059 }
19060 if canonical_chunk_rows == 0
19061 || canonical_chunk_rows > out_f
19062 || out_f % canonical_chunk_rows != 0
19063 {
19064 return Err(format!(
19065 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19066 )
19067 .into());
19068 }
19069 let input = x.slice(0..x.len());
19070 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19071 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19072 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19073 self.linear_device_into(
19074 &input,
19075 &weights,
19076 &mut destination,
19077 1,
19078 in_f,
19079 canonical_chunk_rows,
19080 )?;
19081 }
19082 Ok(())
19083 }
19084
19085 pub fn linear_t1_into(
19088 &self,
19089 x: &cudarc::driver::CudaView<'_, f32>,
19090 w: &cudarc::driver::CudaView<'_, f32>,
19091 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19092 in_f: usize,
19093 out_f: usize,
19094 ) -> Result<(), Box<dyn std::error::Error>> {
19095 self.linear_device_into(x, w, y, 1, in_f, out_f)
19096 }
19097
19098 pub fn linear_decode_exact(
19105 &self,
19106 x: &CudaSlice<f32>,
19107 w: &CudaSlice<f32>,
19108 m_tokens: usize,
19109 in_f: usize,
19110 out_f: usize,
19111 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19112 if m_tokens == 1 {
19113 return self.linear(x, w, 1, in_f, out_f);
19114 }
19115 let xv = self.view(x, m_tokens * in_f);
19116 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19117 for t in 0..m_tokens {
19118 let row = xv.slice(t * in_f..(t + 1) * in_f);
19119 let mut xr = self.alloc_uninit::<f32>(in_f)?;
19120 self.copy_view_into(&mut xr, 0, &row, in_f)?;
19121 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19122 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19123 }
19124 Ok(y)
19125 }
19126
19127 pub fn linear(
19128 &self,
19129 x: &CudaSlice<f32>,
19130 w: &CudaSlice<f32>,
19131 m_tokens: usize,
19132 in_f: usize,
19133 out_f: usize,
19134 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19135 self.linear_device(x, w, m_tokens, in_f, out_f)
19136 }
19137
19138 fn linear_device<I>(
19139 &self,
19140 x: &I,
19141 w: &I,
19142 m_tokens: usize,
19143 in_f: usize,
19144 out_f: usize,
19145 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19146 where
19147 I: cudarc::driver::DevicePtr<f32>,
19148 {
19149 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)?;
19151 Ok(c)
19152 }
19153
19154 fn linear_device_into<I, O>(
19155 &self,
19156 x: &I,
19157 w: &I,
19158 c: &mut O,
19159 m_tokens: usize,
19160 in_f: usize,
19161 out_f: usize,
19162 ) -> Result<(), Box<dyn std::error::Error>>
19163 where
19164 I: cudarc::driver::DevicePtr<f32>,
19165 O: cudarc::driver::DevicePtrMut<f32>,
19166 {
19167 use cudarc::cublaslt::{Matmul, MatmulConfig};
19168 let cfg = MatmulConfig {
19169 transa: true,
19170 transb: false,
19171 transc: false,
19172 m: out_f as u64,
19173 n: m_tokens as u64,
19174 k: in_f as u64,
19175 alpha: 1.0,
19176 lda: in_f as i64,
19177 ldb: in_f as i64,
19178 beta: 0.0,
19179 ldc: out_f as i64,
19180 stride_a: None,
19181 stride_b: None,
19182 stride_c: None,
19183 stride_bias: None,
19184 batch_size: None,
19185 };
19186 let blas = self.gpu.blas();
19187 unsafe {
19188 blas.matmul(cfg, w, x, c, None, None)?;
19189 }
19190 Ok(())
19191 }
19192
19193 pub fn sdpa_naive(
19202 &self,
19203 q: &CudaSlice<f32>,
19204 k: &CudaSlice<f32>,
19205 v: &CudaSlice<f32>,
19206 o: &mut CudaSlice<f32>,
19207 head_dim: usize,
19208 n_head: usize,
19209 n_head_kv: usize,
19210 t: usize,
19211 t_kv: usize,
19212 scale: f32,
19213 causal: bool,
19214 ) -> Result<(), Box<dyn std::error::Error>> {
19215 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19216 return self.sdpa_naive_gmem(
19217 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19218 );
19219 }
19220 let f = self.func("sdpa_naive_f32");
19221 let cfg = LaunchConfig {
19222 grid_dim: (n_head as u32, t as u32, 1),
19223 block_dim: (128, 1, 1),
19224 shared_mem_bytes: (t_kv * 4) as u32,
19225 };
19226 let (hd, nh, nhkv, ti, tkvi, cz) = (
19227 head_dim as i32,
19228 n_head as i32,
19229 n_head_kv as i32,
19230 t as i32,
19231 t_kv as i32,
19232 causal as i32,
19233 );
19234 let __s_b = self.gpu.stream();
19235 let mut b = __s_b.launch_builder(&f);
19236 b.arg(q)
19237 .arg(k)
19238 .arg(v)
19239 .arg(o)
19240 .arg(&hd)
19241 .arg(&nh)
19242 .arg(&nhkv)
19243 .arg(&ti)
19244 .arg(&tkvi)
19245 .arg(&scale)
19246 .arg(&cz);
19247 unsafe {
19248 b.launch(cfg)?;
19249 }
19250 Ok(())
19251 }
19252
19253 #[allow(clippy::too_many_arguments)]
19262 pub fn sdpa_naive_gmem(
19263 &self,
19264 q: &CudaSlice<f32>,
19265 k: &CudaSlice<f32>,
19266 v: &CudaSlice<f32>,
19267 o: &mut CudaSlice<f32>,
19268 head_dim: usize,
19269 n_head: usize,
19270 n_head_kv: usize,
19271 t: usize,
19272 t_kv: usize,
19273 scale: f32,
19274 causal: bool,
19275 ) -> Result<(), Box<dyn std::error::Error>> {
19276 let ws_len = n_head
19277 .checked_mul(t)
19278 .and_then(|x| x.checked_mul(t_kv))
19279 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19280 let ws_bytes = ws_len
19281 .checked_mul(std::mem::size_of::<f32>())
19282 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19283 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19284 return Err(format!(
19285 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19286 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19287 needs a tiled/flash kernel, not the naive oracle"
19288 )
19289 .into());
19290 }
19291 let mut scores = self.uninit(ws_len)?;
19292 let f = self.func("sdpa_naive_gmem_f32");
19293 let cfg = LaunchConfig {
19294 grid_dim: (n_head as u32, t as u32, 1),
19295 block_dim: (128, 1, 1),
19296 shared_mem_bytes: 0,
19297 };
19298 let (hd, nh, nhkv, ti, tkvi, cz) = (
19299 head_dim as i32,
19300 n_head as i32,
19301 n_head_kv as i32,
19302 t as i32,
19303 t_kv as i32,
19304 causal as i32,
19305 );
19306 let __s_b = self.gpu.stream();
19307 let mut b = __s_b.launch_builder(&f);
19308 b.arg(q)
19309 .arg(k)
19310 .arg(v)
19311 .arg(o)
19312 .arg(&mut scores)
19313 .arg(&hd)
19314 .arg(&nh)
19315 .arg(&nhkv)
19316 .arg(&ti)
19317 .arg(&tkvi)
19318 .arg(&scale)
19319 .arg(&cz);
19320 unsafe {
19321 b.launch(cfg)?;
19322 }
19323 Ok(())
19324 }
19325
19326 #[allow(clippy::too_many_arguments)]
19331 pub fn sdpa_naive_island(
19332 &self,
19333 q: &CudaSlice<f32>,
19334 k: &CudaSlice<f32>,
19335 v: &CudaSlice<f32>,
19336 o: &mut CudaSlice<f32>,
19337 span_id: &CudaSlice<i32>,
19338 head_dim: usize,
19339 n_head: usize,
19340 n_head_kv: usize,
19341 t: usize,
19342 t_kv: usize,
19343 scale: f32,
19344 window: usize,
19345 ) -> Result<(), Box<dyn std::error::Error>> {
19346 let f = self.func("sdpa_naive_island_f32");
19347 let cfg = LaunchConfig {
19348 grid_dim: (n_head as u32, t as u32, 1),
19349 block_dim: (128, 1, 1),
19350 shared_mem_bytes: (t_kv * 4) as u32,
19351 };
19352 let (hd, nh, nhkv, ti, tkvi, wi) = (
19353 head_dim as i32,
19354 n_head as i32,
19355 n_head_kv as i32,
19356 t as i32,
19357 t_kv as i32,
19358 window as i32,
19359 );
19360 let __s_b = self.gpu.stream();
19361 let mut b = __s_b.launch_builder(&f);
19362 b.arg(q)
19363 .arg(k)
19364 .arg(v)
19365 .arg(o)
19366 .arg(span_id)
19367 .arg(&hd)
19368 .arg(&nh)
19369 .arg(&nhkv)
19370 .arg(&ti)
19371 .arg(&tkvi)
19372 .arg(&scale)
19373 .arg(&wi);
19374 unsafe {
19375 b.launch(cfg)?;
19376 }
19377 Ok(())
19378 }
19379
19380 #[allow(clippy::too_many_arguments)]
19382 pub fn sdpa_naive_w(
19383 &self,
19384 q: &CudaSlice<f32>,
19385 k: &CudaSlice<f32>,
19386 v: &CudaSlice<f32>,
19387 o: &mut CudaSlice<f32>,
19388 head_dim: usize,
19389 n_head: usize,
19390 n_head_kv: usize,
19391 t: usize,
19392 t_kv: usize,
19393 scale: f32,
19394 causal: bool,
19395 window: usize,
19396 ) -> Result<(), Box<dyn std::error::Error>> {
19397 let f = self.func("sdpa_naive_w_f32");
19398 let cfg = LaunchConfig {
19399 grid_dim: (n_head as u32, t as u32, 1),
19400 block_dim: (128, 1, 1),
19401 shared_mem_bytes: (t_kv * 4) as u32,
19402 };
19403 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19404 head_dim as i32,
19405 n_head as i32,
19406 n_head_kv as i32,
19407 t as i32,
19408 t_kv as i32,
19409 causal as i32,
19410 window as i32,
19411 );
19412 let __s_b = self.gpu.stream();
19413 let mut b = __s_b.launch_builder(&f);
19414 b.arg(q)
19415 .arg(k)
19416 .arg(v)
19417 .arg(o)
19418 .arg(&hd)
19419 .arg(&nh)
19420 .arg(&nhkv)
19421 .arg(&ti)
19422 .arg(&tkvi)
19423 .arg(&scale)
19424 .arg(&cz)
19425 .arg(&wi);
19426 unsafe {
19427 b.launch(cfg)?;
19428 }
19429 Ok(())
19430 }
19431
19432 #[allow(clippy::too_many_arguments)]
19442 pub fn sdpa_naive_w_lo(
19443 &self,
19444 q: &CudaSlice<f32>,
19445 k: &CudaSlice<f32>,
19446 v: &CudaSlice<f32>,
19447 o: &mut CudaSlice<f32>,
19448 head_dim: usize,
19449 n_head: usize,
19450 n_head_kv: usize,
19451 t: usize,
19452 t_kv: usize,
19453 scale: f32,
19454 causal: bool,
19455 window: usize,
19456 ) -> Result<(), Box<dyn std::error::Error>> {
19457 let kv_lo = if window > 0 {
19458 (t_kv - t + 1).saturating_sub(window)
19459 } else {
19460 0
19461 };
19462 let smem = (t_kv - kv_lo) * 4;
19463 if smem > 48 * 1024 {
19464 return Err(format!(
19465 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19466 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19467 a window this wide needs the multi-pass long-ctx kernel"
19468 )
19469 .into());
19470 }
19471 let f = self.func("sdpa_naive_w_lo_f32");
19472 let cfg = LaunchConfig {
19473 grid_dim: (n_head as u32, t as u32, 1),
19474 block_dim: (128, 1, 1),
19475 shared_mem_bytes: smem as u32,
19476 };
19477 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19478 head_dim as i32,
19479 n_head as i32,
19480 n_head_kv as i32,
19481 t as i32,
19482 t_kv as i32,
19483 causal as i32,
19484 window as i32,
19485 kv_lo as i32,
19486 );
19487 let __s_b = self.gpu.stream();
19488 let mut b = __s_b.launch_builder(&f);
19489 b.arg(q)
19490 .arg(k)
19491 .arg(v)
19492 .arg(o)
19493 .arg(&hd)
19494 .arg(&nh)
19495 .arg(&nhkv)
19496 .arg(&ti)
19497 .arg(&tkvi)
19498 .arg(&scale)
19499 .arg(&cz)
19500 .arg(&wi)
19501 .arg(&lo);
19502 unsafe {
19503 b.launch(cfg)?;
19504 }
19505 Ok(())
19506 }
19507
19508 pub fn sdpa_naive_view(
19510 &self,
19511 q: &CudaSlice<f32>,
19512 k: &cudarc::driver::CudaView<f32>,
19513 v: &cudarc::driver::CudaView<f32>,
19514 o: &mut CudaSlice<f32>,
19515 head_dim: usize,
19516 n_head: usize,
19517 n_head_kv: usize,
19518 t: usize,
19519 t_kv: usize,
19520 scale: f32,
19521 causal: bool,
19522 ) -> Result<(), Box<dyn std::error::Error>> {
19523 let f = self.func("sdpa_naive_f32");
19524 let cfg = LaunchConfig {
19525 grid_dim: (n_head as u32, t as u32, 1),
19526 block_dim: (128, 1, 1),
19527 shared_mem_bytes: (t_kv * 4) as u32,
19528 };
19529 let (hd, nh, nhkv, ti, tkvi, cz) = (
19530 head_dim as i32,
19531 n_head as i32,
19532 n_head_kv as i32,
19533 t as i32,
19534 t_kv as i32,
19535 causal as i32,
19536 );
19537 let __s_b = self.gpu.stream();
19538 let mut b = __s_b.launch_builder(&f);
19539 b.arg(q)
19540 .arg(k)
19541 .arg(v)
19542 .arg(o)
19543 .arg(&hd)
19544 .arg(&nh)
19545 .arg(&nhkv)
19546 .arg(&ti)
19547 .arg(&tkvi)
19548 .arg(&scale)
19549 .arg(&cz);
19550 unsafe {
19551 b.launch(cfg)?;
19552 }
19553 Ok(())
19554 }
19555
19556 #[allow(clippy::too_many_arguments)]
19564 pub fn fa_dequant_kv_view_f32(
19565 &self,
19566 k: &cudarc::driver::CudaView<u8>,
19567 v: &cudarc::driver::CudaView<u8>,
19568 kf: &mut CudaSlice<f32>,
19569 vf: &mut CudaSlice<f32>,
19570 kv_dim_k: usize,
19571 kv_dim_v: usize,
19572 t_kv: usize,
19573 k_tok_bytes: usize,
19574 v_tok_bytes: usize,
19575 g: bool,
19576 ) -> Result<(), Box<dyn std::error::Error>> {
19577 let f = if g {
19578 self.func_g("fa_dequant_kv_ws_f32")
19579 } else {
19580 self.func("fa_dequant_kv_ws_f32")
19581 };
19582 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19583 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19584 let cfg = LaunchConfig {
19585 grid_dim: (nblk.max(1), 1, 1),
19586 block_dim: (256, 1, 1),
19587 shared_mem_bytes: 0,
19588 };
19589 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19590 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19591 let __s_b = self.gpu.stream();
19592 let mut b = __s_b.launch_builder(&f);
19593 b.arg(k)
19594 .arg(v)
19595 .arg(&mut *kf)
19596 .arg(&mut *vf)
19597 .arg(&kdk)
19598 .arg(&kdv)
19599 .arg(&tkvi)
19600 .arg(&ktb)
19601 .arg(&vtb);
19602 unsafe {
19603 b.launch(cfg)?;
19604 }
19605 Ok(())
19606 }
19607
19608 #[allow(clippy::too_many_arguments)]
19609 pub fn sdpa_naive_quantized_view(
19610 &self,
19611 q: &CudaSlice<f32>,
19612 k: &cudarc::driver::CudaView<u8>,
19613 v: &cudarc::driver::CudaView<u8>,
19614 o: &mut CudaSlice<f32>,
19615 head_dim: usize,
19616 n_head: usize,
19617 n_head_kv: usize,
19618 t: usize,
19619 t_kv: usize,
19620 scale: f32,
19621 causal: bool,
19622 k_tok_bytes: usize,
19623 v_tok_bytes: usize,
19624 ) -> Result<(), Box<dyn std::error::Error>> {
19625 let kv_dim = n_head_kv * head_dim;
19626 let mut kf = self.uninit(t_kv * kv_dim)?;
19627 let mut vf = self.uninit(t_kv * kv_dim)?;
19628 let f = self.func("fa_dequant_kv_ws_f32");
19629 let total = (2 * t_kv * kv_dim) as u64;
19630 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19631 let cfg = LaunchConfig {
19632 grid_dim: (nblk.max(1), 1, 1),
19633 block_dim: (256, 1, 1),
19634 shared_mem_bytes: 0,
19635 };
19636 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19637 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19638 let __s_b = self.gpu.stream();
19639 let mut b = __s_b.launch_builder(&f);
19640 b.arg(k)
19641 .arg(v)
19642 .arg(&mut kf)
19643 .arg(&mut vf)
19644 .arg(&kv_dim_i)
19645 .arg(&kv_dim_i)
19646 .arg(&t_kv_i)
19647 .arg(&k_tok_bytes_i)
19648 .arg(&v_tok_bytes_i);
19649 unsafe { b.launch(cfg)? };
19650 self.sdpa_naive(
19651 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19652 )
19653 }
19654
19655 #[allow(clippy::too_many_arguments)]
19667 pub fn sdpa_naive_w_quantized_view(
19668 &self,
19669 q: &CudaSlice<f32>,
19670 k: &cudarc::driver::CudaView<u8>,
19671 v: &cudarc::driver::CudaView<u8>,
19672 o: &mut CudaSlice<f32>,
19673 head_dim: usize,
19674 n_head: usize,
19675 n_head_kv: usize,
19676 t: usize,
19677 t_kv: usize,
19678 scale: f32,
19679 causal: bool,
19680 window: usize,
19681 k_tok_bytes: usize,
19682 v_tok_bytes: usize,
19683 ) -> Result<(), Box<dyn std::error::Error>> {
19684 let kv_dim = n_head_kv * head_dim;
19685 let mut kf = self.uninit(t_kv * kv_dim)?;
19686 let mut vf = self.uninit(t_kv * kv_dim)?;
19687 let f = self.func("fa_dequant_kv_ws_f32");
19688 let total = (2 * t_kv * kv_dim) as u64;
19689 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19690 let cfg = LaunchConfig {
19691 grid_dim: (nblk.max(1), 1, 1),
19692 block_dim: (256, 1, 1),
19693 shared_mem_bytes: 0,
19694 };
19695 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19696 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19697 let __s_b = self.gpu.stream();
19698 let mut b = __s_b.launch_builder(&f);
19699 b.arg(k)
19700 .arg(v)
19701 .arg(&mut kf)
19702 .arg(&mut vf)
19703 .arg(&kv_dim_i)
19704 .arg(&kv_dim_i)
19705 .arg(&t_kv_i)
19706 .arg(&k_tok_bytes_i)
19707 .arg(&v_tok_bytes_i);
19708 unsafe { b.launch(cfg)? };
19709 self.sdpa_naive_w(
19710 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19711 )
19712 }
19713
19714 pub fn fa_prefill(
19718 &self,
19719 q: &CudaSlice<f32>,
19720 k: &CudaSlice<f32>,
19721 v: &CudaSlice<f32>,
19722 o: &mut CudaSlice<f32>,
19723 head_dim: usize,
19724 n_head: usize,
19725 n_head_kv: usize,
19726 t: usize,
19727 t_kv: usize,
19728 scale: f32,
19729 causal: bool,
19730 ) -> Result<(), Box<dyn std::error::Error>> {
19731 if portable_mma_gated() {
19732 return self.sdpa_naive(
19733 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19734 );
19735 }
19736 let fa3_on = head_dim == 256
19744 && causal
19745 && t == t_kv
19746 && match std::env::var("MEMRA_FA3").as_deref() {
19747 Ok("0") => false,
19748 Ok("1") => {
19752 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19753 true
19754 }
19755 _ => cfg!(memra_hopper_mma),
19756 };
19757 if fa3_on {
19758 let n = t * n_head * head_dim;
19759 let nkv = t * n_head_kv * head_dim;
19760 let mut q16 = self.alloc_u8_uninit(n * 2)?;
19761 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19762 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19763 self.f32_to_bf16_into(q, &mut q16, n)?;
19764 self.f32_to_bf16_into(k, &mut k16, nkv)?;
19765 self.f32_to_bf16_into(v, &mut v16, nkv)?;
19766 let rc = {
19767 use cudarc::driver::{DevicePtr, DevicePtrMut};
19768 let stream = self.gpu.stream();
19769 let (qp, _g1) = q16.device_ptr(&stream);
19770 let (kp, _g2) = k16.device_ptr(&stream);
19771 let (vp, _g3) = v16.device_ptr(&stream);
19772 let (op, _g4) = o.device_ptr_mut(&stream);
19773 unsafe {
19774 memra_fa3_prefill(
19775 qp as *const core::ffi::c_void,
19776 kp as *const core::ffi::c_void,
19777 vp as *const core::ffi::c_void,
19778 op as *mut f32,
19779 t as i32,
19780 n_head as i32,
19781 n_head_kv as i32,
19782 head_dim as i32,
19783 scale,
19784 stream.cu_stream() as *mut core::ffi::c_void,
19785 )
19786 }
19787 };
19788 if rc != 0 {
19789 return Err(format!("memra_fa3_prefill rc={rc}").into());
19790 }
19791 return Ok(());
19792 }
19793 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19798 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19799 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19800 const BLOCK_Q: usize = 64;
19801 const BKX: usize = 32;
19802 let f = self.func("fa_prefill_bf16_p1");
19803 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19804 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19805 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19806 f.set_attribute(
19807 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19808 shmem as i32,
19809 )?;
19810 let cfg = LaunchConfig {
19811 grid_dim: (
19812 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19813 n_head as u32,
19814 1,
19815 ),
19816 block_dim: (32, 4, 1),
19817 shared_mem_bytes: shmem,
19818 };
19819 let (hd, nh, nhkv, ti, tkvi, cz) = (
19820 head_dim as i32,
19821 n_head as i32,
19822 n_head_kv as i32,
19823 t as i32,
19824 t_kv as i32,
19825 causal as i32,
19826 );
19827 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19828 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19829 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19830 let __s_b = self.gpu.stream();
19831 let mut b = __s_b.launch_builder(&f);
19832 b.arg(&qb)
19833 .arg(&kb)
19834 .arg(&vb)
19835 .arg(o)
19836 .arg(&hd)
19837 .arg(&nh)
19838 .arg(&nhkv)
19839 .arg(&ti)
19840 .arg(&tkvi)
19841 .arg(&scale)
19842 .arg(&cz);
19843 unsafe {
19844 b.launch(cfg)?;
19845 }
19846 return Ok(());
19847 }
19848 const BK: usize = 32;
19854 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19857 let (block_q, warps, w2_sfx): (usize, u32, &str) =
19858 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19859 let hd_sfx = fa_hd_suffix(head_dim)?;
19863 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19864 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
19869 let (kb16, vb16) = if bf16kv {
19870 let n = t_kv * n_head_kv * head_dim;
19871 let mut kb = self.alloc_u8_uninit(n * 2)?;
19872 let mut vb = self.alloc_u8_uninit(n * 2)?;
19873 let fcv = self.func("f32_to_bf16_bulk");
19874 let ni = n as i64;
19875 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
19876 let __s_b = self.gpu.stream();
19877 let mut b = __s_b.launch_builder(&fcv);
19878 b.arg(k).arg(&mut kb).arg(&ni);
19879 unsafe {
19880 b.launch(cfgc)?;
19881 }
19882 let __s_b = self.gpu.stream();
19883 let mut b = __s_b.launch_builder(&fcv);
19884 b.arg(v).arg(&mut vb).arg(&ni);
19885 unsafe {
19886 b.launch(cfgc)?;
19887 }
19888 (Some(kb), Some(vb))
19889 } else {
19890 (None, None)
19891 };
19892 let f = self.func(&if bf16kv {
19893 format!("fa_prefill_bf16kv_pp{hd_sfx}")
19894 } else {
19895 format!(
19896 "fa_prefill_f32{}{}{hd_sfx}",
19897 if floor { "" } else { "_pp" },
19898 if floor { "" } else { w2_sfx }
19899 )
19900 });
19901 let kv_stages = if bf16kv { 2 } else { 1 };
19904 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19905 + 4 * (block_q * BK + 2 * block_q)) as u32;
19906 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19907 f.set_attribute(
19908 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19909 shmem as i32,
19910 )?;
19911 let cfg = LaunchConfig {
19912 grid_dim: (
19913 (t as u32 + block_q as u32 - 1) / block_q as u32,
19914 n_head as u32,
19915 1,
19916 ),
19917 block_dim: (32, warps, 1),
19918 shared_mem_bytes: shmem,
19919 };
19920 let (hd, nh, nhkv, ti, tkvi, cz) = (
19921 head_dim as i32,
19922 n_head as i32,
19923 n_head_kv as i32,
19924 t as i32,
19925 t_kv as i32,
19926 causal as i32,
19927 );
19928 let __s_b = self.gpu.stream();
19929 let mut b = __s_b.launch_builder(&f);
19930 b.arg(q);
19931 match (&kb16, &vb16) {
19932 (Some(kb), Some(vb)) => {
19933 b.arg(kb).arg(vb);
19934 }
19935 _ => {
19936 b.arg(k).arg(v);
19937 }
19938 }
19939 b.arg(o)
19940 .arg(&hd)
19941 .arg(&nh)
19942 .arg(&nhkv)
19943 .arg(&ti)
19944 .arg(&tkvi)
19945 .arg(&scale)
19946 .arg(&cz);
19947 unsafe {
19948 b.launch(cfg)?;
19949 }
19950 Ok(())
19951 }
19952
19953 #[allow(clippy::too_many_arguments)]
19957 pub fn fa_prefill_w(
19958 &self,
19959 q: &CudaSlice<f32>,
19960 k: &CudaSlice<f32>,
19961 v: &CudaSlice<f32>,
19962 o: &mut CudaSlice<f32>,
19963 head_dim: usize,
19964 n_head: usize,
19965 n_head_kv: usize,
19966 t: usize,
19967 t_kv: usize,
19968 scale: f32,
19969 causal: bool,
19970 window: usize,
19971 ) -> Result<(), Box<dyn std::error::Error>> {
19972 if portable_mma_gated() {
19975 return self.sdpa_naive_w(
19976 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19977 );
19978 }
19979 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19983 let faw_f32 =
19984 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
19985 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19986 self.fa_prefill_w_arm(
19987 q,
19988 k,
19989 v,
19990 o,
19991 head_dim,
19992 n_head,
19993 n_head_kv,
19994 t,
19995 t_kv,
19996 scale,
19997 causal,
19998 window,
19999 floor || faw_f32,
20000 floor,
20001 )
20002 }
20003
20004 #[allow(clippy::too_many_arguments)]
20007 pub fn fa_prefill_w_pre(
20008 &self,
20009 qb: &CudaSlice<u8>,
20010 kb: &CudaSlice<u8>,
20011 vb: &CudaSlice<u8>,
20012 o: &mut CudaSlice<f32>,
20013 head_dim: usize,
20014 n_head: usize,
20015 n_head_kv: usize,
20016 t: usize,
20017 t_kv: usize,
20018 scale: f32,
20019 causal: bool,
20020 window: usize,
20021 v_f16: bool,
20022 ) -> Result<(), Box<dyn std::error::Error>> {
20023 const BLOCK_Q: usize = 64;
20024 const BK: usize = 32;
20025 debug_assert_eq!(head_dim, 256);
20026 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20027 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20028 if hp {
20029 const BLOCK_QH: usize = 32;
20030 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20033 let vh: &CudaSlice<u8> = if v_f16 {
20034 vb
20035 } else {
20036 let n = t_kv * n_head_kv * head_dim;
20037 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20038 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20039 }
20040 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20041 vguard.as_ref().unwrap()
20042 };
20043 let f = self.func("fa_prefill_w_bf16_p1h2");
20044 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20045 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20046 f.set_attribute(
20047 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20048 shmem as i32,
20049 )?;
20050 let cfg = LaunchConfig {
20051 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20052 block_dim: (32, 4, 1),
20053 shared_mem_bytes: shmem,
20054 };
20055 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20056 head_dim as i32,
20057 n_head as i32,
20058 n_head_kv as i32,
20059 t as i32,
20060 t_kv as i32,
20061 causal as i32,
20062 window as i32,
20063 );
20064 let __s_b = self.gpu.stream();
20065 let mut b = __s_b.launch_builder(&f);
20066 b.arg(qb)
20067 .arg(kb)
20068 .arg(vh)
20069 .arg(o)
20070 .arg(&hd)
20071 .arg(&nh)
20072 .arg(&nhkv)
20073 .arg(&ti)
20074 .arg(&tkvi)
20075 .arg(&scale)
20076 .arg(&cz)
20077 .arg(&wi);
20078 unsafe {
20079 b.launch(cfg)?;
20080 }
20081 return Ok(());
20082 }
20083 let f = self.func("fa_prefill_w_bf16_p1");
20084 let shmem =
20085 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20086 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20087 f.set_attribute(
20088 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20089 shmem as i32,
20090 )?;
20091 let cfg = LaunchConfig {
20092 grid_dim: (
20093 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20094 n_head as u32,
20095 1,
20096 ),
20097 block_dim: (32, 4, 1),
20098 shared_mem_bytes: shmem,
20099 };
20100 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20101 head_dim as i32,
20102 n_head as i32,
20103 n_head_kv as i32,
20104 t as i32,
20105 t_kv as i32,
20106 causal as i32,
20107 window as i32,
20108 );
20109 let __s_b = self.gpu.stream();
20110 let mut b = __s_b.launch_builder(&f);
20111 b.arg(qb)
20112 .arg(kb)
20113 .arg(vb)
20114 .arg(o)
20115 .arg(&hd)
20116 .arg(&nh)
20117 .arg(&nhkv)
20118 .arg(&ti)
20119 .arg(&tkvi)
20120 .arg(&scale)
20121 .arg(&cz)
20122 .arg(&wi);
20123 unsafe {
20124 b.launch(cfg)?;
20125 }
20126 Ok(())
20127 }
20128
20129 #[allow(clippy::too_many_arguments)]
20131 pub fn fa_prefill_w_arm(
20132 &self,
20133 q: &CudaSlice<f32>,
20134 k: &CudaSlice<f32>,
20135 v: &CudaSlice<f32>,
20136 o: &mut CudaSlice<f32>,
20137 head_dim: usize,
20138 n_head: usize,
20139 n_head_kv: usize,
20140 t: usize,
20141 t_kv: usize,
20142 scale: f32,
20143 causal: bool,
20144 window: usize,
20145 f32_stage: bool,
20146 floor: bool,
20147 ) -> Result<(), Box<dyn std::error::Error>> {
20148 const BLOCK_Q: usize = 64;
20149 const BK: usize = 32;
20150 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20151 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20155 let p1 = !floor
20156 && !f32_stage
20157 && *P1_ON.get_or_init(|| {
20158 std::env::var("MEMRA_FAW_P1")
20159 .map(|v| v != "0")
20160 .unwrap_or(true)
20161 });
20162 let hp =
20163 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20164 if hp {
20165 const BLOCK_QH: usize = 32;
20166 let f = self.func("fa_prefill_w_bf16_p1h2");
20167 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20168 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20169 f.set_attribute(
20170 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20171 shmem as i32,
20172 )?;
20173 let cfg = LaunchConfig {
20174 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20175 block_dim: (32, 4, 1),
20176 shared_mem_bytes: shmem,
20177 };
20178 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20179 head_dim as i32,
20180 n_head as i32,
20181 n_head_kv as i32,
20182 t as i32,
20183 t_kv as i32,
20184 causal as i32,
20185 window as i32,
20186 );
20187 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20188 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20189 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20190 let __s_b = self.gpu.stream();
20191 let mut b = __s_b.launch_builder(&f);
20192 b.arg(&qb)
20193 .arg(&kb)
20194 .arg(&vh)
20195 .arg(o)
20196 .arg(&hd)
20197 .arg(&nh)
20198 .arg(&nhkv)
20199 .arg(&ti)
20200 .arg(&tkvi)
20201 .arg(&scale)
20202 .arg(&cz)
20203 .arg(&wi);
20204 unsafe {
20205 b.launch(cfg)?;
20206 }
20207 return Ok(());
20208 }
20209 if p1 {
20210 let f = self.func("fa_prefill_w_bf16_p1");
20211 let shmem =
20212 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20213 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20214 f.set_attribute(
20215 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20216 shmem as i32,
20217 )?;
20218 let cfg = LaunchConfig {
20219 grid_dim: (
20220 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20221 n_head as u32,
20222 1,
20223 ),
20224 block_dim: (32, 4, 1),
20225 shared_mem_bytes: shmem,
20226 };
20227 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20228 head_dim as i32,
20229 n_head as i32,
20230 n_head_kv as i32,
20231 t as i32,
20232 t_kv as i32,
20233 causal as i32,
20234 window as i32,
20235 );
20236 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20237 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20238 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20239 let __s_b = self.gpu.stream();
20240 let mut b = __s_b.launch_builder(&f);
20241 b.arg(&qb)
20242 .arg(&kb)
20243 .arg(&vb)
20244 .arg(o)
20245 .arg(&hd)
20246 .arg(&nh)
20247 .arg(&nhkv)
20248 .arg(&ti)
20249 .arg(&tkvi)
20250 .arg(&scale)
20251 .arg(&cz)
20252 .arg(&wi);
20253 unsafe {
20254 b.launch(cfg)?;
20255 }
20256 return Ok(());
20257 }
20258 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20261 let g4 = !floor
20262 && !f32_stage
20263 && n_head_kv == 1
20264 && n_head % 4 == 0
20265 && *G4_ON.get_or_init(|| {
20266 std::env::var("MEMRA_FAW_G4")
20267 .map(|v| v != "0")
20268 .unwrap_or(true)
20269 });
20270 if g4 {
20271 const SP_M: usize = 16;
20272 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20275 let o2 = *O2_ON.get_or_init(|| {
20276 std::env::var("MEMRA_FAW_O2")
20277 .map(|v| v != "0")
20278 .unwrap_or(true)
20279 });
20280 let f = self.func(if o2 {
20281 "fa_prefill_w_bf16_g4o2"
20282 } else {
20283 "fa_prefill_w_bf16_g4"
20284 });
20285 let shmem = if o2 {
20286 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20287 } else {
20288 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20289 as u32
20290 };
20291 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20292 f.set_attribute(
20293 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20294 shmem as i32,
20295 )?;
20296 let cfg = LaunchConfig {
20297 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20298 block_dim: (32, 4, 1),
20299 shared_mem_bytes: shmem,
20300 };
20301 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20302 head_dim as i32,
20303 n_head as i32,
20304 n_head_kv as i32,
20305 t as i32,
20306 t_kv as i32,
20307 causal as i32,
20308 window as i32,
20309 );
20310 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20311 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20312 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20313 let __s_b = self.gpu.stream();
20314 let mut b = __s_b.launch_builder(&f);
20315 b.arg(&qb)
20316 .arg(&kb)
20317 .arg(&vb)
20318 .arg(o)
20319 .arg(&hd)
20320 .arg(&nh)
20321 .arg(&nhkv)
20322 .arg(&ti)
20323 .arg(&tkvi)
20324 .arg(&scale)
20325 .arg(&cz)
20326 .arg(&wi);
20327 unsafe {
20328 b.launch(cfg)?;
20329 }
20330 return Ok(());
20331 }
20332 let f = self.func(if floor {
20333 "fa_prefill_w_f32"
20334 } else if f32_stage {
20335 "fa_prefill_w_f32_pp"
20336 } else {
20337 "fa_prefill_w_bf16_pp"
20338 });
20339 let shmem =
20340 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20341 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20342 f.set_attribute(
20343 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20344 shmem as i32,
20345 )?;
20346 let cfg = LaunchConfig {
20347 grid_dim: (
20348 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20349 n_head as u32,
20350 1,
20351 ),
20352 block_dim: (32, 4, 1),
20353 shared_mem_bytes: shmem,
20354 };
20355 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20356 head_dim as i32,
20357 n_head as i32,
20358 n_head_kv as i32,
20359 t as i32,
20360 t_kv as i32,
20361 causal as i32,
20362 window as i32,
20363 );
20364 if f32_stage {
20365 let __s_b = self.gpu.stream();
20366 let mut b = __s_b.launch_builder(&f);
20367 b.arg(q)
20368 .arg(k)
20369 .arg(v)
20370 .arg(o)
20371 .arg(&hd)
20372 .arg(&nh)
20373 .arg(&nhkv)
20374 .arg(&ti)
20375 .arg(&tkvi)
20376 .arg(&scale)
20377 .arg(&cz)
20378 .arg(&wi);
20379 unsafe {
20380 b.launch(cfg)?;
20381 }
20382 } else {
20383 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20384 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20385 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20386 let __s_b = self.gpu.stream();
20387 let mut b = __s_b.launch_builder(&f);
20388 b.arg(&qb)
20389 .arg(&kb)
20390 .arg(&vb)
20391 .arg(o)
20392 .arg(&hd)
20393 .arg(&nh)
20394 .arg(&nhkv)
20395 .arg(&ti)
20396 .arg(&tkvi)
20397 .arg(&scale)
20398 .arg(&cz)
20399 .arg(&wi);
20400 unsafe {
20401 b.launch(cfg)?;
20402 }
20403 }
20404 Ok(())
20405 }
20406
20407 #[allow(clippy::too_many_arguments)]
20411 pub fn fa_prefill_hd512(
20412 &self,
20413 q: &CudaSlice<f32>,
20414 k: &CudaSlice<f32>,
20415 v: &CudaSlice<f32>,
20416 o: &mut CudaSlice<f32>,
20417 head_dim: usize,
20418 n_head: usize,
20419 n_head_kv: usize,
20420 t: usize,
20421 t_kv: usize,
20422 scale: f32,
20423 causal: bool,
20424 ) -> Result<(), Box<dyn std::error::Error>> {
20425 if portable_mma_gated() {
20427 return self.sdpa_naive(
20428 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20429 );
20430 }
20431 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20437 let f32_stage =
20438 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20439 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20443 let sp = !f32_stage
20444 && *SP_ON.get_or_init(|| {
20445 std::env::var("MEMRA_FA512_SP")
20446 .map(|v| v != "0")
20447 .unwrap_or(true)
20448 });
20449 self.fa_prefill_hd512_arm(
20450 q,
20451 k,
20452 v,
20453 o,
20454 head_dim,
20455 n_head,
20456 n_head_kv,
20457 t,
20458 t_kv,
20459 scale,
20460 causal,
20461 f32_stage,
20462 sp,
20463 sp && fa_f16pv_on(),
20464 )
20465 }
20466
20467 #[allow(clippy::too_many_arguments)]
20469 pub fn fa_prefill_hd512_pre(
20470 &self,
20471 qb: &CudaSlice<u8>,
20472 kb: &CudaSlice<u8>,
20473 vb: &CudaSlice<u8>,
20474 o: &mut CudaSlice<f32>,
20475 head_dim: usize,
20476 n_head: usize,
20477 n_head_kv: usize,
20478 t: usize,
20479 t_kv: usize,
20480 scale: f32,
20481 causal: bool,
20482 v_f16: bool,
20483 ) -> Result<(), Box<dyn std::error::Error>> {
20484 debug_assert_eq!(head_dim, 512);
20485 const SP_M: usize = 16;
20486 const BKS: usize = 32;
20487 let f16pv = fa_f16pv_on();
20491 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20492 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20493 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20494 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20495 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20496 let n = t_kv * n_head_kv * head_dim;
20498 let need = n * 2;
20499 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20500 *vguard = Some(self.alloc_uninit::<u8>(need)?);
20501 }
20502 let dst = vguard.as_mut().unwrap();
20503 self.bf16_to_f16_into(vb, n, dst)?;
20504 vguard.as_ref().unwrap()
20505 } else {
20506 vb
20507 };
20508 let f = self.func(if hp {
20509 "fa_prefill_bf16_hd512_sp16h2"
20510 } else {
20511 match (f16pv, nw) {
20512 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20513 (true, _) => "fa_prefill_bf16_hd512_sp16",
20514 _ => "fa_prefill_bf16_hd512_sp",
20515 }
20516 });
20517 let (nwarp, npart) = if hp {
20518 (4usize, 4usize)
20519 } else if nw > 2 {
20520 (nw, nw)
20521 } else {
20522 (2, 1)
20523 };
20524 let shmem = if hp {
20526 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20527 as u32
20528 } else {
20529 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20530 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20531 };
20532 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20533 f.set_attribute(
20534 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20535 shmem as i32,
20536 )?;
20537 let grid_y = if hp {
20538 (n_head / 2) as u32
20539 } else {
20540 n_head as u32
20541 };
20542 let cfg = LaunchConfig {
20543 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20544 block_dim: (32, nwarp as u32, 1),
20545 shared_mem_bytes: shmem,
20546 };
20547 let (hd, nh, nhkv, ti, tkvi, cz) = (
20548 head_dim as i32,
20549 n_head as i32,
20550 n_head_kv as i32,
20551 t as i32,
20552 t_kv as i32,
20553 causal as i32,
20554 );
20555 let __s_b = self.gpu.stream();
20556 let mut b = __s_b.launch_builder(&f);
20557 b.arg(qb)
20558 .arg(kb)
20559 .arg(vref)
20560 .arg(o)
20561 .arg(&hd)
20562 .arg(&nh)
20563 .arg(&nhkv)
20564 .arg(&ti)
20565 .arg(&tkvi)
20566 .arg(&scale)
20567 .arg(&cz);
20568 unsafe {
20569 b.launch(cfg)?;
20570 }
20571 Ok(())
20572 }
20573
20574 #[allow(clippy::too_many_arguments)]
20577 pub fn fa_prefill_hd512_arm(
20578 &self,
20579 q: &CudaSlice<f32>,
20580 k: &CudaSlice<f32>,
20581 v: &CudaSlice<f32>,
20582 o: &mut CudaSlice<f32>,
20583 head_dim: usize,
20584 n_head: usize,
20585 n_head_kv: usize,
20586 t: usize,
20587 t_kv: usize,
20588 scale: f32,
20589 causal: bool,
20590 f32_stage: bool,
20591 sp: bool,
20592 f16pv: bool,
20593 ) -> Result<(), Box<dyn std::error::Error>> {
20594 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20595 if sp && !f32_stage {
20596 const SP_M: usize = 16;
20600 const BKS: usize = 32;
20601 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20602 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20603 let f = self.func(if hp {
20604 "fa_prefill_bf16_hd512_sp16h2"
20605 } else {
20606 match (f16pv, nw) {
20607 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20608 (true, _) => "fa_prefill_bf16_hd512_sp16",
20609 _ => "fa_prefill_bf16_hd512_sp",
20610 }
20611 });
20612 let (nwarp, npart) = if hp {
20613 (4usize, 4usize)
20614 } else if nw > 2 {
20615 (nw, nw)
20616 } else {
20617 (2, 1)
20618 };
20619 let shmem = if hp {
20620 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20621 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20622 } else {
20623 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20624 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20625 };
20626 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20627 f.set_attribute(
20628 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20629 shmem as i32,
20630 )?;
20631 let grid_y = if hp {
20632 (n_head / 2) as u32
20633 } else {
20634 n_head as u32
20635 };
20636 let cfg = LaunchConfig {
20637 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20638 block_dim: (32, nwarp as u32, 1),
20639 shared_mem_bytes: shmem,
20640 };
20641 let (hd, nh, nhkv, ti, tkvi, cz) = (
20642 head_dim as i32,
20643 n_head as i32,
20644 n_head_kv as i32,
20645 t as i32,
20646 t_kv as i32,
20647 causal as i32,
20648 );
20649 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20650 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20651 let vb = if f16pv {
20652 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20653 } else {
20654 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20655 };
20656 let __s_b = self.gpu.stream();
20657 let mut b = __s_b.launch_builder(&f);
20658 b.arg(&qb)
20659 .arg(&kb)
20660 .arg(&vb)
20661 .arg(o)
20662 .arg(&hd)
20663 .arg(&nh)
20664 .arg(&nhkv)
20665 .arg(&ti)
20666 .arg(&tkvi)
20667 .arg(&scale)
20668 .arg(&cz);
20669 unsafe {
20670 b.launch(cfg)?;
20671 }
20672 return Ok(());
20673 }
20674 const BLOCK_Q: usize = 32;
20675 const BK: usize = 32;
20676 const HALF: usize = 256;
20677 let f = self.func(if f32_stage {
20678 "fa_prefill_f32_hd512"
20679 } else {
20680 "fa_prefill_bf16_hd512"
20681 });
20682 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20684 + 4 * BLOCK_Q) as u32;
20685 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20686 f.set_attribute(
20687 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20688 shmem as i32,
20689 )?;
20690 let cfg = LaunchConfig {
20691 grid_dim: (
20692 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20693 n_head as u32,
20694 2,
20695 ),
20696 block_dim: (32, 2, 1),
20697 shared_mem_bytes: shmem,
20698 };
20699 let (hd, nh, nhkv, ti, tkvi, cz) = (
20700 head_dim as i32,
20701 n_head as i32,
20702 n_head_kv as i32,
20703 t as i32,
20704 t_kv as i32,
20705 causal as i32,
20706 );
20707 if f32_stage {
20708 let __s_b = self.gpu.stream();
20709 let mut b = __s_b.launch_builder(&f);
20710 b.arg(q)
20711 .arg(k)
20712 .arg(v)
20713 .arg(o)
20714 .arg(&hd)
20715 .arg(&nh)
20716 .arg(&nhkv)
20717 .arg(&ti)
20718 .arg(&tkvi)
20719 .arg(&scale)
20720 .arg(&cz);
20721 unsafe {
20722 b.launch(cfg)?;
20723 }
20724 } else {
20725 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20726 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20727 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20728 let __s_b = self.gpu.stream();
20729 let mut b = __s_b.launch_builder(&f);
20730 b.arg(&qb)
20731 .arg(&kb)
20732 .arg(&vb)
20733 .arg(o)
20734 .arg(&hd)
20735 .arg(&nh)
20736 .arg(&nhkv)
20737 .arg(&ti)
20738 .arg(&tkvi)
20739 .arg(&scale)
20740 .arg(&cz);
20741 unsafe {
20742 b.launch(cfg)?;
20743 }
20744 }
20745 Ok(())
20746 }
20747
20748 #[allow(clippy::too_many_arguments)]
20752 pub fn rope_neox2_bf16e(
20753 &self,
20754 q: &mut CudaSlice<f32>,
20755 k: &mut CudaSlice<f32>,
20756 qb: &mut CudaSlice<u8>,
20757 kb: &mut CudaSlice<u8>,
20758 pos: &CudaSlice<i32>,
20759 head_dim: usize,
20760 n_dims: usize,
20761 nh_q: usize,
20762 nh_k: usize,
20763 n_tokens: usize,
20764 base: f32,
20765 freq_scale: f32,
20766 ff: Option<&CudaSlice<f32>>,
20767 ) -> Result<(), Box<dyn std::error::Error>> {
20768 let f = self.func("rope_neox2_bf16e_f32");
20769 let rows = ((nh_q + nh_k) * n_tokens) as u32;
20770 let cfg = LaunchConfig {
20771 grid_dim: (rows, 1, 1),
20772 block_dim: ((head_dim / 2) as u32, 1, 1),
20773 shared_mem_bytes: 0,
20774 };
20775 let theta_scale = base.powf(-2.0 / n_dims as f32);
20776 let (hd, nd, nhq, nhk, nt) = (
20777 head_dim as i32,
20778 n_dims as i32,
20779 nh_q as i32,
20780 nh_k as i32,
20781 n_tokens as i32,
20782 );
20783 let __s_b = self.gpu.stream();
20784 let mut b = __s_b.launch_builder(&f);
20785 match ff {
20786 Some(t) => {
20787 b.arg(&mut *q)
20788 .arg(&mut *k)
20789 .arg(&mut *qb)
20790 .arg(&mut *kb)
20791 .arg(pos)
20792 .arg(&hd)
20793 .arg(&nd)
20794 .arg(&nhq)
20795 .arg(&nhk)
20796 .arg(&nt)
20797 .arg(&theta_scale)
20798 .arg(&freq_scale)
20799 .arg(t);
20800 unsafe {
20801 b.launch(cfg)?;
20802 }
20803 }
20804 None => {
20805 let null: u64 = 0;
20806 b.arg(&mut *q)
20807 .arg(&mut *k)
20808 .arg(&mut *qb)
20809 .arg(&mut *kb)
20810 .arg(pos)
20811 .arg(&hd)
20812 .arg(&nd)
20813 .arg(&nhq)
20814 .arg(&nhk)
20815 .arg(&nt)
20816 .arg(&theta_scale)
20817 .arg(&freq_scale)
20818 .arg(&null);
20819 unsafe {
20820 b.launch(cfg)?;
20821 }
20822 }
20823 }
20824 Ok(())
20825 }
20826
20827 pub fn f32_to_bf16(
20830 &self,
20831 x: &CudaSlice<f32>,
20832 n: usize,
20833 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20834 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
20835 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20836 let f = self.func("f32_to_bf16_flat");
20837 let n_i = n as i64;
20838 let cfg = LaunchConfig {
20839 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20840 block_dim: (256, 1, 1),
20841 shared_mem_bytes: 0,
20842 };
20843 let __s_b = self.gpu.stream();
20844 let mut b = __s_b.launch_builder(&f);
20845 b.arg(x).arg(&mut y).arg(&n_i);
20846 unsafe {
20847 b.launch(cfg)?;
20848 }
20849 Ok(y)
20850 }
20851
20852 pub fn f32_to_f16(
20853 &self,
20854 x: &CudaSlice<f32>,
20855 n: usize,
20856 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20857 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20858 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20859 let f = self.func("f32_to_f16_flat");
20860 let n_i = n as i64;
20861 let cfg = LaunchConfig {
20862 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20863 block_dim: (256, 1, 1),
20864 shared_mem_bytes: 0,
20865 };
20866 let __s_b = self.gpu.stream();
20867 let mut b = __s_b.launch_builder(&f);
20868 b.arg(x).arg(&mut y).arg(&n_i);
20869 unsafe {
20870 b.launch(cfg)?;
20871 }
20872 Ok(y)
20873 }
20874
20875 pub fn bf16_to_f16(
20877 &self,
20878 xb: &CudaSlice<u8>,
20879 n: usize,
20880 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20881 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20882 self.bf16_to_f16_into(xb, n, &mut y)?;
20883 Ok(y)
20884 }
20885
20886 pub fn bf16_to_f16_into(
20888 &self,
20889 xb: &CudaSlice<u8>,
20890 n: usize,
20891 y: &mut CudaSlice<u8>,
20892 ) -> Result<(), Box<dyn std::error::Error>> {
20893 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
20894 assert!(y.len() >= n * 2);
20895 let f = self.func("bf16_to_f16_flat");
20896 let n2 = (n / 2) as i64;
20897 let cfg = LaunchConfig {
20898 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
20899 block_dim: (256, 1, 1),
20900 shared_mem_bytes: 0,
20901 };
20902 let __s_b = self.gpu.stream();
20903 let mut b = __s_b.launch_builder(&f);
20904 b.arg(xb).arg(y).arg(&n2);
20905 unsafe {
20906 b.launch(cfg)?;
20907 }
20908 Ok(())
20909 }
20910
20911 #[allow(clippy::too_many_arguments)]
20916 pub fn fa_prefill_vl8(
20917 &self,
20918 seqs: &[FaSeqVl],
20919 head_dim: usize,
20920 n_head: usize,
20921 n_head_kv: usize,
20922 scale: f32,
20923 ) -> Result<(), Box<dyn std::error::Error>> {
20924 const BK: usize = 32;
20925 let b = seqs.len();
20926 assert!(b >= 1 && b <= 8);
20927 let mut packed = [FaSeqVl::default(); 8];
20928 packed[..b].copy_from_slice(seqs);
20929 let v = FaVl8(packed);
20930 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20931 let ept = (n_head_kv * head_dim) as i32;
20932 {
20933 let f = self.func("fa_mirror_vl");
20934 let max_n = (max_t as i64) * ept as i64;
20935 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
20936 for which in 0..2i32 {
20937 let cfg = LaunchConfig {
20938 grid_dim: (blocks, 1, b as u32),
20939 block_dim: (256, 1, 1),
20940 shared_mem_bytes: 0,
20941 };
20942 let __s_lb = self.gpu.stream();
20943 let mut lb = __s_lb.launch_builder(&f);
20944 lb.arg(&v).arg(&ept).arg(&which);
20945 unsafe {
20946 lb.launch(cfg)?;
20947 }
20948 }
20949 }
20950 let hd_sfx = fa_hd_suffix(head_dim)?;
20951 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
20952 let block_q = 64usize;
20953 let kv_stages = 2usize;
20954 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20955 + 4 * (block_q * BK + 2 * block_q)) as u32;
20956 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20957 f.set_attribute(
20958 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20959 shmem as i32,
20960 )?;
20961 let cfg = LaunchConfig {
20962 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
20963 block_dim: (32, 4, 1),
20964 shared_mem_bytes: shmem,
20965 };
20966 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20967 let __s_lb = self.gpu.stream();
20968 let mut lb = __s_lb.launch_builder(&f);
20969 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
20970 unsafe {
20971 lb.launch(cfg)?;
20972 }
20973 Ok(())
20974 }
20975
20976 #[allow(clippy::too_many_arguments)]
20980 pub fn attn_pre_vl8(
20981 &self,
20982 seqs: &[AttnPreVl],
20983 wq: &CudaSlice<f32>,
20984 wk: &CudaSlice<f32>,
20985 head_dim: usize,
20986 rope_dims: usize,
20987 n_head: usize,
20988 n_head_kv: usize,
20989 eps: f32,
20990 freq_base: f32,
20991 freq_scale: f32,
20992 kv_dim_k: usize,
20993 kv_dim_v: usize,
20994 k_tok_bytes: usize,
20995 v_tok_bytes: usize,
20996 ) -> Result<(), Box<dyn std::error::Error>> {
20997 let b = seqs.len();
20998 assert!(b >= 1 && b <= 8);
20999 let mut packed = [AttnPreVl::default(); 8];
21000 packed[..b].copy_from_slice(seqs);
21001 let v = AttnPreVl8(packed);
21002 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21003 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21004 {
21005 let f = self.func("q_gate_split_vl");
21006 let n = max_t * (n_head * head_dim) as u32;
21007 let cfg = LaunchConfig {
21008 grid_dim: (n.div_ceil(256), 1, b as u32),
21009 block_dim: (256, 1, 1),
21010 shared_mem_bytes: 0,
21011 };
21012 let __s_lb = self.gpu.stream();
21013 let mut lb = __s_lb.launch_builder(&f);
21014 lb.arg(&v).arg(&hd).arg(&nh);
21015 unsafe {
21016 lb.launch(cfg)?;
21017 }
21018 }
21019 {
21020 let f = self.func("attn_rms_vl");
21021 let cfg = LaunchConfig {
21022 grid_dim: (max_t * n_head as u32, 2, b as u32),
21023 block_dim: (rms_block(), 1, 1),
21024 shared_mem_bytes: 0,
21025 };
21026 let __s_lb = self.gpu.stream();
21027 let mut lb = __s_lb.launch_builder(&f);
21028 lb.arg(&v)
21029 .arg(wq)
21030 .arg(wk)
21031 .arg(&hd)
21032 .arg(&nh)
21033 .arg(&nhkv)
21034 .arg(&eps);
21035 unsafe {
21036 lb.launch(cfg)?;
21037 }
21038 }
21039 {
21040 let f = self.func("attn_rope_vl");
21041 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21042 let nd = rope_dims as i32;
21043 let cfg = LaunchConfig {
21044 grid_dim: (max_t * n_head as u32, 2, b as u32),
21045 block_dim: ((head_dim / 2) as u32, 1, 1),
21046 shared_mem_bytes: 0,
21047 };
21048 let __s_lb = self.gpu.stream();
21049 let mut lb = __s_lb.launch_builder(&f);
21050 lb.arg(&v)
21051 .arg(&hd)
21052 .arg(&nd)
21053 .arg(&nh)
21054 .arg(&nhkv)
21055 .arg(&theta_scale)
21056 .arg(&freq_scale);
21057 unsafe {
21058 lb.launch(cfg)?;
21059 }
21060 }
21061 {
21062 let f = self.func("append_kv_vl");
21063 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21064 let cfg = LaunchConfig {
21065 grid_dim: (nblk, max_t, b as u32),
21066 block_dim: (32, 1, 1),
21067 shared_mem_bytes: 0,
21068 };
21069 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21070 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21071 let __s_lb = self.gpu.stream();
21072 let mut lb = __s_lb.launch_builder(&f);
21073 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21074 unsafe {
21075 lb.launch(cfg)?;
21076 }
21077 }
21078 Ok(())
21079 }
21080
21081 pub fn fa_prefill_view(
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 k_tok_bytes: usize,
21099 v_tok_bytes: usize,
21100 g: bool,
21101 ) -> Result<(), Box<dyn std::error::Error>> {
21102 if portable_mma_gated() {
21103 return self.sdpa_naive_quantized_view(
21104 q,
21105 k,
21106 v,
21107 o,
21108 head_dim,
21109 n_head,
21110 n_head_kv,
21111 t,
21112 t_kv,
21113 scale,
21114 causal,
21115 k_tok_bytes,
21116 v_tok_bytes,
21117 );
21118 }
21119 const BLOCK_Q: usize = 64;
21120 const BK: usize = 32;
21121 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21124 let f = if g {
21125 self.func_g(&name)
21126 } else {
21127 self.func(&name)
21128 };
21129 let shmem =
21130 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21131 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21132 f.set_attribute(
21133 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21134 shmem as i32,
21135 )?;
21136 let cfg = LaunchConfig {
21137 grid_dim: (
21138 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21139 n_head as u32,
21140 1,
21141 ),
21142 block_dim: (32, 4, 1),
21143 shared_mem_bytes: shmem,
21144 };
21145 let (hd, nh, nhkv, ti, tkvi, cz) = (
21146 head_dim as i32,
21147 n_head as i32,
21148 n_head_kv as i32,
21149 t as i32,
21150 t_kv as i32,
21151 causal as i32,
21152 );
21153 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21154 let __s_b = self.gpu.stream();
21155 let mut b = __s_b.launch_builder(&f);
21156 b.arg(q)
21157 .arg(k)
21158 .arg(v)
21159 .arg(o)
21160 .arg(&hd)
21161 .arg(&nh)
21162 .arg(&nhkv)
21163 .arg(&ti)
21164 .arg(&tkvi)
21165 .arg(&scale)
21166 .arg(&cz)
21167 .arg(&ktb)
21168 .arg(&vtb);
21169 unsafe {
21170 b.launch(cfg)?;
21171 }
21172 Ok(())
21173 }
21174
21175 #[allow(clippy::too_many_arguments)]
21185 pub fn fa_prefill_view_ws(
21186 &self,
21187 q: &CudaSlice<f32>,
21188 k: &cudarc::driver::CudaView<u8>,
21189 v: &cudarc::driver::CudaView<u8>,
21190 o: &mut CudaSlice<f32>,
21191 head_dim: usize,
21192 n_head: usize,
21193 n_head_kv: usize,
21194 t: usize,
21195 t_kv: usize,
21196 scale: f32,
21197 causal: bool,
21198 k_tok_bytes: usize,
21199 v_tok_bytes: usize,
21200 g: bool,
21201 ) -> Result<(), Box<dyn std::error::Error>> {
21202 if portable_mma_gated() {
21203 return self.sdpa_naive_quantized_view(
21204 q,
21205 k,
21206 v,
21207 o,
21208 head_dim,
21209 n_head,
21210 n_head_kv,
21211 t,
21212 t_kv,
21213 scale,
21214 causal,
21215 k_tok_bytes,
21216 v_tok_bytes,
21217 );
21218 }
21219 const BLOCK_Q: usize = 64;
21220 const BK: usize = 32;
21221 let kv_dim_k = n_head_kv * head_dim;
21222 let kv_dim_v = n_head_kv * head_dim;
21223 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21225 let mut guard = self.prime_deqw_ws.lock().unwrap();
21227 let need_grow = match guard.as_ref() {
21228 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21229 None => true,
21230 };
21231 if need_grow {
21232 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21233 let (ck, cv) = guard
21234 .as_ref()
21235 .map(|(a, b)| (a.len(), b.len()))
21236 .unwrap_or((0, 0));
21237 *guard = Some((
21238 self.alloc_u8(grow(ck, k_ws_bytes))?,
21239 self.alloc_u8(grow(cv, v_ws_bytes))?,
21240 ));
21241 }
21242 let (kw, vw) = guard.as_mut().unwrap();
21243 {
21245 let f = if g {
21247 self.func_g("fa_dequant_kv_ws_bf16")
21248 } else {
21249 self.func("fa_dequant_kv_ws_bf16")
21250 };
21251 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21252 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21253 let cfg = LaunchConfig {
21254 grid_dim: (nblk.max(1), 1, 1),
21255 block_dim: (256, 1, 1),
21256 shared_mem_bytes: 0,
21257 };
21258 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21259 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21260 let __s_b = self.gpu.stream();
21261 let mut b = __s_b.launch_builder(&f);
21262 b.arg(k)
21263 .arg(v)
21264 .arg(&mut *kw)
21265 .arg(&mut *vw)
21266 .arg(&kdk)
21267 .arg(&kdv)
21268 .arg(&tkvi)
21269 .arg(&ktb)
21270 .arg(&vtb);
21271 unsafe {
21272 b.launch(cfg)?;
21273 }
21274 }
21275 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21283 .map(|v| v != "0")
21284 .unwrap_or(true);
21285 {
21286 let hd_sfx = fa_hd_suffix(head_dim)?;
21287 let f = self.func(&format!(
21288 "fa_prefill_qw{}{hd_sfx}",
21289 if db { "_db" } else { "" }
21290 ));
21291 let shmem = if db {
21292 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21294 } else {
21295 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21296 };
21297 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21298 f.set_attribute(
21299 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21300 shmem as i32,
21301 )?;
21302 let cfg = LaunchConfig {
21303 grid_dim: (
21304 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21305 n_head as u32,
21306 1,
21307 ),
21308 block_dim: (32, 4, 1),
21309 shared_mem_bytes: shmem,
21310 };
21311 let (hd, nh, nhkv, ti, tkvi, cz) = (
21312 head_dim as i32,
21313 n_head as i32,
21314 n_head_kv as i32,
21315 t as i32,
21316 t_kv as i32,
21317 causal as i32,
21318 );
21319 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21320 let __s_b = self.gpu.stream();
21321 let mut b = __s_b.launch_builder(&f);
21322 b.arg(q)
21323 .arg(&*kw)
21324 .arg(&*vw)
21325 .arg(o)
21326 .arg(&hd)
21327 .arg(&nh)
21328 .arg(&nhkv)
21329 .arg(&ti)
21330 .arg(&tkvi)
21331 .arg(&scale)
21332 .arg(&cz)
21333 .arg(&kdk)
21334 .arg(&kdv);
21335 unsafe {
21336 b.launch(cfg)?;
21337 }
21338 }
21339 Ok(())
21340 }
21341
21342 #[allow(clippy::too_many_arguments)]
21358 pub fn fa_prefill_view_ws_w_hd128(
21359 &self,
21360 q: &CudaSlice<f32>,
21361 k: &cudarc::driver::CudaView<u8>,
21362 v: &cudarc::driver::CudaView<u8>,
21363 o: &mut CudaSlice<f32>,
21364 head_dim: usize,
21365 n_head: usize,
21366 n_head_kv: usize,
21367 t: usize,
21368 t_kv: usize,
21369 scale: f32,
21370 causal: bool,
21371 window: usize,
21372 k_tok_bytes: usize,
21373 v_tok_bytes: usize,
21374 ) -> Result<(), Box<dyn std::error::Error>> {
21375 assert_eq!(
21376 head_dim, 128,
21377 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21378 );
21379 if portable_mma_gated() {
21380 return self.sdpa_naive_w_quantized_view(
21381 q,
21382 k,
21383 v,
21384 o,
21385 head_dim,
21386 n_head,
21387 n_head_kv,
21388 t,
21389 t_kv,
21390 scale,
21391 causal,
21392 window,
21393 k_tok_bytes,
21394 v_tok_bytes,
21395 );
21396 }
21397 const BLOCK_Q: usize = 64;
21398 const BK: usize = 32;
21399 let kv_dim_k = n_head_kv * head_dim;
21400 let kv_dim_v = n_head_kv * head_dim;
21401 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21403 let mut guard = self.prime_deqw_ws.lock().unwrap();
21404 let need_grow = match guard.as_ref() {
21405 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21406 None => true,
21407 };
21408 if need_grow {
21409 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21410 let (ck, cv) = guard
21411 .as_ref()
21412 .map(|(a, b)| (a.len(), b.len()))
21413 .unwrap_or((0, 0));
21414 *guard = Some((
21415 self.alloc_u8(grow(ck, k_ws_bytes))?,
21416 self.alloc_u8(grow(cv, v_ws_bytes))?,
21417 ));
21418 }
21419 let (kw, vw) = guard.as_mut().unwrap();
21420 {
21423 let f = self.func("fa_dequant_kv_ws_bf16");
21424 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21425 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21426 let cfg = LaunchConfig {
21427 grid_dim: (nblk.max(1), 1, 1),
21428 block_dim: (256, 1, 1),
21429 shared_mem_bytes: 0,
21430 };
21431 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21432 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21433 let __s_b = self.gpu.stream();
21434 let mut b = __s_b.launch_builder(&f);
21435 b.arg(k)
21436 .arg(v)
21437 .arg(&mut *kw)
21438 .arg(&mut *vw)
21439 .arg(&kdk)
21440 .arg(&kdv)
21441 .arg(&tkvi)
21442 .arg(&ktb)
21443 .arg(&vtb);
21444 unsafe {
21445 b.launch(cfg)?;
21446 }
21447 }
21448 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21450 .map(|v| v != "0")
21451 .unwrap_or(true);
21452 {
21453 let f = self.func(if db {
21454 "fa_prefill_qw_db_w_hd128"
21455 } else {
21456 "fa_prefill_qw_w_hd128"
21457 });
21458 let shmem = if db {
21459 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21460 } else {
21461 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21462 };
21463 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21464 f.set_attribute(
21465 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21466 shmem as i32,
21467 )?;
21468 let cfg = LaunchConfig {
21469 grid_dim: (
21470 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21471 n_head as u32,
21472 1,
21473 ),
21474 block_dim: (32, 4, 1),
21475 shared_mem_bytes: shmem,
21476 };
21477 let (hd, nh, nhkv, ti, tkvi, cz) = (
21478 head_dim as i32,
21479 n_head as i32,
21480 n_head_kv as i32,
21481 t as i32,
21482 t_kv as i32,
21483 causal as i32,
21484 );
21485 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21486 let __s_b = self.gpu.stream();
21487 let mut b = __s_b.launch_builder(&f);
21488 b.arg(q)
21489 .arg(&*kw)
21490 .arg(&*vw)
21491 .arg(o)
21492 .arg(&hd)
21493 .arg(&nh)
21494 .arg(&nhkv)
21495 .arg(&ti)
21496 .arg(&tkvi)
21497 .arg(&scale)
21498 .arg(&cz)
21499 .arg(&kdk)
21500 .arg(&kdv)
21501 .arg(&wnd);
21502 unsafe {
21503 b.launch(cfg)?;
21504 }
21505 }
21506 Ok(())
21507 }
21508
21509 pub fn fa_decode(
21513 &self,
21514 q: &CudaSlice<f32>,
21515 k: &cudarc::driver::CudaView<u8>,
21516 v: &cudarc::driver::CudaView<u8>,
21517 o: &mut CudaSlice<f32>,
21518 head_dim: usize,
21519 n_head: usize,
21520 n_head_kv: usize,
21521 t_kv: usize,
21522 scale: f32,
21523 k_tok_bytes: usize,
21524 v_tok_bytes: usize,
21525 ) -> Result<(), Box<dyn std::error::Error>> {
21526 self.fa_decode_kvmod(
21527 q,
21528 k,
21529 v,
21530 o,
21531 head_dim,
21532 n_head,
21533 n_head_kv,
21534 t_kv,
21535 scale,
21536 k_tok_bytes,
21537 v_tok_bytes,
21538 false,
21539 )
21540 }
21541
21542 #[allow(clippy::too_many_arguments)]
21546 #[allow(clippy::too_many_arguments)]
21550 #[allow(clippy::too_many_arguments)]
21551 fn fa_decode_scalar_unified(
21552 &self,
21553 q: &cudarc::driver::CudaView<f32>,
21554 k: &cudarc::driver::CudaView<u8>,
21555 v: &cudarc::driver::CudaView<u8>,
21556 o: &mut cudarc::driver::CudaViewMut<f32>,
21557 head_dim: usize,
21558 n_head: usize,
21559 n_head_kv: usize,
21560 t_kv_host: usize,
21561 t_kv_dev: Option<&CudaSlice<i32>>,
21562 scale: f32,
21563 n_splits: usize,
21564 split_keys: usize,
21565 k_tok_bytes: usize,
21566 v_tok_bytes: usize,
21567 g: bool,
21568 part_o: &mut CudaSlice<f32>,
21569 part_m: &mut CudaSlice<f32>,
21570 part_l: &mut CudaSlice<f32>,
21571 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21572 ) -> Result<(), Box<dyn std::error::Error>> {
21573 let f = if g {
21574 self.func_g("fa_decode_f32")
21575 } else {
21576 self.fa_func("fa_decode_f32", head_dim)
21577 };
21578 let cfg = LaunchConfig {
21579 grid_dim: (n_head as u32, n_splits as u32, 1),
21580 block_dim: (head_dim as u32, 1, 1),
21581 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21582 };
21583 let (hd, nh, nhkv, nsp) = (
21584 head_dim as i32,
21585 n_head as i32,
21586 n_head_kv as i32,
21587 n_splits as i32,
21588 );
21589 let (ktb, vtb, tkvi, ski) = (
21590 k_tok_bytes as i64,
21591 v_tok_bytes as i64,
21592 t_kv_host as i32,
21593 split_keys as i32,
21594 );
21595 let __s_b = self.gpu.stream();
21596 let mut b = __s_b.launch_builder(&f);
21597 match t_kv_dev {
21598 Some(d) => {
21599 b.arg(q)
21600 .arg(k)
21601 .arg(v)
21602 .arg(&mut *part_o)
21603 .arg(&mut *part_m)
21604 .arg(&mut *part_l)
21605 .arg(&hd)
21606 .arg(&nh)
21607 .arg(&nhkv)
21608 .arg(&tkvi)
21609 .arg(d)
21610 .arg(&scale)
21611 .arg(&nsp)
21612 .arg(&ski)
21613 .arg(&ktb)
21614 .arg(&vtb);
21615 unsafe {
21616 b.launch(cfg)?;
21617 }
21618 }
21619 None => {
21620 let null: u64 = 0;
21621 b.arg(q)
21622 .arg(k)
21623 .arg(v)
21624 .arg(&mut *part_o)
21625 .arg(&mut *part_m)
21626 .arg(&mut *part_l)
21627 .arg(&hd)
21628 .arg(&nh)
21629 .arg(&nhkv)
21630 .arg(&tkvi)
21631 .arg(&null)
21632 .arg(&scale)
21633 .arg(&nsp)
21634 .arg(&ski)
21635 .arg(&ktb)
21636 .arg(&vtb);
21637 unsafe {
21638 b.launch(cfg)?;
21639 }
21640 }
21641 }
21642 let cfg2 = LaunchConfig {
21643 grid_dim: (n_head as u32, 1, 1),
21644 block_dim: (head_dim as u32, 1, 1),
21645 shared_mem_bytes: 0,
21646 };
21647 if let Some((oq, od)) = q8_out {
21648 let fc = if g {
21650 self.func_g("fa_decode_combine_q8_1")
21651 } else {
21652 self.fa_func("fa_decode_combine_q8_1", head_dim)
21653 };
21654 let __s_b2 = self.gpu.stream();
21655 let mut b2 = __s_b2.launch_builder(&fc);
21656 b2.arg(&*part_o)
21657 .arg(&*part_m)
21658 .arg(&*part_l)
21659 .arg(oq)
21660 .arg(od)
21661 .arg(&hd)
21662 .arg(&nh)
21663 .arg(&nsp);
21664 unsafe {
21665 b2.launch(cfg2)?;
21666 }
21667 return Ok(());
21668 }
21669 let fc = if g {
21670 self.func_g("fa_decode_combine_f32")
21671 } else {
21672 self.fa_func("fa_decode_combine_f32", head_dim)
21673 };
21674 let __s_b2 = self.gpu.stream();
21675 let mut b2 = __s_b2.launch_builder(&fc);
21676 b2.arg(&*part_o)
21677 .arg(&*part_m)
21678 .arg(&*part_l)
21679 .arg(o)
21680 .arg(&hd)
21681 .arg(&nh)
21682 .arg(&nsp);
21683 unsafe {
21684 b2.launch(cfg2)?;
21685 }
21686 Ok(())
21687 }
21688
21689 pub fn fa_decode_kvmod(
21690 &self,
21691 q: &CudaSlice<f32>,
21692 k: &cudarc::driver::CudaView<u8>,
21693 v: &cudarc::driver::CudaView<u8>,
21694 o: &mut CudaSlice<f32>,
21695 head_dim: usize,
21696 n_head: usize,
21697 n_head_kv: usize,
21698 t_kv: usize,
21699 scale: f32,
21700 k_tok_bytes: usize,
21701 v_tok_bytes: usize,
21702 g: bool,
21703 ) -> Result<(), Box<dyn std::error::Error>> {
21704 let q_view = q.as_view();
21705 let mut o_view = o.as_view_mut();
21706 self.fa_decode_kvmod_view(
21707 &q_view,
21708 k,
21709 v,
21710 &mut o_view,
21711 head_dim,
21712 n_head,
21713 n_head_kv,
21714 t_kv,
21715 scale,
21716 k_tok_bytes,
21717 v_tok_bytes,
21718 g,
21719 )
21720 }
21721
21722 #[allow(clippy::too_many_arguments)]
21727 pub fn fa_decode_kvmod_view(
21728 &self,
21729 q: &cudarc::driver::CudaView<f32>,
21730 k: &cudarc::driver::CudaView<u8>,
21731 v: &cudarc::driver::CudaView<u8>,
21732 o: &mut cudarc::driver::CudaViewMut<f32>,
21733 head_dim: usize,
21734 n_head: usize,
21735 n_head_kv: usize,
21736 t_kv: usize,
21737 scale: f32,
21738 k_tok_bytes: usize,
21739 v_tok_bytes: usize,
21740 g: bool,
21741 ) -> Result<(), Box<dyn std::error::Error>> {
21742 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21763 if g && head_dim == 256 && !fa_v4_at(t_kv) {
21767 fa_vec = false;
21768 }
21769 let sp = fa_split_keys(t_kv, n_head_kv);
21770 let n_splits = if fa_vec {
21771 ((t_kv + sp - 1) / sp).max(1)
21772 } else {
21773 ((t_kv + 255) / 256).max(1)
21774 };
21775 let o_len = n_head * n_splits * head_dim;
21776 let ml_len = n_head * n_splits;
21777 let mut part_guard = self.fa_part_pool.lock().unwrap();
21778 if part_guard
21779 .as_ref()
21780 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21781 .unwrap_or(true)
21782 {
21783 let old = part_guard.take();
21794 let (co, cm) = old
21795 .as_ref()
21796 .map(|pp| (pp.0.len(), pp.1.len()))
21797 .unwrap_or((0, 0));
21798 if let Some(old) = old {
21799 self.fa_part_retired.lock().unwrap().push(old);
21800 }
21801 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21802 eprintln!(
21803 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21804 co, o_len, cm, ml_len
21805 );
21806 }
21807 *part_guard = Some((
21808 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21809 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21810 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21811 ));
21812 }
21813 let pg = part_guard.as_mut().unwrap();
21814 self.gpu
21815 .stream()
21816 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21817 self.gpu
21818 .stream()
21819 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21820 self.gpu
21821 .stream()
21822 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21823 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21824 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21825 let (hd, nh, nhkv, tkvi, nsp) = (
21826 head_dim as i32,
21827 n_head as i32,
21828 n_head_kv as i32,
21829 t_kv as i32,
21830 n_splits as i32,
21831 );
21832 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21833 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
21837 let fa512_min = fa512_min_tkv();
21842 let deep = fa_vec
21845 && head_dim == 256
21846 && fa_v4_at(t_kv)
21847 && !g
21848 && fa_deep_at(t_kv)
21849 && !matches!(fa_v4_mode(), "noB3" | "stage");
21850 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21851 let gqa = (n_head / n_head_kv).max(1) as u32;
21854 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21855 (
21856 fv,
21857 LaunchConfig {
21858 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21859 block_dim: (32, gqa, 1),
21860 shared_mem_bytes: 0,
21861 },
21862 )
21863 } else if fa_vec && head_dim <= 256 {
21864 let gqa = (n_head / n_head_kv).max(1) as u32;
21865 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21876 let smem_tkv = *SMEM_TKV.get_or_init(|| {
21877 std::env::var("MEMRA_FA_SMEM_TKV")
21878 .ok()
21879 .and_then(|v| v.parse().ok())
21880 .unwrap_or_else(|| {
21881 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21882 })
21883 });
21884 if fa_v4_at(t_kv) && head_dim == 256 {
21885 let v4name = match fa_v4_mode() {
21889 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
21892 _ => "fa_decode_vec_q_v4",
21893 };
21894 let fv = if g {
21895 self.func_g(v4name)
21896 } else {
21897 self.func(v4name)
21898 };
21899 let shmem = (if deep { 12160 } else { 11520 }
21902 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
21903 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21904 fv.set_attribute(
21905 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21906 shmem as i32,
21907 )?;
21908 (
21909 fv,
21910 LaunchConfig {
21911 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21912 block_dim: (32, gqa, 1),
21913 shared_mem_bytes: shmem,
21914 },
21915 )
21916 } else if fa_v3_active(head_dim) {
21917 let fv = if g {
21920 self.func_g("fa_decode_vec_q_v3")
21921 } else {
21922 self.func("fa_decode_vec_q_v3")
21923 };
21924 let shmem = (32 * head_dim * 2) as u32; (
21926 fv,
21927 LaunchConfig {
21928 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21929 block_dim: (32, gqa, 1),
21930 shared_mem_bytes: shmem,
21931 },
21932 )
21933 } else if fa_v2_on() {
21934 let fv = if g {
21938 self.func_g("fa_decode_vec_q_v2")
21939 } else {
21940 self.func("fa_decode_vec_q_v2")
21941 };
21942 let shmem = (2 * 32 * head_dim * 2) as u32; (
21944 fv,
21945 LaunchConfig {
21946 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21947 block_dim: (32, gqa, 1),
21948 shared_mem_bytes: shmem,
21949 },
21950 )
21951 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
21952 {
21953 let fv = if g {
21957 self.func_g("fa_decode_vec_q_smem")
21958 } else {
21959 self.func("fa_decode_vec_q_smem")
21960 };
21961 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
21963 fv.set_attribute(
21964 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21965 shmem as i32,
21966 )?;
21967 (
21968 fv,
21969 LaunchConfig {
21970 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21971 block_dim: (32, gqa, 1),
21972 shared_mem_bytes: shmem,
21973 },
21974 )
21975 } else {
21976 let fv = if g {
21979 self.func_g("fa_decode_vec_q")
21980 } else {
21981 self.func("fa_decode_vec_q")
21982 };
21983 (
21984 fv,
21985 LaunchConfig {
21986 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21987 block_dim: (32, gqa, 1),
21988 shared_mem_bytes: 0,
21989 },
21990 )
21991 }
21992 } else {
21993 return self.fa_decode_scalar_unified(
21996 q,
21997 k,
21998 v,
21999 o,
22000 head_dim,
22001 n_head,
22002 n_head_kv,
22003 t_kv,
22004 None,
22005 scale,
22006 n_splits,
22007 if fa_vec { sp } else { 256 },
22008 k_tok_bytes,
22009 v_tok_bytes,
22010 g,
22011 part_o,
22012 part_m,
22013 part_l,
22014 None,
22015 );
22016 };
22017 let __s_b = self.gpu.stream();
22018 let mut b = __s_b.launch_builder(&f);
22019 b.arg(q)
22020 .arg(k)
22021 .arg(v)
22022 .arg(&mut *part_o)
22023 .arg(&mut *part_m)
22024 .arg(&mut *part_l)
22025 .arg(&hd)
22026 .arg(&nh)
22027 .arg(&nhkv)
22028 .arg(&tkvi)
22029 .arg(&scale)
22030 .arg(&nsp)
22031 .arg(&ktb)
22032 .arg(&vtb);
22033 unsafe {
22034 b.launch(cfg)?;
22035 }
22036 let (fc, cfg2) = (
22039 if g {
22040 self.func_g("fa_decode_combine_f32")
22041 } else {
22042 self.fa_func("fa_decode_combine_f32", head_dim)
22043 },
22044 LaunchConfig {
22045 grid_dim: (n_head as u32, 1, 1),
22046 block_dim: (head_dim as u32, 1, 1),
22047 shared_mem_bytes: 0,
22048 },
22049 );
22050 let __s_b2 = self.gpu.stream();
22051 let mut b2 = __s_b2.launch_builder(&fc);
22052 b2.arg(&*part_o)
22053 .arg(&*part_m)
22054 .arg(&*part_l)
22055 .arg(o)
22056 .arg(&hd)
22057 .arg(&nh)
22058 .arg(&nsp);
22059 unsafe {
22060 b2.launch(cfg2)?;
22061 }
22062 Ok(())
22063 }
22064
22065 #[allow(clippy::too_many_arguments)]
22076 pub fn fa_decode_batch_seqs_v4(
22077 &self,
22078 q: &CudaSlice<f32>,
22079 kv_ptrs: &cudarc::driver::CudaView<u64>,
22080 pos_seq: &CudaSlice<i32>,
22081 o: &mut CudaSlice<f32>,
22082 head_dim: usize,
22083 n_head: usize,
22084 n_head_kv: usize,
22085 b_n: usize,
22086 t_kv_max: usize,
22087 scale: f32,
22088 split_keys: usize,
22089 k_tok_bytes: usize,
22090 v_tok_bytes: usize,
22091 ) -> Result<(), Box<dyn std::error::Error>> {
22092 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22093 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22094 let o_len = b_n * n_head * n_splits_max * head_dim;
22095 let ml_len = b_n * n_head * n_splits_max;
22096 let mut part_guard = self.fa_part_pool.lock().unwrap();
22097 if part_guard
22098 .as_ref()
22099 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22100 .unwrap_or(true)
22101 {
22102 let old = part_guard.take();
22113 let (co, cm) = old
22114 .as_ref()
22115 .map(|pp| (pp.0.len(), pp.1.len()))
22116 .unwrap_or((0, 0));
22117 if let Some(old) = old {
22118 self.fa_part_retired.lock().unwrap().push(old);
22119 }
22120 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22121 eprintln!(
22122 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22123 co, o_len, cm, ml_len
22124 );
22125 }
22126 *part_guard = Some((
22127 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22128 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22129 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22130 ));
22131 }
22132 let pg = part_guard.as_mut().unwrap();
22133 self.gpu
22134 .stream()
22135 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22136 self.gpu
22137 .stream()
22138 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22139 self.gpu
22140 .stream()
22141 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22142 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22143 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22144 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22145 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22146 let gqa = (n_head / n_head_kv).max(1) as u32;
22147 let f = self.func("fa_decode_vec_q_seqs_v4");
22148 let shmem = (11520 + 32 * head_dim * 2) as u32;
22150 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22151 f.set_attribute(
22152 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22153 shmem as i32,
22154 )?;
22155 let cfg = LaunchConfig {
22156 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22157 block_dim: (32, gqa, 1),
22158 shared_mem_bytes: shmem,
22159 };
22160 {
22161 let __s_b = self.gpu.stream();
22162 let mut b = __s_b.launch_builder(&f);
22163 b.arg(q)
22164 .arg(kv_ptrs)
22165 .arg(pos_seq)
22166 .arg(&mut *part_o)
22167 .arg(&mut *part_m)
22168 .arg(&mut *part_l)
22169 .arg(&hd)
22170 .arg(&nh)
22171 .arg(&nhkv)
22172 .arg(&scale)
22173 .arg(&nspm)
22174 .arg(&spk)
22175 .arg(&ktb)
22176 .arg(&vtb);
22177 unsafe {
22178 b.launch(cfg)?;
22179 }
22180 }
22181 let fc = self.func("fa_decode_combine_seqs");
22182 let cfg2 = LaunchConfig {
22183 grid_dim: (n_head as u32, b_n as u32, 1),
22184 block_dim: (head_dim as u32, 1, 1),
22185 shared_mem_bytes: 0,
22186 };
22187 let __s_b2 = self.gpu.stream();
22188 let mut b2 = __s_b2.launch_builder(&fc);
22189 b2.arg(&*part_o)
22190 .arg(&*part_m)
22191 .arg(&*part_l)
22192 .arg(o)
22193 .arg(&hd)
22194 .arg(&nh)
22195 .arg(pos_seq)
22196 .arg(&nspm)
22197 .arg(&spk);
22198 unsafe {
22199 b2.launch(cfg2)?;
22200 }
22201 Ok(())
22202 }
22203
22204 #[allow(clippy::too_many_arguments)]
22211 pub fn append_kv_quantized_seqs(
22212 &self,
22213 k_rows: &CudaSlice<f32>,
22214 v_rows: &CudaSlice<f32>,
22215 kv_ptrs: &cudarc::driver::CudaView<u64>,
22216 pos_seq: &CudaSlice<i32>,
22217 b_n: usize,
22218 kv_dim_k: usize,
22219 kv_dim_v: usize,
22220 k_tok_bytes: usize,
22221 v_tok_bytes: usize,
22222 ) -> Result<(), Box<dyn std::error::Error>> {
22223 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22224 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22225 let cfg = LaunchConfig {
22226 grid_dim: (nblk, b_n as u32, 1),
22227 block_dim: (32, 1, 1),
22228 shared_mem_bytes: 0,
22229 };
22230 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22231 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22232 let __s_b = self.gpu.stream();
22233 let mut b = __s_b.launch_builder(&f);
22234 b.arg(k_rows)
22235 .arg(v_rows)
22236 .arg(kv_ptrs)
22237 .arg(pos_seq)
22238 .arg(&kdk)
22239 .arg(&kdv)
22240 .arg(&ktb)
22241 .arg(&vtb);
22242 unsafe {
22243 b.launch(cfg)?;
22244 }
22245 Ok(())
22246 }
22247
22248 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22254 std::env::var("MEMRA_NO_FA_VEC").is_err()
22255 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22256 && base_len + 1 >= fa_vec_min_tkv()
22257 && head_dim <= 256
22258 && head_dim % 32 == 0
22259 }
22260
22261 #[allow(clippy::too_many_arguments)]
22270 pub fn fa_decode_rows(
22271 &self,
22272 q: &CudaSlice<f32>,
22273 k: &cudarc::driver::CudaView<u8>,
22274 v: &cudarc::driver::CudaView<u8>,
22275 o: &mut CudaSlice<f32>,
22276 head_dim: usize,
22277 n_head: usize,
22278 n_head_kv: usize,
22279 base_len: usize,
22280 t: usize,
22281 scale: f32,
22282 k_tok_bytes: usize,
22283 v_tok_bytes: usize,
22284 base_dev: Option<(&CudaSlice<i32>, i32)>,
22288 kv_shared: bool,
22291 g: bool,
22295 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22298 ) -> Result<(), Box<dyn std::error::Error>> {
22299 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22300 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
22307 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22308 let v = *SP512.get_or_init(|| {
22311 std::env::var("MEMRA_FA_SP512")
22312 .ok()
22313 .and_then(|x| x.parse().ok())
22314 .unwrap_or(0)
22315 });
22316 sp = if v >= 8 {
22317 v
22318 } else {
22319 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22320 };
22321 }
22322 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22323 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22324 let gqa = (n_head / n_head_kv).max(1) as u32;
22325 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22336 groups.push((0, t, sp));
22337 } else {
22338 let mut r0 = 0usize;
22339 while r0 < t {
22340 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22341 let mut r1 = r0 + 1;
22342 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22343 r1 += 1;
22344 }
22345 groups.push((r0, r1 - r0, sp_g));
22346 r0 = r1;
22347 }
22348 }
22349 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22353 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22354 std::env::var("MEMRA_FA_SMEM_TKV")
22355 .ok()
22356 .and_then(|v| v.parse().ok())
22357 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22358 });
22359 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22360 let v3 = fa_v3_active(head_dim);
22361 let smem_rows =
22362 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22363 let _ = kv_shared;
22368 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22371 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22385 let tb512 = head_dim == 512
22387 && sp <= 32
22388 && n_head / n_head_kv.max(1) <= 16
22389 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22390 let fname = if tb512 {
22391 "fa_decode_vec_q_rows_v4_512_tb"
22392 } else if i2 {
22393 "fa_decode_vec_q_rows_dpl16_i2"
22394 } else if head_dim == 512 {
22395 "fa_decode_vec_q_rows_dpl16"
22396 }
22397 else if v4 {
22399 "fa_decode_vec_q_rows_v4"
22400 } else if v3 {
22401 "fa_decode_vec_q_rows_v3"
22402 } else if fa_v2_on() {
22403 "fa_decode_vec_q_rows_v2"
22404 } else if smem_rows {
22405 "fa_decode_vec_q_rows_smem"
22406 } else {
22407 "fa_decode_vec_q_rows"
22408 };
22409 let f = if head_dim == 512 {
22410 self.fa_func(fname, head_dim)
22411 } else if g {
22412 self.func_g(if smem_rows {
22420 "fa_decode_vec_q_rows"
22421 } else {
22422 fname
22423 })
22424 } else {
22425 self.func(fname)
22426 };
22427 let shmem = if tb512 {
22428 let gk = Self::gkv_on();
22430 let sh =
22431 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22432 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22433 f.set_attribute(
22434 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22435 sh as i32,
22436 )?;
22437 sh
22438 } else if v4 || v3 || smem_rows || fa_v2_on() {
22439 let sh = (if v4 {
22441 11520 + 32 * head_dim * if g { 1 } else { 2 }
22442 } else if v3 {
22443 32 * head_dim * 2
22444 } else {
22445 2 * 32 * head_dim * 2
22446 }) as u32;
22447 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22448 f.set_attribute(
22449 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22450 sh as i32,
22451 )?;
22452 sh
22453 } else {
22454 0
22455 };
22456 for &(r0, t_g, sp_g) in &groups {
22460 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22461 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22462 let base_i = (base_len + r0) as i32;
22463 let o_len = t_g * n_head * n_splits_g * head_dim;
22464 let ml_len = t_g * n_head * n_splits_g;
22465 let mut part_guard = self.fa_part_pool.lock().unwrap();
22466 if part_guard
22467 .as_ref()
22468 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22469 .unwrap_or(true)
22470 {
22471 let old = part_guard.take();
22482 let (co, cm) = old
22483 .as_ref()
22484 .map(|pp| (pp.0.len(), pp.1.len()))
22485 .unwrap_or((0, 0));
22486 if let Some(old) = old {
22487 self.fa_part_retired.lock().unwrap().push(old);
22488 }
22489 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22490 eprintln!(
22491 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22492 co, o_len, cm, ml_len
22493 );
22494 }
22495 *part_guard = Some((
22496 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22497 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22498 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22499 ));
22500 }
22501 let pg = part_guard.as_mut().unwrap();
22502 self.gpu
22503 .stream()
22504 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22505 self.gpu
22506 .stream()
22507 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22508 self.gpu
22509 .stream()
22510 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22511 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22512 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22513 let qv = self.view(q, t * n_head * head_dim);
22514 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22515 let cfg = LaunchConfig {
22516 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22517 block_dim: (32, gqa, 1),
22518 shared_mem_bytes: shmem,
22519 };
22520 {
22521 let __s_b = self.gpu.stream();
22522 let mut b = __s_b.launch_builder(&f);
22523 if tb512 {
22524 let (bd, plus) =
22526 base_dev.expect("hd512 rows twin requires a device base counter");
22527 let plus_g = plus + r0 as i32;
22528 let nr = t_g as i32;
22529 if Self::pdl_on() && Self::pdl_wb_on() {
22530 use cudarc::driver::{DevicePtr, DevicePtrMut};
22532 let s = &self.gpu.stream();
22533 let (pq, _b0) = q_g.device_ptr(s);
22534 let (pk, _b1) = k.device_ptr(s);
22535 let (pv, _b2) = v.device_ptr(s);
22536 let (po, _b3) = part_o.device_ptr_mut(s);
22537 let (pm, _b4) = part_m.device_ptr_mut(s);
22538 let (pl, _b5) = part_l.device_ptr_mut(s);
22539 let (pb, _b6) = bd.device_ptr(s);
22540 let mut ps = [
22541 &pq as *const _ as *mut std::ffi::c_void,
22542 &pk as *const _ as *mut _,
22543 &pv as *const _ as *mut _,
22544 &po as *const _ as *mut _,
22545 &pm as *const _ as *mut _,
22546 &pl as *const _ as *mut _,
22547 &hd as *const _ as *mut _,
22548 &nh as *const _ as *mut _,
22549 &nhkv as *const _ as *mut _,
22550 &pb as *const _ as *mut _,
22551 &plus_g as *const _ as *mut _,
22552 &scale as *const _ as *mut _,
22553 &nspm as *const _ as *mut _,
22554 &spk as *const _ as *mut _,
22555 &ktb as *const _ as *mut _,
22556 &vtb as *const _ as *mut _,
22557 &nr as *const _ as *mut _,
22558 ];
22559 unsafe {
22560 self.launch_pdl_flash(
22561 Self::gkv_on(),
22562 "fa_decode_vec_q_rows_v4_512_tb",
22563 (n_head_kv as u32, n_splits_g as u32, 1),
22564 (32, gqa, 1),
22565 shmem,
22566 &mut ps,
22567 )?;
22568 }
22569 } else {
22570 let cfg_tb = LaunchConfig {
22571 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22572 block_dim: (32, gqa, 1),
22573 shared_mem_bytes: shmem,
22574 };
22575 b.arg(&q_g)
22576 .arg(k)
22577 .arg(v)
22578 .arg(&mut *part_o)
22579 .arg(&mut *part_m)
22580 .arg(&mut *part_l)
22581 .arg(&hd)
22582 .arg(&nh)
22583 .arg(&nhkv)
22584 .arg(bd)
22585 .arg(&plus_g)
22586 .arg(&scale)
22587 .arg(&nspm)
22588 .arg(&spk)
22589 .arg(&ktb)
22590 .arg(&vtb)
22591 .arg(&nr);
22592 unsafe {
22593 b.launch(cfg_tb)?;
22594 }
22595 }
22596 } else if head_dim == 512 {
22597 let (bd, plus) =
22598 base_dev.expect("hd512 rows twin requires a device base counter");
22599 let plus_g = plus + r0 as i32;
22600 b.arg(&q_g)
22601 .arg(k)
22602 .arg(v)
22603 .arg(&mut *part_o)
22604 .arg(&mut *part_m)
22605 .arg(&mut *part_l)
22606 .arg(&hd)
22607 .arg(&nh)
22608 .arg(&nhkv)
22609 .arg(bd)
22610 .arg(&plus_g)
22611 .arg(&scale)
22612 .arg(&nspm)
22613 .arg(&spk)
22614 .arg(&ktb)
22615 .arg(&vtb);
22616 unsafe {
22617 b.launch(cfg)?;
22618 }
22619 } else {
22620 b.arg(&q_g)
22621 .arg(k)
22622 .arg(v)
22623 .arg(&mut *part_o)
22624 .arg(&mut *part_m)
22625 .arg(&mut *part_l)
22626 .arg(&hd)
22627 .arg(&nh)
22628 .arg(&nhkv)
22629 .arg(&base_i)
22630 .arg(&scale)
22631 .arg(&nspm)
22632 .arg(&spk)
22633 .arg(&ktb)
22634 .arg(&vtb);
22635 unsafe {
22636 b.launch(cfg)?;
22637 }
22638 }
22639 }
22640 let cfg2 = LaunchConfig {
22641 grid_dim: (n_head as u32, t_g as u32, 1),
22642 block_dim: (head_dim as u32, 1, 1),
22643 shared_mem_bytes: 0,
22644 };
22645 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22646 if head_dim == 512 {
22647 let (bd, plus) = base_dev.unwrap();
22650 let plus_g = plus + r0 as i32;
22651 if let Some((oq, od)) = q8_out.as_mut() {
22652 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22654 if Self::pdl_on() && Self::pdl_wb_on() {
22655 use cudarc::driver::{DevicePtr, DevicePtrMut};
22657 let s = &self.gpu.stream();
22658 let (po, _g0) = part_o.device_ptr(s);
22659 let (pm, _g1) = part_m.device_ptr(s);
22660 let (pl, _g2) = part_l.device_ptr(s);
22661 let (pq, _g3) = oq.device_ptr_mut(s);
22662 let (pd, _g4) = od.device_ptr_mut(s);
22663 let (pb, _g5) = bd.device_ptr(s);
22664 let mut ps = [
22665 &po as *const _ as *mut std::ffi::c_void,
22666 &pm as *const _ as *mut _,
22667 &pl as *const _ as *mut _,
22668 &pq as *const _ as *mut _,
22669 &pd as *const _ as *mut _,
22670 &hd as *const _ as *mut _,
22671 &nh as *const _ as *mut _,
22672 &pb as *const _ as *mut _,
22673 &plus_g as *const _ as *mut _,
22674 &nspm as *const _ as *mut _,
22675 &spk as *const _ as *mut _,
22676 ];
22677 unsafe {
22678 self.launch_pdl_flash(
22679 Self::gkv_on(),
22680 "fa_decode_combine_rows_dc_q8_1",
22681 cfg2.grid_dim,
22682 cfg2.block_dim,
22683 0,
22684 &mut ps,
22685 )?;
22686 }
22687 continue;
22688 }
22689 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22690 let __s_b2 = self.gpu.stream();
22691 let mut b2 = __s_b2.launch_builder(&fc);
22692 b2.arg(&*part_o)
22693 .arg(&*part_m)
22694 .arg(&*part_l)
22695 .arg(&mut **oq)
22696 .arg(&mut **od)
22697 .arg(&hd)
22698 .arg(&nh)
22699 .arg(bd)
22700 .arg(&plus_g)
22701 .arg(&nspm)
22702 .arg(&spk);
22703 unsafe {
22704 b2.launch(cfg2)?;
22705 }
22706 continue;
22707 }
22708 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22709 let __s_b2 = self.gpu.stream();
22710 let mut b2 = __s_b2.launch_builder(&fc);
22711 b2.arg(&*part_o)
22712 .arg(&*part_m)
22713 .arg(&*part_l)
22714 .arg(&mut o_g)
22715 .arg(&hd)
22716 .arg(&nh)
22717 .arg(bd)
22718 .arg(&plus_g)
22719 .arg(&nspm)
22720 .arg(&spk);
22721 unsafe {
22722 b2.launch(cfg2)?;
22723 }
22724 } else {
22725 assert!(
22728 q8_out.is_none(),
22729 "rows q8 emit requires the hd512 dc combine"
22730 );
22731 let fc = self.func("fa_decode_combine_rows");
22732 let __s_b2 = self.gpu.stream();
22733 let mut b2 = __s_b2.launch_builder(&fc);
22734 b2.arg(&*part_o)
22735 .arg(&*part_m)
22736 .arg(&*part_l)
22737 .arg(&mut o_g)
22738 .arg(&hd)
22739 .arg(&nh)
22740 .arg(&base_i)
22741 .arg(&nspm)
22742 .arg(&spk);
22743 unsafe {
22744 b2.launch(cfg2)?;
22745 }
22746 }
22747 }
22748 Ok(())
22749 }
22750
22751 #[allow(clippy::too_many_arguments)]
22755 pub fn fa_decode_rows_w(
22756 &self,
22757 q: &CudaSlice<f32>,
22758 k: &cudarc::driver::CudaView<u8>,
22759 v: &cudarc::driver::CudaView<u8>,
22760 o: &mut CudaSlice<f32>,
22761 head_dim: usize,
22762 n_head: usize,
22763 n_head_kv: usize,
22764 base_dev: &CudaSlice<i32>,
22765 base_plus: i32,
22766 t: usize,
22767 scale: f32,
22768 window: usize,
22769 k_tok_bytes: usize,
22770 v_tok_bytes: usize,
22771 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22772 ) -> Result<(), Box<dyn std::error::Error>> {
22773 debug_assert!(head_dim == 256);
22778 let sp = {
22786 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22787 let v = *SPW.get_or_init(|| {
22788 std::env::var("MEMRA_FA_SPW")
22789 .ok()
22790 .and_then(|x| x.parse().ok())
22791 .unwrap_or(0)
22792 });
22793 if v >= 8 {
22794 v
22795 } else {
22796 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22797 }
22798 };
22799 let n_splits_max = (window + sp - 1) / sp;
22800 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22801 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22802 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22803 let gqa = (n_head / n_head_kv).max(1) as u32;
22804 let o_len = t * n_head * n_splits_max * head_dim;
22805 let ml_len = t * n_head * n_splits_max;
22806 let mut part_guard = self.fa_part_pool.lock().unwrap();
22807 if part_guard
22808 .as_ref()
22809 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22810 .unwrap_or(true)
22811 {
22812 let old = part_guard.take();
22823 let (co, cm) = old
22824 .as_ref()
22825 .map(|pp| (pp.0.len(), pp.1.len()))
22826 .unwrap_or((0, 0));
22827 if let Some(old) = old {
22828 self.fa_part_retired.lock().unwrap().push(old);
22829 }
22830 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22831 eprintln!(
22832 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22833 co, o_len, cm, ml_len
22834 );
22835 }
22836 *part_guard = Some((
22837 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22838 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22839 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22840 ));
22841 }
22842 let pg = part_guard.as_mut().unwrap();
22843 self.gpu
22844 .stream()
22845 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22846 self.gpu
22847 .stream()
22848 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22849 self.gpu
22850 .stream()
22851 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22852 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22853 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22859 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22860 std::env::var("MEMRA_FA_SMEM_TKV")
22861 .ok()
22862 .and_then(|v| v.parse().ok())
22863 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22864 });
22865 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22871 let wg = Self::wkv_on();
22876 let sp2 =
22879 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
22880 if sp2 {
22881 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22882 if Self::pdl_on() && Self::pdl_wb_on() {
22883 use cudarc::driver::{DevicePtr, DevicePtrMut};
22885 let s = &self.gpu.stream();
22886 let (pq, _b0) = q.device_ptr(s);
22887 let (pk, _b1) = k.device_ptr(s);
22888 let (pv, _b2) = v.device_ptr(s);
22889 let (po, _b3) = part_o.device_ptr_mut(s);
22890 let (pm, _b4) = part_m.device_ptr_mut(s);
22891 let (pl, _b5) = part_l.device_ptr_mut(s);
22892 let (pb, _b6) = base_dev.device_ptr(s);
22893 let mut ps = [
22894 &pq as *const _ as *mut std::ffi::c_void,
22895 &pk as *const _ as *mut _,
22896 &pv as *const _ as *mut _,
22897 &po as *const _ as *mut _,
22898 &pm as *const _ as *mut _,
22899 &pl as *const _ as *mut _,
22900 &hd as *const _ as *mut _,
22901 &nh as *const _ as *mut _,
22902 &nhkv as *const _ as *mut _,
22903 &pb as *const _ as *mut _,
22904 &base_plus as *const _ as *mut _,
22905 &scale as *const _ as *mut _,
22906 &nspm as *const _ as *mut _,
22907 &spk as *const _ as *mut _,
22908 &ktb as *const _ as *mut _,
22909 &vtb as *const _ as *mut _,
22910 &wini as *const _ as *mut _,
22911 ];
22912 unsafe {
22913 self.launch_pdl_flash(
22914 wg,
22915 "fa_decode_vec_q_rows_v4_w_sp",
22916 (n_head_kv as u32, n_splits_max as u32, t as u32),
22917 (32, gqa + 1, 1),
22918 sh,
22919 &mut ps,
22920 )?;
22921 }
22922 } else {
22923 let f = if wg {
22924 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
22925 } else {
22926 self.func("fa_decode_vec_q_rows_v4_w_sp")
22927 };
22928 f.set_attribute(
22929 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22930 sh as i32,
22931 )?;
22932 let cfg = LaunchConfig {
22933 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22934 block_dim: (32, gqa + 1, 1),
22935 shared_mem_bytes: sh,
22936 };
22937 let __s_b = self.gpu.stream();
22938 let mut b = __s_b.launch_builder(&f);
22939 b.arg(q)
22940 .arg(k)
22941 .arg(v)
22942 .arg(&mut *part_o)
22943 .arg(&mut *part_m)
22944 .arg(&mut *part_l)
22945 .arg(&hd)
22946 .arg(&nh)
22947 .arg(&nhkv)
22948 .arg(base_dev)
22949 .arg(&base_plus)
22950 .arg(&scale)
22951 .arg(&nspm)
22952 .arg(&spk)
22953 .arg(&ktb)
22954 .arg(&vtb)
22955 .arg(&wini);
22956 unsafe {
22957 b.launch(cfg)?;
22958 }
22959 }
22960 } else {
22961 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
22962 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22964 use cudarc::driver::{DevicePtr, DevicePtrMut};
22965 let s = &self.gpu.stream();
22966 let (pq, _b0) = q.device_ptr(s);
22967 let (pk, _b1) = k.device_ptr(s);
22968 let (pv, _b2) = v.device_ptr(s);
22969 let (po, _b3) = part_o.device_ptr_mut(s);
22970 let (pm, _b4) = part_m.device_ptr_mut(s);
22971 let (pl, _b5) = part_l.device_ptr_mut(s);
22972 let (pb, _b6) = base_dev.device_ptr(s);
22973 let mut ps = [
22974 &pq as *const _ as *mut std::ffi::c_void,
22975 &pk as *const _ as *mut _,
22976 &pv as *const _ as *mut _,
22977 &po as *const _ as *mut _,
22978 &pm as *const _ as *mut _,
22979 &pl as *const _ as *mut _,
22980 &hd as *const _ as *mut _,
22981 &nh as *const _ as *mut _,
22982 &nhkv as *const _ as *mut _,
22983 &pb as *const _ as *mut _,
22984 &base_plus as *const _ as *mut _,
22985 &scale as *const _ as *mut _,
22986 &nspm as *const _ as *mut _,
22987 &spk as *const _ as *mut _,
22988 &ktb as *const _ as *mut _,
22989 &vtb as *const _ as *mut _,
22990 &wini as *const _ as *mut _,
22991 ];
22992 unsafe {
22993 self.launch_pdl_flash(
22994 wg,
22995 "fa_decode_vec_q_rows_v4_w",
22996 (n_head_kv as u32, n_splits_max as u32, t as u32),
22997 (32, gqa, 1),
22998 sh,
22999 &mut ps,
23000 )?;
23001 }
23002 } else {
23003 let pick = |name: &str| {
23004 if wg {
23005 self.func_g(name)
23006 } else {
23007 self.func(name)
23008 }
23009 };
23010 let (f, sh) = if fa_v4_at(window) {
23011 let f = pick("fa_decode_vec_q_rows_v4_w");
23012 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23013 } else if smem_tkv > 0 && window >= smem_tkv {
23014 (
23017 pick("fa_decode_vec_q_rows_smem_w"),
23018 (2 * 32 * head_dim * 2) as u32,
23019 )
23020 } else {
23021 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23022 };
23023 f.set_attribute(
23024 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23025 sh as i32,
23026 )?;
23027 let cfg = LaunchConfig {
23028 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23029 block_dim: (32, gqa, 1),
23030 shared_mem_bytes: sh,
23031 };
23032 let __s_b = self.gpu.stream();
23033 let mut b = __s_b.launch_builder(&f);
23034 b.arg(q)
23035 .arg(k)
23036 .arg(v)
23037 .arg(&mut *part_o)
23038 .arg(&mut *part_m)
23039 .arg(&mut *part_l)
23040 .arg(&hd)
23041 .arg(&nh)
23042 .arg(&nhkv)
23043 .arg(base_dev)
23044 .arg(&base_plus)
23045 .arg(&scale)
23046 .arg(&nspm)
23047 .arg(&spk)
23048 .arg(&ktb)
23049 .arg(&vtb)
23050 .arg(&wini);
23051 unsafe {
23052 b.launch(cfg)?;
23053 }
23054 }
23055 }
23056 let cfg2 = LaunchConfig {
23057 grid_dim: (n_head as u32, t as u32, 1),
23058 block_dim: (head_dim as u32, 1, 1),
23059 shared_mem_bytes: 0,
23060 };
23061 if let Some((oq, od)) = q8_out {
23062 if Self::pdl_on() && Self::pdl_wb_on() {
23065 use cudarc::driver::{DevicePtr, DevicePtrMut};
23067 let s = &self.gpu.stream();
23068 let (po, _g0) = part_o.device_ptr(s);
23069 let (pm, _g1) = part_m.device_ptr(s);
23070 let (pl, _g2) = part_l.device_ptr(s);
23071 let (pq, _g3) = oq.device_ptr_mut(s);
23072 let (pd, _g4) = od.device_ptr_mut(s);
23073 let mut ps = [
23074 &po as *const _ as *mut std::ffi::c_void,
23075 &pm as *const _ as *mut _,
23076 &pl as *const _ as *mut _,
23077 &pq as *const _ as *mut _,
23078 &pd as *const _ as *mut _,
23079 &hd as *const _ as *mut _,
23080 &nh as *const _ as *mut _,
23081 &nspm as *const _ as *mut _,
23082 &spk as *const _ as *mut _,
23083 &wini as *const _ as *mut _,
23084 ];
23085 unsafe {
23086 self.launch_pdl_flash(
23087 wg,
23088 "fa_decode_combine_rows_w_q8_1",
23089 cfg2.grid_dim,
23090 cfg2.block_dim,
23091 0,
23092 &mut ps,
23093 )?;
23094 }
23095 return Ok(());
23096 }
23097 let fc = if wg {
23098 self.func_g("fa_decode_combine_rows_w_q8_1")
23099 } else {
23100 self.func("fa_decode_combine_rows_w_q8_1")
23101 };
23102 let __s_b2 = self.gpu.stream();
23103 let mut b2 = __s_b2.launch_builder(&fc);
23104 b2.arg(&*part_o)
23105 .arg(&*part_m)
23106 .arg(&*part_l)
23107 .arg(oq)
23108 .arg(od)
23109 .arg(&hd)
23110 .arg(&nh)
23111 .arg(&nspm)
23112 .arg(&spk)
23113 .arg(&wini);
23114 unsafe {
23115 b2.launch(cfg2)?;
23116 }
23117 return Ok(());
23118 }
23119 let fc = if wg {
23120 self.func_g("fa_decode_combine_rows_w")
23121 } else {
23122 self.func("fa_decode_combine_rows_w")
23123 };
23124 let __s_b2 = self.gpu.stream();
23125 let mut b2 = __s_b2.launch_builder(&fc);
23126 b2.arg(&*part_o)
23127 .arg(&*part_m)
23128 .arg(&*part_l)
23129 .arg(o)
23130 .arg(&hd)
23131 .arg(&nh)
23132 .arg(&nspm)
23133 .arg(&spk)
23134 .arg(&wini);
23135 unsafe {
23136 b2.launch(cfg2)?;
23137 }
23138 Ok(())
23139 }
23140
23141 #[allow(clippy::too_many_arguments)]
23147 pub fn fa_decode_rows_dc(
23148 &self,
23149 q: &CudaSlice<f32>,
23150 k: &cudarc::driver::CudaView<u8>,
23151 v: &cudarc::driver::CudaView<u8>,
23152 o: &mut CudaSlice<f32>,
23153 head_dim: usize,
23154 n_head: usize,
23155 n_head_kv: usize,
23156 base_dev: &CudaSlice<i32>,
23157 t_kv_upper: usize,
23158 t: usize,
23159 scale: f32,
23160 k_tok_bytes: usize,
23161 v_tok_bytes: usize,
23162 base_plus: i32,
23163 g: bool,
23164 ) -> Result<(), Box<dyn std::error::Error>> {
23165 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23166 assert!(
23167 v4 || fa_v3_active(head_dim),
23168 "stream fa rows requires the v3 or v4 lane"
23169 );
23170 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23171 if v4 {
23172 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23173 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23174 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23175 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23176 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23177 let gqa = (n_head / n_head_kv).max(1) as u32;
23178 let o_len = t * n_head * n_splits_max * head_dim;
23179 let ml_len = t * n_head * n_splits_max;
23180 let mut part_guard = self.fa_part_pool.lock().unwrap();
23181 if part_guard
23182 .as_ref()
23183 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23184 .unwrap_or(true)
23185 {
23186 let old = part_guard.take();
23197 let (co, cm) = old
23198 .as_ref()
23199 .map(|pp| (pp.0.len(), pp.1.len()))
23200 .unwrap_or((0, 0));
23201 if let Some(old) = old {
23202 self.fa_part_retired.lock().unwrap().push(old);
23203 }
23204 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23205 eprintln!(
23206 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23207 co, o_len, cm, ml_len
23208 );
23209 }
23210 *part_guard = Some((
23211 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23212 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23213 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23214 ));
23215 }
23216 let pg = part_guard.as_mut().unwrap();
23217 self.gpu
23218 .stream()
23219 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23220 self.gpu
23221 .stream()
23222 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23223 self.gpu
23224 .stream()
23225 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23226 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23227 let f = if g {
23228 self.func_g("fa_decode_vec_q_rows_v4_dc")
23229 } else {
23230 self.func("fa_decode_vec_q_rows_v4_dc")
23231 };
23232 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23233 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23234 f.set_attribute(
23235 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23236 sh as i32,
23237 )?;
23238 let cfg = LaunchConfig {
23239 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23240 block_dim: (32, gqa, 1),
23241 shared_mem_bytes: sh,
23242 };
23243 let __s_b = self.gpu.stream();
23244 let mut b = __s_b.launch_builder(&f);
23245 b.arg(q)
23246 .arg(k)
23247 .arg(v)
23248 .arg(&mut *part_o)
23249 .arg(&mut *part_m)
23250 .arg(&mut *part_l)
23251 .arg(&hd)
23252 .arg(&nh)
23253 .arg(&nhkv)
23254 .arg(base_dev)
23255 .arg(&base_plus)
23256 .arg(&scale)
23257 .arg(&nspm)
23258 .arg(&spk)
23259 .arg(&ktb)
23260 .arg(&vtb);
23261 unsafe {
23262 b.launch(cfg)?;
23263 }
23264 let fc = self.func("fa_decode_combine_rows_dc");
23265 let cfg2 = LaunchConfig {
23266 grid_dim: (n_head as u32, t as u32, 1),
23267 block_dim: (head_dim as u32, 1, 1),
23268 shared_mem_bytes: 0,
23269 };
23270 let __s_b2 = self.gpu.stream();
23271 let mut b2 = __s_b2.launch_builder(&fc);
23272 b2.arg(&*part_o)
23273 .arg(&*part_m)
23274 .arg(&*part_l)
23275 .arg(o)
23276 .arg(&hd)
23277 .arg(&nh)
23278 .arg(base_dev)
23279 .arg(&base_plus)
23280 .arg(&nspm)
23281 .arg(&spk);
23282 unsafe {
23283 b2.launch(cfg2)?;
23284 }
23285 return Ok(());
23286 }
23287 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23288 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23289 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23290 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23291 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23292 let gqa = (n_head / n_head_kv).max(1) as u32;
23293 let o_len = t * n_head * n_splits_max * head_dim;
23294 let ml_len = t * n_head * n_splits_max;
23295 let mut part_guard = self.fa_part_pool.lock().unwrap();
23296 if part_guard
23297 .as_ref()
23298 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23299 .unwrap_or(true)
23300 {
23301 let old = part_guard.take();
23312 let (co, cm) = old
23313 .as_ref()
23314 .map(|pp| (pp.0.len(), pp.1.len()))
23315 .unwrap_or((0, 0));
23316 if let Some(old) = old {
23317 self.fa_part_retired.lock().unwrap().push(old);
23318 }
23319 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23320 eprintln!(
23321 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23322 co, o_len, cm, ml_len
23323 );
23324 }
23325 *part_guard = Some((
23326 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23327 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23328 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23329 ));
23330 }
23331 let pg = part_guard.as_mut().unwrap();
23332 self.gpu
23333 .stream()
23334 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23335 self.gpu
23336 .stream()
23337 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23338 self.gpu
23339 .stream()
23340 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23341 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23342 let f = self.func("fa_decode_vec_q_rows_v3_dc");
23343 let sh = (32 * head_dim * 2) as u32;
23344 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23345 f.set_attribute(
23346 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23347 sh as i32,
23348 )?;
23349 let cfg = LaunchConfig {
23350 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23351 block_dim: (32, gqa, 1),
23352 shared_mem_bytes: sh,
23353 };
23354 let __s_b = self.gpu.stream();
23355 let mut b = __s_b.launch_builder(&f);
23356 b.arg(q)
23357 .arg(k)
23358 .arg(v)
23359 .arg(&mut *part_o)
23360 .arg(&mut *part_m)
23361 .arg(&mut *part_l)
23362 .arg(&hd)
23363 .arg(&nh)
23364 .arg(&nhkv)
23365 .arg(base_dev)
23366 .arg(&scale)
23367 .arg(&nspm)
23368 .arg(&spk)
23369 .arg(&ktb)
23370 .arg(&vtb);
23371 unsafe {
23372 b.launch(cfg)?;
23373 }
23374 let fc = self.func("fa_decode_combine_rows_dc");
23375 let cfg2 = LaunchConfig {
23376 grid_dim: (n_head as u32, t as u32, 1),
23377 block_dim: (head_dim as u32, 1, 1),
23378 shared_mem_bytes: 0,
23379 };
23380 let plus0 = 0i32;
23381 let __s_b2 = self.gpu.stream();
23382 let mut b2 = __s_b2.launch_builder(&fc);
23383 b2.arg(&*part_o)
23384 .arg(&*part_m)
23385 .arg(&*part_l)
23386 .arg(o)
23387 .arg(&hd)
23388 .arg(&nh)
23389 .arg(base_dev)
23390 .arg(&plus0)
23391 .arg(&nspm)
23392 .arg(&spk);
23393 unsafe {
23394 b2.launch(cfg2)?;
23395 }
23396 Ok(())
23397 }
23398
23399 pub fn fa_decode_dc(
23410 &self,
23411 q: &CudaSlice<f32>,
23412 k: &cudarc::driver::CudaView<u8>,
23413 v: &cudarc::driver::CudaView<u8>,
23414 o: &mut CudaSlice<f32>,
23415 head_dim: usize,
23416 n_head: usize,
23417 n_head_kv: usize,
23418 t_kv_dev: &CudaSlice<i32>,
23419 bucket_max: usize,
23420 scale: f32,
23421 k_tok_bytes: usize,
23422 v_tok_bytes: usize,
23423 g: bool,
23424 ) -> Result<(), Box<dyn std::error::Error>> {
23425 self.fa_decode_dc_q8(
23426 q,
23427 k,
23428 v,
23429 o,
23430 head_dim,
23431 n_head,
23432 n_head_kv,
23433 t_kv_dev,
23434 bucket_max,
23435 scale,
23436 k_tok_bytes,
23437 v_tok_bytes,
23438 g,
23439 None,
23440 )
23441 }
23442
23443 #[allow(clippy::too_many_arguments)]
23446 pub fn fa_decode_dc_q8(
23447 &self,
23448 q: &CudaSlice<f32>,
23449 k: &cudarc::driver::CudaView<u8>,
23450 v: &cudarc::driver::CudaView<u8>,
23451 o: &mut CudaSlice<f32>,
23452 head_dim: usize,
23453 n_head: usize,
23454 n_head_kv: usize,
23455 t_kv_dev: &CudaSlice<i32>,
23456 bucket_max: usize,
23457 scale: f32,
23458 k_tok_bytes: usize,
23459 v_tok_bytes: usize,
23460 g: bool,
23461 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23462 ) -> Result<(), Box<dyn std::error::Error>> {
23463 let mut fa_vec =
23471 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23472 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23473 fa_vec = false;
23474 } let sp = fa_split_keys(bucket_max, n_head_kv);
23476 let n_splits = if fa_vec {
23477 ((bucket_max + sp - 1) / sp).max(1)
23478 } else {
23479 ((bucket_max + 255) / 256).max(1)
23480 };
23481 let o_len = n_head * n_splits * head_dim;
23482 let ml_len = n_head * n_splits;
23483 let mut part_guard = self.fa_part_pool.lock().unwrap();
23484 if part_guard
23485 .as_ref()
23486 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23487 .unwrap_or(true)
23488 {
23489 let old = part_guard.take();
23500 let (co, cm) = old
23501 .as_ref()
23502 .map(|pp| (pp.0.len(), pp.1.len()))
23503 .unwrap_or((0, 0));
23504 if let Some(old) = old {
23505 self.fa_part_retired.lock().unwrap().push(old);
23506 }
23507 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23508 eprintln!(
23509 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23510 co, o_len, cm, ml_len
23511 );
23512 }
23513 *part_guard = Some((
23514 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23515 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23516 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23517 ));
23518 }
23519 let pg = part_guard.as_mut().unwrap();
23520 self.gpu
23521 .stream()
23522 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23523 self.gpu
23524 .stream()
23525 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23526 self.gpu
23527 .stream()
23528 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23529 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23530 let (hd, nh, nhkv, nsp) = (
23531 head_dim as i32,
23532 n_head as i32,
23533 n_head_kv as i32,
23534 n_splits as i32,
23535 );
23536 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23537 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23538 let deep = fa_vec
23541 && head_dim == 256
23542 && fa_v4_at(bucket_max)
23543 && !g
23544 && fa_deep_at(bucket_max)
23545 && !matches!(fa_v4_mode(), "noB3" | "stage");
23546 let (f, cfg) = if fa_vec
23547 && head_dim == 512
23548 && bucket_max >= {
23549 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23550 *FA512_MIN_DC.get_or_init(|| {
23551 std::env::var("MEMRA_FA512_MIN")
23552 .ok()
23553 .and_then(|v| v.parse().ok())
23554 .unwrap_or(512)
23555 })
23556 } {
23557 let gqa = (n_head / n_head_kv).max(1) as u32;
23559 (
23560 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23561 LaunchConfig {
23562 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23563 block_dim: (32, gqa, 1),
23564 shared_mem_bytes: 0,
23565 },
23566 )
23567 } else if fa_vec && head_dim == 512 {
23568 let q_view = q.as_view();
23571 let mut o_view = o.as_view_mut();
23572 return self.fa_decode_scalar_unified(
23573 &q_view,
23574 k,
23575 v,
23576 &mut o_view,
23577 head_dim,
23578 n_head,
23579 n_head_kv,
23580 0,
23581 Some(t_kv_dev),
23582 scale,
23583 n_splits,
23584 sp,
23585 k_tok_bytes,
23586 v_tok_bytes,
23587 g,
23588 &mut *part_o,
23589 &mut *part_m,
23590 &mut *part_l,
23591 q8_out,
23592 );
23593 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23594 let gqa = (n_head / n_head_kv).max(1) as u32;
23597 let fv = if g {
23598 self.func_g("fa_decode_vec_q_v4_dc")
23599 } else if deep {
23600 self.func("fa_decode_vec_q_v4_deep_dc")
23601 } else {
23602 self.func("fa_decode_vec_q_v4_dc")
23603 };
23604 let shmem =
23605 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23606 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23607 fv.set_attribute(
23608 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23609 shmem as i32,
23610 )?;
23611 (
23612 fv,
23613 LaunchConfig {
23614 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23615 block_dim: (32, gqa, 1),
23616 shared_mem_bytes: shmem,
23617 },
23618 )
23619 } else if fa_vec && fa_v3_active(head_dim) {
23620 let gqa = (n_head / n_head_kv).max(1) as u32;
23623 let fv = if g {
23624 self.func_g("fa_decode_vec_q_v3_dc")
23625 } else {
23626 self.func("fa_decode_vec_q_v3_dc")
23627 };
23628 let shmem = (32 * head_dim * 2) as u32; (
23630 fv,
23631 LaunchConfig {
23632 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23633 block_dim: (32, gqa, 1),
23634 shared_mem_bytes: shmem,
23635 },
23636 )
23637 } else if fa_vec && fa_v2_on() {
23638 let gqa = (n_head / n_head_kv).max(1) as u32;
23642 let fv = if g {
23643 self.func_g("fa_decode_vec_q_v2_dc")
23644 } else {
23645 self.func("fa_decode_vec_q_v2_dc")
23646 };
23647 let shmem = (2 * 32 * head_dim * 2) as u32; (
23649 fv,
23650 LaunchConfig {
23651 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23652 block_dim: (32, gqa, 1),
23653 shared_mem_bytes: shmem,
23654 },
23655 )
23656 } else if fa_vec {
23657 let gqa = (n_head / n_head_kv).max(1) as u32;
23658 let fv = if g {
23660 self.func_g("fa_decode_vec_q_dc")
23661 } else {
23662 self.func("fa_decode_vec_q_dc")
23663 };
23664 (
23665 fv,
23666 LaunchConfig {
23667 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23668 block_dim: (32, gqa, 1),
23669 shared_mem_bytes: 0,
23670 },
23671 )
23672 } else {
23673 let q_view = q.as_view();
23674 let mut o_view = o.as_view_mut();
23675 return self.fa_decode_scalar_unified(
23676 &q_view,
23677 k,
23678 v,
23679 &mut o_view,
23680 head_dim,
23681 n_head,
23682 n_head_kv,
23683 0,
23684 Some(t_kv_dev),
23685 scale,
23686 n_splits,
23687 if fa_vec { sp } else { 256 },
23688 k_tok_bytes,
23689 v_tok_bytes,
23690 g,
23691 &mut *part_o,
23692 &mut *part_m,
23693 &mut *part_l,
23694 q8_out,
23695 );
23696 };
23697 let ski = sp as i32; let __s_b = self.gpu.stream();
23699 let mut b = __s_b.launch_builder(&f);
23700 b.arg(q)
23701 .arg(k)
23702 .arg(v)
23703 .arg(&mut *part_o)
23704 .arg(&mut *part_m)
23705 .arg(&mut *part_l)
23706 .arg(&hd)
23707 .arg(&nh)
23708 .arg(&nhkv)
23709 .arg(t_kv_dev)
23710 .arg(&scale)
23711 .arg(&nsp)
23712 .arg(&ski)
23713 .arg(&ktb)
23714 .arg(&vtb);
23715 unsafe {
23716 b.launch(cfg)?;
23717 }
23718 let cfg2 = LaunchConfig {
23719 grid_dim: (n_head as u32, 1, 1),
23720 block_dim: (head_dim as u32, 1, 1),
23721 shared_mem_bytes: 0,
23722 };
23723 if let Some((oq, od)) = q8_out {
23724 let fc = if g {
23725 self.func_g("fa_decode_combine_q8_1")
23726 } else {
23727 self.fa_func("fa_decode_combine_q8_1", head_dim)
23728 };
23729 let __s_b2 = self.gpu.stream();
23730 let mut b2 = __s_b2.launch_builder(&fc);
23731 b2.arg(&*part_o)
23732 .arg(&*part_m)
23733 .arg(&*part_l)
23734 .arg(oq)
23735 .arg(od)
23736 .arg(&hd)
23737 .arg(&nh)
23738 .arg(&nsp);
23739 unsafe {
23740 b2.launch(cfg2)?;
23741 }
23742 return Ok(());
23743 }
23744 let fc = if g {
23745 self.func_g("fa_decode_combine_f32")
23746 } else {
23747 self.fa_func("fa_decode_combine_f32", head_dim)
23748 };
23749 let __s_b2 = self.gpu.stream();
23750 let mut b2 = __s_b2.launch_builder(&fc);
23751 b2.arg(&*part_o)
23752 .arg(&*part_m)
23753 .arg(&*part_l)
23754 .arg(o)
23755 .arg(&hd)
23756 .arg(&nh)
23757 .arg(&nsp);
23758 unsafe {
23759 b2.launch(cfg2)?;
23760 }
23761 Ok(())
23762 }
23763
23764 #[allow(clippy::too_many_arguments)]
23768 pub fn append_kv_quantized_dcw(
23769 &self,
23770 k_row: &CudaSlice<f32>,
23771 v_row: &CudaSlice<f32>,
23772 kc: &mut CudaSlice<u8>,
23773 vc: &mut CudaSlice<u8>,
23774 len_dev: &CudaSlice<i32>,
23775 base_dev: Option<&CudaSlice<i32>>,
23776 kv_dim_k: usize,
23777 kv_dim_v: usize,
23778 k_tok_bytes: usize,
23779 v_tok_bytes: usize,
23780 ) -> Result<(), Box<dyn std::error::Error>> {
23781 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23782 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23783 let cfg = LaunchConfig {
23784 grid_dim: (nblk, 1, 1),
23785 block_dim: (32, 1, 1),
23786 shared_mem_bytes: 0,
23787 };
23788 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23789 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23790 let null: u64 = 0;
23791 let __s_b = self.gpu.stream();
23792 let mut b = __s_b.launch_builder(&f);
23793 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23794 match base_dev {
23795 Some(base) => {
23796 b.arg(base);
23797 }
23798 None => {
23799 b.arg(&null);
23800 }
23801 }
23802 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23803 unsafe {
23804 b.launch(cfg)?;
23805 }
23806 Ok(())
23807 }
23808
23809 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23811 let f = self.func("inc_i32");
23812 let cfg = LaunchConfig {
23813 grid_dim: (1, 1, 1),
23814 block_dim: (1, 1, 1),
23815 shared_mem_bytes: 0,
23816 };
23817 let __s_b = self.gpu.stream();
23818 let mut b = __s_b.launch_builder(&f);
23819 b.arg(counter);
23820 unsafe {
23821 b.launch(cfg)?;
23822 }
23823 Ok(())
23824 }
23825
23826 #[allow(clippy::too_many_arguments)]
23835 fn fa_part_pool_grow(
23839 &self,
23840 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
23841 o_len: usize,
23842 ml_len: usize,
23843 ) -> Result<(), Box<dyn std::error::Error>> {
23844 if part_guard
23845 .as_ref()
23846 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23847 .unwrap_or(true)
23848 {
23849 let old = part_guard.take();
23850 let (co, cm) = old
23851 .as_ref()
23852 .map(|pp| (pp.0.len(), pp.1.len()))
23853 .unwrap_or((0, 0));
23854 if let Some(old) = old {
23855 self.fa_part_retired.lock().unwrap().push(old);
23856 }
23857 *part_guard = Some((
23858 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23859 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23860 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23861 ));
23862 }
23863 Ok(())
23864 }
23865
23866 pub fn fa_dcw_pool_ensure(
23869 &self,
23870 head_dim: usize,
23871 n_head: usize,
23872 n_head_kv: usize,
23873 bucket_max: usize,
23874 ) -> Result<(), Box<dyn std::error::Error>> {
23875 let sp = fa_split_keys(bucket_max, n_head_kv);
23876 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23877 let o_len = n_head * n_splits * head_dim;
23878 let ml_len = n_head * n_splits;
23879 let mut part_guard = self.fa_part_pool.lock().unwrap();
23880 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
23881 }
23882
23883 #[allow(clippy::too_many_arguments)]
23891 pub fn fa_decode_dcw2(
23892 &self,
23893 q2: &CudaSlice<f32>,
23894 k_ring: &cudarc::driver::CudaView<u8>,
23895 v_ring: &cudarc::driver::CudaView<u8>,
23896 o2: &mut CudaSlice<f32>,
23897 head_dim: usize,
23898 n_head: usize,
23899 n_head_kv: usize,
23900 len_dev: &CudaSlice<i32>,
23901 base_dev: Option<&CudaSlice<i32>>,
23902 window: usize,
23903 bucket_max: usize,
23904 scale: f32,
23905 k_tok_bytes: usize,
23906 v_tok_bytes: usize,
23907 gate2: &CudaSlice<f32>,
23908 ) -> Result<(), Box<dyn std::error::Error>> {
23909 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23910 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23911 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
23912 }
23913 let sp = fa_split_keys(bucket_max, n_head_kv);
23914 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23915 let o_len = 2 * n_head * n_splits * head_dim;
23917 let ml_len = 2 * n_head * n_splits;
23918 let mut part_guard = self.fa_part_pool.lock().unwrap();
23919 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23920 let pg = part_guard.as_mut().unwrap();
23921 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23922 let (hd, nh, nhkv, nsp) = (
23923 head_dim as i32,
23924 n_head as i32,
23925 n_head_kv as i32,
23926 n_splits as i32,
23927 );
23928 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23929 let (ski, win) = (sp as i32, window as i32);
23930 let gqa = (n_head / n_head_kv).max(1) as u32;
23931 let smem = (32 * head_dim * 2) as u32;
23932 let f = self.func("fa_decode_vec_q_v3_dcw2");
23933 let cfg = LaunchConfig {
23934 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23935 block_dim: (32, gqa, 1),
23936 shared_mem_bytes: smem,
23937 };
23938 let null: u64 = 0;
23939 {
23940 let __s_b = self.gpu.stream();
23941 let mut b = __s_b.launch_builder(&f);
23942 b.arg(q2)
23943 .arg(k_ring)
23944 .arg(v_ring)
23945 .arg(&mut *part_o)
23946 .arg(&mut *part_m)
23947 .arg(&mut *part_l)
23948 .arg(&hd)
23949 .arg(&nh)
23950 .arg(&nhkv)
23951 .arg(len_dev);
23952 match base_dev {
23953 Some(base) => {
23954 b.arg(base);
23955 }
23956 None => {
23957 b.arg(&null);
23958 }
23959 }
23960 b.arg(&win)
23961 .arg(&scale)
23962 .arg(&nsp)
23963 .arg(&ski)
23964 .arg(&ktb)
23965 .arg(&vtb);
23966 unsafe {
23967 b.launch(cfg)?;
23968 }
23969 }
23970 let fc = {
23974 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23975 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23976 self.func("fa_decode_combine_gate_f32_s")
23977 } else {
23978 self.func("fa_decode_combine_gate_f32")
23979 }
23980 };
23981 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23982 let nh2 = (2 * n_head) as i32;
23983 let cfg2 = LaunchConfig {
23984 grid_dim: ((2 * n_head) as u32, 1, 1),
23985 block_dim: (head_dim as u32, 1, 1),
23986 shared_mem_bytes: if combine_shared {
23987 (2 * n_splits * 4) as u32
23988 } else {
23989 0
23990 },
23991 };
23992 let __s_b2 = self.gpu.stream();
23993 let mut b2 = __s_b2.launch_builder(&fc);
23994 b2.arg(&*part_o)
23995 .arg(&*part_m)
23996 .arg(&*part_l)
23997 .arg(gate2)
23998 .arg(o2)
23999 .arg(&hd)
24000 .arg(&nh2)
24001 .arg(&nsp);
24002 unsafe {
24003 b2.launch(cfg2)?;
24004 }
24005 Ok(())
24006 }
24007
24008 #[allow(clippy::too_many_arguments)]
24017 pub fn fa_decode_dcw_rows(
24018 &self,
24019 q_rows: &CudaSlice<f32>,
24020 tab: &CudaSlice<u64>,
24021 o_rows: &mut CudaSlice<f32>,
24022 t: usize,
24023 head_dim: usize,
24024 n_head: usize,
24025 n_head_kv: usize,
24026 window: usize,
24027 max_ns: usize,
24028 scale: f32,
24029 k_tok_bytes: usize,
24030 v_tok_bytes: usize,
24031 gate_rows: &CudaSlice<f32>,
24032 ) -> Result<(), Box<dyn std::error::Error>> {
24033 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24034 || head_dim > 256
24035 || head_dim % 32 != 0
24036 || !fa_v3_on()
24037 {
24038 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24039 }
24040 if fa_sm_count() < 128
24041 || std::env::var("MEMRA_FA_SPLIT").is_ok()
24042 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24043 || std::env::var("MEMRA_FA_SP16").is_ok()
24044 {
24045 return Err(
24046 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24047 (or a <128-SM rig) keep the per-row path"
24048 .into(),
24049 );
24050 }
24051 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24052 return Err("fa_decode_dcw_rows geometry".into());
24053 }
24054 let o_len = t * n_head * max_ns * head_dim;
24055 let ml_len = t * n_head * max_ns;
24056 let mut part_guard = self.fa_part_pool.lock().unwrap();
24057 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24058 let pg = part_guard.as_mut().unwrap();
24059 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24060 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24061 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24062 let (win, mns) = (window as i32, max_ns as i32);
24063 let gqa = (n_head / n_head_kv).max(1) as u32;
24064 let smem = (32 * head_dim * 2) as u32;
24065 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24066 let cfg = LaunchConfig {
24067 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24068 block_dim: (32, gqa, 1),
24069 shared_mem_bytes: smem,
24070 };
24071 {
24072 let __s_b = self.gpu.stream();
24073 let mut b = __s_b.launch_builder(&f);
24074 b.arg(q_rows)
24075 .arg(tab)
24076 .arg(&mut *part_o)
24077 .arg(&mut *part_m)
24078 .arg(&mut *part_l)
24079 .arg(&hd)
24080 .arg(&nh)
24081 .arg(&nhkv)
24082 .arg(&win)
24083 .arg(&scale)
24084 .arg(&mns)
24085 .arg(&ktb)
24086 .arg(&vtb);
24087 unsafe {
24088 b.launch(cfg)?;
24089 }
24090 }
24091 let fc = {
24095 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24096 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24097 self.func("fa_decode_combine_gate_f32_s")
24098 } else {
24099 self.func("fa_decode_combine_gate_f32")
24100 }
24101 };
24102 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24103 let nht = (t * n_head) as i32;
24104 let cfg2 = LaunchConfig {
24105 grid_dim: ((t * n_head) as u32, 1, 1),
24106 block_dim: (head_dim as u32, 1, 1),
24107 shared_mem_bytes: if combine_shared {
24108 (2 * max_ns * 4) as u32
24109 } else {
24110 0
24111 },
24112 };
24113 let __s_b2 = self.gpu.stream();
24114 let mut b2 = __s_b2.launch_builder(&fc);
24115 b2.arg(&*part_o)
24116 .arg(&*part_m)
24117 .arg(&*part_l)
24118 .arg(gate_rows)
24119 .arg(o_rows)
24120 .arg(&hd)
24121 .arg(&nht)
24122 .arg(&mns);
24123 unsafe {
24124 b2.launch(cfg2)?;
24125 }
24126 Ok(())
24127 }
24128
24129 pub fn fa_decode_dcw(
24130 &self,
24131 q: &CudaSlice<f32>,
24132 k_ring: &cudarc::driver::CudaView<u8>,
24133 v_ring: &cudarc::driver::CudaView<u8>,
24134 o: &mut CudaSlice<f32>,
24135 head_dim: usize,
24136 n_head: usize,
24137 n_head_kv: usize,
24138 len_dev: &CudaSlice<i32>,
24139 base_dev: Option<&CudaSlice<i32>>,
24140 window: usize,
24141 bucket_max: usize,
24142 scale: f32,
24143 k_tok_bytes: usize,
24144 v_tok_bytes: usize,
24145 fused_gate: Option<&CudaSlice<f32>>,
24149 ) -> Result<(), Box<dyn std::error::Error>> {
24150 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24151 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24152 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"
24153 .into());
24154 }
24155 let sp = fa_split_keys(bucket_max, n_head_kv);
24156 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24157 let o_len = n_head * n_splits * head_dim;
24158 let ml_len = n_head * n_splits;
24159 let mut part_guard = self.fa_part_pool.lock().unwrap();
24160 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24161 let pg = part_guard.as_mut().unwrap();
24162 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24167 let memset_on = *MEMSET_ON
24172 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24173 || crate::tp::token_graph_building();
24174 if memset_on {
24175 self.gpu
24176 .stream()
24177 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24178 self.gpu
24179 .stream()
24180 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24181 self.gpu
24182 .stream()
24183 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24184 }
24185 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24186 let (hd, nh, nhkv, nsp) = (
24187 head_dim as i32,
24188 n_head as i32,
24189 n_head_kv as i32,
24190 n_splits as i32,
24191 );
24192 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24193 let (ski, win) = (sp as i32, window as i32);
24194 let gqa = (n_head / n_head_kv).max(1) as u32;
24195 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24199 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24200 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24201 Ok("2") => 2,
24202 Ok("1") => 1,
24203 _ => 0,
24204 });
24205 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24210 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24211 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24212 std::sync::Mutex::new(None);
24213 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24217 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24218 && (n_head / n_head_kv) % 2 == 0
24219 && (n_head / n_head_kv) >= 2;
24220 let f = if fprof {
24221 self.func("fa_decode_vec_q_v3_dcw_prof")
24222 } else if hs2 {
24223 self.func("fa_decode_vec_q_v3_dcw_hs2")
24224 } else if hoist == 2 {
24225 self.func("fa_decode_vec_q_v3_dcw_hc")
24227 } else if hoist == 1 {
24228 self.func("fa_decode_vec_q_v3_dcw_h")
24230 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24231 self.func("fa_decode_vec_q_v3_dcw_u8")
24232 } else {
24233 self.func("fa_decode_vec_q_v3_dcw")
24234 };
24235 let cfg = LaunchConfig {
24236 grid_dim: if hs2 {
24237 ((2 * n_head_kv) as u32, n_splits as u32, 1)
24238 } else {
24239 (n_head_kv as u32, n_splits as u32, 1)
24240 },
24241 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24242 shared_mem_bytes: smem,
24243 };
24244 let null: u64 = 0;
24245 let __s_b = self.gpu.stream();
24246 let mut b = __s_b.launch_builder(&f);
24247 b.arg(q)
24248 .arg(k_ring)
24249 .arg(v_ring)
24250 .arg(&mut *part_o)
24251 .arg(&mut *part_m)
24252 .arg(&mut *part_l)
24253 .arg(&hd)
24254 .arg(&nh)
24255 .arg(&nhkv)
24256 .arg(len_dev);
24257 match base_dev {
24258 Some(base) => {
24259 b.arg(base);
24260 }
24261 None => {
24262 b.arg(&null);
24263 }
24264 }
24265 b.arg(&win)
24266 .arg(&scale)
24267 .arg(&nsp)
24268 .arg(&ski)
24269 .arg(&ktb)
24270 .arg(&vtb);
24271 if fprof {
24272 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24273 if guard
24274 .as_ref()
24275 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24276 {
24277 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24278 }
24279 let (_, buf) = guard.as_mut().expect("armed above");
24280 b.arg(&*buf);
24281 unsafe {
24282 b.launch(cfg)?;
24283 }
24284 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24285 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24286 if n % 430 == 0 {
24287 self.stream().synchronize()?;
24288 let h = self.dtoh_u64(buf)?;
24289 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24290 let tot: u64 = h[..6].iter().sum();
24291 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24292 for (i, name) in phases.iter().enumerate() {
24293 let pct = if tot > 0 {
24294 h[i] as f64 / tot as f64 * 100.0
24295 } else {
24296 0.0
24297 };
24298 line.push_str(&format!(" {name}={pct:.1}%"));
24299 }
24300 if h[6] > 0 {
24301 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24302 }
24303 eprintln!("{line}");
24304 }
24305 } else {
24306 unsafe {
24307 b.launch(cfg)?;
24308 }
24309 }
24310 let mut combine_shared = false;
24311 let fc = if fused_gate.is_some() {
24312 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24315 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24316 combine_shared = true;
24317 self.func("fa_decode_combine_gate_f32_s")
24318 } else {
24319 self.func("fa_decode_combine_gate_f32")
24320 }
24321 } else {
24322 self.fa_func("fa_decode_combine_f32", head_dim)
24323 };
24324 let cfg2 = LaunchConfig {
24325 grid_dim: (n_head as u32, 1, 1),
24326 block_dim: (head_dim as u32, 1, 1),
24327 shared_mem_bytes: if combine_shared {
24328 (2 * n_splits * 4) as u32
24329 } else {
24330 0
24331 },
24332 };
24333 let __s_b2 = self.gpu.stream();
24334 let mut b2 = __s_b2.launch_builder(&fc);
24335 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24336 if let Some(gate_row) = fused_gate {
24337 b2.arg(gate_row);
24338 }
24339 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24340 unsafe {
24341 b2.launch(cfg2)?;
24342 }
24343 Ok(())
24344 }
24345
24346 pub fn fa_geom_eager(
24352 &self,
24353 t_kv: usize,
24354 head_dim: usize,
24355 n_head_kv: usize,
24356 g: bool,
24357 ) -> (bool, usize) {
24358 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24362 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24368 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24369 if g && head_dim == 256 && !fa_v4_at(t_kv) {
24375 fa_vec = false;
24376 }
24377 let sp = fa_split_keys(t_kv, n_head_kv);
24378 let n_splits = if fa_vec {
24379 ((t_kv + sp - 1) / sp).max(1)
24380 } else {
24381 ((t_kv + 255) / 256).max(1)
24382 };
24383 (fa_vec, n_splits)
24384 }
24385
24386 pub fn fa_bucket_key(
24392 &self,
24393 t_kv: usize,
24394 head_dim: usize,
24395 n_head_kv: usize,
24396 g: bool,
24397 ) -> (bool, usize) {
24398 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24399 }
24400
24401 pub fn capture_graph_retained<F>(
24413 &self,
24414 step: F,
24415 ) -> Result<
24416 (
24417 cudarc::driver::CudaGraph,
24418 Vec<Box<dyn std::any::Any + Send>>,
24419 ),
24420 Box<dyn std::error::Error>,
24421 >
24422 where
24423 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24424 {
24425 use cudarc::driver::sys::CUgraphInstantiate_flags;
24426 self.capture_graph_retained_flags(
24427 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24428 step,
24429 )
24430 }
24431
24432 pub fn capture_graph_retained_flags<F>(
24437 &self,
24438 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24439 mut step: F,
24440 ) -> Result<
24441 (
24442 cudarc::driver::CudaGraph,
24443 Vec<Box<dyn std::any::Any + Send>>,
24444 ),
24445 Box<dyn std::error::Error>,
24446 >
24447 where
24448 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24449 {
24450 use cudarc::driver::sys::CUstreamCaptureMode;
24451 self.capture_keep.lock().unwrap().clear();
24459 let was_tracking = self.gpu.ctx.is_event_tracking();
24460 if was_tracking {
24461 unsafe {
24462 self.gpu.ctx.disable_event_tracking();
24463 }
24464 }
24465 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24466 self.capture_keep_on
24467 .store(true, std::sync::atomic::Ordering::Relaxed);
24468 let w = (|| {
24469 step(self)?;
24470 step(self)
24471 })();
24472 self.capture_keep_on
24473 .store(false, std::sync::atomic::Ordering::Relaxed);
24474 w?;
24475 self.gpu.stream().synchronize()?;
24476 self.gpu
24477 .stream()
24478 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24479 let r = step(self);
24480 let g = self.gpu.stream().end_capture(flags);
24481 r?;
24482 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24483 graph.upload()?;
24484 Ok(graph)
24485 };
24486 let result = run();
24487 self.capture_keep_on
24488 .store(false, std::sync::atomic::Ordering::Relaxed);
24489 if was_tracking {
24490 unsafe {
24491 self.gpu.ctx.enable_event_tracking();
24492 }
24493 }
24494 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24495 Ok((result?, keeper))
24496 }
24497
24498 pub fn capture_graph_retained_nowarm<F>(
24504 &self,
24505 mut step: F,
24506 ) -> Result<
24507 (
24508 cudarc::driver::CudaGraph,
24509 Vec<Box<dyn std::any::Any + Send>>,
24510 ),
24511 Box<dyn std::error::Error>,
24512 >
24513 where
24514 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24515 {
24516 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24517 let was_tracking = self.gpu.ctx.is_event_tracking();
24518 if was_tracking {
24519 unsafe {
24520 self.gpu.ctx.disable_event_tracking();
24521 }
24522 }
24523 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24524 self.gpu.stream().synchronize()?;
24525 self.gpu
24526 .stream()
24527 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24528 let r = step(self);
24529 let g = self.gpu.stream().end_capture(
24530 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24531 );
24532 r?;
24533 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24534 graph.upload()?;
24535 Ok(graph)
24536 };
24537 let result = run();
24538 if was_tracking {
24539 unsafe {
24540 self.gpu.ctx.enable_event_tracking();
24541 }
24542 }
24543 Ok((result?, Vec::new()))
24544 }
24545
24546 pub fn capture_graph<F>(
24547 &self,
24548 mut step: F,
24549 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24550 where
24551 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24552 {
24553 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24554 let was_tracking = self.gpu.ctx.is_event_tracking();
24562 if was_tracking {
24563 unsafe {
24564 self.gpu.ctx.disable_event_tracking();
24565 }
24566 }
24567 let iflag = {
24574 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24575 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24576 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24579 Ok("priority") => {
24580 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24581 }
24582 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24583 })
24584 };
24585 let ct = {
24592 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24593 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24594 };
24595 let warmups = {
24618 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24619 *W.get_or_init(|| {
24620 std::env::var("MEMRA_GRAPH_WARMUPS")
24621 .ok()
24622 .and_then(|v| v.parse().ok())
24623 .filter(|n| *n >= 1)
24624 .unwrap_or(1)
24625 })
24626 };
24627 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24628 let t_w = std::time::Instant::now();
24629 for _ in 0..warmups {
24631 step(self)?;
24632 }
24633 self.gpu.stream().synchronize()?;
24634 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24635 let t_c = std::time::Instant::now();
24637 self.gpu
24638 .stream()
24639 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24640 let r = step(self);
24643 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24644 let t_i = std::time::Instant::now();
24645 let g = self.gpu.stream().end_capture(iflag);
24646 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24647 r?;
24648 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24649 let t_u = std::time::Instant::now();
24650 graph.upload()?;
24651 if ct {
24652 println!(
24653 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
24654 instantiate {ms_inst:.2} ms upload {:.2} ms",
24655 t_u.elapsed().as_secs_f64() * 1e3
24656 );
24657 }
24658 Ok(graph)
24659 };
24660 let result = run();
24661 if was_tracking {
24662 unsafe {
24663 self.gpu.ctx.enable_event_tracking();
24664 }
24665 }
24666 result
24667 }
24668
24669 pub fn gdn_scan_s128_view(
24671 &self,
24672 q: &CudaSlice<f32>,
24673 k: &CudaSlice<f32>,
24674 v: &CudaSlice<f32>,
24675 g: &CudaSlice<f32>,
24676 beta: &CudaSlice<f32>,
24677 state_in: &cudarc::driver::CudaView<f32>,
24678 state_out: &mut cudarc::driver::CudaViewMut<f32>,
24679 o: &mut CudaSlice<f32>,
24680 n_head: usize,
24681 t: usize,
24682 scale: f32,
24683 ) -> Result<(), Box<dyn std::error::Error>> {
24684 let f = self.func("gdn_scan_s128");
24685 const S_V: u32 = 128;
24686 const WARP: u32 = 32;
24687 const COLS: u32 = 4;
24688 let cfg = LaunchConfig {
24689 grid_dim: (n_head as u32, 1, S_V / COLS),
24690 block_dim: (WARP, COLS, 1),
24691 shared_mem_bytes: 0,
24692 };
24693 let (h, ti) = (n_head as i32, t as i32);
24694 let __s_b = self.gpu.stream();
24695 let mut b = __s_b.launch_builder(&f);
24696 b.arg(q)
24697 .arg(k)
24698 .arg(v)
24699 .arg(g)
24700 .arg(beta)
24701 .arg(state_in)
24702 .arg(state_out)
24703 .arg(o)
24704 .arg(&h)
24705 .arg(&ti)
24706 .arg(&scale);
24707 unsafe {
24708 b.launch(cfg)?;
24709 }
24710 Ok(())
24711 }
24712
24713 pub fn ssm_conv1d_view(
24715 &self,
24716 x: &cudarc::driver::CudaView<f32>,
24717 w: &CudaSlice<f32>,
24718 y: &mut CudaSlice<f32>,
24719 conv_dim: usize,
24720 t: usize,
24721 d_conv: usize,
24722 silu: bool,
24723 ) -> Result<(), Box<dyn std::error::Error>> {
24724 let f = self.func("ssm_conv1d_silu_f32");
24725 let cfg = LaunchConfig {
24727 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24728 block_dim: (256, 1, 1),
24729 shared_mem_bytes: 0,
24730 };
24731 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24732 let __s_b = self.gpu.stream();
24733 let mut b = __s_b.launch_builder(&f);
24734 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24735 unsafe {
24736 b.launch(cfg)?;
24737 }
24738 Ok(())
24739 }
24740
24741 pub fn ssm_conv1d_tm(
24748 &self,
24749 qkv_tm: &CudaSlice<f32>,
24750 w: &CudaSlice<f32>,
24751 y: &mut CudaSlice<f32>,
24752 conv_dim: usize,
24753 t: usize,
24754 d_conv: usize,
24755 ) -> Result<(), Box<dyn std::error::Error>> {
24756 let f = self.func("ssm_conv1d_tm_f32");
24757 let cfg = LaunchConfig {
24758 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24759 block_dim: (256, 1, 1),
24760 shared_mem_bytes: 0,
24761 };
24762 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24763 let __s_b = self.gpu.stream();
24764 let mut b = __s_b.launch_builder(&f);
24765 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24766 unsafe {
24767 b.launch(cfg)?;
24768 }
24769 Ok(())
24770 }
24771
24772 pub fn ssm_conv1d_tm_state(
24780 &self,
24781 qkv_tm: &CudaSlice<f32>,
24782 conv_state: &mut CudaSlice<f32>,
24783 w: &CudaSlice<f32>,
24784 y: &mut CudaSlice<f32>,
24785 conv_dim: usize,
24786 t: usize,
24787 d_conv: usize,
24788 ) -> Result<(), Box<dyn std::error::Error>> {
24789 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24790 }
24791
24792 #[allow(clippy::too_many_arguments)]
24795 pub fn ssm_conv1d_tm_state_pad(
24796 &self,
24797 qkv_tm: &CudaSlice<f32>,
24798 conv_state: &mut CudaSlice<f32>,
24799 w: &CudaSlice<f32>,
24800 y: &mut CudaSlice<f32>,
24801 conv_dim: usize,
24802 t: usize,
24803 d_conv: usize,
24804 pad_len: Option<&CudaSlice<i32>>,
24805 ) -> Result<(), Box<dyn std::error::Error>> {
24806 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24807 let ring_old = if t < d_conv - 1 {
24811 Some(self.clone_dtod(conv_state)?)
24812 } else {
24813 None
24814 };
24815 {
24816 let f = self.func("ssm_conv1d_tm_state_f32");
24817 let cfg = LaunchConfig {
24818 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24819 block_dim: (256, 1, 1),
24820 shared_mem_bytes: 0,
24821 };
24822 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24823 let __s_b = self.gpu.stream();
24824 let mut b = __s_b.launch_builder(&f);
24825 b.arg(qkv_tm)
24826 .arg(&*conv_state)
24827 .arg(w)
24828 .arg(y)
24829 .arg(&cd)
24830 .arg(&ti)
24831 .arg(&dc);
24832 unsafe {
24833 b.launch(cfg)?;
24834 }
24835 }
24836 match (ring_old, pad_len) {
24837 (None, Some(len_d)) => {
24838 let f = self.func("ssm_conv_ring_update_dev_f32");
24839 let n = conv_dim * (d_conv - 1);
24840 let cfg = LaunchConfig::for_num_elems(n as u32);
24841 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24842 let __s_b = self.gpu.stream();
24843 let mut b = __s_b.launch_builder(&f);
24844 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24845 unsafe {
24846 b.launch(cfg)?;
24847 }
24848 }
24849 (None, None) => {
24850 let f = self.func("ssm_conv_ring_update_f32");
24851 let n = conv_dim * (d_conv - 1);
24852 let cfg = LaunchConfig::for_num_elems(n as u32);
24853 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24854 let __s_b = self.gpu.stream();
24855 let mut b = __s_b.launch_builder(&f);
24856 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24857 unsafe {
24858 b.launch(cfg)?;
24859 }
24860 }
24861 (Some(old), _) => {
24862 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
24863 }
24864 }
24865 Ok(())
24866 }
24867
24868 pub fn ssm_conv1d_tm_state_pad_v(
24870 &self,
24871 qkv_tm: &cudarc::driver::CudaView<f32>,
24872 conv_state: &mut CudaSlice<f32>,
24873 w: &CudaSlice<f32>,
24874 y: &mut CudaSlice<f32>,
24875 conv_dim: usize,
24876 t: usize,
24877 d_conv: usize,
24878 pad_len: Option<&CudaSlice<i32>>,
24879 ) -> Result<(), Box<dyn std::error::Error>> {
24880 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24881 let ring_old = if t < d_conv - 1 {
24885 Some(self.clone_dtod(conv_state)?)
24886 } else {
24887 None
24888 };
24889 {
24890 let f = self.func("ssm_conv1d_tm_state_f32");
24891 let cfg = LaunchConfig {
24892 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24893 block_dim: (256, 1, 1),
24894 shared_mem_bytes: 0,
24895 };
24896 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24897 let __s_b = self.gpu.stream();
24898 let mut b = __s_b.launch_builder(&f);
24899 b.arg(qkv_tm)
24900 .arg(&*conv_state)
24901 .arg(w)
24902 .arg(y)
24903 .arg(&cd)
24904 .arg(&ti)
24905 .arg(&dc);
24906 unsafe {
24907 b.launch(cfg)?;
24908 }
24909 }
24910 match (ring_old, pad_len) {
24911 (None, Some(len_d)) => {
24912 let f = self.func("ssm_conv_ring_update_dev_f32");
24913 let n = conv_dim * (d_conv - 1);
24914 let cfg = LaunchConfig::for_num_elems(n as u32);
24915 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24916 let __s_b = self.gpu.stream();
24917 let mut b = __s_b.launch_builder(&f);
24918 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24919 unsafe {
24920 b.launch(cfg)?;
24921 }
24922 }
24923 (None, None) => {
24924 let f = self.func("ssm_conv_ring_update_f32");
24925 let n = conv_dim * (d_conv - 1);
24926 let cfg = LaunchConfig::for_num_elems(n as u32);
24927 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24928 let __s_b = self.gpu.stream();
24929 let mut b = __s_b.launch_builder(&f);
24930 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24931 unsafe {
24932 b.launch(cfg)?;
24933 }
24934 }
24935 (Some(_), _) => unreachable!(
24936 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
24937 ),
24938 }
24939 Ok(())
24940 }
24941
24942 pub fn ssm_conv_ring_rebuild(
24947 &self,
24948 qkv_tm: &CudaSlice<f32>,
24949 ring_old: &CudaSlice<f32>,
24950 conv_state: &mut CudaSlice<f32>,
24951 conv_dim: usize,
24952 tc: usize,
24953 d_conv: usize,
24954 ) -> Result<(), Box<dyn std::error::Error>> {
24955 let f = self.func("ssm_conv_ring_rebuild_f32");
24956 let n = conv_dim * (d_conv - 1);
24957 let cfg = LaunchConfig::for_num_elems(n as u32);
24958 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
24959 let __s_b = self.gpu.stream();
24960 let mut b = __s_b.launch_builder(&f);
24961 b.arg(qkv_tm)
24962 .arg(ring_old)
24963 .arg(conv_state)
24964 .arg(&cd)
24965 .arg(&ti)
24966 .arg(&dc);
24967 unsafe {
24968 b.launch(cfg)?;
24969 }
24970 Ok(())
24971 }
24972
24973 #[allow(clippy::too_many_arguments)]
24978 pub fn gdn_prep_decode(
24979 &self,
24980 conv_out: &CudaSlice<f32>,
24981 beta_raw: &CudaSlice<f32>,
24982 alpha: &CudaSlice<f32>,
24983 dt_bias: &CudaSlice<f32>,
24984 a: &CudaSlice<f32>,
24985 q_l2: &mut CudaSlice<f32>,
24986 k_l2: &mut CudaSlice<f32>,
24987 v_g: &mut CudaSlice<f32>,
24988 beta: &mut CudaSlice<f32>,
24989 g_log: &mut CudaSlice<f32>,
24990 d_state: usize,
24991 num_v: usize,
24992 num_k: usize,
24993 key_dim: usize,
24994 eps: f32,
24995 ) -> Result<(), Box<dyn std::error::Error>> {
24996 let f = self.func("gdn_prep_decode_f32");
24997 let cfg = LaunchConfig {
24998 grid_dim: (num_v as u32, 1, 1),
24999 block_dim: (32, 4, 1),
25000 shared_mem_bytes: 0,
25001 };
25002 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25003 let __s_b = self.gpu.stream();
25004 let mut b = __s_b.launch_builder(&f);
25005 b.arg(conv_out)
25006 .arg(beta_raw)
25007 .arg(alpha)
25008 .arg(dt_bias)
25009 .arg(a)
25010 .arg(q_l2)
25011 .arg(k_l2)
25012 .arg(v_g)
25013 .arg(beta)
25014 .arg(g_log)
25015 .arg(&ds)
25016 .arg(&nv)
25017 .arg(&nk)
25018 .arg(&kd)
25019 .arg(&eps);
25020 unsafe {
25021 b.launch(cfg)?;
25022 }
25023 Ok(())
25024 }
25025
25026 #[allow(clippy::too_many_arguments)]
25030 pub fn ssm_conv1d_gdn(
25031 &self,
25032 qkv_tm: &CudaSlice<f32>,
25033 w: &CudaSlice<f32>,
25034 q_g: &mut CudaSlice<f32>,
25035 k_g: &mut CudaSlice<f32>,
25036 v_g: &mut CudaSlice<f32>,
25037 conv_dim: usize,
25038 t: usize,
25039 d_conv: usize,
25040 d_state: usize,
25041 num_v: usize,
25042 num_k: usize,
25043 key_dim: usize,
25044 ) -> Result<(), Box<dyn std::error::Error>> {
25045 let f = self.func("ssm_conv1d_gdn_f32");
25046 let cfg = LaunchConfig {
25047 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25048 block_dim: (256, 1, 1),
25049 shared_mem_bytes: 0,
25050 };
25051 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25052 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25053 let __s_b = self.gpu.stream();
25054 let mut b = __s_b.launch_builder(&f);
25055 b.arg(qkv_tm)
25056 .arg(w)
25057 .arg(q_g)
25058 .arg(k_g)
25059 .arg(v_g)
25060 .arg(&cd)
25061 .arg(&ti)
25062 .arg(&dc)
25063 .arg(&ds)
25064 .arg(&nv)
25065 .arg(&nk)
25066 .arg(&kd);
25067 unsafe {
25068 b.launch(cfg)?;
25069 }
25070 Ok(())
25071 }
25072
25073 pub fn ssm_conv1d(
25074 &self,
25075 x: &CudaSlice<f32>,
25076 w: &CudaSlice<f32>,
25077 y: &mut CudaSlice<f32>,
25078 conv_dim: usize,
25079 t: usize,
25080 d_conv: usize,
25081 silu: bool,
25082 ) -> Result<(), Box<dyn std::error::Error>> {
25083 let f = self.func("ssm_conv1d_silu_f32");
25084 let cfg = LaunchConfig {
25085 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25086 block_dim: (256, 1, 1),
25087 shared_mem_bytes: 0,
25088 };
25089 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25090 let __s_b = self.gpu.stream();
25091 let mut b = __s_b.launch_builder(&f);
25092 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25093 unsafe {
25094 b.launch(cfg)?;
25095 }
25096 Ok(())
25097 }
25098
25099 pub fn gdn_scan_s128(
25102 &self,
25103 q: &CudaSlice<f32>,
25104 k: &CudaSlice<f32>,
25105 v: &CudaSlice<f32>,
25106 g: &CudaSlice<f32>,
25107 beta: &CudaSlice<f32>,
25108 state_in: &CudaSlice<f32>,
25109 state_out: &mut CudaSlice<f32>,
25110 o: &mut CudaSlice<f32>,
25111 n_head: usize,
25112 t: usize,
25113 scale: f32,
25114 ) -> Result<(), Box<dyn std::error::Error>> {
25115 let f = self.func("gdn_scan_s128");
25116 const S_V: u32 = 128;
25117 const WARP: u32 = 32;
25118 const COLS_PER_BLOCK: u32 = 4;
25119 let cfg = LaunchConfig {
25120 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25121 block_dim: (WARP, COLS_PER_BLOCK, 1),
25122 shared_mem_bytes: 0,
25123 };
25124 let (h, ti) = (n_head as i32, t as i32);
25125 let __s_b = self.gpu.stream();
25126 let mut b = __s_b.launch_builder(&f);
25127 b.arg(q)
25128 .arg(k)
25129 .arg(v)
25130 .arg(g)
25131 .arg(beta)
25132 .arg(state_in)
25133 .arg(state_out)
25134 .arg(o)
25135 .arg(&h)
25136 .arg(&ti)
25137 .arg(&scale);
25138 unsafe {
25139 b.launch(cfg)?;
25140 }
25141 Ok(())
25142 }
25143
25144 #[allow(clippy::too_many_arguments)]
25149 pub fn ssm_conv1d_fused_decode_b(
25150 &self,
25151 qkv_cols: &CudaSlice<f32>,
25152 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25153 w: &CudaSlice<f32>,
25154 conv_outs: &mut CudaSlice<f32>,
25155 conv_dim: usize,
25156 d_conv: usize,
25157 b_n: usize,
25158 ) -> Result<(), Box<dyn std::error::Error>> {
25159 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25160 let cfg = LaunchConfig {
25161 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25162 block_dim: (256, 1, 1),
25163 shared_mem_bytes: 0,
25164 };
25165 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25166 let __s_b = self.gpu.stream();
25167 let mut b = __s_b.launch_builder(&f);
25168 b.arg(qkv_cols)
25169 .arg(conv_state_ptrs)
25170 .arg(w)
25171 .arg(conv_outs)
25172 .arg(&cd)
25173 .arg(&dc);
25174 unsafe {
25175 b.launch(cfg)?;
25176 }
25177 Ok(())
25178 }
25179
25180 #[allow(clippy::too_many_arguments)]
25181 pub fn gdn_prep_decode_b(
25182 &self,
25183 conv_outs: &CudaSlice<f32>,
25184 beta_raws: &CudaSlice<f32>,
25185 alphas: &CudaSlice<f32>,
25186 dt_bias: &CudaSlice<f32>,
25187 a: &CudaSlice<f32>,
25188 q_l2: &mut CudaSlice<f32>,
25189 k_l2: &mut CudaSlice<f32>,
25190 v_g: &mut CudaSlice<f32>,
25191 beta: &mut CudaSlice<f32>,
25192 g_log: &mut CudaSlice<f32>,
25193 d_state: usize,
25194 num_v: usize,
25195 num_k: usize,
25196 key_dim: usize,
25197 eps: f32,
25198 conv_dim: usize,
25199 b_n: usize,
25200 ) -> Result<(), Box<dyn std::error::Error>> {
25201 let f = self.func("gdn_prep_decode_b_f32");
25202 let cfg = LaunchConfig {
25203 grid_dim: (num_v as u32, 1, b_n as u32),
25204 block_dim: (32, 4, 1),
25205 shared_mem_bytes: 0,
25206 };
25207 let (ds, nv, nk, kd, cd) = (
25208 d_state as i32,
25209 num_v as i32,
25210 num_k as i32,
25211 key_dim as i32,
25212 conv_dim as i32,
25213 );
25214 let __s_b = self.gpu.stream();
25215 let mut b = __s_b.launch_builder(&f);
25216 b.arg(conv_outs)
25217 .arg(beta_raws)
25218 .arg(alphas)
25219 .arg(dt_bias)
25220 .arg(a)
25221 .arg(q_l2)
25222 .arg(k_l2)
25223 .arg(v_g)
25224 .arg(beta)
25225 .arg(g_log)
25226 .arg(&ds)
25227 .arg(&nv)
25228 .arg(&nk)
25229 .arg(&kd)
25230 .arg(&eps)
25231 .arg(&cd);
25232 unsafe {
25233 b.launch(cfg)?;
25234 }
25235 Ok(())
25236 }
25237
25238 #[allow(clippy::too_many_arguments)]
25239 pub fn gdn_scan_s128_batched(
25240 &self,
25241 q: &CudaSlice<f32>,
25242 k: &CudaSlice<f32>,
25243 v: &CudaSlice<f32>,
25244 g: &CudaSlice<f32>,
25245 beta: &CudaSlice<f32>,
25246 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25247 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25248 o: &mut CudaSlice<f32>,
25249 n_head: usize,
25250 b_n: usize,
25251 scale: f32,
25252 ) -> Result<(), Box<dyn std::error::Error>> {
25253 let f = self.func("gdn_scan_s128_b");
25254 const S_V: u32 = 128;
25255 const WARP: u32 = 32;
25256 const COLS_PER_BLOCK: u32 = 4;
25257 let cfg = LaunchConfig {
25258 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25259 block_dim: (WARP, COLS_PER_BLOCK, 1),
25260 shared_mem_bytes: 0,
25261 };
25262 let h = n_head as i32;
25263 let __s_b = self.gpu.stream();
25264 let mut b = __s_b.launch_builder(&f);
25265 b.arg(q)
25266 .arg(k)
25267 .arg(v)
25268 .arg(g)
25269 .arg(beta)
25270 .arg(state_in_ptrs)
25271 .arg(state_out_ptrs)
25272 .arg(o)
25273 .arg(&h)
25274 .arg(&scale);
25275 unsafe {
25276 b.launch(cfg)?;
25277 }
25278 Ok(())
25279 }
25280
25281 #[allow(clippy::too_many_arguments)]
25287 pub fn ssm_conv1d_fused_decode_b_view(
25288 &self,
25289 qkv_cols: &cudarc::driver::CudaView<f32>,
25290 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25291 w: &CudaSlice<f32>,
25292 conv_outs: &mut CudaSlice<f32>,
25293 conv_dim: usize,
25294 d_conv: usize,
25295 b_n: usize,
25296 ) -> Result<(), Box<dyn std::error::Error>> {
25297 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25298 let cfg = LaunchConfig {
25299 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25300 block_dim: (256, 1, 1),
25301 shared_mem_bytes: 0,
25302 };
25303 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25304 let __s_b = self.gpu.stream();
25305 let mut b = __s_b.launch_builder(&f);
25306 b.arg(qkv_cols)
25307 .arg(conv_state_ptrs)
25308 .arg(w)
25309 .arg(conv_outs)
25310 .arg(&cd)
25311 .arg(&dc);
25312 unsafe {
25313 b.launch(cfg)?;
25314 }
25315 Ok(())
25316 }
25317
25318 #[allow(clippy::too_many_arguments)]
25319 pub fn gdn_prep_decode_b_view(
25320 &self,
25321 conv_outs: &CudaSlice<f32>,
25322 beta_raws: &cudarc::driver::CudaView<f32>,
25323 alphas: &cudarc::driver::CudaView<f32>,
25324 dt_bias: &CudaSlice<f32>,
25325 a: &CudaSlice<f32>,
25326 q_l2: &mut CudaSlice<f32>,
25327 k_l2: &mut CudaSlice<f32>,
25328 v_g: &mut CudaSlice<f32>,
25329 beta: &mut CudaSlice<f32>,
25330 g_log: &mut CudaSlice<f32>,
25331 d_state: usize,
25332 num_v: usize,
25333 num_k: usize,
25334 key_dim: usize,
25335 eps: f32,
25336 conv_dim: usize,
25337 b_n: usize,
25338 ) -> Result<(), Box<dyn std::error::Error>> {
25339 let f = self.func("gdn_prep_decode_b_f32");
25340 let cfg = LaunchConfig {
25341 grid_dim: (num_v as u32, 1, b_n as u32),
25342 block_dim: (32, 4, 1),
25343 shared_mem_bytes: 0,
25344 };
25345 let (ds, nv, nk, kd, cd) = (
25346 d_state as i32,
25347 num_v as i32,
25348 num_k as i32,
25349 key_dim as i32,
25350 conv_dim as i32,
25351 );
25352 let __s_b = self.gpu.stream();
25353 let mut b = __s_b.launch_builder(&f);
25354 b.arg(conv_outs)
25355 .arg(beta_raws)
25356 .arg(alphas)
25357 .arg(dt_bias)
25358 .arg(a)
25359 .arg(q_l2)
25360 .arg(k_l2)
25361 .arg(v_g)
25362 .arg(beta)
25363 .arg(g_log)
25364 .arg(&ds)
25365 .arg(&nv)
25366 .arg(&nk)
25367 .arg(&kd)
25368 .arg(&eps)
25369 .arg(&cd);
25370 unsafe {
25371 b.launch(cfg)?;
25372 }
25373 Ok(())
25374 }
25375
25376 #[allow(clippy::too_many_arguments)]
25377 pub fn gdn_scan_s128_batched_view(
25378 &self,
25379 q: &CudaSlice<f32>,
25380 k: &CudaSlice<f32>,
25381 v: &CudaSlice<f32>,
25382 g: &CudaSlice<f32>,
25383 beta: &CudaSlice<f32>,
25384 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25385 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25386 o: &mut cudarc::driver::CudaViewMut<f32>,
25387 n_head: usize,
25388 b_n: usize,
25389 scale: f32,
25390 ) -> Result<(), Box<dyn std::error::Error>> {
25391 let f = self.func("gdn_scan_s128_b");
25392 const S_V: u32 = 128;
25393 const WARP: u32 = 32;
25394 const COLS_PER_BLOCK: u32 = 4;
25395 let cfg = LaunchConfig {
25396 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25397 block_dim: (WARP, COLS_PER_BLOCK, 1),
25398 shared_mem_bytes: 0,
25399 };
25400 let h = n_head as i32;
25401 let __s_b = self.gpu.stream();
25402 let mut b = __s_b.launch_builder(&f);
25403 b.arg(q)
25404 .arg(k)
25405 .arg(v)
25406 .arg(g)
25407 .arg(beta)
25408 .arg(state_in_ptrs)
25409 .arg(state_out_ptrs)
25410 .arg(o)
25411 .arg(&h)
25412 .arg(&scale);
25413 unsafe {
25414 b.launch(cfg)?;
25415 }
25416 Ok(())
25417 }
25418
25419 pub fn gdn_chunked_enabled() -> bool {
25428 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25429 *E.get_or_init(|| {
25430 std::env::var("MEMRA_GDN_CHUNKED")
25431 .map(|v| v != "0")
25432 .unwrap_or(true)
25433 })
25434 }
25435
25436 pub fn gdn_chunk_size() -> usize {
25441 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25442 *C.get_or_init(|| {
25443 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25444 .ok()
25445 .and_then(|v| v.parse().ok())
25446 .unwrap_or(32);
25447 c.clamp(32, 128) / 32 * 32
25448 })
25449 }
25450
25451 #[allow(clippy::too_many_arguments)]
25456 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25459 #[allow(clippy::too_many_arguments)]
25460 pub fn gdn_chunk_k123(
25461 &self,
25462 q: &CudaSlice<f32>,
25463 k: &CudaSlice<f32>,
25464 v: &CudaSlice<f32>,
25465 g: &CudaSlice<f32>,
25466 beta: &CudaSlice<f32>,
25467 wb16: Option<&mut CudaSlice<u8>>,
25468 n_head: usize,
25469 t: usize,
25470 c: usize,
25471 hk: usize,
25472 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25473 ) -> Result<
25474 (
25475 CudaSlice<f32>,
25476 CudaSlice<f32>,
25477 CudaSlice<f32>,
25478 CudaSlice<f32>,
25479 ),
25480 Box<dyn std::error::Error>,
25481 > {
25482 const D: usize = 128;
25483 let h = n_head;
25484 let nc = (t + c - 1) / c;
25485 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25486 let mut gcum = self.uninit(t * h)?;
25487 let mut a = self.uninit(nc * h * c * c)?;
25488 let mut p = self.uninit(nc * h * c * c)?;
25489 let mut u = self.uninit(nc * h * c * D)?;
25490 let mut w = self.uninit(nc * h * c * D)?;
25491 {
25492 let f = self.func("gdn_chunk_cumgate_f32");
25494 let cfg = LaunchConfig {
25495 grid_dim: (nc as u32, h as u32, 1),
25496 block_dim: (32, 1, 1),
25497 shared_mem_bytes: 0,
25498 };
25499 let __s_b = self.gpu.stream();
25500 let mut b = __s_b.launch_builder(&f);
25501 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25502 unsafe {
25503 b.launch(cfg)?;
25504 }
25505 }
25506 if let Some((qb, kb, pb)) = k2w {
25507 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25510 let f = self.func("gdn_k2_wgmma");
25511 let cfg = LaunchConfig {
25512 grid_dim: (nc as u32, h as u32, 1),
25513 block_dim: (128, 1, 1),
25514 shared_mem_bytes: 0,
25515 };
25516 let hki = hk as i32;
25517 let __s_b = self.gpu.stream();
25518 let mut b = __s_b.launch_builder(&f);
25519 b.arg(qb)
25520 .arg(kb)
25521 .arg(&gcum)
25522 .arg(beta)
25523 .arg(&mut a)
25524 .arg(&mut *pb)
25525 .arg(&hi)
25526 .arg(&ti)
25527 .arg(&ci)
25528 .arg(&hki);
25529 unsafe {
25530 b.launch(cfg)?;
25531 }
25532 } else if c <= 64 && !portable_mma_gated() {
25533 let f = self.func("gdn_chunk_attn_f32");
25535 f.set_attribute(
25536 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25537 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25538 )?;
25539 let jt = ((c + 31) / 32) as u32;
25540 let cfg = LaunchConfig {
25541 grid_dim: (nc as u32, h as u32, jt),
25542 block_dim: (256, 1, 1),
25543 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25544 };
25545 let hki = hk as i32;
25546 let __s_b = self.gpu.stream();
25547 let mut b = __s_b.launch_builder(&f);
25548 b.arg(q)
25549 .arg(k)
25550 .arg(&gcum)
25551 .arg(beta)
25552 .arg(&mut a)
25553 .arg(&mut p)
25554 .arg(&hi)
25555 .arg(&ti)
25556 .arg(&ci)
25557 .arg(&hki);
25558 unsafe {
25559 b.launch(cfg)?;
25560 }
25561 } else {
25562 assert!(
25564 hk == h,
25565 "generic K2 is broadcast-only (de-broadcast rides C==32)"
25566 );
25567 let f = self.func("gdn_chunk_attn_g_f32");
25568 let cfg = LaunchConfig {
25569 grid_dim: (nc as u32, h as u32, 1),
25570 block_dim: (32, 8, 1),
25571 shared_mem_bytes: 0,
25572 };
25573 let __s_b = self.gpu.stream();
25574 let mut b = __s_b.launch_builder(&f);
25575 b.arg(q)
25576 .arg(k)
25577 .arg(&gcum)
25578 .arg(beta)
25579 .arg(&mut a)
25580 .arg(&mut p)
25581 .arg(&hi)
25582 .arg(&ti)
25583 .arg(&ci);
25584 unsafe {
25585 b.launch(cfg)?;
25586 }
25587 }
25588 {
25589 let cfg = LaunchConfig {
25591 grid_dim: (nc as u32, h as u32, 1),
25592 block_dim: (256, 1, 1),
25593 shared_mem_bytes: 0,
25594 };
25595 match c {
25596 32 | 64 => {
25597 let f = self.func(if c == 32 {
25598 "gdn_chunk_solve32_f32"
25599 } else {
25600 "gdn_chunk_solve64_f32"
25601 });
25602 let wb: u64 = match wb16 {
25604 Some(d) => self.addr_u8(d),
25605 None => 0,
25606 };
25607 let hki = hk as i32;
25608 let __s_b = self.gpu.stream();
25609 let mut b = __s_b.launch_builder(&f);
25610 b.arg(v)
25611 .arg(k)
25612 .arg(&a)
25613 .arg(&gcum)
25614 .arg(&mut u)
25615 .arg(&mut w)
25616 .arg(&wb)
25617 .arg(&hi)
25618 .arg(&ti)
25619 .arg(&hki);
25620 unsafe {
25621 b.launch(cfg)?;
25622 }
25623 }
25624 _ => {
25625 assert!(hk == h, "generic K3 is broadcast-only");
25626 let f = self.func("gdn_chunk_solve_f32");
25627 let __s_b = self.gpu.stream();
25628 let mut b = __s_b.launch_builder(&f);
25629 b.arg(v)
25630 .arg(k)
25631 .arg(&a)
25632 .arg(&gcum)
25633 .arg(&mut u)
25634 .arg(&mut w)
25635 .arg(&hi)
25636 .arg(&ti)
25637 .arg(&ci);
25638 unsafe {
25639 b.launch(cfg)?;
25640 }
25641 }
25642 }
25643 }
25644 Ok((gcum, p, u, w))
25645 }
25646
25647 pub fn gdn_db_on() -> bool {
25651 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25652 }
25653
25654 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25663 !portable_mma_gated()
25664 && c == 32
25665 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25666 Ok("1") => true,
25667 Ok("0") => false,
25668 _ => gdn_mma_default_on(),
25669 }
25670 }
25671
25672 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25679 cfg!(memra_hopper_mma)
25680 && self.gdn_mma_enabled(c)
25681 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25682 }
25683
25684 #[allow(clippy::too_many_arguments)]
25689 pub fn ssm_conv1d_gdn_state_pad(
25690 &self,
25691 qkv_tm: &cudarc::driver::CudaView<f32>,
25692 conv_state: &mut CudaSlice<f32>,
25693 w: &CudaSlice<f32>,
25694 q_g: &mut CudaSlice<f32>,
25695 k_g: &mut CudaSlice<f32>,
25696 v_g: &mut CudaSlice<f32>,
25697 conv_dim: usize,
25698 t: usize,
25699 d_conv: usize,
25700 d_state: usize,
25701 num_v: usize,
25702 num_k: usize,
25703 key_dim: usize,
25704 hk: usize,
25705 pad_len: Option<&CudaSlice<i32>>,
25706 ) -> Result<(), Box<dyn std::error::Error>> {
25707 assert!(
25708 t >= d_conv - 1,
25709 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25710 );
25711 {
25712 let f = self.func("ssm_conv1d_gdn_state_f32");
25713 let cfg = LaunchConfig {
25714 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25715 block_dim: (256, 1, 1),
25716 shared_mem_bytes: 0,
25717 };
25718 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25719 let (ds, nv, nk, kd, hki) = (
25720 d_state as i32,
25721 num_v as i32,
25722 num_k as i32,
25723 key_dim as i32,
25724 hk as i32,
25725 );
25726 let __s_b = self.gpu.stream();
25727 let mut b = __s_b.launch_builder(&f);
25728 b.arg(qkv_tm)
25729 .arg(&*conv_state)
25730 .arg(w)
25731 .arg(q_g)
25732 .arg(k_g)
25733 .arg(v_g)
25734 .arg(&cd)
25735 .arg(&ti)
25736 .arg(&dc)
25737 .arg(&ds)
25738 .arg(&nv)
25739 .arg(&nk)
25740 .arg(&kd)
25741 .arg(&hki);
25742 unsafe {
25743 b.launch(cfg)?;
25744 }
25745 }
25746 match pad_len {
25747 Some(len_d) => {
25748 let f = self.func("ssm_conv_ring_update_dev_f32");
25749 let n = conv_dim * (d_conv - 1);
25750 let cfg = LaunchConfig::for_num_elems(n as u32);
25751 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25752 let __s_b = self.gpu.stream();
25753 let mut b = __s_b.launch_builder(&f);
25754 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25755 unsafe {
25756 b.launch(cfg)?;
25757 }
25758 }
25759 None => {
25760 let f = self.func("ssm_conv_ring_update_f32");
25761 let n = conv_dim * (d_conv - 1);
25762 let cfg = LaunchConfig::for_num_elems(n as u32);
25763 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25764 let __s_b = self.gpu.stream();
25765 let mut b = __s_b.launch_builder(&f);
25766 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25767 unsafe {
25768 b.launch(cfg)?;
25769 }
25770 }
25771 }
25772 Ok(())
25773 }
25774
25775 pub fn gdn_chunk_alloc(
25779 &self,
25780 n_head: usize,
25781 t: usize,
25782 c: usize,
25783 hk: usize,
25784 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25785 const D: usize = 128;
25786 assert!(
25787 c == 32,
25788 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25789 );
25790 let h = n_head;
25791 let nc = (t + c - 1) / c;
25792 Ok(GdnChunkBufs {
25793 gcum: self.uninit(t * h)?,
25794 a: self.uninit(nc * h * c * c)?,
25795 p: self.uninit(nc * h * c * c)?,
25796 u: self.uninit(nc * h * c * D)?,
25797 w: self.uninit(nc * h * c * D)?,
25798 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25799 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25800 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25801 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25802 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25803 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25804 o: self.uninit(D * h * t)?,
25805 t,
25806 nc,
25807 })
25808 }
25809
25810 pub fn f32_to_bf16_v(
25812 &self,
25813 x: &cudarc::driver::CudaView<f32>,
25814 dst: &mut CudaSlice<u8>,
25815 n: usize,
25816 ) -> Result<(), Box<dyn std::error::Error>> {
25817 let f = self.func("f32_to_bf16_bulk");
25818 let ni = n as i64;
25819 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25820 let __s_b = self.gpu.stream();
25821 let mut b = __s_b.launch_builder(&f);
25822 b.arg(x).arg(dst).arg(&ni);
25823 unsafe {
25824 b.launch(cfg)?;
25825 }
25826 Ok(())
25827 }
25828
25829 pub fn f32_to_bf16_into(
25831 &self,
25832 x: &CudaSlice<f32>,
25833 dst: &mut CudaSlice<u8>,
25834 n: usize,
25835 ) -> Result<(), Box<dyn std::error::Error>> {
25836 let f = self.func("f32_to_bf16_bulk");
25837 let ni = n as i64;
25838 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25839 let __s_b = self.gpu.stream();
25840 let mut b = __s_b.launch_builder(&f);
25841 b.arg(x).arg(dst).arg(&ni);
25842 unsafe {
25843 b.launch(cfg)?;
25844 }
25845 Ok(())
25846 }
25847
25848 pub fn gdn_chunk_k123_vl8(
25851 &self,
25852 seqs: &[GdnSeqVl],
25853 n_head: usize,
25854 hk: usize,
25855 wq: Option<&GdnWVl8>,
25856 ) -> Result<(), Box<dyn std::error::Error>> {
25857 let b = seqs.len();
25858 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
25859 let mut packed = [GdnSeqVl::default(); 8];
25860 packed[..b].copy_from_slice(seqs);
25861 let v = GdnVl8(packed);
25862 let (hi, ci) = (n_head as i32, 32i32);
25863 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25864 {
25865 let f = self.func("gdn_chunk_cumgate_vl");
25866 let cfg = LaunchConfig {
25867 grid_dim: (max_nc, n_head as u32, b as u32),
25868 block_dim: (32, 1, 1),
25869 shared_mem_bytes: 0,
25870 };
25871 let __s_lb = self.gpu.stream();
25872 let mut lb = __s_lb.launch_builder(&f);
25873 lb.arg(&v).arg(&hi).arg(&ci);
25874 unsafe {
25875 lb.launch(cfg)?;
25876 }
25877 }
25878 let hki = hk as i32;
25879 if let Some(w) = wq {
25880 let f = self.func("gdn_k2_wgmma_vl");
25882 let cfg = LaunchConfig {
25883 grid_dim: (max_nc, n_head as u32, b as u32),
25884 block_dim: (128, 1, 1),
25885 shared_mem_bytes: 0,
25886 };
25887 let __s_lb = self.gpu.stream();
25888 let mut lb = __s_lb.launch_builder(&f);
25889 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
25890 unsafe {
25891 lb.launch(cfg)?;
25892 }
25893 } else {
25894 let f = self.func("gdn_chunk_attn_vl");
25895 f.set_attribute(
25896 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25897 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25898 )?;
25899 let cfg = LaunchConfig {
25900 grid_dim: (max_nc, n_head as u32, b as u32),
25901 block_dim: (256, 1, 1),
25902 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25903 };
25904 let __s_lb = self.gpu.stream();
25905 let mut lb = __s_lb.launch_builder(&f);
25906 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25907 unsafe {
25908 lb.launch(cfg)?;
25909 }
25910 }
25911 {
25912 let f = self.func("gdn_chunk_solve32_vl");
25913 let cfg = LaunchConfig {
25914 grid_dim: (max_nc, n_head as u32, b as u32),
25915 block_dim: (256, 1, 1),
25916 shared_mem_bytes: 0,
25917 };
25918 let __s_lb = self.gpu.stream();
25919 let mut lb = __s_lb.launch_builder(&f);
25920 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25921 unsafe {
25922 lb.launch(cfg)?;
25923 }
25924 }
25925 Ok(())
25926 }
25927
25928 #[allow(clippy::too_many_arguments)]
25932 pub fn gdn_prep_vl8(
25933 &self,
25934 seqs: &[GdnPrepVl],
25935 conv_w: &CudaSlice<f32>,
25936 dt_bias: &CudaSlice<f32>,
25937 a: &CudaSlice<f32>,
25938 conv_dim: usize,
25939 d_conv: usize,
25940 d_state: usize,
25941 num_v: usize,
25942 num_k: usize,
25943 key_dim: usize,
25944 hk: usize,
25945 eps: f32,
25946 ) -> Result<(), Box<dyn std::error::Error>> {
25947 let b = seqs.len();
25948 assert!(b >= 1 && b <= 8);
25949 let mut packed = [GdnPrepVl::default(); 8];
25950 packed[..b].copy_from_slice(seqs);
25951 let v = GdnPrepVl8(packed);
25952 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25953 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
25954 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
25955 assert!(
25956 conv_fuse || hk == num_v,
25957 "de-broadcast requires the fused conv"
25958 );
25959 if conv_fuse {
25960 let f = self.func("ssm_conv1d_gdn_state_vl");
25961 let cfg = LaunchConfig {
25962 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25963 block_dim: (256, 1, 1),
25964 shared_mem_bytes: 0,
25965 };
25966 let (dsi, nvi, nki, kdi, hki) = (
25967 d_state as i32,
25968 num_v as i32,
25969 num_k as i32,
25970 key_dim as i32,
25971 hk as i32,
25972 );
25973 let __s_lb = self.gpu.stream();
25974 let mut lb = __s_lb.launch_builder(&f);
25975 lb.arg(&v)
25976 .arg(conv_w)
25977 .arg(&cdi)
25978 .arg(&dci)
25979 .arg(&dsi)
25980 .arg(&nvi)
25981 .arg(&nki)
25982 .arg(&kdi)
25983 .arg(&hki);
25984 unsafe {
25985 lb.launch(cfg)?;
25986 }
25987 } else {
25988 let f = self.func("ssm_conv1d_tm_state_vl");
25989 let cfg = LaunchConfig {
25990 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25991 block_dim: (256, 1, 1),
25992 shared_mem_bytes: 0,
25993 };
25994 let __s_lb = self.gpu.stream();
25995 let mut lb = __s_lb.launch_builder(&f);
25996 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
25997 unsafe {
25998 lb.launch(cfg)?;
25999 }
26000 }
26001 {
26002 let f = self.func("ssm_conv_ring_update_vl");
26003 let n = (conv_dim * (d_conv - 1)) as u32;
26004 let cfg = LaunchConfig {
26005 grid_dim: (n.div_ceil(256), 1, b as u32),
26006 block_dim: (256, 1, 1),
26007 shared_mem_bytes: 0,
26008 };
26009 let __s_lb = self.gpu.stream();
26010 let mut lb = __s_lb.launch_builder(&f);
26011 lb.arg(&v).arg(&cdi).arg(&dci);
26012 unsafe {
26013 lb.launch(cfg)?;
26014 }
26015 }
26016 if !conv_fuse {
26017 let f = self.func("qkv_to_gdn_repack_vl");
26018 let n = max_t * (num_v * d_state) as u32;
26019 let cfg = LaunchConfig {
26020 grid_dim: (n.div_ceil(256), 1, b as u32),
26021 block_dim: (256, 1, 1),
26022 shared_mem_bytes: 0,
26023 };
26024 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26025 let __s_lb = self.gpu.stream();
26026 let mut lb = __s_lb.launch_builder(&f);
26027 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26028 unsafe {
26029 lb.launch(cfg)?;
26030 }
26031 }
26032 if Self::l2_v2_on(d_state) {
26033 let f = self.func("gdn_l2_v2_vl");
26034 let cfg = LaunchConfig {
26035 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26036 block_dim: (256, 1, 1),
26037 shared_mem_bytes: 0,
26038 };
26039 let (dsi, nvi) = (d_state as i32, hk as i32);
26040 let __s_lb = self.gpu.stream();
26041 let mut lb = __s_lb.launch_builder(&f);
26042 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26043 unsafe {
26044 lb.launch(cfg)?;
26045 }
26046 } else {
26047 let f = self.func("gdn_l2_vl");
26048 let cfg = LaunchConfig {
26049 grid_dim: (max_t * hk as u32, 2, b as u32),
26050 block_dim: (256, 1, 1),
26051 shared_mem_bytes: 0,
26052 };
26053 let (dsi, nvi) = (d_state as i32, hk as i32);
26054 let __s_lb = self.gpu.stream();
26055 let mut lb = __s_lb.launch_builder(&f);
26056 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26057 unsafe {
26058 lb.launch(cfg)?;
26059 }
26060 }
26061 {
26062 let f = self.func("gdn_gate_prep_vl");
26063 let n = max_t * num_v as u32;
26064 let cfg = LaunchConfig {
26065 grid_dim: (n.div_ceil(256), 1, b as u32),
26066 block_dim: (256, 1, 1),
26067 shared_mem_bytes: 0,
26068 };
26069 let nvi = num_v as i32;
26070 let __s_lb = self.gpu.stream();
26071 let mut lb = __s_lb.launch_builder(&f);
26072 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26073 unsafe {
26074 lb.launch(cfg)?;
26075 }
26076 }
26077 Ok(())
26078 }
26079
26080 pub fn gdn_mirror_vl8(
26082 &self,
26083 seqs: &[GdnSeqVl],
26084 n_head: usize,
26085 which: i32,
26086 hk: usize,
26087 ) -> Result<(), Box<dyn std::error::Error>> {
26088 let b = seqs.len();
26089 assert!(b >= 1 && b <= 8);
26090 let mut packed = [GdnSeqVl::default(); 8];
26091 packed[..b].copy_from_slice(seqs);
26092 let v = GdnVl8(packed);
26093 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26094 let max_n = seqs
26095 .iter()
26096 .map(|s| {
26097 if which == 0 {
26098 s.t as i64 * ept as i64
26099 } else {
26100 s.nc as i64 * ept as i64 * 32
26101 }
26102 })
26103 .max()
26104 .unwrap();
26105 let f = self.func("gdn_mirror_vl");
26106 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26107 let cfg = LaunchConfig {
26108 grid_dim: (blocks, 1, b as u32),
26109 block_dim: (256, 1, 1),
26110 shared_mem_bytes: 0,
26111 };
26112 let __s_lb = self.gpu.stream();
26113 let mut lb = __s_lb.launch_builder(&f);
26114 lb.arg(&v).arg(&ept).arg(&which);
26115 unsafe {
26116 lb.launch(cfg)?;
26117 }
26118 Ok(())
26119 }
26120
26121 pub fn gdn_tail_vl8(
26123 &self,
26124 seqs: &[GdnPrepVl],
26125 norm_w: &CudaSlice<f32>,
26126 d_state: usize,
26127 num_v: usize,
26128 eps: f32,
26129 ) -> Result<(), Box<dyn std::error::Error>> {
26130 let b = seqs.len();
26131 assert!(b >= 1 && b <= 8);
26132 let mut packed = [GdnPrepVl::default(); 8];
26133 packed[..b].copy_from_slice(seqs);
26134 let v = GdnPrepVl8(packed);
26135 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26136 let f = self.func("gated_rmsnorm_f16out_vl");
26137 let cfg = LaunchConfig {
26139 grid_dim: (max_t * num_v as u32, 1, b as u32),
26140 block_dim: (128, 1, 1),
26141 shared_mem_bytes: 0,
26142 };
26143 let (dsi, nvi) = (d_state as i32, num_v as i32);
26144 let __s_lb = self.gpu.stream();
26145 let mut lb = __s_lb.launch_builder(&f);
26146 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26147 unsafe {
26148 lb.launch(cfg)?;
26149 }
26150 Ok(())
26151 }
26152
26153 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26156 use cudarc::driver::DevicePtr;
26157 let s = self.gpu.stream();
26158 let (p, _g) = x.device_ptr(&s);
26159 p as u64
26160 }
26161 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26162 use cudarc::driver::DevicePtrMut;
26163 let s = self.gpu.stream();
26164 let (p, _g) = x.device_ptr_mut(&s);
26165 p as u64
26166 }
26167 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26168 use cudarc::driver::DevicePtr;
26169 let s = self.gpu.stream();
26170 let (p, _g) = x.device_ptr(&s);
26171 p as u64
26172 }
26173 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26174 use cudarc::driver::DevicePtr;
26175 let s = self.gpu.stream();
26176 let (p, _g) = x.device_ptr(&s);
26177 p as u64
26178 }
26179
26180 pub fn gdn_chunk_vl8(
26184 &self,
26185 seqs: &[GdnSeqVl],
26186 n_head: usize,
26187 scale: f32,
26188 hk: usize,
26189 wq: Option<&GdnWVl8>,
26190 ) -> Result<(), Box<dyn std::error::Error>> {
26191 const NSPLIT: u32 = 4;
26192 let b = seqs.len();
26193 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26194 let mut packed = [GdnSeqVl::default(); 8];
26195 packed[..b].copy_from_slice(seqs);
26196 let v = GdnVl8(packed);
26197 let (hi, ci) = (n_head as i32, 32i32);
26198 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26199 let hki = hk as i32;
26200 if let Some(w) = wq {
26201 let f = self.func("gdn_k45_wgmma_vl");
26203 let cfg = LaunchConfig {
26204 grid_dim: (n_head as u32, NSPLIT, b as u32),
26205 block_dim: (256, 1, 1),
26206 shared_mem_bytes: 0,
26207 };
26208 let __s_lb = self.gpu.stream();
26209 let mut lb = __s_lb.launch_builder(&f);
26210 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26211 unsafe {
26212 lb.launch(cfg)?;
26213 }
26214 let _ = max_nc;
26215 return Ok(());
26216 }
26217 {
26218 let f = self.func("gdn_chunk_state_mma_vl");
26219 let cfg = LaunchConfig {
26220 grid_dim: (n_head as u32, NSPLIT, b as u32),
26221 block_dim: (256, 1, 1),
26222 shared_mem_bytes: 0,
26223 };
26224 let __s_lb = self.gpu.stream();
26225 let mut lb = __s_lb.launch_builder(&f);
26226 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26227 unsafe {
26228 lb.launch(cfg)?;
26229 }
26230 }
26231 {
26232 let f = self.func("gdn_chunk_output_mma_vl");
26233 let cfg = LaunchConfig {
26234 grid_dim: (max_nc, n_head as u32, b as u32),
26235 block_dim: (256, 1, 1),
26236 shared_mem_bytes: 0,
26237 };
26238 let __s_lb = self.gpu.stream();
26239 let mut lb = __s_lb.launch_builder(&f);
26240 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26241 unsafe {
26242 lb.launch(cfg)?;
26243 }
26244 }
26245 Ok(())
26246 }
26247 pub fn gdn_scan_chunked(
26248 &self,
26249 q: &CudaSlice<f32>,
26250 k: &CudaSlice<f32>,
26251 v: &CudaSlice<f32>,
26252 g: &CudaSlice<f32>,
26253 beta: &CudaSlice<f32>,
26254 kb16_pre: Option<&CudaSlice<u8>>,
26255 qb16_pre: Option<&CudaSlice<u8>>,
26256 state_in: &CudaSlice<f32>,
26257 state_out: &mut CudaSlice<f32>,
26258 o: &mut CudaSlice<f32>,
26259 n_head: usize,
26260 t: usize,
26261 scale: f32,
26262 c: usize,
26263 hk: usize,
26264 ) -> Result<(), Box<dyn std::error::Error>> {
26265 const D: usize = 128;
26266 const NSPLIT: u32 = 4;
26267 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26268 let h = n_head;
26269 let nc = (t + c - 1) / c;
26270 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26271 let gdn_mma_pre = !portable_mma_gated()
26276 && c == 32
26277 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26278 Ok("1") => true,
26279 Ok("0") => false,
26280 _ => gdn_mma_default_on(),
26281 };
26282 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26283 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26284 } else {
26285 None
26286 };
26287 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26292 && gdn_mma_pre
26293 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26294 let nk = t * hk * D;
26295 let mut kb16_local: Option<CudaSlice<u8>> = None;
26296 if gdn_mma_pre && kb16_pre.is_none() {
26297 let mut kb = self.alloc_u8_uninit(nk * 2)?;
26298 let f = self.func("f32_to_bf16_bulk");
26299 let n2 = nk as i64;
26300 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26301 let __s_b = self.gpu.stream();
26302 let mut b = __s_b.launch_builder(&f);
26303 b.arg(k).arg(&mut kb).arg(&n2);
26304 unsafe {
26305 b.launch(cfg2)?;
26306 }
26307 kb16_local = Some(kb);
26308 }
26309 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26310 if let Some(kb) = kb16_pre {
26311 assert!(kb.len() >= nk * 2, "kb16_pre too small");
26312 }
26313 let mut qb16: Option<CudaSlice<u8>> = None;
26314 let mut pb16: Option<CudaSlice<u8>> = None;
26315 if gdn_wgmma_pre {
26316 if qb16_pre.is_none() {
26319 let mut qb = self.alloc_u8_uninit(nk * 2)?;
26320 let f = self.func("f32_to_bf16_bulk");
26321 let n2 = nk as i64;
26322 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26323 let __s_b = self.gpu.stream();
26324 let mut b = __s_b.launch_builder(&f);
26325 b.arg(q).arg(&mut qb).arg(&n2);
26326 unsafe {
26327 b.launch(cfg2)?;
26328 }
26329 qb16 = Some(qb);
26330 } else if let Some(qb) = qb16_pre {
26331 assert!(qb.len() >= nk * 2, "qb16_pre too small");
26332 }
26333 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26334 }
26335 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26336 let k2w = if gdn_wgmma_pre {
26337 Some((
26338 *qb16_ref0.as_ref().unwrap(),
26339 *kb16_ref0.as_ref().unwrap(),
26340 pb16.as_mut().unwrap(),
26341 ))
26342 } else {
26343 None
26344 };
26345 let (gcum, p, u, w) =
26346 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26347 let _ = &w;
26348 let mut y = self.uninit(nc * h * c * D)?;
26349 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
26363 && c == 32
26364 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26365 Ok("1") => true,
26366 Ok("0") => false,
26367 _ => gdn_mma_default_on(),
26368 };
26369 if gdn_mma {
26370 let wb16 = wb16_pre
26371 .take()
26372 .expect("mma path pre-allocates wb16 (K3 store fold)");
26373 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26374 if gdn_wgmma_pre {
26386 let qb16 = qb16_ref0.unwrap();
26388 let pb16 = pb16.as_ref().unwrap();
26389 {
26390 let f = self.func("gdn_k45_wgmma");
26391 let cfg = LaunchConfig {
26392 grid_dim: (h as u32, 4, 1),
26393 block_dim: (256, 1, 1),
26394 shared_mem_bytes: 0,
26395 };
26396 let hki = hk as i32;
26397 let __s_b = self.gpu.stream();
26398 let mut b = __s_b.launch_builder(&f);
26399 b.arg(kb16_ref)
26400 .arg(&gcum)
26401 .arg(beta)
26402 .arg(&u)
26403 .arg(&wb16)
26404 .arg(qb16)
26405 .arg(pb16)
26406 .arg(o)
26407 .arg(&scale)
26408 .arg(state_in)
26409 .arg(&mut *state_out)
26410 .arg(&hi)
26411 .arg(&ti)
26412 .arg(&ci)
26413 .arg(&hki);
26414 unsafe {
26415 b.launch(cfg)?;
26416 }
26417 }
26418 return Ok(());
26419 }
26420 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26424 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26425 {
26426 let f = self.func("gdn_chunk_state_mma");
26427 let cfg = LaunchConfig {
26428 grid_dim: (h as u32, NSPLIT, 1),
26429 block_dim: (256, 1, 1),
26430 shared_mem_bytes: 0,
26431 };
26432 let hki = hk as i32;
26433 let __s_b = self.gpu.stream();
26434 let mut b = __s_b.launch_builder(&f);
26435 b.arg(kb16_ref)
26436 .arg(&gcum)
26437 .arg(beta)
26438 .arg(&u)
26439 .arg(&wb16)
26440 .arg(&mut y16)
26441 .arg(&mut ssnap16)
26442 .arg(state_in)
26443 .arg(&mut *state_out)
26444 .arg(&hi)
26445 .arg(&ti)
26446 .arg(&ci)
26447 .arg(&hki);
26448 unsafe {
26449 b.launch(cfg)?;
26450 }
26451 }
26452 {
26453 let f = self.func("gdn_chunk_output_mma");
26455 let jt = ((c + 31) / 32) as u32;
26456 let cfg = LaunchConfig {
26457 grid_dim: (nc as u32, h as u32, jt),
26458 block_dim: (256, 1, 1),
26459 shared_mem_bytes: 0,
26460 };
26461 let hki = hk as i32;
26462 let __s_b = self.gpu.stream();
26463 let mut b = __s_b.launch_builder(&f);
26464 b.arg(q)
26465 .arg(&gcum)
26466 .arg(&p)
26467 .arg(&y16)
26468 .arg(&ssnap16)
26469 .arg(o)
26470 .arg(&hi)
26471 .arg(&ti)
26472 .arg(&ci)
26473 .arg(&scale)
26474 .arg(&hki);
26475 unsafe {
26476 b.launch(cfg)?;
26477 }
26478 }
26479 return Ok(());
26480 }
26481 {
26482 let f = self.func("gdn_chunk_state_f32");
26484 let cfg = LaunchConfig {
26485 grid_dim: (h as u32, NSPLIT, 1),
26486 block_dim: (256, 1, 1),
26487 shared_mem_bytes: 0,
26488 };
26489 let __s_b = self.gpu.stream();
26490 let mut b = __s_b.launch_builder(&f);
26491 b.arg(k)
26492 .arg(&gcum)
26493 .arg(beta)
26494 .arg(&u)
26495 .arg(&w)
26496 .arg(&mut y)
26497 .arg(&mut ssnap)
26498 .arg(state_in)
26499 .arg(&mut *state_out)
26500 .arg(&hi)
26501 .arg(&ti)
26502 .arg(&ci);
26503 unsafe {
26504 b.launch(cfg)?;
26505 }
26506 }
26507 {
26508 let f = self.func("gdn_chunk_output_f32");
26510 let jt = ((c + 31) / 32) as u32;
26511 let cfg = LaunchConfig {
26512 grid_dim: (nc as u32, h as u32, jt),
26513 block_dim: (256, 1, 1),
26514 shared_mem_bytes: 0,
26515 };
26516 let __s_b = self.gpu.stream();
26517 let mut b = __s_b.launch_builder(&f);
26518 b.arg(q)
26519 .arg(&gcum)
26520 .arg(&p)
26521 .arg(&y)
26522 .arg(&ssnap)
26523 .arg(o)
26524 .arg(&hi)
26525 .arg(&ti)
26526 .arg(&ci)
26527 .arg(&scale);
26528 unsafe {
26529 b.launch(cfg)?;
26530 }
26531 }
26532 Ok(())
26533 }
26534
26535 #[allow(clippy::too_many_arguments)]
26544 #[allow(clippy::too_many_arguments)]
26545 pub fn gdn_scan_prefill(
26546 &self,
26547 q: &CudaSlice<f32>,
26548 k: &CudaSlice<f32>,
26549 v: &CudaSlice<f32>,
26550 g: &CudaSlice<f32>,
26551 beta: &CudaSlice<f32>,
26552 kb16_pre: Option<&CudaSlice<u8>>,
26553 qb16_pre: Option<&CudaSlice<u8>>,
26554 state_in: &CudaSlice<f32>,
26555 state_out: &mut CudaSlice<f32>,
26556 o: &mut CudaSlice<f32>,
26557 n_head: usize,
26558 t: usize,
26559 scale: f32,
26560 hk: usize,
26561 ) -> Result<(), Box<dyn std::error::Error>> {
26562 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26563 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26564 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26565 }
26566 if Self::gdn_chunked_enabled() && t >= 16 {
26567 self.gdn_scan_chunked(
26568 q,
26569 k,
26570 v,
26571 g,
26572 beta,
26573 kb16_pre,
26574 qb16_pre,
26575 state_in,
26576 state_out,
26577 o,
26578 n_head,
26579 t,
26580 scale,
26581 Self::gdn_chunk_size(),
26582 hk,
26583 )
26584 } else {
26585 assert!(
26586 hk == n_head,
26587 "s128 scan is broadcast-only (prep guarantees by predicate)"
26588 );
26589 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26590 }
26591 }
26592
26593 #[allow(clippy::too_many_arguments)]
26595 fn gdn_scan_diff(
26596 &self,
26597 q: &CudaSlice<f32>,
26598 k: &CudaSlice<f32>,
26599 v: &CudaSlice<f32>,
26600 g: &CudaSlice<f32>,
26601 beta: &CudaSlice<f32>,
26602 state_in: &CudaSlice<f32>,
26603 state_out: &mut CudaSlice<f32>,
26604 o: &mut CudaSlice<f32>,
26605 n_head: usize,
26606 t: usize,
26607 scale: f32,
26608 ) -> Result<(), Box<dyn std::error::Error>> {
26609 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26610 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26611 let mut o_c = self.uninit(o.len())?;
26612 let mut st_c = self.uninit(state_out.len())?;
26613 self.gdn_scan_chunked(
26614 q,
26615 k,
26616 v,
26617 g,
26618 beta,
26619 None,
26620 None,
26621 state_in,
26622 &mut st_c,
26623 &mut o_c,
26624 n_head,
26625 t,
26626 scale,
26627 Self::gdn_chunk_size(),
26628 n_head,
26629 )?;
26630 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26631 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26632 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26633 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26634 let mut max_abs = 0f32;
26635 let mut max_rel = 0f32;
26636 let mut sum_rel = 0f64;
26637 for (x, y) in a.iter().zip(b) {
26638 let ad = (x - y).abs();
26639 let rel = ad / x.abs().max(y.abs()).max(1e-3);
26640 if ad > max_abs {
26641 max_abs = ad;
26642 }
26643 if rel > max_rel {
26644 max_rel = rel;
26645 }
26646 sum_rel += rel as f64;
26647 }
26648 (max_abs, max_rel, sum_rel / a.len() as f64)
26649 };
26650 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26651 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26652 println!(
26653 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26654 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26655 Self::gdn_chunk_size()
26656 );
26657 Ok(())
26658 }
26659
26660 pub fn gdn_glog(
26662 &self,
26663 alpha: &CudaSlice<f32>,
26664 dt_bias: &CudaSlice<f32>,
26665 a: &CudaSlice<f32>,
26666 g_log: &mut CudaSlice<f32>,
26667 n_head: usize,
26668 t: usize,
26669 ) -> Result<(), Box<dyn std::error::Error>> {
26670 let f = self.func("gdn_glog_f32");
26671 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26672 let (h, ti) = (n_head as i32, t as i32);
26673 let __s_b = self.gpu.stream();
26674 let mut b = __s_b.launch_builder(&f);
26675 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26676 unsafe {
26677 b.launch(cfg)?;
26678 }
26679 Ok(())
26680 }
26681
26682 pub fn sigmoid_v(
26685 &self,
26686 x: &cudarc::driver::CudaView<f32>,
26687 y: &mut CudaSlice<f32>,
26688 n: usize,
26689 ) -> Result<(), Box<dyn std::error::Error>> {
26690 let f = self.func("sigmoid_f32");
26691 let cfg = LaunchConfig::for_num_elems(n as u32);
26692 let ni = n as i32;
26693 let __s_b = self.gpu.stream();
26694 let mut b = __s_b.launch_builder(&f);
26695 b.arg(x).arg(y).arg(&ni);
26696 unsafe {
26697 b.launch(cfg)?;
26698 }
26699 Ok(())
26700 }
26701
26702 pub fn gdn_glog_v(
26703 &self,
26704 alpha: &cudarc::driver::CudaView<f32>,
26705 dt_bias: &CudaSlice<f32>,
26706 a: &CudaSlice<f32>,
26707 g_log: &mut CudaSlice<f32>,
26708 n_head: usize,
26709 t: usize,
26710 ) -> Result<(), Box<dyn std::error::Error>> {
26711 let f = self.func("gdn_glog_f32");
26712 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26713 let (h, ti) = (n_head as i32, t as i32);
26714 let __s_b = self.gpu.stream();
26715 let mut b = __s_b.launch_builder(&f);
26716 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26717 unsafe {
26718 b.launch(cfg)?;
26719 }
26720 Ok(())
26721 }
26722
26723 pub fn sigmoid(
26724 &self,
26725 x: &CudaSlice<f32>,
26726 y: &mut CudaSlice<f32>,
26727 n: usize,
26728 ) -> Result<(), Box<dyn std::error::Error>> {
26729 let f = self.func("sigmoid_f32");
26730 let cfg = LaunchConfig::for_num_elems(n as u32);
26731 let ni = n as i32;
26732 let __s_b = self.gpu.stream();
26733 let mut b = __s_b.launch_builder(&f);
26734 b.arg(x).arg(y).arg(&ni);
26735 unsafe {
26736 b.launch(cfg)?;
26737 }
26738 Ok(())
26739 }
26740
26741 pub fn sig_mul_f16out(
26744 &self,
26745 a: &CudaSlice<f32>,
26746 g: &CudaSlice<f32>,
26747 dst: &mut CudaSlice<f32>,
26748 dst16: &mut CudaSlice<u8>,
26749 n: usize,
26750 ) -> Result<(), Box<dyn std::error::Error>> {
26751 let f = self.func("sig_mul_f16out_f32");
26752 let cfg = LaunchConfig::for_num_elems(n as u32);
26753 let ni = n as i32;
26754 let __s_b = self.gpu.stream();
26755 let mut b = __s_b.launch_builder(&f);
26756 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26757 unsafe {
26758 b.launch(cfg)?;
26759 }
26760 Ok(())
26761 }
26762
26763 #[allow(clippy::too_many_arguments)]
26772 pub fn attn_head_gate(
26773 &self,
26774 a: &CudaSlice<f32>,
26775 g: &CudaSlice<f32>,
26776 dst: &mut CudaSlice<f32>,
26777 dst16: Option<&mut CudaSlice<u8>>,
26778 head_dim: usize,
26779 n_head: usize,
26780 t: usize,
26781 ) -> Result<(), Box<dyn std::error::Error>> {
26782 let f = self.func("attn_head_gate_f32");
26783 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26784 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26785 let d16: u64 = match dst16 {
26787 Some(d) => self.addr_u8(d),
26788 None => 0,
26789 };
26790 let __s_b = self.gpu.stream();
26791 let mut b = __s_b.launch_builder(&f);
26792 b.arg(a)
26793 .arg(g)
26794 .arg(dst)
26795 .arg(&d16)
26796 .arg(&hd)
26797 .arg(&nh)
26798 .arg(&ti);
26799 unsafe {
26800 b.launch(cfg)?;
26801 }
26802 Ok(())
26803 }
26804
26805 #[allow(clippy::too_many_arguments)]
26814 pub fn swiglu_clamped_mul_scaled(
26815 &self,
26816 gate: &CudaSlice<f32>,
26817 up: &CudaSlice<f32>,
26818 gs: f32,
26819 us: f32,
26820 limit: f32,
26821 dst: &mut CudaSlice<f32>,
26822 n: usize,
26823 ) -> Result<(), Box<dyn std::error::Error>> {
26824 debug_assert!(
26825 limit > 1e-6,
26826 "swiglu_clamped needs a live limit; use silu_mul_scaled"
26827 );
26828 let f = self.func("swiglu_clamped_mul_scaled_f32");
26829 let cfg = LaunchConfig::for_num_elems(n as u32);
26830 let ni = n as i32;
26831 let __s_b = self.gpu.stream();
26832 let mut b = __s_b.launch_builder(&f);
26833 b.arg(gate)
26834 .arg(up)
26835 .arg(&gs)
26836 .arg(&us)
26837 .arg(&limit)
26838 .arg(dst)
26839 .arg(&ni);
26840 unsafe {
26841 b.launch(cfg)?;
26842 }
26843 Ok(())
26844 }
26845
26846 pub fn gated_rmsnorm(
26848 &self,
26849 o: &CudaSlice<f32>,
26850 w: &CudaSlice<f32>,
26851 z: &CudaSlice<f32>,
26852 dst: &mut CudaSlice<f32>,
26853 ncols: usize,
26854 nrows: usize,
26855 eps: f32,
26856 ) -> Result<(), Box<dyn std::error::Error>> {
26857 let f = self.func("gated_rmsnorm_f32");
26858 let cfg = LaunchConfig {
26859 grid_dim: (nrows as u32, 1, 1),
26860 block_dim: (128, 1, 1),
26861 shared_mem_bytes: 0,
26862 };
26863 let (nc, e) = (ncols as i32, eps);
26864 let __s_b = self.gpu.stream();
26865 let mut b = __s_b.launch_builder(&f);
26866 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26867 unsafe {
26868 b.launch(cfg)?;
26869 }
26870 Ok(())
26871 }
26872
26873 pub fn gated_rmsnorm_f16out(
26876 &self,
26877 o: &CudaSlice<f32>,
26878 w: &CudaSlice<f32>,
26879 z: &CudaSlice<f32>,
26880 dst: &mut CudaSlice<f32>,
26881 dst16: &mut CudaSlice<u8>,
26882 ncols: usize,
26883 nrows: usize,
26884 eps: f32,
26885 ) -> Result<(), Box<dyn std::error::Error>> {
26886 let f = self.func("gated_rmsnorm_f16out_f32");
26887 let cfg = LaunchConfig {
26889 grid_dim: (nrows as u32, 1, 1),
26890 block_dim: (128, 1, 1),
26891 shared_mem_bytes: 0,
26892 };
26893 let (nc, e) = (ncols as i32, eps);
26894 let __s_b = self.gpu.stream();
26895 let mut b = __s_b.launch_builder(&f);
26896 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26897 unsafe {
26898 b.launch(cfg)?;
26899 }
26900 Ok(())
26901 }
26902
26903 #[allow(clippy::too_many_arguments)]
26907 pub fn add_rms_norm_zq8(
26908 &self,
26909 a: &CudaSlice<f32>,
26910 b_in: &CudaSlice<f32>,
26911 w: &CudaSlice<f32>,
26912 res: &mut CudaSlice<f32>,
26913 z: &mut CudaSlice<f32>,
26914 ncols: usize,
26915 nrows: usize,
26916 eps: f32,
26917 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26918 assert!(ncols % 32 == 0);
26919 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
26920 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26921 let f = self.func("add_rms_norm_zq8");
26922 let cfg = LaunchConfig {
26923 grid_dim: (nrows as u32, 1, 1),
26924 block_dim: (1024, 1, 1),
26925 shared_mem_bytes: 0,
26926 };
26927 let (nc, ep) = (ncols as i32, eps);
26928 let __s_b = self.gpu.stream();
26929 let mut b = __s_b.launch_builder(&f);
26930 b.arg(a)
26931 .arg(b_in)
26932 .arg(w)
26933 .arg(res)
26934 .arg(z)
26935 .arg(&mut q)
26936 .arg(&mut d)
26937 .arg(&nc)
26938 .arg(&ep);
26939 unsafe {
26940 b.launch(cfg)?;
26941 }
26942 Ok((q, d))
26943 }
26944
26945 pub fn gated_rmsnorm_zv(
26950 &self,
26951 o: &CudaSlice<f32>,
26952 w: &CudaSlice<f32>,
26953 z: &cudarc::driver::CudaView<f32>,
26954 dst: &mut CudaSlice<f32>,
26955 ncols: usize,
26956 nrows: usize,
26957 eps: f32,
26958 ) -> Result<(), Box<dyn std::error::Error>> {
26959 let f = self.func("gated_rmsnorm_f32");
26960 let cfg = LaunchConfig {
26961 grid_dim: (nrows as u32, 1, 1),
26962 block_dim: (128, 1, 1),
26963 shared_mem_bytes: 0,
26964 };
26965 let (nc, e) = (ncols as i32, eps);
26966 let __s_b = self.gpu.stream();
26967 let mut b = __s_b.launch_builder(&f);
26968 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
26969 unsafe {
26970 b.launch(cfg)?;
26971 }
26972 Ok(())
26973 }
26974
26975 pub fn gated_rmsnorm_f16out_zv(
26976 &self,
26977 o: &CudaSlice<f32>,
26978 w: &CudaSlice<f32>,
26979 z: &cudarc::driver::CudaView<f32>,
26980 dst: &mut CudaSlice<f32>,
26981 dst16: &mut CudaSlice<u8>,
26982 ncols: usize,
26983 nrows: usize,
26984 eps: f32,
26985 ) -> Result<(), Box<dyn std::error::Error>> {
26986 let f = self.func("gated_rmsnorm_f16out_f32");
26987 let cfg = LaunchConfig {
26989 grid_dim: (nrows as u32, 1, 1),
26990 block_dim: (128, 1, 1),
26991 shared_mem_bytes: 0,
26992 };
26993 let (nc, e) = (ncols as i32, eps);
26994 let __s_b = self.gpu.stream();
26995 let mut b = __s_b.launch_builder(&f);
26996 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26997 unsafe {
26998 b.launch(cfg)?;
26999 }
27000 Ok(())
27001 }
27002
27003 pub fn gated_rmsnorm_q8_1(
27004 &self,
27005 o: &CudaSlice<f32>,
27006 w: &CudaSlice<f32>,
27007 z: &CudaSlice<f32>,
27008 ncols: usize,
27009 nrows: usize,
27010 eps: f32,
27011 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27012 assert!(ncols % 32 == 0);
27013 let f = self.func("gated_rmsnorm_q8_1");
27014 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27015 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27016 let cfg = LaunchConfig {
27017 grid_dim: (nrows as u32, 1, 1),
27018 block_dim: (128, 1, 1),
27019 shared_mem_bytes: 0,
27020 };
27021 let (nc, ep) = (ncols as i32, eps);
27022 let __s_b = self.gpu.stream();
27023 let mut b = __s_b.launch_builder(&f);
27024 b.arg(o)
27025 .arg(w)
27026 .arg(z)
27027 .arg(&mut out_q)
27028 .arg(&mut out_d)
27029 .arg(&nc)
27030 .arg(&ep);
27031 unsafe {
27032 b.launch(cfg)?;
27033 }
27034 Ok((out_q, out_d))
27035 }
27036
27037 pub fn transpose(
27039 &self,
27040 inp: &CudaSlice<f32>,
27041 rows: usize,
27042 cols: usize,
27043 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27044 let f = self.func("transpose_f32");
27045 let mut out = self.zeros(rows * cols)?;
27046 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27047 let (r, c) = (rows as i32, cols as i32);
27048 let __s_b = self.gpu.stream();
27049 let mut b = __s_b.launch_builder(&f);
27050 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27051 unsafe {
27052 b.launch(cfg)?;
27053 }
27054 Ok(out)
27055 }
27056
27057 pub fn repeat_heads(
27059 &self,
27060 inp: &CudaSlice<f32>,
27061 out: &mut CudaSlice<f32>,
27062 head_dim: usize,
27063 n_in: usize,
27064 n_out: usize,
27065 t: usize,
27066 ) -> Result<(), Box<dyn std::error::Error>> {
27067 let f = self.func("repeat_heads_f32");
27068 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27069 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27070 let __s_b = self.gpu.stream();
27071 let mut b = __s_b.launch_builder(&f);
27072 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27073 unsafe {
27074 b.launch(cfg)?;
27075 }
27076 Ok(())
27077 }
27078
27079 pub fn q_gate_split(
27086 &self,
27087 qf: &CudaSlice<f32>,
27088 q_out: &mut CudaSlice<f32>,
27089 gate_out: &mut CudaSlice<f32>,
27090 head_dim: usize,
27091 n_head: usize,
27092 t: usize,
27093 ) -> Result<(), Box<dyn std::error::Error>> {
27094 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27095 let out_need = head_dim * n_head * t;
27096 if q_out.len() < out_need || gate_out.len() < out_need {
27097 return Err(format!(
27098 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27099 q_out.len(),
27100 gate_out.len()
27101 )
27102 .into());
27103 }
27104 let f = self.func("q_gate_split_f32");
27105 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27106 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27107 let __s_b = self.gpu.stream();
27108 let mut b = __s_b.launch_builder(&f);
27109 b.arg(qf)
27110 .arg(q_out)
27111 .arg(gate_out)
27112 .arg(&hd)
27113 .arg(&nh)
27114 .arg(&ti);
27115 unsafe {
27116 b.launch(cfg)?;
27117 }
27118 Ok(())
27119 }
27120
27121 pub fn qkv_to_gdn_repack(
27125 &self,
27126 conv_out: &CudaSlice<f32>,
27127 q_g: &mut CudaSlice<f32>,
27128 k_g: &mut CudaSlice<f32>,
27129 v_g: &mut CudaSlice<f32>,
27130 d_state: usize,
27131 num_v: usize,
27132 num_k: usize,
27133 key_dim: usize,
27134 t: usize,
27135 ) -> Result<(), Box<dyn std::error::Error>> {
27136 let f = self.func("qkv_to_gdn_repack_f32");
27137 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27138 let (ds, nv, nk, kd, ti) = (
27139 d_state as i32,
27140 num_v as i32,
27141 num_k as i32,
27142 key_dim as i32,
27143 t as i32,
27144 );
27145 let __s_b = self.gpu.stream();
27146 let mut b = __s_b.launch_builder(&f);
27147 b.arg(conv_out)
27148 .arg(q_g)
27149 .arg(k_g)
27150 .arg(v_g)
27151 .arg(&ds)
27152 .arg(&nv)
27153 .arg(&nk)
27154 .arg(&kd)
27155 .arg(&ti);
27156 unsafe {
27157 b.launch(cfg)?;
27158 }
27159 Ok(())
27160 }
27161
27162 pub fn conv_left_pad(
27165 &self,
27166 src: &CudaSlice<f32>,
27167 dst: &mut CudaSlice<f32>,
27168 conv_dim: usize,
27169 t: usize,
27170 pad: usize,
27171 ) -> Result<(), Box<dyn std::error::Error>> {
27172 let f = self.func("conv_left_pad_f32");
27173 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27174 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27175 let __s_b = self.gpu.stream();
27176 let mut b = __s_b.launch_builder(&f);
27177 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27178 unsafe {
27179 b.launch(cfg)?;
27180 }
27181 Ok(())
27182 }
27183
27184 pub fn conv_assemble_and_roll(
27188 &self,
27189 qkv_col: &CudaSlice<f32>,
27190 conv_state: &mut CudaSlice<f32>,
27191 conv_in: &mut CudaSlice<f32>,
27192 conv_dim: usize,
27193 pad: usize,
27194 ) -> Result<(), Box<dyn std::error::Error>> {
27195 let f = self.func("conv_assemble_and_roll_f32");
27196 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27197 let (cd, p) = (conv_dim as i32, pad as i32);
27198 let __s_b = self.gpu.stream();
27199 let mut b = __s_b.launch_builder(&f);
27200 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27201 unsafe {
27202 b.launch(cfg)?;
27203 }
27204 Ok(())
27205 }
27206
27207 pub fn ssm_conv1d_fused_decode(
27213 &self,
27214 qkv_col: &CudaSlice<f32>,
27215 conv_state: &mut CudaSlice<f32>,
27216 w: &CudaSlice<f32>,
27217 conv_out: &mut CudaSlice<f32>,
27218 conv_dim: usize,
27219 d_conv: usize,
27220 ) -> Result<(), Box<dyn std::error::Error>> {
27221 let f = self.func("ssm_conv1d_fused_decode_f32");
27222 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27223 let (cd, dc) = (conv_dim as i32, d_conv as i32);
27224 let __s_b = self.gpu.stream();
27225 let mut b = __s_b.launch_builder(&f);
27226 b.arg(qkv_col)
27227 .arg(conv_state)
27228 .arg(w)
27229 .arg(conv_out)
27230 .arg(&cd)
27231 .arg(&dc);
27232 unsafe {
27233 b.launch(cfg)?;
27234 }
27235 Ok(())
27236 }
27237
27238 pub fn slice_range(
27241 &self,
27242 src: &CudaSlice<f32>,
27243 start: usize,
27244 len: usize,
27245 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27246 let host = self.gpu.stream().clone_dtoh(src)?;
27247 self.gpu.stream().synchronize()?;
27248 Ok(self.htod(&host[start..start + len])?)
27249 }
27250}
27251
27252#[cfg(test)]
27253mod target_dispatch_tests {
27254 use super::legacy_quant_gemm_allowed;
27255
27256 #[test]
27257 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27258 assert!(legacy_quant_gemm_allowed(false, false, false));
27260 assert!(!legacy_quant_gemm_allowed(false, false, true));
27261 assert!(!legacy_quant_gemm_allowed(true, false, false));
27263 assert!(!legacy_quant_gemm_allowed(true, false, true));
27264 assert!(legacy_quant_gemm_allowed(true, true, false));
27266 assert!(!legacy_quant_gemm_allowed(true, true, true));
27267 }
27268
27269 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27270 #[test]
27271 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27272 assert!(!legacy_quant_gemm_allowed(
27273 cfg!(memra_portable_cuda),
27274 cfg!(memra_hopper_mma),
27275 false
27276 ));
27277 }
27278
27279 #[cfg(memra_hopper_mma)]
27280 #[test]
27281 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27282 assert!(legacy_quant_gemm_allowed(
27283 cfg!(memra_portable_cuda),
27284 cfg!(memra_hopper_mma),
27285 false
27286 ));
27287 assert!(super::portable_mma_gated() == false);
27288 }
27289}
27290
27291impl memra_kv::KvDev for Engine {
27294 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27295 Engine::zeros(self, n)
27296 }
27297 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27298 Engine::uninit(self, n)
27299 }
27300 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27301 Engine::alloc_u8(self, n)
27302 }
27303 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27304 Engine::htod_i32(self, v)
27305 }
27306 fn clone_dtod(
27307 &self,
27308 src: &CudaSlice<f32>,
27309 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27310 Engine::clone_dtod(self, src)
27311 }
27312 fn copy_into(
27313 &self,
27314 dst: &mut CudaSlice<f32>,
27315 off: usize,
27316 src: &CudaSlice<f32>,
27317 len: usize,
27318 ) -> Result<(), Box<dyn std::error::Error>> {
27319 Engine::copy_into(self, dst, off, src, len)
27320 }
27321 fn set_i32_one(
27322 &self,
27323 d: &mut CudaSlice<i32>,
27324 v: i32,
27325 ) -> Result<(), Box<dyn std::error::Error>> {
27326 Engine::set_i32_one(self, d, v)
27327 }
27328}
27329
27330#[cfg(test)]
27331mod fused_gate_bounds_tests {
27332 use super::*;
27333
27334 #[test]
27347 #[ignore = "requires a CUDA GPU"]
27348 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27349 let e = Engine::new(0).unwrap();
27350 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27351 let fused = 2 * head_dim * n_head * t;
27352 let out_n = head_dim * n_head * t;
27353
27354 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27356 let mut q = e.uninit(out_n).unwrap();
27357 let mut gate = e.uninit(out_n).unwrap();
27358 let err = e
27359 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27360 .expect_err("half-width wq must be refused, not read past")
27361 .to_string();
27362 assert!(err.contains("NO fused gate"), "{err}");
27363 assert!(err.contains(&format!("{fused}")), "{err}");
27364
27365 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27368 let wide = e.htod(&host).unwrap();
27369 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27370 .expect("full-width wq splits");
27371 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27372 for tok in 0..t {
27373 for hh in 0..n_head {
27374 for d in 0..head_dim {
27375 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27376 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27377 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27378 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27379 }
27380 }
27381 }
27382
27383 let mut small = e.uninit(out_n - 1).unwrap();
27385 assert!(
27386 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27387 .is_err()
27388 );
27389 }
27390}
27391
27392#[cfg(test)]
27396mod fused_rope_width_tests {
27397 use super::Engine;
27398
27399 #[test]
27402 fn full_width_is_accepted() {
27403 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27404 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27405 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27406 }
27407
27408 #[test]
27421 fn gemma4_official_artifact_widths_pass() {
27422 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27423 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27424 }
27425
27426 #[test]
27429 fn partial_rotary_is_refused_with_the_geometry_named() {
27430 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27432 .expect_err("partial rotary must refuse");
27433 let msg = err.to_string();
27434 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27435 assert!(msg.contains("n_rot 64"), "{msg}");
27436 assert!(msg.contains("head_dim 256"), "{msg}");
27437 assert!(
27438 msg.contains("64..256"),
27439 "names the band it would corrupt: {msg}"
27440 );
27441 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27443 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27445 }
27446}