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
617fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
622 match (sig_expf, fast && n_used <= 8) {
623 (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
624 (true, false) => "moe_router_sigmoid_topk_f32_dexp",
625 (false, true) => "moe_router_sigmoid_topk_f32_fast",
626 (false, false) => "moe_router_sigmoid_topk_f32",
627 }
628}
629
630#[cfg(test)]
631mod sigmoid_topk_dispatch_tests {
632 #[test]
633 fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
634 use super::sigmoid_topk_kernel;
635
636 assert_eq!(
637 sigmoid_topk_kernel(false, true, 8),
638 "moe_router_sigmoid_topk_f32_fast"
639 );
640 assert_eq!(
641 sigmoid_topk_kernel(true, true, 8),
642 "moe_router_sigmoid_topk_f32_dexp_fast"
643 );
644 assert_eq!(
645 sigmoid_topk_kernel(false, true, 9),
646 "moe_router_sigmoid_topk_f32"
647 );
648 assert_eq!(
649 sigmoid_topk_kernel(true, true, 9),
650 "moe_router_sigmoid_topk_f32_dexp"
651 );
652 }
653}
654
655pub(crate) fn rms_block() -> u32 {
656 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
657 *V.get_or_init(|| {
658 std::env::var("MEMRA_RMS_BLOCK")
659 .ok()
660 .and_then(|v| v.parse().ok())
661 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
662 })
663}
664
665pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
666 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
667 if let Some(forced) = *S.get_or_init(|| {
668 std::env::var("MEMRA_FA_SPLIT")
669 .ok()
670 .and_then(|v| v.parse().ok())
671 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
672 }) {
673 return forced;
674 }
675 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
693 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
694 {
695 return if t_kv <= 8192 {
696 16
697 } else if t_kv <= 16384 {
698 64
699 } else {
700 128
701 };
702 }
703 let big_rig = fa_sm_count() >= 128;
704 if big_rig {
705 let _ = n_head_kv;
706 if t_kv <= 2048 {
707 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
716 if let Some(sp) = *SHORT.get_or_init(|| {
717 std::env::var("MEMRA_FA_SP_SHORT")
718 .ok()
719 .and_then(|v| v.parse().ok())
720 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
721 }) {
722 return sp;
723 }
724 16
725 } else if t_kv <= 16384 {
726 64
727 } else {
728 128
729 }
730 } else if n_head_kv <= 4 {
731 if t_kv <= 512 {
752 8
753 } else if t_kv <= 16384 {
754 64
755 } else {
756 128
757 }
758 } else {
759 if t_kv <= 8192 {
760 32
761 } else if t_kv <= 16384 {
762 64
763 } else {
764 128
765 }
766 }
767}
768
769pub(crate) fn fa_sm_count() -> i32 {
772 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
773 *N.get_or_init(|| {
774 cudarc::driver::result::init().ok();
775 cudarc::driver::result::device::get(0)
776 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
777 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
778 .unwrap_or(82)
779 })
780}
781
782fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
786 match head_dim {
787 256 => Ok(""),
788 128 => Ok("_hd128"),
789 d => Err(format!(
790 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
791 callers must gate to sdpa_naive"
792 )
793 .into()),
794 }
795}
796
797pub const QT_Q8_0: i32 = 0;
799pub const QT_Q4_K: i32 = 1;
800pub const QT_Q6_K: i32 = 2;
801pub const QT_Q5_K: i32 = 3;
802pub const QT_Q3_K: i32 = 4;
803pub const QT_IQ4_XS: i32 = 5;
804pub const QT_IQ3_S: i32 = 6;
805pub const QT_NVFP4: i32 = 7;
806pub const QT_F8_E4M3: i32 = 10;
812pub const QT_NVFP4_RP: i32 = 9;
815pub const QT_F32: i32 = 8;
817pub const QT_BF16: i32 = 11;
818pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
823pub const QT_F8_E4M3_BLK: i32 = 14;
839
840pub struct Engine {
842 pub gpu: memra_runtime::Gpu,
843 module: Arc<CudaModule>,
844 hybrid: Arc<CudaModule>,
845 qmatvec: Arc<CudaModule>,
846 flash: Arc<CudaModule>,
847 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
851 gemm: Arc<CudaModule>,
852 router: Arc<CudaModule>,
853 sample: Arc<CudaModule>,
855 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
859 w8_mirrors: Mutex<std::collections::HashMap<u64, CudaSlice<u8>>>,
864 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
867 moe_cache_layout: Mutex<Option<Vec<usize>>>,
871 capture_keep_on: std::sync::atomic::AtomicBool,
877 verify_exact: std::sync::atomic::AtomicBool,
882 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
883 pub copy_stream: Arc<CudaStream>,
885 #[cfg(memra_cutlass)]
892 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
893 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
897 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
900 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
904 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
908 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
910 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
911 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
916 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
921 router_stage: Mutex<Option<PinnedStage>>,
925}
926
927fn fa_v2_on() -> bool {
937 std::env::var("MEMRA_FA_V2")
943 .map(|v| v != "0")
944 .unwrap_or(true)
945}
946
947pub(crate) fn fa_v3_on() -> bool {
955 std::env::var("MEMRA_FA_V3")
959 .map(|v| v != "0")
960 .unwrap_or(true)
961}
962
963fn fa_v4_mode() -> &'static str {
968 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
969 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
970}
971fn fa_v4_on() -> bool {
972 fa_v4_mode() != "0"
973} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
982 std::sync::atomic::AtomicUsize::new(1024);
983pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
984 std::sync::atomic::AtomicUsize::new(usize::MAX);
985pub fn fa_v4_at_pub(t_kv: usize) -> bool {
986 fa_v4_at(t_kv)
987}
988fn fa_v4_at(t_kv: usize) -> bool {
989 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
990 let mx = *M.get_or_init(|| {
991 std::env::var("MEMRA_FA_V4_MAX")
992 .ok()
993 .and_then(|v| v.parse().ok())
994 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
995 });
996 fa_v4_on() && t_kv < mx
997}
998pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1012fn fa_deep_at(t_kv: usize) -> bool {
1013 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1014 return false;
1015 }
1016 let min = std::env::var("MEMRA_FA_DEEP_MIN")
1017 .ok()
1018 .and_then(|v| v.parse().ok())
1019 .unwrap_or(FA_DEEP_MIN_DEFAULT);
1020 t_kv >= min
1021}
1022pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1024 fa_deep_at(t_kv)
1025}
1026
1027fn fa_v3_active(head_dim: usize) -> bool {
1028 fa_v3_on()
1031 && head_dim % 128 == 0
1032 && kv_cache_formats() == ("q8_0", "q5_1")
1033 && !Engine::kv_fp8_on()
1034}
1035
1036pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1044 std::env::var("MEMRA_NO_FA_VEC").is_err()
1045 && t_kv >= fa_vec_min_tkv()
1046 && head_dim == 256
1047 && fa_v4_at(t_kv)
1048 && !matches!(fa_v4_mode(), "noB3" | "stage")
1049 && !Engine::kv_fp8_on()
1050}
1051pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1053 fa_split_keys(t_kv, n_head_kv)
1054}
1055
1056struct PinnedStage {
1061 ptr: *mut u8,
1062 cap: usize,
1063}
1064unsafe impl Send for PinnedStage {}
1065impl PinnedStage {
1066 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1067 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1068 Ok(PinnedStage { ptr, cap })
1069 }
1070}
1071impl Drop for PinnedStage {
1072 fn drop(&mut self) {
1073 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1074 }
1075}
1076
1077pub const ARGMAX_NB: usize = 256;
1080
1081pub(crate) use memra_fa3_vl as fa3_vl_raw;
1083
1084unsafe extern "C" {
1085 fn memra_fa3_prefill(
1087 q16: *const core::ffi::c_void,
1088 k16: *const core::ffi::c_void,
1089 v16: *const core::ffi::c_void,
1090 o: *mut f32,
1091 t: i32,
1092 h: i32,
1093 hkv: i32,
1094 d: i32,
1095 scale: f32,
1096 stream: *mut core::ffi::c_void,
1097 ) -> i32;
1098 pub(crate) fn memra_fa3_vl(
1100 q16s: *const *const core::ffi::c_void,
1101 k16s: *const *const core::ffi::c_void,
1102 v16s: *const *const core::ffi::c_void,
1103 os: *const *mut f32,
1104 ts: *const i32,
1105 b: i32,
1106 h: i32,
1107 hkv: i32,
1108 d: i32,
1109 scale: f32,
1110 stream: *mut core::ffi::c_void,
1111 ) -> i32;
1112}
1113
1114#[repr(C)]
1119#[derive(Clone, Copy)]
1120pub struct WPtr8(pub [u64; 8]);
1121unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1122
1123#[repr(C)]
1128#[derive(Clone, Copy, Default)]
1129pub struct GdnSeqVl {
1130 pub kb16: u64,
1131 pub gcum: u64,
1132 pub beta: u64,
1133 pub u: u64,
1134 pub wb16: u64,
1135 pub y: u64,
1136 pub ssnap: u64,
1137 pub state_in: u64,
1138 pub state_out: u64,
1139 pub q: u64,
1140 pub p: u64,
1141 pub o: u64,
1142 pub k: u64,
1143 pub v: u64,
1144 pub g: u64,
1145 pub a: u64,
1146 pub w: u64,
1147 pub t: i32,
1148 pub nc: i32,
1149}
1150unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1151#[repr(C)]
1152#[derive(Clone, Copy)]
1153pub struct GdnVl8(pub [GdnSeqVl; 8]);
1154unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1155
1156#[repr(C)]
1159#[derive(Clone, Copy, Default)]
1160pub struct GdnWVl {
1161 pub qb16: u64,
1162 pub pb16: u64,
1163}
1164unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1165#[repr(C)]
1166#[derive(Clone, Copy)]
1167pub struct GdnWVl8(pub [GdnWVl; 8]);
1168unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1169
1170#[repr(C)]
1172#[derive(Clone, Copy, Default)]
1173pub struct GdnPrepVl {
1174 pub qkv: u64,
1175 pub conv_state: u64,
1176 pub conv_out: u64,
1177 pub q_g: u64,
1178 pub k_g: u64,
1179 pub v_g: u64,
1180 pub q_l2: u64,
1181 pub k_l2: u64,
1182 pub beta_raw: u64,
1183 pub alpha: u64,
1184 pub beta: u64,
1185 pub g_log: u64,
1186 pub o: u64,
1187 pub z: u64,
1188 pub gn: u64,
1189 pub gn16: u64,
1190 pub kb16: u64,
1191 pub qb16: u64,
1192 pub t: i32,
1193 pub pad: i32,
1194}
1195unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1196#[repr(C)]
1197#[derive(Clone, Copy)]
1198pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1199unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1200
1201#[repr(C)]
1203#[derive(Clone, Copy, Default)]
1204pub struct FaSeqVl {
1205 pub q: u64,
1206 pub k16: u64,
1207 pub v16: u64,
1208 pub o: u64,
1209 pub kf: u64,
1210 pub vf: u64,
1211 pub t: i32,
1212 pub pad: i32,
1213}
1214unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1215#[repr(C)]
1216#[derive(Clone, Copy)]
1217pub struct FaVl8(pub [FaSeqVl; 8]);
1218unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1219
1220#[repr(C)]
1222#[derive(Clone, Copy, Default)]
1223pub struct AttnPreVl {
1224 pub qf: u64,
1225 pub kf: u64,
1226 pub vf: u64,
1227 pub q: u64,
1228 pub gate: u64,
1229 pub qn: u64,
1230 pub kn: u64,
1231 pub kc: u64,
1232 pub vc: u64,
1233 pub t: i32,
1234 pub pad: i32,
1235}
1236unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1237#[repr(C)]
1238#[derive(Clone, Copy)]
1239pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1240unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1241
1242pub struct GdnChunkBufs {
1245 pub gcum: CudaSlice<f32>,
1246 pub a: CudaSlice<f32>,
1247 pub p: CudaSlice<f32>,
1248 pub u: CudaSlice<f32>,
1249 pub w: CudaSlice<f32>,
1250 pub kb16: CudaSlice<u8>,
1251 pub wb16: CudaSlice<u8>,
1252 pub y16: CudaSlice<u8>,
1253 pub ssnap16: CudaSlice<u8>,
1254 pub qb16: CudaSlice<u8>,
1255 pub pb16: CudaSlice<u8>,
1256 pub o: CudaSlice<f32>,
1257 pub t: usize,
1258 pub nc: usize,
1259}
1260
1261#[repr(C)]
1263#[derive(Clone, Copy)]
1264pub struct F32x8(pub [f32; 8]);
1265unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1266
1267pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1271
1272#[must_use = "dropping immediately ends the exact scope"]
1277pub struct ExactScope<'a> {
1278 flag: &'a std::sync::atomic::AtomicBool,
1279 prev: bool,
1280}
1281
1282impl<'a> ExactScope<'a> {
1283 pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
1284 let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
1285 flag.store(on, std::sync::atomic::Ordering::Relaxed);
1286 ExactScope { flag, prev }
1287 }
1288}
1289
1290impl Drop for ExactScope<'_> {
1291 fn drop(&mut self) {
1292 self.flag
1293 .store(self.prev, std::sync::atomic::Ordering::Relaxed);
1294 }
1295}
1296
1297#[cfg(test)]
1298mod exact_scope_tests {
1299 use std::sync::atomic::{AtomicBool, Ordering};
1300
1301 #[test]
1302 fn error_path_restores_verify_exact() {
1303 let flag = AtomicBool::new(false);
1308 let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
1309 let _scope = super::ExactScope::set(flag, true);
1310 assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
1311 Err("draft forward failed")? };
1313 assert!(failing(&flag).is_err());
1314 assert!(
1315 !flag.load(Ordering::Relaxed),
1316 "error propagation must restore the pre-scope value"
1317 );
1318 let flag = AtomicBool::new(true);
1320 {
1321 let _scope = super::ExactScope::set(&flag, true);
1322 }
1323 assert!(flag.load(Ordering::Relaxed));
1324 let flag = AtomicBool::new(false);
1326 let scope = super::ExactScope::set(&flag, true);
1327 drop(scope);
1328 assert!(!flag.load(Ordering::Relaxed));
1329 }
1330}
1331
1332impl Engine {
1333 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1334 let gpu = memra_runtime::Gpu::new(ordinal)?;
1335 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1339 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1340 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1341 .and_then(|d| unsafe {
1342 Ok((
1343 cudarc::driver::result::device::get_attribute(
1344 d,
1345 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1346 )?,
1347 cudarc::driver::result::device::get_attribute(
1348 d,
1349 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1350 )?,
1351 ))
1352 })
1353 .unwrap_or((0, 0));
1354 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1355 let ok = matches!(
1356 (built, maj, min),
1357 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1358 );
1359 if !ok {
1360 return Err(format!(
1361 "memra was built for sm_{built} but device {ordinal} reports compute \
1362 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1363 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1364 )
1365 .into());
1366 }
1367 }
1368 unsafe {
1373 use cudarc::driver::sys;
1374 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1375 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1376 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1377 let mut thresh: u64 = u64::MAX;
1378 let _ = sys::cuMemPoolSetAttribute(
1379 pool,
1380 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1381 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1382 );
1383 }
1384 }
1385 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1386 let hybrid = gpu
1387 .ctx
1388 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1389 let qmatvec = gpu
1390 .ctx
1391 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1392 let flash = gpu
1393 .ctx
1394 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1395 let gemm = gpu
1396 .ctx
1397 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1398 let router = gpu
1399 .ctx
1400 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1401 let sample = gpu
1402 .ctx
1403 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1404 let copy_stream = gpu.ctx.new_stream()?;
1405 if std::env::var("MEMRA_EVT")
1421 .map(|v| v == "1")
1422 .unwrap_or(false)
1423 {
1424 } else {
1426 unsafe {
1427 gpu.ctx.disable_event_tracking();
1428 }
1429 }
1430 Ok(Self {
1431 gpu,
1432 module,
1433 hybrid,
1434 qmatvec,
1435 flash,
1436 flash_g: std::sync::OnceLock::new(),
1437 gemm,
1438 router,
1439 sample,
1440 moe_cache: Mutex::new(None),
1441 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
1442 w8_act: Mutex::new(std::collections::HashMap::new()),
1443 moe_cache_layout: Mutex::new(None),
1444 copy_stream,
1445 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1446 verify_exact: std::sync::atomic::AtomicBool::new(false),
1447 capture_keep: Mutex::new(Vec::new()),
1448 argmax_partials: Mutex::new(None),
1449 prime_deqw_ws: Mutex::new(None),
1450 router_stage: Mutex::new(None),
1451 fp8_scratch: Mutex::new(None),
1452 fa_vf16_scratch: Mutex::new(None),
1453 fa_part_pool: Mutex::new(None),
1454 fa_part_retired: Mutex::new(Vec::new()),
1455 fn_cache: Mutex::new(Default::default()),
1456 f16_scratch: Mutex::new(None),
1457 #[cfg(memra_cutlass)]
1458 cutlass_scratch: Mutex::new(None),
1459 })
1460 }
1461
1462 pub fn ctx(&self) -> &Arc<CudaContext> {
1463 &self.gpu.ctx
1464 }
1465
1466 pub fn pool_cached_bytes(&self) -> usize {
1484 let (reserved, used) = self.pool_reserved_used();
1485 reserved.saturating_sub(used)
1486 }
1487
1488 pub fn device_graph_mem_reserved(&self) -> usize {
1498 use cudarc::driver::sys as cus;
1499 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1500 return 0;
1501 };
1502 let mut bytes: u64 = 0;
1503 let rc = unsafe {
1504 cus::cuDeviceGetGraphMemAttribute(
1505 dev,
1506 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1507 &mut bytes as *mut u64 as *mut std::ffi::c_void,
1508 )
1509 };
1510 if rc == cus::cudaError_enum::CUDA_SUCCESS {
1511 bytes as usize
1512 } else {
1513 0
1514 }
1515 }
1516
1517 pub fn pool_reserved_used(&self) -> (usize, usize) {
1524 use cudarc::driver::sys;
1525 unsafe {
1526 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1527 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1528 != sys::CUresult::CUDA_SUCCESS
1529 {
1530 return (0, 0);
1531 }
1532 let (mut reserved, mut used) = (0u64, 0u64);
1533 if sys::cuMemPoolGetAttribute(
1534 pool,
1535 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1536 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1537 ) != sys::CUresult::CUDA_SUCCESS
1538 {
1539 return (0, 0);
1540 }
1541 if sys::cuMemPoolGetAttribute(
1542 pool,
1543 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1544 &mut used as *mut u64 as *mut core::ffi::c_void,
1545 ) != sys::CUresult::CUDA_SUCCESS
1546 {
1547 return (0, 0);
1548 }
1549 (reserved as usize, used as usize)
1550 }
1551 }
1552
1553 pub fn stream(&self) -> Arc<CudaStream> {
1556 self.gpu.stream()
1557 }
1558 pub fn gkv_on() -> bool {
1561 memra_kv::gkv_on()
1562 }
1563
1564 pub fn wkv_on() -> bool {
1576 memra_kv::wkv_on()
1577 }
1578
1579 pub fn kv_fp8_on() -> bool {
1585 memra_kv::kv_fp8_on()
1586 }
1587
1588 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1591 if head_dim == 512 && Self::gkv_on() {
1592 self.func_g(name)
1593 } else {
1594 self.func(name)
1595 }
1596 }
1597
1598 fn func_g(&self, name: &str) -> CudaFunction {
1602 let m = self.flash_g.get_or_init(|| {
1603 self.gpu
1604 .ctx
1605 .load_module(cudarc::nvrtc::Ptx::from_binary(
1606 FLASH_FATBIN_KF8VF8.to_vec(),
1607 ))
1608 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1609 });
1610 let key = format!("g:{name}");
1611 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1612 return f.clone();
1613 }
1614 let f = match m.load_function(name) {
1615 Ok(f) => f,
1616 Err(_) => self.func(name),
1617 };
1618 self.fn_cache.lock().unwrap().insert(key, f.clone());
1619 f
1620 }
1621
1622 fn func(&self, name: &str) -> CudaFunction {
1623 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1626 return f.clone();
1627 }
1628 let f = self
1629 .module
1630 .load_function(name)
1631 .or_else(|_| self.hybrid.load_function(name))
1632 .or_else(|_| self.qmatvec.load_function(name))
1633 .or_else(|_| self.flash.load_function(name))
1634 .or_else(|_| self.gemm.load_function(name))
1635 .or_else(|_| self.router.load_function(name))
1636 .or_else(|_| self.sample.load_function(name))
1637 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1638 self.fn_cache
1639 .lock()
1640 .unwrap()
1641 .insert(name.to_string(), f.clone());
1642 f
1643 }
1644
1645 pub fn scatter_trim_logits(
1648 &self,
1649 src: &CudaSlice<f32>,
1650 d2t: &CudaSlice<u32>,
1651 dst: &mut CudaSlice<f32>,
1652 d_vocab: usize,
1653 n_vocab: usize,
1654 ) -> Result<(), Box<dyn std::error::Error>> {
1655 let f1 = self.func("scatter_trim_logits_f32");
1656 let f2 = self.func("scatter_trim_logits_pass2_f32");
1657 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1658 let cfg1 = LaunchConfig {
1659 grid_dim: (256, 1, 1),
1660 block_dim: (256, 1, 1),
1661 shared_mem_bytes: 0,
1662 };
1663 let __s_b1 = self.gpu.stream();
1664 let mut b1 = __s_b1.launch_builder(&f1);
1665 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1666 unsafe {
1667 b1.launch(cfg1)?;
1668 }
1669 let cfg2 = LaunchConfig {
1670 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1671 block_dim: (256, 1, 1),
1672 shared_mem_bytes: 0,
1673 };
1674 let __s_b2 = self.gpu.stream();
1675 let mut b2 = __s_b2.launch_builder(&f2);
1676 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1677 unsafe {
1678 b2.launch(cfg2)?;
1679 }
1680 Ok(())
1681 }
1682
1683 #[allow(clippy::too_many_arguments)]
1689 pub fn filter_stats(
1690 &self,
1691 x: &CudaSlice<f32>,
1692 row_stride: usize,
1693 rows: &CudaSlice<i32>,
1694 out_th: &mut CudaSlice<f32>,
1695 out_z: &mut CudaSlice<f32>,
1696 out_max: &mut CudaSlice<f32>,
1697 n: usize,
1698 nrow: usize,
1699 temp: f32,
1700 top_k: i32,
1701 top_p: f32,
1702 min_p: f32,
1703 ) -> Result<(), Box<dyn std::error::Error>> {
1704 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1731 let coop_on =
1732 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1733 if coop_on && self.sm_count() >= 16 {
1734 let cap = self.sm_count() as usize / 16;
1735 let mut done = 0usize;
1736 while done < nrow {
1737 let chunk = cap.min(nrow - done);
1738 self.filter_stats_coop_chunk(
1739 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1740 top_p, min_p,
1741 )?;
1742 done += chunk;
1743 }
1744 return Ok(());
1745 }
1746 self.filter_stats_plain_program(
1747 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1748 )
1749 }
1750
1751 #[allow(clippy::too_many_arguments)]
1756 pub fn filter_stats_coop_chunk(
1757 &self,
1758 x: &CudaSlice<f32>,
1759 row_stride: usize,
1760 rows: &CudaSlice<i32>,
1761 row0: usize,
1762 out_th: &mut CudaSlice<f32>,
1763 out_z: &mut CudaSlice<f32>,
1764 out_max: &mut CudaSlice<f32>,
1765 n: usize,
1766 chunk: usize,
1767 temp: f32,
1768 top_k: i32,
1769 top_p: f32,
1770 min_p: f32,
1771 ) -> Result<(), Box<dyn std::error::Error>> {
1772 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1773 let f = self.func("filter_stats_coop_f32");
1774 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1775 let cfg = LaunchConfig {
1776 grid_dim: (16, chunk as u32, 1),
1777 block_dim: (512, 1, 1),
1778 shared_mem_bytes: 0,
1779 };
1780 let rows_v = rows.slice(row0..row0 + chunk);
1781 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1782 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1783 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1784 let __s_b = self.gpu.stream();
1785 let mut b = __s_b.launch_builder(&f);
1786 b.arg(x)
1787 .arg(&rs)
1788 .arg(&rows_v)
1789 .arg(&mut th_v)
1790 .arg(&mut z_v)
1791 .arg(&mut mx_v)
1792 .arg(&mut ws)
1793 .arg(&ni)
1794 .arg(&nr)
1795 .arg(&temp)
1796 .arg(&top_k)
1797 .arg(&top_p)
1798 .arg(&min_p);
1799 unsafe {
1800 b.launch_cooperative(cfg)?;
1801 }
1802 Ok(())
1803 }
1804
1805 #[allow(clippy::too_many_arguments)]
1809 pub fn filter_stats_plain_program(
1810 &self,
1811 x: &CudaSlice<f32>,
1812 row_stride: usize,
1813 rows: &CudaSlice<i32>,
1814 out_th: &mut CudaSlice<f32>,
1815 out_z: &mut CudaSlice<f32>,
1816 out_max: &mut CudaSlice<f32>,
1817 n: usize,
1818 nrow: usize,
1819 temp: f32,
1820 top_k: i32,
1821 top_p: f32,
1822 min_p: f32,
1823 ) -> Result<(), Box<dyn std::error::Error>> {
1824 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1825 let f = self.func("filter_stats_f32");
1826 let cfg = LaunchConfig {
1827 grid_dim: (nrow as u32, 1, 1),
1828 block_dim: (1024, 1, 1),
1829 shared_mem_bytes: 0,
1830 };
1831 let __s_b = self.gpu.stream();
1832 let mut b = __s_b.launch_builder(&f);
1833 b.arg(x)
1834 .arg(&rs)
1835 .arg(rows)
1836 .arg(&mut *out_th)
1837 .arg(&mut *out_z)
1838 .arg(&mut *out_max)
1839 .arg(&ni)
1840 .arg(&nr)
1841 .arg(&temp)
1842 .arg(&top_k)
1843 .arg(&top_p)
1844 .arg(&min_p);
1845 unsafe {
1846 b.launch(cfg)?;
1847 }
1848 Ok(())
1849 }
1850
1851 #[allow(clippy::too_many_arguments)]
1853 pub fn softmax_gather_filtered(
1854 &self,
1855 x: &CudaSlice<f32>,
1856 row_stride: usize,
1857 ids: &CudaSlice<u32>,
1858 rows: &CudaSlice<i32>,
1859 th: &CudaSlice<f32>,
1860 z: &CudaSlice<f32>,
1861 out: &mut CudaSlice<f32>,
1862 n: usize,
1863 npair: usize,
1864 temp: f32,
1865 ) -> Result<(), Box<dyn std::error::Error>> {
1866 let f = self.func("softmax_gather_filtered_f32");
1867 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1868 let cfg = LaunchConfig {
1869 grid_dim: (npair as u32, 1, 1),
1870 block_dim: (256, 1, 1),
1871 shared_mem_bytes: 0,
1872 };
1873 let __s_b = self.gpu.stream();
1874 let mut b = __s_b.launch_builder(&f);
1875 b.arg(x)
1876 .arg(&rs)
1877 .arg(ids)
1878 .arg(rows)
1879 .arg(th)
1880 .arg(z)
1881 .arg(&mut *out)
1882 .arg(&ni)
1883 .arg(&np)
1884 .arg(&temp);
1885 unsafe {
1886 b.launch(cfg)?;
1887 }
1888 Ok(())
1889 }
1890
1891 #[allow(clippy::too_many_arguments)]
1893 pub fn residual_sample_filtered(
1894 &self,
1895 p: &CudaSlice<f32>,
1896 q: Option<&CudaSlice<f32>>,
1897 n: usize,
1898 temp: f32,
1899 seed: u64,
1900 stream_pos: u32,
1901 p_stats: (f32, f32, f32),
1902 q_stats: (f32, f32, f32),
1903 out_tok: &mut CudaSlice<u32>,
1904 ) -> Result<(), Box<dyn std::error::Error>> {
1905 let f = self.func("residual_sample_filtered_f32");
1906 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1907 let has_q: i32 = q.is_some() as i32;
1908 let qbuf = q.unwrap_or(p);
1909 let (pm, pth, pz) = p_stats;
1910 let (qm, qth, qz) = q_stats;
1911 let cfg = LaunchConfig {
1912 grid_dim: (1, 1, 1),
1913 block_dim: (1024, 1, 1),
1914 shared_mem_bytes: 0,
1915 };
1916 let __s_b = self.gpu.stream();
1917 let mut b = __s_b.launch_builder(&f);
1918 b.arg(p)
1919 .arg(qbuf)
1920 .arg(&has_q)
1921 .arg(&ni)
1922 .arg(&temp)
1923 .arg(&slo)
1924 .arg(&shi)
1925 .arg(&stream_pos)
1926 .arg(&pm)
1927 .arg(&pth)
1928 .arg(&pz)
1929 .arg(&qm)
1930 .arg(&qth)
1931 .arg(&qz)
1932 .arg(&mut *out_tok);
1933 unsafe {
1934 b.launch(cfg)?;
1935 }
1936 Ok(())
1937 }
1938
1939 #[allow(clippy::too_many_arguments)]
1945 pub fn residual_sample_sparse_q(
1946 &self,
1947 p: &CudaSlice<f32>,
1948 cand_ids: &CudaSlice<u32>,
1949 q_probs: &CudaSlice<f32>,
1950 n_cand: usize,
1951 n: usize,
1952 temp: f32,
1953 seed: u64,
1954 stream_pos: u32,
1955 p_stats: (f32, f32, f32),
1956 out_tok: &mut CudaSlice<u32>,
1957 ) -> Result<(), Box<dyn std::error::Error>> {
1958 assert!(
1959 n_cand >= 1 && n_cand <= 32,
1960 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1961 );
1962 let f = self.func("residual_sample_sparse_q_f32");
1963 let (ni, nc) = (n as i32, n_cand as i32);
1964 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1965 let (pm, pth, pz) = p_stats;
1966 let cfg = LaunchConfig {
1967 grid_dim: (1, 1, 1),
1968 block_dim: (1024, 1, 1),
1969 shared_mem_bytes: 0,
1970 };
1971 let __s_b = self.gpu.stream();
1972 let mut b = __s_b.launch_builder(&f);
1973 b.arg(p)
1974 .arg(cand_ids)
1975 .arg(q_probs)
1976 .arg(&nc)
1977 .arg(&ni)
1978 .arg(&temp)
1979 .arg(&slo)
1980 .arg(&shi)
1981 .arg(&stream_pos)
1982 .arg(&pm)
1983 .arg(&pth)
1984 .arg(&pz)
1985 .arg(&mut *out_tok);
1986 unsafe {
1987 b.launch(cfg)?;
1988 }
1989 Ok(())
1990 }
1991
1992 #[allow(clippy::too_many_arguments)]
1994 pub fn gumbel_perturb_filtered(
1995 &self,
1996 x: &CudaSlice<f32>,
1997 y: &mut CudaSlice<f32>,
1998 n: usize,
1999 seed: u64,
2000 stream_pos: u32,
2001 temp: f32,
2002 row_max: f32,
2003 th: f32,
2004 ) -> Result<(), Box<dyn std::error::Error>> {
2005 let f = self.func("gumbel_perturb_filtered_f32");
2006 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2007 let cfg = LaunchConfig {
2008 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2009 block_dim: (256, 1, 1),
2010 shared_mem_bytes: 0,
2011 };
2012 let __s_b = self.gpu.stream();
2013 let mut b = __s_b.launch_builder(&f);
2014 b.arg(x)
2015 .arg(&mut *y)
2016 .arg(&ni)
2017 .arg(&slo)
2018 .arg(&shi)
2019 .arg(&stream_pos)
2020 .arg(&temp)
2021 .arg(&row_max)
2022 .arg(&th);
2023 unsafe {
2024 b.launch(cfg)?;
2025 }
2026 Ok(())
2027 }
2028
2029 #[allow(clippy::too_many_arguments)]
2033 pub fn penalize_logits(
2034 &self,
2035 x: &mut CudaSlice<f32>,
2036 hist: &CudaSlice<u32>,
2037 n_hist: usize,
2038 rep: f32,
2039 freq: f32,
2040 present: f32,
2041 n: usize,
2042 ) -> Result<(), Box<dyn std::error::Error>> {
2043 if n_hist == 0 {
2044 return Ok(());
2045 }
2046 let f = self.func("penalize_logits_f32");
2047 let (nh, ni) = (n_hist as i32, n as i32);
2048 let cfg = LaunchConfig {
2049 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2050 block_dim: (128, 1, 1),
2051 shared_mem_bytes: 0,
2052 };
2053 let __s_b = self.gpu.stream();
2054 let mut b = __s_b.launch_builder(&f);
2055 b.arg(&mut *x)
2056 .arg(hist)
2057 .arg(&nh)
2058 .arg(&rep)
2059 .arg(&freq)
2060 .arg(&present)
2061 .arg(&ni);
2062 unsafe {
2063 b.launch(cfg)?;
2064 }
2065 Ok(())
2066 }
2067
2068 #[allow(clippy::too_many_arguments)]
2070 pub fn penalize_logits_rows(
2071 &self,
2072 x: &mut CudaSlice<f32>,
2073 hist: &CudaSlice<u32>,
2074 n_hist: usize,
2075 rep: f32,
2076 freq: f32,
2077 present: f32,
2078 n: usize,
2079 nrow: usize,
2080 ) -> Result<(), Box<dyn std::error::Error>> {
2081 if n_hist == 0 || nrow == 0 {
2082 return Ok(());
2083 }
2084 let f = self.func("penalize_logits_rows_f32");
2085 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2086 let cfg = LaunchConfig {
2087 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2088 block_dim: (128, 1, 1),
2089 shared_mem_bytes: 0,
2090 };
2091 let __s_b = self.gpu.stream();
2092 let mut b = __s_b.launch_builder(&f);
2093 b.arg(&mut *x)
2094 .arg(hist)
2095 .arg(&nh)
2096 .arg(&rep)
2097 .arg(&freq)
2098 .arg(&present)
2099 .arg(&ni)
2100 .arg(&nr);
2101 unsafe {
2102 b.launch(cfg)?;
2103 }
2104 Ok(())
2105 }
2106
2107 #[allow(clippy::too_many_arguments)]
2113 pub fn penalize_logits_sparse_rows(
2114 &self,
2115 x: &mut CudaSlice<f32>,
2116 ids: &[u32],
2117 counts: &[u32],
2118 offsets: &[i32],
2119 rows: &[i32],
2120 reps: &[f32],
2121 freqs: &[f32],
2122 presents: &[f32],
2123 n: usize,
2124 ) -> Result<(), Box<dyn std::error::Error>> {
2125 let nrow = rows.len();
2126 if nrow == 0 {
2127 return Ok(());
2128 }
2129 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2130 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2131 let entry_count =
2132 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2133 if ids.len() != counts.len()
2134 || offsets.len() != nrow + 1
2135 || reps.len() != nrow
2136 || freqs.len() != nrow
2137 || presents.len() != nrow
2138 || offsets.first().copied() != Some(0)
2139 || offsets.last().copied() != Some(entry_count)
2140 {
2141 return Err("sparse penalty row metadata shape mismatch".into());
2142 }
2143 if counts.contains(&0) {
2144 return Err("sparse penalty counts must be positive".into());
2145 }
2146 let mut max_len = 0usize;
2147 for pair in offsets.windows(2) {
2148 if pair[0] < 0 || pair[1] < pair[0] {
2149 return Err("sparse penalty offsets must be monotonic".into());
2150 }
2151 max_len = max_len.max((pair[1] - pair[0]) as usize);
2152 }
2153 if max_len == 0 {
2154 return Ok(());
2155 }
2156
2157 let mut seen = std::collections::HashSet::with_capacity(ids.len());
2158 for (r, &row) in rows.iter().enumerate() {
2159 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2160 return Err("sparse penalty row index exceeds logits shape".into());
2161 }
2162 let begin = offsets[r] as usize;
2163 let end = offsets[r + 1] as usize;
2164 for &id in &ids[begin..end] {
2165 if id as usize >= n {
2166 return Err("sparse penalty token id exceeds logits row".into());
2167 }
2168 if !seen.insert((row, id)) {
2169 return Err("sparse penalty entries must be unique per logits row".into());
2170 }
2171 }
2172 }
2173
2174 unsafe {
2176 self.penalize_logits_sparse_rows_unchecked(
2177 x, ids, counts, offsets, rows, reps, freqs, presents, n,
2178 )
2179 }
2180 }
2181
2182 #[allow(clippy::too_many_arguments)]
2191 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2192 &self,
2193 x: &mut CudaSlice<f32>,
2194 ids: &[u32],
2195 counts: &[u32],
2196 offsets: &[i32],
2197 rows: &[i32],
2198 reps: &[f32],
2199 freqs: &[f32],
2200 presents: &[f32],
2201 n: usize,
2202 ) -> Result<(), Box<dyn std::error::Error>> {
2203 let nrow = rows.len();
2204 if nrow == 0 {
2205 return Ok(());
2206 }
2207 let max_len = offsets
2208 .windows(2)
2209 .map(|pair| (pair[1] - pair[0]) as usize)
2210 .max()
2211 .unwrap_or(0);
2212 if max_len == 0 {
2213 return Ok(());
2214 }
2215 let ids_d = self.htod_u32_v(ids)?;
2216 let counts_d = self.htod_u32_v(counts)?;
2217 let offsets_d = self.htod_i32(offsets)?;
2218 let rows_d = self.htod_i32(rows)?;
2219 let reps_d = self.htod(reps)?;
2220 let freqs_d = self.htod(freqs)?;
2221 let presents_d = self.htod(presents)?;
2222 let f = self.func("penalize_logits_sparse_rows_f32");
2223 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2224 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2225 let cfg = LaunchConfig {
2226 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2227 block_dim: (128, 1, 1),
2228 shared_mem_bytes: 0,
2229 };
2230 let __s_b = self.gpu.stream();
2231 let mut b = __s_b.launch_builder(&f);
2232 b.arg(&mut *x)
2233 .arg(&ids_d)
2234 .arg(&counts_d)
2235 .arg(&offsets_d)
2236 .arg(&rows_d)
2237 .arg(&reps_d)
2238 .arg(&freqs_d)
2239 .arg(&presents_d)
2240 .arg(&ni)
2241 .arg(&nr);
2242 unsafe {
2243 b.launch(cfg)?;
2244 }
2245 Ok(())
2246 }
2247
2248 #[allow(clippy::too_many_arguments)]
2256 pub fn penalize_logits_rows_inc(
2257 &self,
2258 x: &mut CudaSlice<f32>,
2259 hist: &CudaSlice<u32>,
2260 n_hist0: usize,
2261 rep: f32,
2262 freq: f32,
2263 present: f32,
2264 n: usize,
2265 nrow: usize,
2266 win: usize,
2267 ) -> Result<(), Box<dyn std::error::Error>> {
2268 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2269 return Ok(());
2270 }
2271 debug_assert!(
2272 hist.len() >= n_hist0 + nrow - 1,
2273 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2274 );
2275 let f = self.func("penalize_logits_rows_inc_f32");
2276 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2277 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2278 let cfg = LaunchConfig {
2279 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2280 block_dim: (128, 1, 1),
2281 shared_mem_bytes: 0,
2282 };
2283 let __s_b = self.gpu.stream();
2284 let mut b = __s_b.launch_builder(&f);
2285 b.arg(&mut *x)
2286 .arg(hist)
2287 .arg(&nh)
2288 .arg(&rep)
2289 .arg(&freq)
2290 .arg(&present)
2291 .arg(&ni)
2292 .arg(&nr)
2293 .arg(&wi);
2294 unsafe {
2295 b.launch(cfg)?;
2296 }
2297 Ok(())
2298 }
2299
2300 pub fn wpf_level() -> u32 {
2308 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2309 *ON.get_or_init(|| {
2310 std::env::var("MEMRA_WPF")
2311 .ok()
2312 .and_then(|v| v.parse().ok())
2313 .unwrap_or(1)
2314 })
2315 }
2316
2317 pub fn set_verify_exact(&self, on: bool) {
2333 self.verify_exact
2334 .store(on, std::sync::atomic::Ordering::Relaxed);
2335 }
2336 pub(crate) fn verify_exact_on(&self) -> bool {
2337 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2338 }
2339
2340 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2346 ExactScope::set(&self.verify_exact, on)
2347 }
2348
2349 pub fn qkv_append_on() -> bool {
2352 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2353 *ON.get_or_init(|| {
2354 std::env::var("MEMRA_QKV_APPEND")
2355 .map(|v| v != "0")
2356 .unwrap_or(true)
2357 })
2358 }
2359
2360 pub fn pdl_wb_on() -> bool {
2363 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2364 *ON.get_or_init(|| {
2365 std::env::var("MEMRA_PDL_WB")
2366 .map(|v| v != "0")
2367 .unwrap_or(true)
2368 })
2369 }
2370
2371 pub fn norm_ilp_on() -> bool {
2379 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2380 *ON.get_or_init(|| {
2381 std::env::var("MEMRA_NORM_ILP")
2382 .map(|v| v != "0")
2383 .unwrap_or(true)
2384 })
2385 }
2386
2387 pub fn tk_ffn_dual_on() -> bool {
2395 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2396 *ON.get_or_init(|| {
2397 std::env::var("MEMRA_TK_FFN_DUAL")
2398 .map(|v| v != "0")
2399 .unwrap_or(true)
2400 })
2401 }
2402
2403 pub fn pdl_mmvq_on() -> bool {
2407 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2408 *ON.get_or_init(|| {
2409 std::env::var("MEMRA_PDL_MMVQ")
2410 .map(|v| v != "0")
2411 .unwrap_or(true)
2412 })
2413 }
2414
2415 pub fn pdl_on() -> bool {
2416 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2417 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2418 }
2419
2420 pub fn pdl_nvfp4q8_on() -> bool {
2426 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2427 *ON.get_or_init(|| {
2428 std::env::var("MEMRA_PDL_NVFP4")
2429 .map(|v| v != "0")
2430 .unwrap_or(true)
2431 })
2432 }
2433
2434 fn q40_mr1_on() -> bool {
2440 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2441 match *Q40MR.get_or_init(|| {
2442 std::env::var("MEMRA_Q40_MR")
2443 .ok()
2444 .and_then(|v| v.parse().ok())
2445 }) {
2446 Some(v) => v == 1,
2447 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2448 }
2449 }
2450
2451 fn pdl_func_flash(
2456 &self,
2457 g: bool,
2458 name: &'static str,
2459 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2460 use cudarc::driver::sys as cu;
2461 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2468 std::sync::Mutex::new(None);
2469 static FNS: std::sync::Mutex<
2470 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2471 > = std::sync::Mutex::new(None);
2472 let ctx_key = self.ctx().cu_ctx() as usize;
2473 if let Some(&f) = FNS
2474 .lock()
2475 .unwrap()
2476 .get_or_insert_with(Default::default)
2477 .get(&(ctx_key, g, name))
2478 {
2479 return Ok(f as cu::CUfunction);
2480 }
2481 let module = {
2482 let mut mods = MODS.lock().unwrap();
2483 let map = mods.get_or_insert_with(Default::default);
2484 match map.get(&(ctx_key, g)) {
2485 Some(&m) => m,
2486 None => {
2487 let m = self.pdl_load_module_in_ctx(if g {
2488 FLASH_FATBIN_KF8VF8
2489 } else {
2490 FLASH_FATBIN
2491 })?;
2492 map.insert((ctx_key, g), m);
2493 m
2494 }
2495 }
2496 };
2497 let cname = std::ffi::CString::new(name)?;
2498 let mut f: cu::CUfunction = std::ptr::null_mut();
2499 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2500 if r != cu::CUresult::CUDA_SUCCESS {
2501 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2502 }
2503 FNS.lock()
2504 .unwrap()
2505 .get_or_insert_with(Default::default)
2506 .insert((ctx_key, g, name), f as usize);
2507 Ok(f)
2508 }
2509
2510 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2515 use cudarc::driver::sys as cu;
2516 let mut prev: cu::CUcontext = std::ptr::null_mut();
2517 unsafe {
2518 cu::cuCtxGetCurrent(&mut prev).result()?;
2519 }
2520 self.ctx().bind_to_thread()?;
2521 let mut m: cu::CUmodule = std::ptr::null_mut();
2522 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2523 let restore = if prev.is_null() {
2524 cu::CUresult::CUDA_SUCCESS
2525 } else {
2526 unsafe { cu::cuCtxSetCurrent(prev) }
2527 };
2528 if r != cu::CUresult::CUDA_SUCCESS {
2529 return Err(format!("pdl module load: {r:?}").into());
2530 }
2531 if restore != cu::CUresult::CUDA_SUCCESS {
2532 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2533 }
2534 Ok(m as usize)
2535 }
2536
2537 pub fn raw_kernel_function(
2540 &self,
2541 name: &'static str,
2542 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2543 self.pdl_func(name)
2544 }
2545
2546 fn pdl_func(
2547 &self,
2548 name: &'static str,
2549 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2550 use cudarc::driver::sys as cu;
2551 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2554 std::sync::Mutex::new(None);
2555 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2558 std::sync::Mutex::new(None);
2559 static FNS: std::sync::Mutex<
2560 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2561 > = std::sync::Mutex::new(None);
2562 let ctx_key = self.ctx().cu_ctx() as usize;
2563 if let Some(&f) = FNS
2564 .lock()
2565 .unwrap()
2566 .get_or_insert_with(Default::default)
2567 .get(&(ctx_key, name))
2568 {
2569 return Ok(f as cu::CUfunction);
2570 }
2571 let module = {
2572 let mut mods = MODULES.lock().unwrap();
2573 let map = mods.get_or_insert_with(Default::default);
2574 match map.get(&ctx_key) {
2575 Some(&m) => m,
2576 None => {
2577 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2578 map.insert(ctx_key, m);
2579 m
2580 }
2581 }
2582 };
2583 let cname = std::ffi::CString::new(name)?;
2584 let mut f: cu::CUfunction = std::ptr::null_mut();
2585 let mut r =
2586 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2587 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2588 let qmodule = {
2589 let mut mods = QMODULES.lock().unwrap();
2590 let map = mods.get_or_insert_with(Default::default);
2591 match map.get(&ctx_key) {
2592 Some(&m) => m,
2593 None => {
2594 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2595 map.insert(ctx_key, m);
2596 m
2597 }
2598 }
2599 };
2600 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2601 }
2602 if r != cu::CUresult::CUDA_SUCCESS {
2603 return Err(format!("pdl_func {name}: {r:?}").into());
2604 }
2605 FNS.lock()
2606 .unwrap()
2607 .get_or_insert_with(Default::default)
2608 .insert((ctx_key, name), f as usize);
2609 Ok(f)
2610 }
2611
2612 unsafe fn launch_pdl_flash(
2624 &self,
2625 g: bool,
2626 name: &'static str,
2627 grid: (u32, u32, u32),
2628 block: (u32, u32, u32),
2629 smem: u32,
2630 params: &mut [*mut std::ffi::c_void],
2631 ) -> Result<(), Box<dyn std::error::Error>> {
2632 use cudarc::driver::sys as cu;
2633 let f = self.pdl_func_flash(g, name)?;
2634 if smem > 0 {
2635 let r =
2637 unsafe {
2638 cu::cuFuncSetAttribute(f,
2639 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2640 smem as i32)
2641 };
2642 if r != cu::CUresult::CUDA_SUCCESS {
2643 return Err(format!("pdl smem attr {name}: {r:?}").into());
2644 }
2645 }
2646 let mut attr = cu::CUlaunchAttribute {
2647 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2648 pad: [0; 4],
2649 value: cu::CUlaunchAttributeValue {
2650 programmaticStreamSerializationAllowed: 1,
2651 },
2652 };
2653 let cfg = cu::CUlaunchConfig {
2654 gridDimX: grid.0,
2655 gridDimY: grid.1,
2656 gridDimZ: grid.2,
2657 blockDimX: block.0,
2658 blockDimY: block.1,
2659 blockDimZ: block.2,
2660 sharedMemBytes: smem,
2661 hStream: self.gpu.stream().cu_stream(),
2662 attrs: &mut attr,
2663 numAttrs: 1,
2664 };
2665 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2666 if r != cu::CUresult::CUDA_SUCCESS {
2667 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2668 }
2669 Ok(())
2670 }
2671
2672 unsafe fn launch_pdl(
2673 &self,
2674 name: &'static str,
2675 grid: (u32, u32, u32),
2676 block: (u32, u32, u32),
2677 params: &mut [*mut std::ffi::c_void],
2678 ) -> Result<(), Box<dyn std::error::Error>> {
2679 use cudarc::driver::sys as cu;
2680 let f = self.pdl_func(name)?;
2681 let mut attr = cu::CUlaunchAttribute {
2682 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2683 pad: [0; 4],
2684 value: cu::CUlaunchAttributeValue {
2685 programmaticStreamSerializationAllowed: 1,
2686 },
2687 };
2688 let cfg = cu::CUlaunchConfig {
2689 gridDimX: grid.0,
2690 gridDimY: grid.1,
2691 gridDimZ: grid.2,
2692 blockDimX: block.0,
2693 blockDimY: block.1,
2694 blockDimZ: block.2,
2695 sharedMemBytes: 0,
2696 hStream: self.gpu.stream().cu_stream(),
2697 attrs: &mut attr,
2698 numAttrs: 1,
2699 };
2700 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2701 if r != cu::CUresult::CUDA_SUCCESS {
2702 return Err(format!("launch_pdl {name}: {r:?}").into());
2703 }
2704 Ok(())
2705 }
2706
2707 pub fn prefetch_weight_l2(
2710 &self,
2711 w: &crate::model::GpuTensor,
2712 ) -> Result<(), Box<dyn std::error::Error>> {
2713 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2714 let p = rp4.as_ref().unwrap_or(bytes);
2715 self.prefetch_l2(p, p.len())?;
2716 }
2717 Ok(())
2718 }
2719
2720 pub fn gather_row_bf16(
2723 &self,
2724 table: &CudaSlice<u8>,
2725 tok: &CudaSlice<u32>,
2726 idx: usize,
2727 dst: &mut CudaSlice<f32>,
2728 ncols: usize,
2729 ) -> Result<(), Box<dyn std::error::Error>> {
2730 let f = self.func("gather_row_bf16_f32");
2731 let cfg = LaunchConfig {
2732 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2733 block_dim: (256, 1, 1),
2734 shared_mem_bytes: 0,
2735 };
2736 let (nc, ix) = (ncols as i32, idx as i32);
2737 let __s_b = self.gpu.stream();
2738 let mut b = __s_b.launch_builder(&f);
2739 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2740 unsafe {
2741 b.launch(cfg)?;
2742 }
2743 Ok(())
2744 }
2745
2746 #[allow(clippy::too_many_arguments)]
2752 pub fn dflash2_dynconv(
2753 &self,
2754 x: &CudaSlice<f32>,
2755 dyn_: &CudaSlice<f32>,
2756 base: &CudaSlice<f32>,
2757 out: &mut CudaSlice<f32>,
2758 rows: usize,
2759 hidden: usize,
2760 group_size: usize,
2761 ksize: usize,
2762 half: usize,
2763 ) -> Result<(), Box<dyn std::error::Error>> {
2764 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2765 let f = self.func("dflash2_dynconv_f32");
2766 let n = rows * hidden;
2767 let cfg = LaunchConfig {
2768 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2769 block_dim: (256, 1, 1),
2770 shared_mem_bytes: 0,
2771 };
2772 let (ri, hi, gi, ki, hf) = (
2773 rows as i32,
2774 hidden as i32,
2775 group_size as i32,
2776 ksize as i32,
2777 half as i32,
2778 );
2779 let __s_b = self.gpu.stream();
2780 let mut b = __s_b.launch_builder(&f);
2781 b.arg(x)
2782 .arg(dyn_)
2783 .arg(base)
2784 .arg(out)
2785 .arg(&ri)
2786 .arg(&hi)
2787 .arg(&gi)
2788 .arg(&ki)
2789 .arg(&hf);
2790 unsafe {
2791 b.launch(cfg)?;
2792 }
2793 Ok(())
2794 }
2795
2796 pub fn topk_rows(
2800 &self,
2801 logits: &CudaSlice<f32>,
2802 n_rows: usize,
2803 n_cols: usize,
2804 k: usize,
2805 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2806 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2807 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2808 let f = self.func("topk_rows_f32");
2809 let nth = 256usize;
2810 let mut vals = self.uninit(n_rows * k)?;
2811 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2812 let cfg = LaunchConfig {
2813 grid_dim: (n_rows as u32, 1, 1),
2814 block_dim: (nth as u32, 1, 1),
2815 shared_mem_bytes: (nth * k * 8) as u32,
2816 };
2817 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2818 let __s_b = self.gpu.stream();
2819 let mut b = __s_b.launch_builder(&f);
2820 b.arg(logits)
2821 .arg(&nr)
2822 .arg(&nc)
2823 .arg(&ki)
2824 .arg(&mut vals)
2825 .arg(&mut idxs);
2826 unsafe {
2827 b.launch(cfg)?;
2828 }
2829 Ok((vals, idxs))
2830 }
2831
2832 pub fn add_row_inplace(
2834 &self,
2835 logits: &mut CudaSlice<f32>,
2836 bias: &CudaSlice<f32>,
2837 n: usize,
2838 row_off: usize,
2839 ) -> Result<(), Box<dyn std::error::Error>> {
2840 let f = self.func("add_row_inplace_f32");
2841 let cfg = LaunchConfig {
2842 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2843 block_dim: (256, 1, 1),
2844 shared_mem_bytes: 0,
2845 };
2846 let (ni, off) = (n as i32, row_off as i64);
2847 let __s_b = self.gpu.stream();
2848 let mut b = __s_b.launch_builder(&f);
2849 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2850 unsafe {
2851 b.launch(cfg)?;
2852 }
2853 Ok(())
2854 }
2855
2856 pub fn prefetch_l2(
2858 &self,
2859 p: &CudaSlice<u8>,
2860 n: usize,
2861 ) -> Result<(), Box<dyn std::error::Error>> {
2862 let f = self.func("prefetch_l2_bytes");
2863 let lines = n.div_ceil(128);
2864 let ni = n as i64;
2865 let cfg = LaunchConfig {
2866 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2867 block_dim: (256, 1, 1),
2868 shared_mem_bytes: 0,
2869 };
2870 let __s_b = self.gpu.stream();
2871 let mut b = __s_b.launch_builder(&f);
2872 b.arg(p).arg(&ni);
2873 unsafe {
2874 b.launch(cfg)?;
2875 }
2876 Ok(())
2877 }
2878
2879 pub fn router_gemv(
2882 &self,
2883 w: &CudaSlice<f32>,
2884 x: &CudaSlice<f32>,
2885 n_embd: usize,
2886 n_experts: usize,
2887 t: usize,
2888 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2889 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2895 Ok("0") => false,
2896 Ok(_) => true,
2897 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2898 };
2899 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2909 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2910 }
2911
2912 pub fn router_gemv_form(
2915 &self,
2916 w: &CudaSlice<f32>,
2917 x: &CudaSlice<f32>,
2918 n_embd: usize,
2919 n_experts: usize,
2920 t: usize,
2921 w8: bool,
2922 batch: bool,
2923 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2924 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2925 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2926 let f = if batch {
2927 self.func("router_gemv_f32_w8_batch")
2928 } else if w8 {
2929 self.func("router_gemv_f32_w8")
2930 } else {
2931 self.func("router_gemv_f32")
2932 };
2933 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2934 let cfg = if batch {
2935 LaunchConfig {
2936 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2937 block_dim: (32, 8, 1),
2938 shared_mem_bytes: 0,
2939 }
2940 } else {
2941 LaunchConfig {
2942 grid_dim: (n_experts as u32, t as u32, 1),
2943 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2944 shared_mem_bytes: 0,
2945 }
2946 };
2947 let __s_b = self.gpu.stream();
2948 let mut b = __s_b.launch_builder(&f);
2949 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2950 unsafe {
2951 b.launch(cfg)?;
2952 }
2953 Ok(y)
2954 }
2955
2956 pub fn router_gemv_into(
2959 &self,
2960 w: &CudaSlice<f32>,
2961 x: &CudaSlice<f32>,
2962 y: &mut CudaSlice<f32>,
2963 n_embd: usize,
2964 n_experts: usize,
2965 t: usize,
2966 ) -> Result<(), Box<dyn std::error::Error>> {
2967 if y.len() < t * n_experts {
2968 return Err("router_gemv_into output too small".into());
2969 }
2970 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2971 Ok("0") => false,
2972 Ok(_) => true,
2973 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2974 };
2975 let f = if w8 {
2976 self.func("router_gemv_f32_w8")
2977 } else {
2978 self.func("router_gemv_f32")
2979 };
2980 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2981 let cfg = LaunchConfig {
2982 grid_dim: (n_experts as u32, t as u32, 1),
2983 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2984 shared_mem_bytes: 0,
2985 };
2986 let __s_b = self.gpu.stream();
2987 let mut b = __s_b.launch_builder(&f);
2988 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2989 unsafe {
2990 b.launch(cfg)?;
2991 }
2992 Ok(())
2993 }
2994
2995 pub fn rows_permute(
2997 &self,
2998 src: &CudaSlice<f32>,
2999 idx: &CudaSlice<i32>,
3000 nrows: usize,
3001 ncols: usize,
3002 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3003 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3004 let f = self.func("rows_permute_f32");
3005 let (nc, nr) = (ncols as i32, nrows as i32);
3006 let cfg = LaunchConfig {
3007 grid_dim: (nrows as u32, 1, 1),
3008 block_dim: (256, 1, 1),
3009 shared_mem_bytes: 0,
3010 };
3011 let __s_b = self.gpu.stream();
3012 let mut b = __s_b.launch_builder(&f);
3013 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3014 unsafe {
3015 b.launch(cfg)?;
3016 }
3017 Ok(dst)
3018 }
3019
3020 pub fn sigmoid_dot_rows(
3025 &self,
3026 x: &CudaSlice<f32>,
3027 w: &CudaSlice<f32>,
3028 n_embd: usize,
3029 t: usize,
3030 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3031 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3034 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3035 let gs = self.linear(x, w, t, n_embd, 1)?;
3036 let mut g = self.uninit(t)?;
3037 self.sigmoid(&gs, &mut g, t)?;
3038 return Ok(g);
3039 }
3040 let mut g = self.alloc_uninit::<f32>(t)?;
3046 let f = self.func("sigmoid_dot_rows_f32");
3047 let (ne, ti) = (n_embd as i32, t as i32);
3048 let cfg = LaunchConfig {
3049 grid_dim: (t as u32, 1, 1),
3050 block_dim: (32, 8, 1),
3051 shared_mem_bytes: 0,
3052 };
3053 let __s_b = self.gpu.stream();
3054 let mut b = __s_b.launch_builder(&f);
3055 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3056 unsafe {
3057 b.launch(cfg)?;
3058 }
3059 Ok(g)
3060 }
3061
3062 pub fn sigmoid_dot_rows_into(
3064 &self,
3065 x: &CudaSlice<f32>,
3066 w: &CudaSlice<f32>,
3067 g: &mut CudaSlice<f32>,
3068 n_embd: usize,
3069 t: usize,
3070 ) -> Result<(), Box<dyn std::error::Error>> {
3071 if g.len() < t {
3072 return Err("sigmoid_dot_rows_into output too small".into());
3073 }
3074 let f = self.func("sigmoid_dot_rows_f32");
3075 let (ne, ti) = (n_embd as i32, t as i32);
3076 let cfg = LaunchConfig {
3077 grid_dim: (t as u32, 1, 1),
3078 block_dim: (32, 8, 1),
3079 shared_mem_bytes: 0,
3080 };
3081 let __s_b = self.gpu.stream();
3082 let mut b = __s_b.launch_builder(&f);
3083 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3084 unsafe {
3085 b.launch(cfg)?;
3086 }
3087 Ok(())
3088 }
3089
3090 pub fn spec_rollback_stream(
3092 &self,
3093 len_ptrs: &CudaSlice<u64>,
3094 pos_start: &CudaSlice<i32>,
3095 acc: &CudaSlice<u32>,
3096 base: usize,
3097 n_rows: usize,
3098 ) -> Result<(), Box<dyn std::error::Error>> {
3099 let f = self.func("spec_rollback_stream");
3100 let (b, nr) = (base as i32, n_rows as i32);
3101 let cfg = LaunchConfig {
3102 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3103 block_dim: (64, 1, 1),
3104 shared_mem_bytes: 0,
3105 };
3106 let __s_bl = self.gpu.stream();
3107 let mut bl = __s_bl.launch_builder(&f);
3108 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3109 unsafe {
3110 bl.launch(cfg)?;
3111 }
3112 Ok(())
3113 }
3114
3115 pub fn plain_tok_ring(
3117 &self,
3118 vam: &CudaSlice<u32>,
3119 pos_start: &CudaSlice<i32>,
3120 base: usize,
3121 ring: &mut CudaSlice<u32>,
3122 ) -> Result<(), Box<dyn std::error::Error>> {
3123 let f = self.func("plain_tok_ring");
3124 let (b, cap) = (base as i32, ring.len() as i32);
3125 let cfg = LaunchConfig {
3126 grid_dim: (1, 1, 1),
3127 block_dim: (32, 1, 1),
3128 shared_mem_bytes: 0,
3129 };
3130 let __s_bl = self.gpu.stream();
3131 let mut bl = __s_bl.launch_builder(&f);
3132 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3133 unsafe {
3134 bl.launch(cfg)?;
3135 }
3136 Ok(())
3137 }
3138
3139 pub fn spec_ring_commit(
3141 &self,
3142 vtok: &CudaSlice<u32>,
3143 acc: &CudaSlice<u32>,
3144 brk: &CudaSlice<u32>,
3145 ring: &mut CudaSlice<u32>,
3146 pend: &mut CudaSlice<u32>,
3147 ) -> Result<(), Box<dyn std::error::Error>> {
3148 let f = self.func("spec_ring_commit");
3149 let cfg = LaunchConfig {
3150 grid_dim: (1, 1, 1),
3151 block_dim: (32, 1, 1),
3152 shared_mem_bytes: 0,
3153 };
3154 let __s_b = self.gpu.stream();
3155 let mut b = __s_b.launch_builder(&f);
3156 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3157 unsafe {
3158 b.launch(cfg)?;
3159 }
3160 Ok(())
3161 }
3162 pub fn i32_copy_add(
3163 &self,
3164 src: &CudaSlice<i32>,
3165 dst: &mut CudaSlice<i32>,
3166 delta: i32,
3167 ) -> Result<(), Box<dyn std::error::Error>> {
3168 let f = self.func("i32_copy_add");
3169 let cfg = LaunchConfig {
3170 grid_dim: (1, 1, 1),
3171 block_dim: (32, 1, 1),
3172 shared_mem_bytes: 0,
3173 };
3174 let __s_b = self.gpu.stream();
3175 let mut b = __s_b.launch_builder(&f);
3176 b.arg(src).arg(dst).arg(&delta);
3177 unsafe {
3178 b.launch(cfg)?;
3179 }
3180 Ok(())
3181 }
3182 pub fn u32_copy(
3183 &self,
3184 src: &CudaSlice<u32>,
3185 dst: &mut CudaSlice<u32>,
3186 ) -> Result<(), Box<dyn std::error::Error>> {
3187 let f = self.func("u32_copy");
3188 let cfg = LaunchConfig {
3189 grid_dim: (1, 1, 1),
3190 block_dim: (32, 1, 1),
3191 shared_mem_bytes: 0,
3192 };
3193 let __s_b = self.gpu.stream();
3194 let mut b = __s_b.launch_builder(&f);
3195 b.arg(src).arg(dst);
3196 unsafe {
3197 b.launch(cfg)?;
3198 }
3199 Ok(())
3200 }
3201
3202 pub fn spec_adapt_k(
3206 &self,
3207 acc: &CudaSlice<u32>,
3208 brk: &mut CudaSlice<u32>,
3209 floor: usize,
3210 cap: usize,
3211 ) -> Result<(), Box<dyn std::error::Error>> {
3212 let f = self.func("spec_adapt_k");
3213 let (fl, cp) = (floor as i32, cap as i32);
3214 let cfg = LaunchConfig {
3215 grid_dim: (1, 1, 1),
3216 block_dim: (32, 1, 1),
3217 shared_mem_bytes: 0,
3218 };
3219 let __s_b = self.gpu.stream();
3220 let mut b = __s_b.launch_builder(&f);
3221 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3222 unsafe {
3223 b.launch(cfg)?;
3224 }
3225 Ok(())
3226 }
3227
3228 pub fn spec_accept_greedy_dc(
3230 &self,
3231 preds: &CudaSlice<u32>,
3232 vtok: &CudaSlice<u32>,
3233 last_pred: &CudaSlice<u32>,
3234 brk: &CudaSlice<u32>,
3235 out: &mut CudaSlice<u32>,
3236 ) -> Result<(), Box<dyn std::error::Error>> {
3237 let f = self.func("spec_accept_greedy_dc");
3238 let cfg = LaunchConfig {
3239 grid_dim: (1, 1, 1),
3240 block_dim: (32, 1, 1),
3241 shared_mem_bytes: 0,
3242 };
3243 let __s_b = self.gpu.stream();
3244 let mut b = __s_b.launch_builder(&f);
3245 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3246 unsafe {
3247 b.launch(cfg)?;
3248 }
3249 Ok(())
3250 }
3251
3252 pub fn pos_iota(
3254 &self,
3255 pos0: &CudaSlice<i32>,
3256 out: &mut CudaSlice<i32>,
3257 t: usize,
3258 ) -> Result<(), Box<dyn std::error::Error>> {
3259 let f = self.func("pos_iota_i32");
3260 let ti = t as i32;
3261 let cfg = LaunchConfig {
3262 grid_dim: (1, 1, 1),
3263 block_dim: (t.max(1) as u32, 1, 1),
3264 shared_mem_bytes: 0,
3265 };
3266 let __s_b = self.gpu.stream();
3267 let mut b = __s_b.launch_builder(&f);
3268 b.arg(pos0).arg(out).arg(&ti);
3269 unsafe {
3270 b.launch(cfg)?;
3271 }
3272 Ok(())
3273 }
3274 #[allow(clippy::too_many_arguments)]
3275 pub fn append_kv_quantized_rows_dc(
3276 &self,
3277 k_rows: &CudaSlice<f32>,
3278 v_rows: &CudaSlice<f32>,
3279 kc: &mut CudaSlice<u8>,
3280 vc: &mut CudaSlice<u8>,
3281 t0_dev: &CudaSlice<i32>,
3282 t: usize,
3283 kv_dim_k: usize,
3284 kv_dim_v: usize,
3285 k_tok_bytes: usize,
3286 v_tok_bytes: usize,
3287 g: bool,
3288 ) -> Result<(), Box<dyn std::error::Error>> {
3289 let f = if g {
3290 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3291 } else {
3292 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3293 };
3294 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3295 let cfg = LaunchConfig {
3296 grid_dim: (nblk, t as u32, 1),
3297 block_dim: (32, 1, 1),
3298 shared_mem_bytes: 0,
3299 };
3300 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3301 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3302 let __s_b = self.gpu.stream();
3303 let mut b = __s_b.launch_builder(&f);
3304 b.arg(k_rows)
3305 .arg(v_rows)
3306 .arg(kc)
3307 .arg(vc)
3308 .arg(t0_dev)
3309 .arg(&kdk)
3310 .arg(&kdv)
3311 .arg(&ktb)
3312 .arg(&vtb);
3313 unsafe {
3314 b.launch(cfg)?;
3315 }
3316 Ok(())
3317 }
3318
3319 #[allow(clippy::too_many_arguments)]
3322 pub fn append_kv_quantized_row_dc_inc(
3323 &self,
3324 k_row: &CudaSlice<f32>,
3325 v_row: &CudaSlice<f32>,
3326 kc: &mut CudaSlice<u8>,
3327 vc: &mut CudaSlice<u8>,
3328 t0_dev: &mut CudaSlice<i32>,
3329 kv_dim_k: usize,
3330 kv_dim_v: usize,
3331 k_tok_bytes: usize,
3332 v_tok_bytes: usize,
3333 g: bool,
3334 ) -> Result<(), Box<dyn std::error::Error>> {
3335 let f = if g {
3336 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3337 } else {
3338 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3339 };
3340 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3341 let cfg = LaunchConfig {
3342 grid_dim: (1, 1, 1),
3343 block_dim: (nthreads, 1, 1),
3344 shared_mem_bytes: 0,
3345 };
3346 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3347 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3348 let __s_b = self.gpu.stream();
3349 let mut b = __s_b.launch_builder(&f);
3350 b.arg(k_row)
3351 .arg(v_row)
3352 .arg(kc)
3353 .arg(vc)
3354 .arg(t0_dev)
3355 .arg(&kdk)
3356 .arg(&kdv)
3357 .arg(&ktb)
3358 .arg(&vtb);
3359 unsafe {
3360 b.launch(cfg)?;
3361 }
3362 Ok(())
3363 }
3364
3365 pub fn pack_tok_p(
3367 &self,
3368 tok: &CudaSlice<u32>,
3369 p: &CudaSlice<f32>,
3370 out: &mut CudaSlice<u32>,
3371 slot: usize,
3372 ) -> Result<(), Box<dyn std::error::Error>> {
3373 let f = self.func("pack_tok_p");
3374 let sl = slot as i32;
3375 let cfg = LaunchConfig {
3376 grid_dim: (1, 1, 1),
3377 block_dim: (32, 1, 1),
3378 shared_mem_bytes: 0,
3379 };
3380 let __s_b = self.gpu.stream();
3381 let mut b = __s_b.launch_builder(&f);
3382 b.arg(tok).arg(p).arg(out).arg(&sl);
3383 unsafe {
3384 b.launch(cfg)?;
3385 }
3386 Ok(())
3387 }
3388 pub fn tok_map_u32(
3389 &self,
3390 tok: &mut CudaSlice<u32>,
3391 map: &CudaSlice<u32>,
3392 ) -> Result<(), Box<dyn std::error::Error>> {
3393 let f = self.func("tok_map_u32");
3394 let cfg = LaunchConfig {
3395 grid_dim: (1, 1, 1),
3396 block_dim: (32, 1, 1),
3397 shared_mem_bytes: 0,
3398 };
3399 let __s_b = self.gpu.stream();
3400 let mut b = __s_b.launch_builder(&f);
3401 b.arg(tok).arg(map);
3402 unsafe {
3403 b.launch(cfg)?;
3404 }
3405 Ok(())
3406 }
3407
3408 #[allow(clippy::too_many_arguments)]
3410 pub fn spec_assemble_verify(
3411 &self,
3412 tokp: &CudaSlice<u32>,
3413 pend: &CudaSlice<u32>,
3414 d2t: Option<&CudaSlice<u32>>,
3415 vtok: &mut CudaSlice<u32>,
3416 brk: &mut CudaSlice<u32>,
3417 p_min: f32,
3418 k: usize,
3419 pmin0: bool,
3420 ) -> Result<(), Box<dyn std::error::Error>> {
3421 let f = self.func("spec_assemble_verify");
3422 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3423 let cfg = LaunchConfig {
3424 grid_dim: (1, 1, 1),
3425 block_dim: (32, 1, 1),
3426 shared_mem_bytes: 0,
3427 };
3428 let __s_b = self.gpu.stream();
3429 let mut b = __s_b.launch_builder(&f);
3430 match d2t {
3431 Some(m) => {
3432 b.arg(tokp)
3433 .arg(pend)
3434 .arg(m)
3435 .arg(vtok)
3436 .arg(brk)
3437 .arg(&p_min)
3438 .arg(&ki)
3439 .arg(&pm);
3440 unsafe {
3441 b.launch(cfg)?;
3442 }
3443 }
3444 None => {
3445 let null: u64 = 0;
3446 b.arg(tokp)
3447 .arg(pend)
3448 .arg(&null)
3449 .arg(vtok)
3450 .arg(brk)
3451 .arg(&p_min)
3452 .arg(&ki)
3453 .arg(&pm);
3454 unsafe {
3455 b.launch(cfg)?;
3456 }
3457 }
3458 }
3459 Ok(())
3460 }
3461
3462 #[allow(clippy::too_many_arguments)]
3464 pub fn ssm_conv_ring_rebuild_dc(
3465 &self,
3466 qkv_tm: &CudaSlice<f32>,
3467 ring_old: &CudaSlice<f32>,
3468 conv_state: &mut CudaSlice<f32>,
3469 conv_dim: usize,
3470 acc: &CudaSlice<u32>,
3471 base: usize,
3472 t_v: usize,
3473 d_conv: usize,
3474 ) -> Result<(), Box<dyn std::error::Error>> {
3475 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3476 let n = conv_dim * (d_conv - 1);
3477 let cfg = LaunchConfig::for_num_elems(n as u32);
3478 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3479 let __s_b = self.gpu.stream();
3480 let mut b = __s_b.launch_builder(&f);
3481 b.arg(qkv_tm)
3482 .arg(ring_old)
3483 .arg(conv_state)
3484 .arg(&cd)
3485 .arg(acc)
3486 .arg(&b0)
3487 .arg(&tv)
3488 .arg(&dc);
3489 unsafe {
3490 b.launch(cfg)?;
3491 }
3492 Ok(())
3493 }
3494 #[allow(clippy::too_many_arguments)]
3495 pub fn gdn_scan_s128_dc(
3496 &self,
3497 q: &CudaSlice<f32>,
3498 k: &CudaSlice<f32>,
3499 v: &CudaSlice<f32>,
3500 g: &CudaSlice<f32>,
3501 beta: &CudaSlice<f32>,
3502 state_in: &CudaSlice<f32>,
3503 state_out: &mut CudaSlice<f32>,
3504 o: &mut CudaSlice<f32>,
3505 n_head: usize,
3506 acc: &CudaSlice<u32>,
3507 base: usize,
3508 t_v: usize,
3509 scale: f32,
3510 ) -> Result<(), Box<dyn std::error::Error>> {
3511 let f = self.func("gdn_scan_s128_dc");
3512 const S_V: u32 = 128;
3513 const WARP: u32 = 32;
3514 const COLS_PER_BLOCK: u32 = 4;
3515 let cfg = LaunchConfig {
3516 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3517 block_dim: (WARP, COLS_PER_BLOCK, 1),
3518 shared_mem_bytes: 0,
3519 };
3520 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3521 let __s_b = self.gpu.stream();
3522 let mut b = __s_b.launch_builder(&f);
3523 b.arg(q)
3524 .arg(k)
3525 .arg(v)
3526 .arg(g)
3527 .arg(beta)
3528 .arg(state_in)
3529 .arg(state_out)
3530 .arg(o)
3531 .arg(&h)
3532 .arg(acc)
3533 .arg(&b0)
3534 .arg(&tv)
3535 .arg(&scale);
3536 unsafe {
3537 b.launch(cfg)?;
3538 }
3539 Ok(())
3540 }
3541
3542 pub fn spec_rollback_kv(
3544 &self,
3545 len_ptrs: &CudaSlice<u64>,
3546 saved: &CudaSlice<i32>,
3547 acc: &CudaSlice<u32>,
3548 base: usize,
3549 n_layer: usize,
3550 ) -> Result<(), Box<dyn std::error::Error>> {
3551 let f = self.func("spec_rollback_kv");
3552 let (b, nl) = (base as i32, n_layer as i32);
3553 let cfg = LaunchConfig {
3554 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3555 block_dim: (64, 1, 1),
3556 shared_mem_bytes: 0,
3557 };
3558 let __s_bl = self.gpu.stream();
3559 let mut bl = __s_bl.launch_builder(&f);
3560 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3561 unsafe {
3562 bl.launch(cfg)?;
3563 }
3564 Ok(())
3565 }
3566
3567 pub fn spec_fork_valid(
3569 &self,
3570 acc: &CudaSlice<u32>,
3571 optimistic_pending: u32,
3572 valid: &mut CudaSlice<u32>,
3573 ) -> Result<(), Box<dyn std::error::Error>> {
3574 let f = self.func("spec_fork_valid");
3575 let cfg = LaunchConfig {
3576 grid_dim: (1, 1, 1),
3577 block_dim: (1, 1, 1),
3578 shared_mem_bytes: 0,
3579 };
3580 let __s_bl = self.gpu.stream();
3581 let mut bl = __s_bl.launch_builder(&f);
3582 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3583 unsafe {
3584 bl.launch(cfg)?;
3585 }
3586 Ok(())
3587 }
3588
3589 pub fn spec_fork_reconcile_kv(
3591 &self,
3592 len_ptrs: &CudaSlice<u64>,
3593 saved: &CudaSlice<i32>,
3594 acc: &CudaSlice<u32>,
3595 valid: &CudaSlice<u32>,
3596 base: usize,
3597 n_layer: usize,
3598 ) -> Result<(), Box<dyn std::error::Error>> {
3599 let f = self.func("spec_fork_reconcile_kv");
3600 let (b, nl) = (base as i32, n_layer as i32);
3601 let cfg = LaunchConfig {
3602 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3603 block_dim: (64, 1, 1),
3604 shared_mem_bytes: 0,
3605 };
3606 let __s_bl = self.gpu.stream();
3607 let mut bl = __s_bl.launch_builder(&f);
3608 bl.arg(len_ptrs)
3609 .arg(saved)
3610 .arg(acc)
3611 .arg(valid)
3612 .arg(&b)
3613 .arg(&nl);
3614 unsafe {
3615 bl.launch(cfg)?;
3616 }
3617 Ok(())
3618 }
3619
3620 pub fn spec_fork_restore_f32(
3622 &self,
3623 snapshot: &CudaSlice<f32>,
3624 state: &mut CudaSlice<f32>,
3625 valid: &CudaSlice<u32>,
3626 ) -> Result<(), Box<dyn std::error::Error>> {
3627 assert_eq!(
3628 snapshot.len(),
3629 state.len(),
3630 "fork recurrent snapshot shape mismatch"
3631 );
3632 let f = self.func("spec_fork_restore_f32");
3633 let n = state.len() as i32;
3634 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3635 let cfg = LaunchConfig {
3636 grid_dim: (blocks, 1, 1),
3637 block_dim: (256, 1, 1),
3638 shared_mem_bytes: 0,
3639 };
3640 let __s_bl = self.gpu.stream();
3641 let mut bl = __s_bl.launch_builder(&f);
3642 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3643 unsafe {
3644 bl.launch(cfg)?;
3645 }
3646 Ok(())
3647 }
3648
3649 pub fn spec_seed_gather(
3652 &self,
3653 vx: &CudaSlice<f32>,
3654 fill_prev: &CudaSlice<f32>,
3655 acc: &CudaSlice<u32>,
3656 h_seed: &mut CudaSlice<f32>,
3657 base: usize,
3658 n_embd: usize,
3659 ) -> Result<(), Box<dyn std::error::Error>> {
3660 let f = self.func("spec_seed_gather");
3661 let (b, ne) = (base as i32, n_embd as i32);
3662 let cfg = LaunchConfig {
3663 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3664 block_dim: (256, 1, 1),
3665 shared_mem_bytes: 0,
3666 };
3667 let __s_bl = self.gpu.stream();
3668 let mut bl = __s_bl.launch_builder(&f);
3669 bl.arg(vx)
3670 .arg(fill_prev)
3671 .arg(acc)
3672 .arg(h_seed)
3673 .arg(&b)
3674 .arg(&ne);
3675 unsafe {
3676 bl.launch(cfg)?;
3677 }
3678 Ok(())
3679 }
3680
3681 pub fn spec_accept_greedy(
3683 &self,
3684 preds: &CudaSlice<u32>,
3685 draft: &CudaSlice<u32>,
3686 last_pred: u32,
3687 base: usize,
3688 k_round: usize,
3689 out: &mut CudaSlice<u32>,
3690 ) -> Result<(), Box<dyn std::error::Error>> {
3691 let f = self.func("spec_accept_greedy");
3692 let (b, k) = (base as i32, k_round as i32);
3693 let cfg = LaunchConfig {
3694 grid_dim: (1, 1, 1),
3695 block_dim: (32, 1, 1),
3696 shared_mem_bytes: 0,
3697 };
3698 let __s_bl = self.gpu.stream();
3699 let mut bl = __s_bl.launch_builder(&f);
3700 bl.arg(preds)
3701 .arg(draft)
3702 .arg(&last_pred)
3703 .arg(&b)
3704 .arg(&k)
3705 .arg(out);
3706 unsafe {
3707 bl.launch(cfg)?;
3708 }
3709 Ok(())
3710 }
3711
3712 pub fn gumbel_perturb(
3719 &self,
3720 x: &CudaSlice<f32>,
3721 y: &mut CudaSlice<f32>,
3722 n: usize,
3723 seed: u64,
3724 stream_pos: u32,
3725 temp: f32,
3726 ) -> Result<(), Box<dyn std::error::Error>> {
3727 let f = self.func("gumbel_perturb_f32");
3728 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3729 let cfg = LaunchConfig {
3730 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3731 block_dim: (256, 1, 1),
3732 shared_mem_bytes: 0,
3733 };
3734 let __s_b = self.gpu.stream();
3735 let mut b = __s_b.launch_builder(&f);
3736 b.arg(x)
3737 .arg(&mut *y)
3738 .arg(&ni)
3739 .arg(&slo)
3740 .arg(&shi)
3741 .arg(&stream_pos)
3742 .arg(&temp);
3743 unsafe {
3744 b.launch(cfg)?;
3745 }
3746 Ok(())
3747 }
3748
3749 pub fn mask_logits_col(
3757 &self,
3758 logits: &mut CudaSlice<f32>,
3759 mask: &CudaSlice<u32>,
3760 col: usize,
3761 n: usize,
3762 mask_words: usize,
3763 ) -> Result<(), Box<dyn std::error::Error>> {
3764 let f = self.func("mask_logits_f32");
3765 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3766 let cfg = LaunchConfig {
3767 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3768 block_dim: (256, 1, 1),
3769 shared_mem_bytes: 0,
3770 };
3771 let __s_b = self.gpu.stream();
3772 let mut b = __s_b.launch_builder(&f);
3773 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3774 unsafe {
3775 b.launch(cfg)?;
3776 }
3777 Ok(())
3778 }
3779
3780 pub fn gumbel_perturb_col(
3787 &self,
3788 x: &CudaSlice<f32>,
3789 col: usize,
3790 y: &mut CudaSlice<f32>,
3791 n: usize,
3792 seed: u64,
3793 stream_pos: u32,
3794 temp: f32,
3795 ) -> Result<(), Box<dyn std::error::Error>> {
3796 let f = self.func("gumbel_perturb_f32");
3797 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3798 let col_view = x.slice(col * n..(col + 1) * n);
3799 let cfg = LaunchConfig {
3800 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3801 block_dim: (256, 1, 1),
3802 shared_mem_bytes: 0,
3803 };
3804 let __s_b = self.gpu.stream();
3805 let mut b = __s_b.launch_builder(&f);
3806 b.arg(&col_view)
3807 .arg(&mut *y)
3808 .arg(&ni)
3809 .arg(&slo)
3810 .arg(&shi)
3811 .arg(&stream_pos)
3812 .arg(&temp);
3813 unsafe {
3814 b.launch(cfg)?;
3815 }
3816 Ok(())
3817 }
3818
3819 #[allow(clippy::too_many_arguments)]
3825 pub fn gumbel_perturb_filtered_col(
3826 &self,
3827 x: &CudaSlice<f32>,
3828 col: usize,
3829 y: &mut CudaSlice<f32>,
3830 n: usize,
3831 seed: u64,
3832 stream_pos: u32,
3833 temp: f32,
3834 stat_max: &CudaSlice<f32>,
3835 stat_th: &CudaSlice<f32>,
3836 stat_idx: usize,
3837 ) -> Result<(), Box<dyn std::error::Error>> {
3838 let f = self.func("gumbel_perturb_filtered_col_f32");
3839 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3840 let (ci, si) = (col as i32, stat_idx as i32);
3841 let cfg = LaunchConfig {
3842 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3843 block_dim: (256, 1, 1),
3844 shared_mem_bytes: 0,
3845 };
3846 let __s_b = self.gpu.stream();
3847 let mut b = __s_b.launch_builder(&f);
3848 b.arg(x)
3849 .arg(&ci)
3850 .arg(&mut *y)
3851 .arg(&ni)
3852 .arg(&slo)
3853 .arg(&shi)
3854 .arg(&stream_pos)
3855 .arg(&temp)
3856 .arg(stat_max)
3857 .arg(stat_th)
3858 .arg(&si);
3859 unsafe {
3860 b.launch(cfg)?;
3861 }
3862 Ok(())
3863 }
3864
3865 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3870 let f = self.func("memra_sctr_inc");
3871 let cfg = LaunchConfig {
3872 grid_dim: (1, 1, 1),
3873 block_dim: (1, 1, 1),
3874 shared_mem_bytes: 0,
3875 };
3876 let __s_b = self.gpu.stream();
3877 let mut b = __s_b.launch_builder(&f);
3878 b.arg(&mut *ctr);
3879 unsafe {
3880 b.launch(cfg)?;
3881 }
3882 Ok(())
3883 }
3884
3885 pub fn gumbel_perturb_ctr(
3890 &self,
3891 x: &CudaSlice<f32>,
3892 y: &mut CudaSlice<f32>,
3893 n: usize,
3894 seed: u64,
3895 ctr: &CudaSlice<u32>,
3896 temp: f32,
3897 ) -> Result<(), Box<dyn std::error::Error>> {
3898 let f = self.func("gumbel_perturb_ctr_f32");
3899 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3900 let cfg = LaunchConfig {
3901 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3902 block_dim: (256, 1, 1),
3903 shared_mem_bytes: 0,
3904 };
3905 let __s_b = self.gpu.stream();
3906 let mut b = __s_b.launch_builder(&f);
3907 b.arg(x)
3908 .arg(&mut *y)
3909 .arg(&ni)
3910 .arg(&slo)
3911 .arg(&shi)
3912 .arg(ctr)
3913 .arg(&temp);
3914 unsafe {
3915 b.launch(cfg)?;
3916 }
3917 Ok(())
3918 }
3919
3920 pub fn softmax_gather(
3924 &self,
3925 x: &CudaSlice<f32>,
3926 row_stride: usize,
3927 ids: &CudaSlice<u32>,
3928 rows: &CudaSlice<i32>,
3929 out: &mut CudaSlice<f32>,
3930 n: usize,
3931 npair: usize,
3932 temp: f32,
3933 ) -> Result<(), Box<dyn std::error::Error>> {
3934 let f = self.func("softmax_gather_f32");
3935 let (ni, rs) = (n as i32, row_stride as i64);
3936 let np = npair as i32;
3937 let cfg = LaunchConfig {
3938 grid_dim: (npair as u32, 1, 1),
3939 block_dim: (256, 1, 1),
3940 shared_mem_bytes: 0,
3941 };
3942 let __s_b = self.gpu.stream();
3943 let mut b = __s_b.launch_builder(&f);
3944 b.arg(x)
3945 .arg(&rs)
3946 .arg(ids)
3947 .arg(rows)
3948 .arg(&mut *out)
3949 .arg(&ni)
3950 .arg(&np)
3951 .arg(&temp);
3952 unsafe {
3953 b.launch(cfg)?;
3954 }
3955 Ok(())
3956 }
3957
3958 pub fn residual_sample(
3962 &self,
3963 p: &CudaSlice<f32>,
3964 q: Option<&CudaSlice<f32>>,
3965 n: usize,
3966 temp: f32,
3967 seed: u64,
3968 stream_pos: u32,
3969 out_tok: &mut CudaSlice<u32>,
3970 ) -> Result<(), Box<dyn std::error::Error>> {
3971 let f = self.func("residual_sample_f32");
3972 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3973 let nth = 1024u32;
3974 let cfg = LaunchConfig {
3975 grid_dim: (1, 1, 1),
3976 block_dim: (nth, 1, 1),
3977 shared_mem_bytes: 0,
3978 };
3979 let has_q: i32 = q.is_some() as i32;
3980 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3982 let mut b = __s_b.launch_builder(&f);
3983 b.arg(p)
3984 .arg(qbuf)
3985 .arg(&has_q)
3986 .arg(&ni)
3987 .arg(&temp)
3988 .arg(&slo)
3989 .arg(&shi)
3990 .arg(&stream_pos)
3991 .arg(&mut *out_tok);
3992 unsafe {
3993 b.launch(cfg)?;
3994 }
3995 Ok(())
3996 }
3997
3998 pub fn with_moe_cache<R>(
4003 &self,
4004 max_block_bytes: usize,
4005 f: impl FnOnce(
4006 &mut crate::moe_cache::MoeSlotCache,
4007 &Engine,
4008 ) -> Result<R, Box<dyn std::error::Error>>,
4009 ) -> Result<R, Box<dyn std::error::Error>> {
4010 let mut guard = self.moe_cache.lock().unwrap();
4011 if guard.is_none() {
4012 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4013 }
4014 let cache = guard.as_mut().unwrap();
4015 f(cache, self)
4016 }
4017
4018 pub fn freeze_moe_cache(&self) {
4021 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4022 cache.freeze();
4023 }
4024 }
4025
4026 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4029 self.moe_cache
4030 .lock()
4031 .unwrap()
4032 .as_ref()
4033 .map(crate::moe_cache::MoeSlotCache::export_residency)
4034 }
4035
4036 pub(crate) fn moe_cache_frozen(&self) -> bool {
4037 self.moe_cache
4038 .lock()
4039 .unwrap()
4040 .as_ref()
4041 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4042 }
4043
4044 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4051 crate::cpu_experts::configured()
4052 && self.moe_cache_frozen()
4053 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4054 }
4055
4056 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4058 assert!(
4059 self.moe_cache.lock().unwrap().is_none(),
4060 "MoE cache layout configured after cache construction"
4061 );
4062 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4063 }
4064
4065 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4066 self.moe_cache_layout.lock().unwrap().clone()
4067 }
4068
4069 pub fn moe_cache_enabled() -> bool {
4071 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4072 }
4073
4074 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4077 let guard = self.moe_cache.lock().unwrap();
4078 guard
4079 .as_ref()
4080 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4081 }
4082
4083 pub fn cpu_expert_stats(
4087 &self,
4088 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4089 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4090 }
4091
4092 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4095 crate::cpu_experts::predictor_stats()
4096 }
4097
4098 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4099 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4100 }
4101
4102 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4105 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4106 }
4107
4108 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4111 let guard = self.moe_cache.lock().unwrap();
4112 guard
4113 .as_ref()
4114 .and_then(|cache| cache.pread_stats())
4115 .map(|stats| {
4116 (
4117 stats.reads,
4118 stats.bytes,
4119 stats.read_errors,
4120 stats.short_reads,
4121 stats.fallbacks,
4122 stats.buffer_waits,
4123 stats.ring_full,
4124 )
4125 })
4126 }
4127
4128 pub fn spill_config_fallbacks(&self) -> u64 {
4130 crate::spill_pread::config_fallbacks()
4131 }
4132
4133 pub fn moe_cache_reset_counters(&self) {
4135 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4136 c.reset_counters();
4137 }
4138 }
4139
4140 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4141 Ok(self.gpu.stream().clone_htod(v)?)
4142 }
4143
4144 pub fn htod_bytes_padded(
4148 &self,
4149 v: &[u8],
4150 pad: usize,
4151 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4152 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4153 {
4154 let mut view = d.slice_mut(0..v.len());
4155 self.gpu.stream().memcpy_htod(v, &mut view)?;
4156 }
4157 Ok(d)
4158 }
4159
4160 pub fn copy_into(
4162 &self,
4163 dst: &mut CudaSlice<f32>,
4164 off: usize,
4165 src: &CudaSlice<f32>,
4166 len: usize,
4167 ) -> Result<(), Box<dyn std::error::Error>> {
4168 let mut view = dst.slice_mut(off..off + len);
4169 self.gpu
4170 .stream()
4171 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4172 Ok(())
4173 }
4174
4175 pub fn copy_u8_into(
4178 &self,
4179 dst: &mut CudaSlice<u8>,
4180 off: usize,
4181 src: &CudaSlice<u8>,
4182 len: usize,
4183 ) -> Result<(), Box<dyn std::error::Error>> {
4184 let mut view = dst.slice_mut(off..off + len);
4185 self.gpu
4186 .stream()
4187 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4188 Ok(())
4189 }
4190
4191 pub fn copy_u8_range_into(
4193 &self,
4194 dst: &mut CudaSlice<u8>,
4195 dst_off: usize,
4196 src: &CudaSlice<u8>,
4197 src_off: usize,
4198 len: usize,
4199 ) -> Result<(), Box<dyn std::error::Error>> {
4200 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
4201 self.gpu
4202 .stream()
4203 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4204 Ok(())
4205 }
4206
4207 pub fn prepare_kv_append(
4211 &self,
4212 kv: &mut crate::cache::KvLayer,
4213 retain_from: usize,
4214 append_rows: usize,
4215 ) -> Result<usize, Box<dyn std::error::Error>> {
4216 let Some(plan) = kv
4217 .ring
4218 .as_ref()
4219 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4220 .transpose()?
4221 else {
4222 return Ok(kv.len);
4223 };
4224 match plan {
4225 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4226 crate::cache::KvRingAppend::Rebase {
4227 src_row,
4228 keep_rows,
4229 new_base,
4230 write_row,
4231 } => {
4232 if keep_rows > 0 {
4233 let k_len = keep_rows * kv.k_tok_bytes;
4234 let v_len = keep_rows * kv.v_tok_bytes;
4235 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4236 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4237 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4238 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4239 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4240 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4241 }
4242 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4243 Ok(write_row)
4244 }
4245 }
4246 }
4247
4248 pub fn htod_u8_into(
4251 &self,
4252 dst: &mut CudaSlice<u8>,
4253 off: usize,
4254 src: &[u8],
4255 ) -> Result<(), Box<dyn std::error::Error>> {
4256 let mut view = dst.slice_mut(off..off + src.len());
4257 self.gpu.stream().memcpy_htod(src, &mut view)?;
4258 Ok(())
4259 }
4260
4261 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4262 b.slice(0..len)
4263 }
4264
4265 pub fn view_u8_range<'a>(
4268 &self,
4269 b: &'a CudaSlice<u8>,
4270 start: usize,
4271 end: usize,
4272 ) -> cudarc::driver::CudaView<'a, u8> {
4273 b.slice(start..end)
4274 }
4275 pub fn view_u8<'a>(
4276 &self,
4277 b: &'a CudaSlice<u8>,
4278 len: usize,
4279 ) -> cudarc::driver::CudaView<'a, u8> {
4280 b.slice(0..len)
4281 }
4282
4283 pub fn append_kv_quantized(
4287 &self,
4288 k_row: &CudaSlice<f32>,
4289 v_row: &CudaSlice<f32>,
4290 kc: &mut CudaSlice<u8>,
4291 vc: &mut CudaSlice<u8>,
4292 t: usize,
4293 kv_dim_k: usize,
4294 kv_dim_v: usize,
4295 k_tok_bytes: usize,
4296 v_tok_bytes: usize,
4297 g: bool,
4298 ) -> Result<(), Box<dyn std::error::Error>> {
4299 let f = if g {
4300 self.func_g("append_quantize_kv_q8_0_q5_1")
4301 } else {
4302 self.func("append_quantize_kv_q8_0_q5_1")
4303 };
4304 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4305 let cfg = LaunchConfig {
4306 grid_dim: (nblk, 1, 1),
4307 block_dim: (32, 1, 1),
4308 shared_mem_bytes: 0,
4309 };
4310 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4311 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4312 let __s_b = self.gpu.stream();
4313 let mut b = __s_b.launch_builder(&f);
4314 b.arg(k_row)
4315 .arg(v_row)
4316 .arg(kc)
4317 .arg(vc)
4318 .arg(&ti)
4319 .arg(&kdk)
4320 .arg(&kdv)
4321 .arg(&ktb)
4322 .arg(&vtb);
4323 unsafe {
4324 b.launch(cfg)?;
4325 }
4326 Ok(())
4327 }
4328
4329 pub fn append_kv_quantized_dc(
4333 &self,
4334 k_row: &CudaSlice<f32>,
4335 v_row: &CudaSlice<f32>,
4336 kc: &mut CudaSlice<u8>,
4337 vc: &mut CudaSlice<u8>,
4338 t_dev: &CudaSlice<i32>,
4339 kv_dim_k: usize,
4340 kv_dim_v: usize,
4341 k_tok_bytes: usize,
4342 v_tok_bytes: usize,
4343 g: bool,
4344 ) -> Result<(), Box<dyn std::error::Error>> {
4345 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4346 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4347 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4348 if Self::pdl_on() && Self::pdl_wb_on() {
4350 use cudarc::driver::{DevicePtr, DevicePtrMut};
4351 let s = &self.gpu.stream();
4352 let (pk, _g0) = k_row.device_ptr(s);
4353 let (pv, _g1) = v_row.device_ptr(s);
4354 let (pkc, _g2) = kc.device_ptr_mut(s);
4355 let (pvc, _g3) = vc.device_ptr_mut(s);
4356 let (pt, _g4) = t_dev.device_ptr(s);
4357 let mut ps = [
4358 &pk as *const _ as *mut std::ffi::c_void,
4359 &pv as *const _ as *mut _,
4360 &pkc as *const _ as *mut _,
4361 &pvc as *const _ as *mut _,
4362 &pt as *const _ as *mut _,
4363 &kdk as *const _ as *mut _,
4364 &kdv as *const _ as *mut _,
4365 &ktb as *const _ as *mut _,
4366 &vtb as *const _ as *mut _,
4367 ];
4368 unsafe {
4369 self.launch_pdl_flash(
4370 g,
4371 "append_quantize_kv_q8_0_q5_1_dc",
4372 (nblk, 1, 1),
4373 (32, 1, 1),
4374 0,
4375 &mut ps,
4376 )?;
4377 }
4378 return Ok(());
4379 }
4380 let f = if g {
4381 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4382 } else {
4383 self.func("append_quantize_kv_q8_0_q5_1_dc")
4384 };
4385 let cfg = LaunchConfig {
4386 grid_dim: (nblk, 1, 1),
4387 block_dim: (32, 1, 1),
4388 shared_mem_bytes: 0,
4389 };
4390 let __s_b = self.gpu.stream();
4391 let mut b = __s_b.launch_builder(&f);
4392 b.arg(k_row)
4393 .arg(v_row)
4394 .arg(kc)
4395 .arg(vc)
4396 .arg(t_dev)
4397 .arg(&kdk)
4398 .arg(&kdv)
4399 .arg(&ktb)
4400 .arg(&vtb);
4401 unsafe {
4402 b.launch(cfg)?;
4403 }
4404 Ok(())
4405 }
4406
4407 #[allow(clippy::too_many_arguments)]
4414 pub fn append_kv_quantized_rows(
4415 &self,
4416 k_rows: &CudaSlice<f32>,
4417 v_rows: &CudaSlice<f32>,
4418 kc: &mut CudaSlice<u8>,
4419 vc: &mut CudaSlice<u8>,
4420 t0: usize,
4421 t: usize,
4422 kv_dim_k: usize,
4423 kv_dim_v: usize,
4424 k_tok_bytes: usize,
4425 v_tok_bytes: usize,
4426 g: bool,
4427 ) -> Result<(), Box<dyn std::error::Error>> {
4428 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4429 for i in 0..t {
4430 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4431 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4432 self.append_kv_quantized_view(
4433 &k_row,
4434 &v_row,
4435 kc,
4436 vc,
4437 t0 + i,
4438 kv_dim_k,
4439 kv_dim_v,
4440 k_tok_bytes,
4441 v_tok_bytes,
4442 g,
4443 )?;
4444 }
4445 return Ok(());
4446 }
4447 let f = if g {
4448 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4449 } else {
4450 self.func("append_quantize_kv_q8_0_q5_1_rows")
4451 };
4452 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4453 let cfg = LaunchConfig {
4454 grid_dim: (nblk, t as u32, 1),
4455 block_dim: (32, 1, 1),
4456 shared_mem_bytes: 0,
4457 };
4458 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4459 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4460 let __s_b = self.gpu.stream();
4461 let mut b = __s_b.launch_builder(&f);
4462 b.arg(k_rows)
4463 .arg(v_rows)
4464 .arg(kc)
4465 .arg(vc)
4466 .arg(&t0i)
4467 .arg(&kdk)
4468 .arg(&kdv)
4469 .arg(&ktb)
4470 .arg(&vtb);
4471 unsafe {
4472 b.launch(cfg)?;
4473 }
4474 Ok(())
4475 }
4476
4477 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4481 let f = self.func("inc_i32");
4482 let cfg = LaunchConfig {
4483 grid_dim: (1, 1, 1),
4484 block_dim: (1, 1, 1),
4485 shared_mem_bytes: 0,
4486 };
4487 let __s_b = self.gpu.stream();
4488 let mut b = __s_b.launch_builder(&f);
4489 b.arg(p);
4490 unsafe {
4491 b.launch(cfg)?;
4492 }
4493 Ok(())
4494 }
4495
4496 pub fn append_kv_quantized_view(
4499 &self,
4500 k_row: &cudarc::driver::CudaView<f32>,
4501 v_row: &cudarc::driver::CudaView<f32>,
4502 kc: &mut CudaSlice<u8>,
4503 vc: &mut CudaSlice<u8>,
4504 t: usize,
4505 kv_dim_k: usize,
4506 kv_dim_v: usize,
4507 k_tok_bytes: usize,
4508 v_tok_bytes: usize,
4509 g: bool,
4510 ) -> Result<(), Box<dyn std::error::Error>> {
4511 let stream = self.gpu.stream();
4512 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4513 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4514 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4515 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4516 let f = if g {
4517 self.func_g("append_quantize_kv_q8_0_q5_1")
4518 } else {
4519 self.func("append_quantize_kv_q8_0_q5_1")
4520 };
4521 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4522 let cfg = LaunchConfig {
4523 grid_dim: (nblk, 1, 1),
4524 block_dim: (32, 1, 1),
4525 shared_mem_bytes: 0,
4526 };
4527 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4528 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4529 let mut b = stream.launch_builder(&f);
4530 b.arg(k_row)
4531 .arg(v_row)
4532 .arg(kc)
4533 .arg(vc)
4534 .arg(&ti)
4535 .arg(&kdk)
4536 .arg(&kdv)
4537 .arg(&ktb)
4538 .arg(&vtb);
4539 unsafe {
4540 b.launch(cfg)?;
4541 }
4542 Ok(())
4543 }
4544
4545 pub fn copy_view_into(
4548 &self,
4549 dst: &mut CudaSlice<f32>,
4550 off: usize,
4551 src: &cudarc::driver::CudaView<f32>,
4552 len: usize,
4553 ) -> Result<(), Box<dyn std::error::Error>> {
4554 let mut view = dst.slice_mut(off..off + len);
4555 self.gpu
4556 .stream()
4557 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4558 Ok(())
4559 }
4560
4561 pub fn clone_dtod(
4565 &self,
4566 src: &CudaSlice<f32>,
4567 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4568 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4569 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4570 Ok(dst)
4571 }
4572
4573 pub fn dtod_copy_view(
4576 &self,
4577 src: &cudarc::driver::CudaView<f32>,
4578 dst: &mut CudaSlice<f32>,
4579 ) -> Result<(), Box<dyn std::error::Error>> {
4580 self.gpu.stream().memcpy_dtod(src, dst)?;
4581 Ok(())
4582 }
4583
4584 pub fn dtod_copy_view_i8(
4586 &self,
4587 src: &cudarc::driver::CudaView<i8>,
4588 dst: &mut CudaSlice<i8>,
4589 ) -> Result<(), Box<dyn std::error::Error>> {
4590 self.gpu.stream().memcpy_dtod(src, dst)?;
4591 Ok(())
4592 }
4593
4594 pub fn dtod_copy_into(
4596 &self,
4597 src: &CudaSlice<f32>,
4598 dst: &mut CudaSlice<f32>,
4599 offset: usize,
4600 ) -> Result<(), Box<dyn std::error::Error>> {
4601 let n = src.len();
4602 let mut dv = dst.slice_mut(offset..offset + n);
4603 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4604 Ok(())
4605 }
4606
4607 pub fn copy_batch_uniform_f32(
4613 &self,
4614 table: &CudaSlice<u64>,
4615 n: usize,
4616 words: usize,
4617 ) -> Result<(), Box<dyn std::error::Error>> {
4618 if n == 0 || words == 0 {
4619 return Ok(());
4620 }
4621 debug_assert!(
4622 table.len() >= 2 * n,
4623 "pointer table must hold n srcs + n dsts"
4624 );
4625 let f = self.func("copy_batch_uniform_f32");
4626 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4629 let (ni, wi) = (n as i32, words as i32);
4630 let cfg = LaunchConfig {
4631 grid_dim: (chunks, n as u32, 1),
4632 block_dim: (256, 1, 1),
4633 shared_mem_bytes: 0,
4634 };
4635 let __s = self.gpu.stream();
4636 let mut b = __s.launch_builder(&f);
4637 b.arg(table).arg(&ni).arg(&wi);
4638 unsafe {
4639 b.launch(cfg)?;
4640 }
4641 Ok(())
4642 }
4643
4644 pub fn htod_u64_into(
4647 &self,
4648 v: &[u64],
4649 dst: &mut CudaSlice<u64>,
4650 ) -> Result<(), Box<dyn std::error::Error>> {
4651 let mut view = dst.slice_mut(0..v.len());
4652 self.gpu.stream().memcpy_htod(v, &mut view)?;
4653 Ok(())
4654 }
4655
4656 pub fn copy_indirect_src_f32(
4661 &self,
4662 src_entry: &cudarc::driver::CudaView<u64>,
4663 dst: &mut CudaSlice<f32>,
4664 dst_off: usize,
4665 words: usize,
4666 ) -> Result<(), Box<dyn std::error::Error>> {
4667 let f = self.func("copy_indirect_src_f32");
4668 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4669 let wi = words as i32;
4670 let cfg = LaunchConfig {
4671 grid_dim: (chunks, 1, 1),
4672 block_dim: (256, 1, 1),
4673 shared_mem_bytes: 0,
4674 };
4675 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4676 let __s = self.gpu.stream();
4677 let mut b = __s.launch_builder(&f);
4678 b.arg(src_entry).arg(&mut dv).arg(&wi);
4679 unsafe {
4680 b.launch(cfg)?;
4681 }
4682 Ok(())
4683 }
4684
4685 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4687 self.alloc_uninit::<i8>(n)
4688 }
4689
4690 pub fn qmatvec(
4692 &self,
4693 w: &CudaSlice<u8>,
4694 x: &CudaSlice<f32>,
4695 m: usize,
4696 in_f: usize,
4697 out_f: usize,
4698 qtype: i32,
4699 row_bytes: usize,
4700 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4701 let f = self.func("qmatvec_f32");
4702 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4704 grid_dim: (out_f as u32, m as u32, 1),
4705 block_dim: (256, 1, 1),
4706 shared_mem_bytes: 0,
4707 };
4708 let (inf, outf, mi, qt, rb) =
4709 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4710 let __s_b = self.gpu.stream();
4711 let mut b = __s_b.launch_builder(&f);
4712 b.arg(w)
4713 .arg(x)
4714 .arg(&mut y)
4715 .arg(&inf)
4716 .arg(&outf)
4717 .arg(&mi)
4718 .arg(&qt)
4719 .arg(&rb);
4720 unsafe {
4721 b.launch(cfg)?;
4722 }
4723 Ok(y)
4724 }
4725
4726 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4728 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4729 self.keep_if_capturing(&s);
4730 Ok(s)
4731 }
4732
4733 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4737 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4738 self.keep_if_capturing(&s);
4739 Ok(s)
4740 }
4741
4742 pub fn memset_zeros_view(
4745 &self,
4746 dst: &mut cudarc::driver::CudaViewMut<f32>,
4747 ) -> Result<(), Box<dyn std::error::Error>> {
4748 self.gpu.stream().memset_zeros(dst)?;
4749 Ok(())
4750 }
4751
4752 pub fn stage_expert(
4758 &self,
4759 host_bytes: &[u8],
4760 scratch: &mut CudaSlice<u8>,
4761 off: usize,
4762 ) -> Result<(), Box<dyn std::error::Error>> {
4763 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4766 }
4767
4768 pub fn moe_router_topk(
4774 &self,
4775 logits: &CudaSlice<f32>,
4776 t: usize,
4777 n_expert: usize,
4778 n_used: usize,
4779 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4780 let f = self.func("moe_router_topk_f32");
4781 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 {
4784 grid_dim: (t as u32, 1, 1),
4785 block_dim: (n_expert as u32, 1, 1),
4786 shared_mem_bytes: 0,
4787 };
4788 let (ne, nu) = (n_expert as i32, n_used as i32);
4789 let __s_b = self.gpu.stream();
4790 let mut b = __s_b.launch_builder(&f);
4791 b.arg(logits)
4792 .arg(&mut sel_idx)
4793 .arg(&mut sel_w)
4794 .arg(&ne)
4795 .arg(&nu);
4796 unsafe {
4797 b.launch(cfg)?;
4798 }
4799 Ok((sel_idx, sel_w))
4800 }
4801
4802 pub fn moe_router_topk_scaled(
4805 &self,
4806 logits: &CudaSlice<f32>,
4807 t: usize,
4808 n_expert: usize,
4809 n_used: usize,
4810 ex_scale: &CudaSlice<f32>,
4811 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4812 let f = self.func("moe_router_topk_scaled_f32");
4817 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4818 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4819 let cfg = LaunchConfig {
4820 grid_dim: (t as u32, 1, 1),
4821 block_dim: (n_expert as u32, 1, 1),
4822 shared_mem_bytes: 0,
4823 };
4824 let (ne, nu) = (n_expert as i32, n_used as i32);
4825 let __s_b = self.gpu.stream();
4826 let mut b = __s_b.launch_builder(&f);
4827 b.arg(logits)
4828 .arg(&mut sel_idx)
4829 .arg(&mut sel_w)
4830 .arg(&ne)
4831 .arg(&nu)
4832 .arg(ex_scale);
4833 unsafe {
4834 b.launch(cfg)?;
4835 }
4836 Ok((sel_idx, sel_w))
4837 }
4838
4839 pub fn moe_router_topk_host(
4847 &self,
4848 logits: &CudaSlice<f32>,
4849 t: usize,
4850 n_expert: usize,
4851 n_used: usize,
4852 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4853 let f = self.func("moe_router_topk_f32");
4854 let n = t * n_used;
4855 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4856 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4857 let cfg = LaunchConfig {
4858 grid_dim: (t as u32, 1, 1),
4859 block_dim: (n_expert as u32, 1, 1),
4860 shared_mem_bytes: 0,
4861 };
4862 let (ne, nu) = (n_expert as i32, n_used as i32);
4863 let __s_b = self.gpu.stream();
4864 let mut b = __s_b.launch_builder(&f);
4865 b.arg(logits)
4866 .arg(&mut sel_idx)
4867 .arg(&mut sel_w)
4868 .arg(&ne)
4869 .arg(&nu);
4870 unsafe {
4871 b.launch(cfg)?;
4872 }
4873 let bytes = n * 8;
4875 let mut guard = self.router_stage.lock().unwrap();
4876 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4877 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4878 }
4879 let stage = guard.as_mut().unwrap();
4880 let (si, sw) = unsafe {
4881 (
4882 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4883 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4884 )
4885 };
4886 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()))
4890 }
4891
4892 #[allow(clippy::too_many_arguments)]
4896 pub fn moe_router_sigmoid_topk(
4897 &self,
4898 logits: &CudaSlice<f32>,
4899 t: usize,
4900 n_expert: usize,
4901 n_used: usize,
4902 active_count: usize,
4903 correction_bias: &CudaSlice<f32>,
4904 active: &CudaSlice<u8>,
4905 scaling_factor: f32,
4906 route_norm: bool,
4907 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4908 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4909 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4910 return Err(format!(
4911 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4912 )
4913 .into());
4914 }
4915 if logits.len() < t * n_expert
4916 || correction_bias.len() != n_expert
4917 || active.len() != n_expert
4918 {
4919 return Err(format!(
4920 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4921 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4922 ).into());
4923 }
4924 let f = self.func(crate::sigmoid_topk_kernel(
4925 crate::sig_expf_dev_on(),
4926 crate::topk_fast_on(),
4927 n_used,
4928 ));
4929 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4930 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4931 let threads = n_expert.div_ceil(32) * 32;
4932 let cfg = LaunchConfig {
4933 grid_dim: (t as u32, 1, 1),
4934 block_dim: (threads as u32, 1, 1),
4935 shared_mem_bytes: 0,
4936 };
4937 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4938 let __s_b = self.gpu.stream();
4939 let mut b = __s_b.launch_builder(&f);
4940 b.arg(logits)
4941 .arg(correction_bias)
4942 .arg(active)
4943 .arg(&mut sel_idx)
4944 .arg(&mut sel_w)
4945 .arg(&ne)
4946 .arg(&nu)
4947 .arg(&scaling_factor)
4948 .arg(&rn);
4949 unsafe {
4950 b.launch(cfg)?;
4951 }
4952 Ok((sel_idx, sel_w))
4953 }
4954
4955 #[allow(clippy::too_many_arguments)]
4958 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4962 if ptr == 0 {
4963 return Err("ring_flag_raw: unarmed flag".into());
4964 }
4965 let f = self.func("memra_ring_flag");
4966 let cfg = LaunchConfig {
4967 grid_dim: (1, 1, 1),
4968 block_dim: (32, 1, 1),
4969 shared_mem_bytes: 0,
4970 };
4971 let __s_b = self.gpu.stream();
4972 let mut b = __s_b.launch_builder(&f);
4973 b.arg(&ptr).arg(&value);
4974 unsafe {
4975 b.launch(cfg)?;
4976 }
4977 Ok(())
4978 }
4979
4980 pub fn moe_sel_w_mirror(
4983 &self,
4984 sel_src: &CudaSlice<i32>,
4985 w_src: &CudaSlice<f32>,
4986 sel_dst: &mut CudaSlice<i32>,
4987 w_dst: &mut CudaSlice<f32>,
4988 n: usize,
4989 ) -> Result<(), Box<dyn std::error::Error>> {
4990 if n == 0
4991 || n > 32
4992 || sel_src.len() < n
4993 || w_src.len() < n
4994 || sel_dst.len() < n
4995 || w_dst.len() < n
4996 {
4997 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4998 }
4999 let f = self.func("moe_sel_w_mirror");
5000 let cfg = LaunchConfig {
5001 grid_dim: (1, 1, 1),
5002 block_dim: (32, 1, 1),
5003 shared_mem_bytes: 0,
5004 };
5005 let ni = n as i32;
5006 let __s_b = self.gpu.stream();
5007 let mut b = __s_b.launch_builder(&f);
5008 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5009 unsafe {
5010 b.launch(cfg)?;
5011 }
5012 Ok(())
5013 }
5014
5015 pub fn moe_router_sigmoid_topk_into(
5016 &self,
5017 logits: &CudaSlice<f32>,
5018 t: usize,
5019 n_expert: usize,
5020 n_used: usize,
5021 active_count: usize,
5022 correction_bias: &CudaSlice<f32>,
5023 active: &CudaSlice<u8>,
5024 scaling_factor: f32,
5025 route_norm: bool,
5026 sel_idx: &mut CudaSlice<i32>,
5027 sel_w: &mut CudaSlice<f32>,
5028 ) -> Result<(), Box<dyn std::error::Error>> {
5029 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5030 if n_expert == 0
5031 || n_expert > 1024
5032 || n_used == 0
5033 || n_used > 32 || n_used > n_expert
5035 || logits.len() < t * n_expert
5036 || correction_bias.len() != n_expert
5037 || active.len() != n_expert
5038 || sel_idx.len() < t * n_used
5039 || sel_w.len() < t * n_used
5040 {
5041 return Err("sigmoid router _into geometry mismatch".into());
5042 }
5043 let f = self.func(crate::sigmoid_topk_kernel(
5044 crate::sig_expf_dev_on(),
5045 crate::topk_fast_on(),
5046 n_used,
5047 ));
5048 let threads = n_expert.div_ceil(32) * 32;
5049 let cfg = LaunchConfig {
5050 grid_dim: (t as u32, 1, 1),
5051 block_dim: (threads as u32, 1, 1),
5052 shared_mem_bytes: 0,
5053 };
5054 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5055 let __s_b = self.gpu.stream();
5056 let mut b = __s_b.launch_builder(&f);
5057 b.arg(logits)
5058 .arg(correction_bias)
5059 .arg(active)
5060 .arg(&mut *sel_idx)
5061 .arg(&mut *sel_w)
5062 .arg(&ne)
5063 .arg(&nu)
5064 .arg(&scaling_factor)
5065 .arg(&rn);
5066 unsafe {
5067 b.launch(cfg)?;
5068 }
5069 Ok(())
5070 }
5071
5072 #[allow(clippy::too_many_arguments)]
5075 pub fn moe_router_sigmoid_topk_host(
5076 &self,
5077 logits: &CudaSlice<f32>,
5078 t: usize,
5079 n_expert: usize,
5080 n_used: usize,
5081 active_count: usize,
5082 correction_bias: &CudaSlice<f32>,
5083 active: &CudaSlice<u8>,
5084 scaling_factor: f32,
5085 route_norm: bool,
5086 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5087 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5088 logits,
5089 t,
5090 n_expert,
5091 n_used,
5092 active_count,
5093 correction_bias,
5094 active,
5095 scaling_factor,
5096 route_norm,
5097 )?;
5098 let n = t * n_used;
5099 let bytes = n * 8;
5100 let mut guard = self.router_stage.lock().unwrap();
5101 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5102 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5103 }
5104 let stage = guard.as_mut().unwrap();
5105 let (si, sw) = unsafe {
5106 (
5107 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5108 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5109 )
5110 };
5111 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5112 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5113 self.gpu.stream().synchronize()?;
5114 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5115 }
5116
5117 pub fn stage_expert_async(
5121 &self,
5122 host_bytes: &[u8],
5123 scratch: &mut CudaSlice<u8>,
5124 off: usize,
5125 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5126 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5127 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5128 Ok(self.copy_stream.record_event(None)?)
5129 }
5130
5131 pub fn compute_wait(
5133 &self,
5134 ev: &cudarc::driver::CudaEvent,
5135 ) -> Result<(), Box<dyn std::error::Error>> {
5136 self.gpu.stream().wait(ev)?;
5137 Ok(())
5138 }
5139
5140 pub fn qmatvec_view(
5145 &self,
5146 w: &CudaSlice<u8>,
5147 range: std::ops::Range<usize>,
5148 x: &cudarc::driver::CudaView<f32>,
5149 m: usize,
5150 in_f: usize,
5151 out_f: usize,
5152 qtype: i32,
5153 row_bytes: usize,
5154 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5155 let f = self.func("qmatvec_f32");
5156 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
5159 grid_dim: (out_f as u32, m as u32, 1),
5160 block_dim: (256, 1, 1),
5161 shared_mem_bytes: 0,
5162 };
5163 let (inf, outf, mi, qt, rb) =
5164 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5165 let __s_b = self.gpu.stream();
5166 let mut b = __s_b.launch_builder(&f);
5167 b.arg(&wv)
5168 .arg(x)
5169 .arg(&mut y)
5170 .arg(&inf)
5171 .arg(&outf)
5172 .arg(&mi)
5173 .arg(&qt)
5174 .arg(&rb);
5175 unsafe {
5176 b.launch(cfg)?;
5177 }
5178 Ok(y)
5179 }
5180
5181 #[allow(clippy::too_many_arguments)]
5188 pub fn moe_gate_up_silu8_q8(
5192 &self,
5193 gp: WPtr8,
5194 up: WPtr8,
5195 aq: &CudaSlice<i8>,
5196 ad: &CudaSlice<f32>,
5197 in_f: usize,
5198 n_ff: usize,
5199 n_used: usize,
5200 qt_g: i32,
5201 qt_u: i32,
5202 rb_g: usize,
5203 rb_u: usize,
5204 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5205 let f = self.func("moe_gate_up_silu8_q8");
5206 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5207 let cfg = LaunchConfig {
5208 grid_dim: (n_ff as u32, n_used as u32, 1),
5209 block_dim: (32, 1, 1),
5210 shared_mem_bytes: 0,
5211 };
5212 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5213 let __s_b = self.gpu.stream();
5214 let mut b = __s_b.launch_builder(&f);
5215 b.arg(&gp)
5216 .arg(&up)
5217 .arg(aq)
5218 .arg(ad)
5219 .arg(&mut act)
5220 .arg(&inf)
5221 .arg(&nff)
5222 .arg(&qt_g)
5223 .arg(&qt_u)
5224 .arg(&rbg)
5225 .arg(&rbu);
5226 unsafe {
5227 b.launch(cfg)?;
5228 }
5229 Ok(act)
5230 }
5231
5232 #[allow(clippy::too_many_arguments)]
5233 pub fn moe_down8_fma_q8(
5234 &self,
5235 dp: WPtr8,
5236 w: F32x8,
5237 aq2: &CudaSlice<i8>,
5238 ad2: &CudaSlice<f32>,
5239 dst: &mut cudarc::driver::CudaViewMut<f32>,
5240 in_f: usize,
5241 out_f: usize,
5242 n_used: usize,
5243 qt: i32,
5244 rb: usize,
5245 ) -> Result<(), Box<dyn std::error::Error>> {
5246 let f = self.func("moe_down8_fma_q8");
5247 let cfg = LaunchConfig {
5248 grid_dim: (out_f as u32, 1, 1),
5249 block_dim: (32, 1, 1),
5250 shared_mem_bytes: 0,
5251 };
5252 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5253 let __s_b = self.gpu.stream();
5254 let mut b = __s_b.launch_builder(&f);
5255 b.arg(&dp)
5256 .arg(&w)
5257 .arg(aq2)
5258 .arg(ad2)
5259 .arg(dst)
5260 .arg(&inf)
5261 .arg(&outf)
5262 .arg(&nu)
5263 .arg(&qt)
5264 .arg(&rbi);
5265 unsafe {
5266 b.launch(cfg)?;
5267 }
5268 Ok(())
5269 }
5270
5271 pub fn qmatvec_expert_q8(
5273 &self,
5274 w: &CudaSlice<u8>,
5275 range: std::ops::Range<usize>,
5276 aq: &CudaSlice<i8>,
5277 ad: &CudaSlice<f32>,
5278 m: usize,
5279 in_f: usize,
5280 out_f: usize,
5281 qtype: i32,
5282 row_bytes: usize,
5283 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5284 let f = self.func("qmatvec_expert_q8");
5285 let wv = w.slice(range);
5286 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5287 const ROWS: u32 = 4; let cfg = LaunchConfig {
5289 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5290 block_dim: (32, ROWS, 1),
5291 shared_mem_bytes: 0,
5292 };
5293 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5294 let __s_b = self.gpu.stream();
5295 let mut b = __s_b.launch_builder(&f);
5296 b.arg(&wv)
5297 .arg(aq)
5298 .arg(ad)
5299 .arg(&mut y)
5300 .arg(&inf)
5301 .arg(&outf)
5302 .arg(&mi)
5303 .arg(&qtype)
5304 .arg(&rbi);
5305 unsafe {
5306 b.launch(cfg)?;
5307 }
5308 Ok(y)
5309 }
5310
5311 pub fn moe_gate_up_silu8(
5312 &self,
5313 gp: WPtr8,
5314 up: WPtr8,
5315 x: &cudarc::driver::CudaView<f32>,
5316 in_f: usize,
5317 n_ff: usize,
5318 n_used: usize,
5319 qt_g: i32,
5320 qt_u: i32,
5321 rb_g: usize,
5322 rb_u: usize,
5323 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5324 let f = self.func("moe_gate_up_silu8_f32");
5325 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5327 grid_dim: (n_ff as u32, n_used as u32, 1),
5328 block_dim: (256, 1, 1),
5329 shared_mem_bytes: 0,
5330 };
5331 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5332 let __s_b = self.gpu.stream();
5333 let mut b = __s_b.launch_builder(&f);
5334 b.arg(&gp)
5335 .arg(&up)
5336 .arg(x)
5337 .arg(&mut act)
5338 .arg(&inf)
5339 .arg(&nff)
5340 .arg(&qt_g)
5341 .arg(&qt_u)
5342 .arg(&rbg)
5343 .arg(&rbu);
5344 unsafe {
5345 b.launch(cfg)?;
5346 }
5347 Ok(act)
5348 }
5349
5350 #[allow(clippy::too_many_arguments)]
5356 pub fn moe_down8_fma_into(
5357 &self,
5358 dp: WPtr8,
5359 w: F32x8,
5360 act: &CudaSlice<f32>,
5361 dst: &mut cudarc::driver::CudaViewMut<f32>,
5362 in_f: usize,
5363 out_f: usize,
5364 n_used: usize,
5365 qt: i32,
5366 rb: usize,
5367 ) -> Result<(), Box<dyn std::error::Error>> {
5368 let f = self.func("moe_down8_fma_f32");
5369 let cfg = LaunchConfig {
5370 grid_dim: (out_f as u32, 1, 1),
5371 block_dim: (256, 1, 1),
5372 shared_mem_bytes: 0,
5373 };
5374 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5375 let __s_b = self.gpu.stream();
5376 let mut b = __s_b.launch_builder(&f);
5377 b.arg(&dp)
5378 .arg(&w)
5379 .arg(act)
5380 .arg(dst)
5381 .arg(&inf)
5382 .arg(&outf)
5383 .arg(&nu)
5384 .arg(&qt)
5385 .arg(&rbv);
5386 unsafe {
5387 b.launch(cfg)?;
5388 }
5389 Ok(())
5390 }
5391
5392 #[allow(clippy::too_many_arguments)]
5397 #[allow(clippy::too_many_arguments)]
5412 #[allow(clippy::too_many_arguments)]
5414 pub fn moe_pairs_matvec_q8(
5415 &self,
5416 table: &CudaSlice<u64>,
5417 proj: i32,
5418 pair_tok: &CudaSlice<i32>,
5419 pair_ex: &CudaSlice<i32>,
5420 aq: &CudaSlice<i8>,
5421 ad: &CudaSlice<f32>,
5422 in_f: usize,
5423 out_f: usize,
5424 n_expert: usize,
5425 n_pairs: usize,
5426 qtype: i32,
5427 row_bytes: usize,
5428 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5429 let f = self.func("moe_pairs_matvec_q8");
5430 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5431 const ROWS: u32 = 4;
5432 let cfg = LaunchConfig {
5433 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5434 block_dim: (32, ROWS, 1),
5435 shared_mem_bytes: 0,
5436 };
5437 let (inf, outf, ne, np, rbi) = (
5438 in_f as i32,
5439 out_f as i32,
5440 n_expert as i32,
5441 n_pairs as i32,
5442 row_bytes as i64,
5443 );
5444 let __s_b = self.gpu.stream();
5445 let mut b = __s_b.launch_builder(&f);
5446 b.arg(table)
5447 .arg(&proj)
5448 .arg(pair_tok)
5449 .arg(pair_ex)
5450 .arg(aq)
5451 .arg(ad)
5452 .arg(&mut y)
5453 .arg(&inf)
5454 .arg(&outf)
5455 .arg(&ne)
5456 .arg(&np)
5457 .arg(&qtype)
5458 .arg(&rbi);
5459 unsafe {
5460 b.launch(cfg)?;
5461 }
5462 Ok(y)
5463 }
5464
5465 #[allow(clippy::too_many_arguments)]
5467 pub fn moe_pairs_matvec_q8_em(
5468 &self,
5469 table: &CudaSlice<u64>,
5470 proj: i32,
5471 ex_ids: &CudaSlice<i32>,
5472 ex_off: &CudaSlice<i32>,
5473 ex_pairs: &CudaSlice<i32>,
5474 pair_tok: &CudaSlice<i32>,
5475 aq: &CudaSlice<i8>,
5476 ad: &CudaSlice<f32>,
5477 in_f: usize,
5478 out_f: usize,
5479 n_expert: usize,
5480 n_active: usize,
5481 n_pairs: usize,
5482 qtype: i32,
5483 row_bytes: usize,
5484 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5485 let f = self.func("moe_pairs_matvec_q8_em");
5486 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5487 const ROWS: u32 = 4;
5488 let cfg = LaunchConfig {
5489 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5490 block_dim: (32, ROWS, 1),
5491 shared_mem_bytes: 0,
5492 };
5493 let (inf, outf, ne, na, rbi) = (
5494 in_f as i32,
5495 out_f as i32,
5496 n_expert as i32,
5497 n_active as i32,
5498 row_bytes as i64,
5499 );
5500 let __s_b = self.gpu.stream();
5501 let mut b = __s_b.launch_builder(&f);
5502 b.arg(table)
5503 .arg(&proj)
5504 .arg(ex_ids)
5505 .arg(ex_off)
5506 .arg(ex_pairs)
5507 .arg(pair_tok)
5508 .arg(aq)
5509 .arg(ad)
5510 .arg(&mut y)
5511 .arg(&inf)
5512 .arg(&outf)
5513 .arg(&ne)
5514 .arg(&na)
5515 .arg(&qtype)
5516 .arg(&rbi);
5517 unsafe {
5518 b.launch(cfg)?;
5519 }
5520 Ok(y)
5521 }
5522
5523 #[allow(clippy::too_many_arguments)]
5526 pub fn moe_pairs_matvec_q8_dec(
5527 &self,
5528 table: &CudaSlice<u64>,
5529 proj: i32,
5530 ex_ids: &CudaSlice<i32>,
5531 ex_off: &CudaSlice<i32>,
5532 ex_pairs: &CudaSlice<i32>,
5533 pair_tok: &CudaSlice<i32>,
5534 aq: &CudaSlice<i8>,
5535 ad: &CudaSlice<f32>,
5536 in_f: usize,
5537 out_f: usize,
5538 n_expert: usize,
5539 n_active: usize,
5540 n_pairs: usize,
5541 qtype: i32,
5542 row_bytes: usize,
5543 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5544 let f = self.func("moe_pairs_matvec_q8_dec");
5545 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5546 const ROWS: u32 = 4;
5547 let cfg = LaunchConfig {
5548 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5549 block_dim: (32, ROWS, 1),
5550 shared_mem_bytes: 0,
5551 };
5552 let (inf, outf, ne, na, rbi) = (
5553 in_f as i32,
5554 out_f as i32,
5555 n_expert as i32,
5556 n_active as i32,
5557 row_bytes as i64,
5558 );
5559 let __s_b = self.gpu.stream();
5560 let mut b = __s_b.launch_builder(&f);
5561 b.arg(table)
5562 .arg(&proj)
5563 .arg(ex_ids)
5564 .arg(ex_off)
5565 .arg(ex_pairs)
5566 .arg(pair_tok)
5567 .arg(aq)
5568 .arg(ad)
5569 .arg(&mut y)
5570 .arg(&inf)
5571 .arg(&outf)
5572 .arg(&ne)
5573 .arg(&na)
5574 .arg(&qtype)
5575 .arg(&rbi);
5576 unsafe {
5577 b.launch(cfg)?;
5578 }
5579 Ok(y)
5580 }
5581
5582 pub fn moe_pairs_gelu_mul(
5583 &self,
5584 gate: &CudaSlice<f32>,
5585 up: &CudaSlice<f32>,
5586 n: usize,
5587 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5588 let f = self.func("moe_pairs_gelu_mul");
5589 let mut act = self.alloc_uninit::<f32>(n)?;
5590 let cfg = LaunchConfig::for_num_elems(n as u32);
5591 let nl = n as i64;
5592 let __s_b = self.gpu.stream();
5593 let mut b = __s_b.launch_builder(&f);
5594 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5595 unsafe {
5596 b.launch(cfg)?;
5597 }
5598 Ok(act)
5599 }
5600
5601 pub fn moe_pairs_silu_mul(
5602 &self,
5603 gate: &CudaSlice<f32>,
5604 up: &CudaSlice<f32>,
5605 n: usize,
5606 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5607 let f = self.func("moe_pairs_silu_mul");
5608 let mut act = self.alloc_uninit::<f32>(n)?;
5609 let cfg = LaunchConfig::for_num_elems(n as u32);
5610 let nl = n as i64;
5611 let __s_b = self.gpu.stream();
5612 let mut b = __s_b.launch_builder(&f);
5613 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5614 unsafe {
5615 b.launch(cfg)?;
5616 }
5617 Ok(act)
5618 }
5619
5620 #[allow(clippy::too_many_arguments)]
5621 pub fn moe_pairs_scatter(
5622 &self,
5623 y_down: &CudaSlice<f32>,
5624 pair_w: &CudaSlice<f32>,
5625 tok_pair_off: &CudaSlice<i32>,
5626 tok_pair_ids: &CudaSlice<i32>,
5627 moe_out: &mut CudaSlice<f32>,
5628 t: usize,
5629 n_embd: usize,
5630 ) -> Result<(), Box<dyn std::error::Error>> {
5631 let f = self.func("moe_pairs_scatter");
5632 let cfg = LaunchConfig {
5633 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5634 block_dim: (256, 1, 1),
5635 shared_mem_bytes: 0,
5636 };
5637 let ne = n_embd as i32;
5638 let __s_b = self.gpu.stream();
5639 let mut b = __s_b.launch_builder(&f);
5640 b.arg(y_down)
5641 .arg(pair_w)
5642 .arg(tok_pair_off)
5643 .arg(tok_pair_ids)
5644 .arg(moe_out)
5645 .arg(&ne);
5646 unsafe {
5647 b.launch(cfg)?;
5648 }
5649 Ok(())
5650 }
5651
5652 #[allow(clippy::too_many_arguments)]
5656 pub fn moe_gate_up_gelu8_dev_q8(
5657 &self,
5658 table: &CudaSlice<u64>,
5659 sel: &cudarc::driver::CudaView<i32>,
5660 aq: &CudaSlice<i8>,
5661 ad: &CudaSlice<f32>,
5662 in_f: usize,
5663 n_ff: usize,
5664 n_used: usize,
5665 n_expert: usize,
5666 qt_g: i32,
5667 qt_u: i32,
5668 rb_g: usize,
5669 rb_u: usize,
5670 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5671 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5672 let (inf, nff, ne, rbg, rbu) = (
5673 in_f as i32,
5674 n_ff as i32,
5675 n_expert as i32,
5676 rb_g as i64,
5677 rb_u as i64,
5678 );
5679 let f = self.func("moe_gate_up_gelu8_dev_q8");
5680 let cfg = LaunchConfig {
5681 grid_dim: (n_ff as u32, n_used as u32, 1),
5682 block_dim: (32, 1, 1),
5683 shared_mem_bytes: 0,
5684 };
5685 let __s_b = self.gpu.stream();
5686 let mut b = __s_b.launch_builder(&f);
5687 b.arg(table)
5688 .arg(sel)
5689 .arg(aq)
5690 .arg(ad)
5691 .arg(&mut act)
5692 .arg(&inf)
5693 .arg(&nff)
5694 .arg(&ne)
5695 .arg(&qt_g)
5696 .arg(&qt_u)
5697 .arg(&rbg)
5698 .arg(&rbu);
5699 unsafe {
5700 b.launch(cfg)?;
5701 }
5702 Ok(act)
5703 }
5704
5705 #[allow(clippy::too_many_arguments)]
5707 pub fn moe_gate_up_gelu8_dev_q8_rows(
5708 &self,
5709 table: &CudaSlice<u64>,
5710 sel: &CudaSlice<i32>,
5711 aq: &CudaSlice<i8>,
5712 ad: &CudaSlice<f32>,
5713 t: usize,
5714 in_f: usize,
5715 n_ff: usize,
5716 n_used: usize,
5717 n_expert: usize,
5718 qt_g: i32,
5719 qt_u: i32,
5720 rb_g: usize,
5721 rb_u: usize,
5722 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5723 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5724 let (inf, nff, ne, rbg, rbu, nu) = (
5725 in_f as i32,
5726 n_ff as i32,
5727 n_expert as i32,
5728 rb_g as i64,
5729 rb_u as i64,
5730 n_used as i32,
5731 );
5732 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5733 let cfg = LaunchConfig {
5734 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5735 block_dim: (32, 1, 1),
5736 shared_mem_bytes: 0,
5737 };
5738 let __s_b = self.gpu.stream();
5739 let mut b = __s_b.launch_builder(&f);
5740 b.arg(table)
5741 .arg(sel)
5742 .arg(aq)
5743 .arg(ad)
5744 .arg(&mut act)
5745 .arg(&inf)
5746 .arg(&nff)
5747 .arg(&ne)
5748 .arg(&qt_g)
5749 .arg(&qt_u)
5750 .arg(&rbg)
5751 .arg(&rbu)
5752 .arg(&nu);
5753 unsafe {
5754 b.launch(cfg)?;
5755 }
5756 Ok(act)
5757 }
5758
5759 #[allow(clippy::too_many_arguments)]
5761 pub fn moe_gate_up_gelu8_dev_q8_csr(
5762 &self,
5763 table: &CudaSlice<u64>,
5764 sel: &CudaSlice<i32>,
5765 aq: &CudaSlice<i8>,
5766 ad: &CudaSlice<f32>,
5767 n_pairs: usize,
5768 in_f: usize,
5769 n_ff: usize,
5770 n_used: usize,
5771 n_expert: usize,
5772 qt_g: i32,
5773 qt_u: i32,
5774 rb_g: usize,
5775 rb_u: usize,
5776 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5777 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5778 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5779 in_f as i32,
5780 n_ff as i32,
5781 n_expert as i32,
5782 rb_g as i64,
5783 rb_u as i64,
5784 n_used as i32,
5785 n_pairs as i32,
5786 );
5787 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5788 let cfg = LaunchConfig {
5789 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5790 block_dim: (32, 1, 1),
5791 shared_mem_bytes: 0,
5792 };
5793 let __s_b = self.gpu.stream();
5794 let mut b = __s_b.launch_builder(&f);
5795 b.arg(table)
5796 .arg(sel)
5797 .arg(aq)
5798 .arg(ad)
5799 .arg(&mut act)
5800 .arg(&inf)
5801 .arg(&nff)
5802 .arg(&ne)
5803 .arg(&qt_g)
5804 .arg(&qt_u)
5805 .arg(&rbg)
5806 .arg(&rbu)
5807 .arg(&nu)
5808 .arg(&npi);
5809 unsafe {
5810 b.launch(cfg)?;
5811 }
5812 Ok(act)
5813 }
5814
5815 #[allow(clippy::too_many_arguments)]
5817 pub fn moe_down8_fma_dev_q8_rows_g(
5818 &self,
5819 table: &CudaSlice<u64>,
5820 sel: &CudaSlice<i32>,
5821 w: &CudaSlice<f32>,
5822 aq2: &CudaSlice<i8>,
5823 ad2: &CudaSlice<f32>,
5824 dst: &mut CudaSlice<f32>,
5825 t: usize,
5826 in_f: usize,
5827 out_f: usize,
5828 n_used: usize,
5829 n_expert: usize,
5830 qt: i32,
5831 rb: usize,
5832 ) -> Result<(), Box<dyn std::error::Error>> {
5833 let (inf, outf, nu, ne, rbi) = (
5834 in_f as i32,
5835 out_f as i32,
5836 n_used as i32,
5837 n_expert as i32,
5838 rb as i64,
5839 );
5840 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5844 let f = self.func(if step_b1_w8 {
5845 "moe_down8_fma_dev_q8_rows_w8"
5846 } else {
5847 "moe_down8_fma_dev_q8_rows_g"
5848 });
5849 let cfg = LaunchConfig {
5850 grid_dim: (out_f as u32, 1, t as u32),
5851 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5852 shared_mem_bytes: 0,
5853 };
5854 let __s_b = self.gpu.stream();
5855 let mut b = __s_b.launch_builder(&f);
5856 b.arg(table)
5857 .arg(sel)
5858 .arg(w)
5859 .arg(aq2)
5860 .arg(ad2)
5861 .arg(dst)
5862 .arg(&inf)
5863 .arg(&outf)
5864 .arg(&nu)
5865 .arg(&ne)
5866 .arg(&qt)
5867 .arg(&rbi);
5868 unsafe {
5869 b.launch(cfg)?;
5870 }
5871 Ok(())
5872 }
5873
5874 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5878 let (out_f, in_f) = (2048usize, 2816usize);
5879 let nblk = in_f / 32;
5880 let mut seed = 0x9E3779B97F4A7C15u64;
5881 let mut rng = move || {
5882 seed = seed
5883 .wrapping_mul(6364136223846793005)
5884 .wrapping_add(1442695040888963407);
5885 (seed >> 33) as u8
5886 };
5887 let mut w = vec![0u8; out_f * nblk * 18];
5888 for b in w.iter_mut() {
5889 *b = rng();
5890 }
5891 for r in 0..out_f {
5892 for g in 0..nblk {
5893 let off = (r * nblk + g) * 18;
5894 w[off] = 0x00;
5895 w[off + 1] = 0x2C; }
5897 }
5898 let qplane = out_f * nblk * 16;
5899 let mut wrp = vec![0u8; w.len()];
5900 for r in 0..out_f {
5901 for g in 0..nblk {
5902 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5903 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5904 .copy_from_slice(&src[0..2]);
5905 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5906 }
5907 }
5908 let w_d = self.htod_bytes(&w)?;
5909 let wrp_d = self.htod_bytes(&wrp)?;
5910 let mut aq = vec![0i8; m * in_f];
5911 for v in aq.iter_mut() {
5912 *v = rng() as i8;
5913 }
5914 let aq_d = self.htod_i8(&aq)?;
5915 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5916 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5917 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5918 const RPB: u32 = 4;
5919 let cfg = LaunchConfig {
5920 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5921 block_dim: (32, RPB, 1),
5922 shared_mem_bytes: 0,
5923 };
5924 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5925 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5926 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5927 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5928 {
5929 let __s_b = self.gpu.stream();
5930 let mut b = __s_b.launch_builder(&fb);
5931 b.arg(&w_d)
5932 .arg(&aq_d)
5933 .arg(&ad_d)
5934 .arg(&mut y0)
5935 .arg(&inf)
5936 .arg(&outf)
5937 .arg(&mi)
5938 .arg(&rb);
5939 unsafe {
5940 b.launch(cfg)?;
5941 }
5942 let __s_b = self.gpu.stream();
5943 let mut b = __s_b.launch_builder(&fr);
5944 b.arg(&wrp_d)
5945 .arg(&aq_d)
5946 .arg(&ad_d)
5947 .arg(&mut y1)
5948 .arg(&inf)
5949 .arg(&outf)
5950 .arg(&mi)
5951 .arg(&qp);
5952 unsafe {
5953 b.launch(cfg)?;
5954 }
5955 }
5956 self.gpu.stream().synchronize()?;
5957 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5958 let nd = h0
5959 .iter()
5960 .zip(&h1)
5961 .filter(|(a, b)| a.to_bits() != b.to_bits())
5962 .count();
5963 if nd != 0 {
5964 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5965 }
5966 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5967 self.gpu.stream().synchronize()?;
5968 let t0 = std::time::Instant::now();
5969 for _ in 0..500 {
5970 if rp {
5971 let __s_b = self.gpu.stream();
5972 let mut b = __s_b.launch_builder(&fr);
5973 b.arg(&wrp_d)
5974 .arg(&aq_d)
5975 .arg(&ad_d)
5976 .arg(&mut y1)
5977 .arg(&inf)
5978 .arg(&outf)
5979 .arg(&mi)
5980 .arg(&qp);
5981 unsafe {
5982 b.launch(cfg)?;
5983 }
5984 } else {
5985 let __s_b = self.gpu.stream();
5986 let mut b = __s_b.launch_builder(&fb);
5987 b.arg(&w_d)
5988 .arg(&aq_d)
5989 .arg(&ad_d)
5990 .arg(&mut y0)
5991 .arg(&inf)
5992 .arg(&outf)
5993 .arg(&mi)
5994 .arg(&rb);
5995 unsafe {
5996 b.launch(cfg)?;
5997 }
5998 }
5999 }
6000 self.gpu.stream().synchronize()?;
6001 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6002 };
6003 let _ = time(false)?;
6004 let _ = time(true)?; Ok((time(false)?, time(true)?))
6006 }
6007
6008 pub fn build_q4_rp4(
6013 &self,
6014 t: &mut crate::model::GpuTensor,
6015 ) -> Result<(), Box<dyn std::error::Error>> {
6016 use crate::model::GpuTensor;
6017 let GpuTensor::Quant {
6018 bytes,
6019 qtype,
6020 row_bytes,
6021 ne,
6022 rp4,
6023 ..
6024 } = t
6025 else {
6026 return Ok(());
6027 };
6028 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6029 return Ok(());
6030 }
6031 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6032 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6033 return Ok(());
6034 }
6035 let nblk = in_f / 32;
6036 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6037 let f = self.func("q4_0_split_rp_build");
6038 let n = (out_f * nblk) as i32;
6039 let cfg = LaunchConfig {
6040 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6041 block_dim: (256, 1, 1),
6042 shared_mem_bytes: 0,
6043 };
6044 let (of, nb) = (out_f as i32, nblk as i32);
6045 let _ = n;
6046 let __s_b = self.gpu.stream();
6047 let mut b = __s_b.launch_builder(&f);
6048 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6049 unsafe {
6050 b.launch(cfg)?;
6051 }
6052 *rp4 = Some(dst);
6053 Ok(())
6054 }
6055
6056 pub fn build_q8_rp4(
6061 &self,
6062 t: &mut crate::model::GpuTensor,
6063 ) -> Result<(), Box<dyn std::error::Error>> {
6064 use crate::model::GpuTensor;
6065 let GpuTensor::Quant {
6066 bytes,
6067 qtype,
6068 row_bytes,
6069 ne,
6070 rp4,
6071 ..
6072 } = t
6073 else {
6074 return Ok(());
6075 };
6076 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6077 return Ok(());
6078 }
6079 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6080 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6081 return Ok(());
6082 }
6083 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6084 Ok(())
6085 }
6086
6087 pub fn build_q8_rp4_raw(
6090 &self,
6091 bytes: &CudaSlice<u8>,
6092 in_f: usize,
6093 out_f: usize,
6094 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6095 assert!(in_f % 32 == 0);
6096 let nblk = in_f / 32;
6097 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6098 let f = self.func("q8_0_split_rp_build");
6099 let cfg = LaunchConfig {
6100 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6101 block_dim: (256, 1, 1),
6102 shared_mem_bytes: 0,
6103 };
6104 let (of, nb) = (out_f as i32, nblk as i32);
6105 let __s_b = self.gpu.stream();
6106 let mut b = __s_b.launch_builder(&f);
6107 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6108 unsafe {
6109 b.launch(cfg)?;
6110 }
6111 Ok(dst)
6112 }
6113
6114 pub fn build_q4k_rp4(
6122 &self,
6123 t: &mut crate::model::GpuTensor,
6124 ) -> Result<(), Box<dyn std::error::Error>> {
6125 use crate::model::GpuTensor;
6126 let GpuTensor::Quant {
6127 bytes,
6128 qtype,
6129 row_bytes,
6130 ne,
6131 rp4,
6132 ..
6133 } = t
6134 else {
6135 return Ok(());
6136 };
6137 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6138 return Ok(());
6139 }
6140 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6141 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6142 return Ok(());
6143 }
6144 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6145 Ok(())
6146 }
6147
6148 pub fn build_q6k_rp4(
6149 &self,
6150 t: &mut crate::model::GpuTensor,
6151 ) -> Result<(), Box<dyn std::error::Error>> {
6152 use crate::model::GpuTensor;
6153 let GpuTensor::Quant {
6154 bytes,
6155 qtype,
6156 row_bytes,
6157 ne,
6158 rp4,
6159 ..
6160 } = t
6161 else {
6162 return Ok(());
6163 };
6164 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6165 return Ok(());
6166 }
6167 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6168 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6169 return Ok(());
6170 }
6171 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6172 Ok(())
6173 }
6174
6175 pub fn build_kq_rp4_raw(
6177 &self,
6178 bytes: &CudaSlice<u8>,
6179 in_f: usize,
6180 out_f: usize,
6181 qtype: i32,
6182 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6183 assert!(in_f % 256 == 0);
6184 let nsbk = in_f / 256;
6185 let (sb_bytes, kname) = match qtype {
6186 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6187 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6188 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6189 };
6190 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6191 let f = self.func(kname);
6192 let cfg = LaunchConfig {
6193 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6194 block_dim: (256, 1, 1),
6195 shared_mem_bytes: 0,
6196 };
6197 let (of, nb) = (out_f as i32, nsbk as i32);
6198 let __s_b = self.gpu.stream();
6199 let mut b = __s_b.launch_builder(&f);
6200 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6201 unsafe {
6202 b.launch(cfg)?;
6203 }
6204 Ok(dst)
6205 }
6206
6207 pub fn kqrp_enabled() -> bool {
6211 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6212 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6213 Ok("0") => false,
6214 Ok(_) => true,
6215 Err(_) => cfg!(memra_hopper_mma),
6216 })
6217 }
6218
6219 pub fn build_q4_rp_swap(
6225 &self,
6226 t: &mut crate::model::GpuTensor,
6227 ) -> Result<bool, Box<dyn std::error::Error>> {
6228 use crate::model::GpuTensor;
6229 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6239 return Ok(false);
6240 }
6241 self.build_q4_rp4(t)?;
6242 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6244 return Ok(false);
6245 };
6246 match rp4.take() {
6247 Some(split) => {
6248 *bytes = split; *rp = true;
6250 Ok(true)
6251 }
6252 None => Ok(false),
6253 }
6254 }
6255
6256 pub fn q4rp_enabled() -> bool {
6258 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6259 *ON.get_or_init(|| {
6260 std::env::var("MEMRA_Q4RP")
6261 .map(|v| v != "0")
6262 .unwrap_or(true)
6263 })
6264 }
6265
6266 pub fn copy_rows_strided(
6269 &self,
6270 src: &CudaSlice<f32>,
6271 dst: &mut CudaSlice<f32>,
6272 row_elems: usize,
6273 n_rows: usize,
6274 src_stride: usize,
6275 src_off: usize,
6276 ) -> Result<(), Box<dyn std::error::Error>> {
6277 let f = self.func("copy_rows_strided_f32");
6278 let cfg = LaunchConfig {
6279 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6280 block_dim: (256, 1, 1),
6281 shared_mem_bytes: 0,
6282 };
6283 let (re, nr) = (row_elems as i32, n_rows as i32);
6284 let (st, off) = (src_stride as i64, src_off as i64);
6285 let __s_b = self.gpu.stream();
6286 let mut b = __s_b.launch_builder(&f);
6287 b.arg(src)
6288 .arg(&mut *dst)
6289 .arg(&re)
6290 .arg(&nr)
6291 .arg(&st)
6292 .arg(&off);
6293 unsafe {
6294 b.launch(cfg)?;
6295 }
6296 Ok(())
6297 }
6298
6299 pub fn place_rows_strided(
6305 &self,
6306 src: &CudaSlice<f32>,
6307 dst: &mut CudaSlice<f32>,
6308 row_elems: usize,
6309 n_rows: usize,
6310 dst_stride: usize,
6311 dst_off: usize,
6312 ) -> Result<(), Box<dyn std::error::Error>> {
6313 if row_elems == 0 || n_rows == 0 {
6314 return Err("strided row placement requires nonzero rows and row width".into());
6315 }
6316 let src_len = n_rows
6317 .checked_mul(row_elems)
6318 .ok_or("strided row placement source size overflow")?;
6319 let dst_len = n_rows
6320 .checked_sub(1)
6321 .and_then(|rows| rows.checked_mul(dst_stride))
6322 .and_then(|base| base.checked_add(dst_off))
6323 .and_then(|base| base.checked_add(row_elems))
6324 .ok_or("strided row placement destination size overflow")?;
6325 let row_end = dst_off
6326 .checked_add(row_elems)
6327 .ok_or("strided row placement row size overflow")?;
6328 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6329 return Err(format!(
6330 "strided row placement geometry mismatch: src={} need_src={src_len} \
6331 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6332 dst_stride={dst_stride} dst_off={dst_off}",
6333 src.len(),
6334 dst.len(),
6335 )
6336 .into());
6337 }
6338 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6339 return Err("strided row placement exceeds CUDA kernel geometry".into());
6340 }
6341 let f = self.func("place_rows_strided_f32");
6342 let cfg = LaunchConfig {
6343 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6344 block_dim: (256, 1, 1),
6345 shared_mem_bytes: 0,
6346 };
6347 let (re, nr) = (row_elems as i32, n_rows as i32);
6348 let (st, off) = (dst_stride as i64, dst_off as i64);
6349 let __s_b = self.gpu.stream();
6350 let mut b = __s_b.launch_builder(&f);
6351 b.arg(src)
6352 .arg(&mut *dst)
6353 .arg(&re)
6354 .arg(&nr)
6355 .arg(&st)
6356 .arg(&off);
6357 unsafe {
6358 b.launch(cfg)?;
6359 }
6360 Ok(())
6361 }
6362
6363 pub fn u32_set_k(
6365 &self,
6366 dst: &mut CudaSlice<u32>,
6367 v: u32,
6368 idx: usize,
6369 ) -> Result<(), Box<dyn std::error::Error>> {
6370 let f = self.func("u32_set_k");
6371 let cfg = LaunchConfig {
6372 grid_dim: (1, 1, 1),
6373 block_dim: (1, 1, 1),
6374 shared_mem_bytes: 0,
6375 };
6376 let ii = idx as i32;
6377 let __s_b = self.gpu.stream();
6378 let mut b = __s_b.launch_builder(&f);
6379 b.arg(dst).arg(&v).arg(&ii);
6380 unsafe {
6381 b.launch(cfg)?;
6382 }
6383 Ok(())
6384 }
6385
6386 pub fn i32_add_k(
6388 &self,
6389 d: &mut CudaSlice<i32>,
6390 v: i32,
6391 ) -> Result<(), Box<dyn std::error::Error>> {
6392 let f = self.func("i32_add_k");
6393 let cfg = LaunchConfig {
6394 grid_dim: (1, 1, 1),
6395 block_dim: (32, 1, 1),
6396 shared_mem_bytes: 0,
6397 };
6398 let __s_b = self.gpu.stream();
6399 let mut b = __s_b.launch_builder(&f);
6400 b.arg(d).arg(&v);
6401 unsafe {
6402 b.launch(cfg)?;
6403 }
6404 Ok(())
6405 }
6406
6407 pub fn i32_iota_from(
6409 &self,
6410 ctr: &CudaSlice<i32>,
6411 dst: &mut CudaSlice<i32>,
6412 n: usize,
6413 ) -> Result<(), Box<dyn std::error::Error>> {
6414 let f = self.func("i32_iota_from");
6415 let cfg = LaunchConfig::for_num_elems(n as u32);
6416 let ni = n as i32;
6417 let __s_b = self.gpu.stream();
6418 let mut b = __s_b.launch_builder(&f);
6419 b.arg(ctr).arg(dst).arg(&ni);
6420 unsafe {
6421 b.launch(cfg)?;
6422 }
6423 Ok(())
6424 }
6425
6426 pub fn u32_map_k(
6428 &self,
6429 buf: &mut CudaSlice<u32>,
6430 map: &CudaSlice<u32>,
6431 idx: usize,
6432 ) -> Result<(), Box<dyn std::error::Error>> {
6433 let f = self.func("u32_map_k");
6434 let cfg = LaunchConfig {
6435 grid_dim: (1, 1, 1),
6436 block_dim: (1, 1, 1),
6437 shared_mem_bytes: 0,
6438 };
6439 let ii = idx as i32;
6440 let __s_b = self.gpu.stream();
6441 let mut b = __s_b.launch_builder(&f);
6442 b.arg(buf).arg(map).arg(&ii);
6443 unsafe {
6444 b.launch(cfg)?;
6445 }
6446 Ok(())
6447 }
6448
6449 #[allow(clippy::too_many_arguments)]
6451 pub fn u32_pack2(
6452 &self,
6453 a: &CudaSlice<u32>,
6454 off_a: usize,
6455 n1: usize,
6456 b_in: &CudaSlice<u32>,
6457 n2: usize,
6458 out: &mut CudaSlice<u32>,
6459 ) -> Result<(), Box<dyn std::error::Error>> {
6460 let f = self.func("u32_pack2");
6461 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6462 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6463 let __s_b = self.gpu.stream();
6464 let mut b = __s_b.launch_builder(&f);
6465 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6466 unsafe {
6467 b.launch(cfg)?;
6468 }
6469 Ok(())
6470 }
6471
6472 pub fn moe_w_exscale(
6474 &self,
6475 w: &mut CudaSlice<f32>,
6476 sel: &CudaSlice<i32>,
6477 s: &CudaSlice<f32>,
6478 n: usize,
6479 ) -> Result<(), Box<dyn std::error::Error>> {
6480 let f = self.func("moe_w_exscale");
6481 let cfg = LaunchConfig::for_num_elems(n as u32);
6482 let ni = n as i32;
6483 let __s_b = self.gpu.stream();
6484 let mut b = __s_b.launch_builder(&f);
6485 b.arg(w).arg(sel).arg(s).arg(&ni);
6486 unsafe {
6487 b.launch(cfg)?;
6488 }
6489 Ok(())
6490 }
6491
6492 pub fn moe_w_scale_by_expert(
6495 &self,
6496 w: &mut CudaSlice<f32>,
6497 sel: &CudaSlice<i32>,
6498 macros: &CudaSlice<f32>,
6499 n_expert: usize,
6500 n: usize,
6501 ) -> Result<(), Box<dyn std::error::Error>> {
6502 let f = self.func("moe_w_scale_by_expert");
6503 let cfg = LaunchConfig {
6504 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6505 block_dim: (64, 1, 1),
6506 shared_mem_bytes: 0,
6507 };
6508 let (ne, nn) = (n_expert as i32, n as i32);
6509 let __s_b = self.gpu.stream();
6510 let mut b = __s_b.launch_builder(&f);
6511 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6512 unsafe {
6513 b.launch(cfg)?;
6514 }
6515 Ok(())
6516 }
6517
6518 pub fn moe_gate_up_silu8_dev_q8(
6519 &self,
6520 table: &CudaSlice<u64>,
6521 sel: &cudarc::driver::CudaView<i32>,
6522 aq: &CudaSlice<i8>,
6523 ad: &CudaSlice<f32>,
6524 in_f: usize,
6525 n_ff: usize,
6526 n_used: usize,
6527 n_expert: usize,
6528 qt_g: i32,
6529 qt_u: i32,
6530 rb_g: usize,
6531 rb_u: usize,
6532 macros: &CudaSlice<f32>,
6533 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6534 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6535 let (mode, wpb) = GU.get_or_init(|| {
6536 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6537 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6538 .ok()
6539 .and_then(|v| v.parse().ok())
6540 .unwrap_or(4u32)
6541 .clamp(1, 16);
6542 (mode, wpb)
6543 });
6544 let (mode, wpb) = (mode.as_str(), *wpb);
6545 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6546 let (inf, nff, ne, rbg, rbu) = (
6547 in_f as i32,
6548 n_ff as i32,
6549 n_expert as i32,
6550 rb_g as i64,
6551 rb_u as i64,
6552 );
6553 let (f, cfg) = match mode {
6554 "1" | "2" | "4" => {
6555 let rpw: u32 = mode.parse().unwrap();
6556 let f = self.func(match rpw {
6557 1 => "moe_gate_up_silu8_dev_q8_r1",
6558 2 => "moe_gate_up_silu8_dev_q8_r2",
6559 _ => "moe_gate_up_silu8_dev_q8_r4",
6560 });
6561 let rows_per_block = (rpw * wpb) as usize;
6562 let gx = n_ff.div_ceil(rows_per_block) as u32;
6563 (
6564 f,
6565 LaunchConfig {
6566 grid_dim: (gx, n_used as u32, 1),
6567 block_dim: (32, wpb, 1),
6568 shared_mem_bytes: 0,
6569 },
6570 )
6571 }
6572 "j8" if n_used <= 32 => (
6573 self.func("moe_gate_up_silu8_dev_q8_j8"),
6574 LaunchConfig {
6575 grid_dim: (n_ff as u32, 1, 1),
6576 block_dim: (32, n_used as u32, 1),
6577 shared_mem_bytes: 0,
6578 },
6579 ),
6580 "vsm2" => {
6582 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6583 let sh = (rb_g + rb_u) as u32;
6584 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6585 f.set_attribute(
6586 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6587 sh as i32,
6588 )?;
6589 (
6590 f,
6591 LaunchConfig {
6592 grid_dim: (n_ff as u32, n_used as u32, 1),
6593 block_dim: (32, 1, 1),
6594 shared_mem_bytes: sh,
6595 },
6596 )
6597 }
6598 "vsm" => {
6599 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6600 let sh = (rb_g + rb_u) as u32;
6601 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6602 f.set_attribute(
6603 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6604 sh as i32,
6605 )?;
6606 (
6607 f,
6608 LaunchConfig {
6609 grid_dim: (n_ff as u32, n_used as u32, 1),
6610 block_dim: (32, 1, 1),
6611 shared_mem_bytes: sh,
6612 },
6613 )
6614 }
6615 "sg" => (
6616 self.func("moe_gate_up_silu8_dev_q8_sg"),
6617 LaunchConfig {
6618 grid_dim: (n_ff as u32, n_used as u32, 1),
6619 block_dim: (32, 1, 1),
6620 shared_mem_bytes: 0,
6621 },
6622 ),
6623 "j8sg" if n_used <= 32 => (
6624 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6625 LaunchConfig {
6626 grid_dim: (n_ff as u32, 1, 1),
6627 block_dim: (32, n_used as u32, 1),
6628 shared_mem_bytes: 0,
6629 },
6630 ),
6631 "u64" if in_f == 2048 => (
6632 self.func("moe_gate_up_silu8_dev_q8_u64"),
6633 LaunchConfig {
6634 grid_dim: (n_ff as u32, n_used as u32, 1),
6635 block_dim: (32, 1, 1),
6636 shared_mem_bytes: 0,
6637 },
6638 ),
6639 "gs4" if in_f == 2048 => (
6640 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6641 LaunchConfig {
6642 grid_dim: (n_ff as u32, n_used as u32, 1),
6643 block_dim: (32, 4, 1),
6644 shared_mem_bytes: 0,
6645 },
6646 ),
6647 "v" | "" => (
6649 self.func("moe_gate_up_silu8_dev_q8_v"),
6650 LaunchConfig {
6651 grid_dim: (n_ff as u32, n_used as u32, 1),
6652 block_dim: (32, 1, 1),
6653 shared_mem_bytes: 0,
6654 },
6655 ),
6656 "s2" => (
6657 self.func("moe_gate_up_silu8_dev_q8_s2"),
6658 LaunchConfig {
6659 grid_dim: (n_ff as u32, n_used as u32, 1),
6660 block_dim: (32, 2, 1),
6661 shared_mem_bytes: 0,
6662 },
6663 ),
6664 "s2z" => {
6665 let rz = wpb.min(16); (
6667 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6668 LaunchConfig {
6669 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6670 block_dim: (32, 2, rz),
6671 shared_mem_bytes: 0,
6672 },
6673 )
6674 }
6675 _ => (
6676 self.func("moe_gate_up_silu8_dev_q8"),
6677 LaunchConfig {
6678 grid_dim: (n_ff as u32, n_used as u32, 1),
6679 block_dim: (32, 1, 1),
6680 shared_mem_bytes: 0,
6681 },
6682 ),
6683 };
6684 let __s_b = self.gpu.stream();
6685 let mut b = __s_b.launch_builder(&f);
6686 b.arg(table)
6687 .arg(sel)
6688 .arg(aq)
6689 .arg(ad)
6690 .arg(&mut act)
6691 .arg(&inf)
6692 .arg(&nff)
6693 .arg(&ne)
6694 .arg(&qt_g)
6695 .arg(&qt_u)
6696 .arg(&rbg)
6697 .arg(&rbu)
6698 .arg(macros);
6699 unsafe {
6700 b.launch(cfg)?;
6701 }
6702 Ok(act)
6703 }
6704
6705 #[allow(clippy::too_many_arguments)]
6706 pub fn moe_down8_fma_dev_q8(
6707 &self,
6708 table: &CudaSlice<u64>,
6709 sel: &cudarc::driver::CudaView<i32>,
6710 w: &cudarc::driver::CudaView<f32>,
6711 aq2: &CudaSlice<i8>,
6712 ad2: &CudaSlice<f32>,
6713 dst: &mut cudarc::driver::CudaViewMut<f32>,
6714 in_f: usize,
6715 out_f: usize,
6716 n_used: usize,
6717 n_expert: usize,
6718 qt: i32,
6719 rb: usize,
6720 ) -> Result<(), Box<dyn std::error::Error>> {
6721 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6722 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6723 let (inf, outf, nu, ne, rbi) = (
6724 in_f as i32,
6725 out_f as i32,
6726 n_used as i32,
6727 n_expert as i32,
6728 rb as i64,
6729 );
6730 let (f, cfg) = match mode.as_str() {
6733 m @ ("1" | "2" | "4") if n_used <= 8 => {
6734 let rpw: usize = m.parse().unwrap();
6735 let f = self.func(match rpw {
6736 1 => "moe_down8_fma_dev_q8_w8r1",
6737 2 => "moe_down8_fma_dev_q8_w8r2",
6738 _ => "moe_down8_fma_dev_q8_w8r4",
6739 });
6740 (
6741 f,
6742 LaunchConfig {
6743 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6744 block_dim: (32, n_used as u32, 1),
6745 shared_mem_bytes: 0,
6746 },
6747 )
6748 }
6749 "h2" if in_f == 512 => (
6750 self.func("moe_down8_fma_dev_q8_h2"),
6751 LaunchConfig {
6752 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6753 block_dim: (32, 1, 1),
6754 shared_mem_bytes: 0,
6755 },
6756 ),
6757 "" if in_f == 704 && n_used <= 8 => (
6760 self.func("moe_down8_fma_dev_q8_w8r2"),
6761 LaunchConfig {
6762 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6763 block_dim: (32, n_used as u32, 1),
6764 shared_mem_bytes: 0,
6765 },
6766 ),
6767 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6771 self.func("moe_down8_fma_dev_q8_w8h2v"),
6772 LaunchConfig {
6773 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6774 block_dim: (32, n_used as u32, 1),
6775 shared_mem_bytes: 0,
6776 },
6777 ),
6778 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6779 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6780 LaunchConfig {
6781 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6782 block_dim: (32, n_used as u32, 1),
6783 shared_mem_bytes: 0,
6784 },
6785 ),
6786 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6787 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6788 LaunchConfig {
6789 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6790 block_dim: (32, n_used as u32, 1),
6791 shared_mem_bytes: 0,
6792 },
6793 ),
6794 "w8h2" if in_f == 512 && n_used <= 8 => (
6795 self.func("moe_down8_fma_dev_q8_w8h2"),
6796 LaunchConfig {
6797 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6798 block_dim: (32, n_used as u32, 1),
6799 shared_mem_bytes: 0,
6800 },
6801 ),
6802 _ => (
6803 self.func("moe_down8_fma_dev_q8"),
6804 LaunchConfig {
6805 grid_dim: (out_f as u32, 1, 1),
6806 block_dim: (32, 1, 1),
6807 shared_mem_bytes: 0,
6808 },
6809 ),
6810 };
6811 let __s_b = self.gpu.stream();
6812 let mut b = __s_b.launch_builder(&f);
6813 b.arg(table)
6814 .arg(sel)
6815 .arg(w)
6816 .arg(aq2)
6817 .arg(ad2)
6818 .arg(dst)
6819 .arg(&inf)
6820 .arg(&outf)
6821 .arg(&nu)
6822 .arg(&ne)
6823 .arg(&qt)
6824 .arg(&rbi);
6825 unsafe {
6826 b.launch(cfg)?;
6827 }
6828 Ok(())
6829 }
6830
6831 #[allow(clippy::too_many_arguments)]
6838 pub fn moe_gate_up_silu8_dev_q8_rows(
6839 &self,
6840 table: &CudaSlice<u64>,
6841 sel: &CudaSlice<i32>,
6842 aq: &CudaSlice<i8>,
6843 ad: &CudaSlice<f32>,
6844 t: usize,
6845 in_f: usize,
6846 n_ff: usize,
6847 n_used: usize,
6848 n_expert: usize,
6849 qt_g: i32,
6850 qt_u: i32,
6851 rb_g: usize,
6852 rb_u: usize,
6853 macros: &CudaSlice<f32>,
6854 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6855 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6856 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6857 let cfg = LaunchConfig {
6858 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6859 block_dim: (32, 1, 1),
6860 shared_mem_bytes: 0,
6861 };
6862 let (inf, nff, ne, nu, rbg, rbu) = (
6863 in_f as i32,
6864 n_ff as i32,
6865 n_expert as i32,
6866 n_used as i32,
6867 rb_g as i64,
6868 rb_u as i64,
6869 );
6870 let __s_b = self.gpu.stream();
6871 let mut b = __s_b.launch_builder(&f);
6872 b.arg(table)
6873 .arg(sel)
6874 .arg(aq)
6875 .arg(ad)
6876 .arg(&mut act)
6877 .arg(&inf)
6878 .arg(&nff)
6879 .arg(&ne)
6880 .arg(&qt_g)
6881 .arg(&qt_u)
6882 .arg(&rbg)
6883 .arg(&rbu)
6884 .arg(&nu)
6885 .arg(macros);
6886 unsafe {
6887 b.launch(cfg)?;
6888 }
6889 Ok(act)
6890 }
6891
6892 #[allow(clippy::too_many_arguments)]
6897 pub fn moe_down8_fma_dev_q8_rows(
6898 &self,
6899 table: &CudaSlice<u64>,
6900 sel: &CudaSlice<i32>,
6901 w: &CudaSlice<f32>,
6902 aq2: &CudaSlice<i8>,
6903 ad2: &CudaSlice<f32>,
6904 dst: &mut CudaSlice<f32>,
6905 t: usize,
6906 in_f: usize,
6907 out_f: usize,
6908 n_used: usize,
6909 n_expert: usize,
6910 qt: i32,
6911 rb: usize,
6912 ) -> Result<(), Box<dyn std::error::Error>> {
6913 assert!(
6914 in_f == 512 && n_used <= 8,
6915 "down rows twin is w8h2v shape-gated"
6916 );
6917 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6918 let cfg = LaunchConfig {
6919 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6920 block_dim: (32, n_used as u32, 1),
6921 shared_mem_bytes: 0,
6922 };
6923 let (inf, outf, nu, ne, rbi) = (
6924 in_f as i32,
6925 out_f as i32,
6926 n_used as i32,
6927 n_expert as i32,
6928 rb as i64,
6929 );
6930 let __s_b = self.gpu.stream();
6931 let mut b = __s_b.launch_builder(&f);
6932 b.arg(table)
6933 .arg(sel)
6934 .arg(w)
6935 .arg(aq2)
6936 .arg(ad2)
6937 .arg(dst)
6938 .arg(&inf)
6939 .arg(&outf)
6940 .arg(&nu)
6941 .arg(&ne)
6942 .arg(&qt)
6943 .arg(&rbi);
6944 unsafe {
6945 b.launch(cfg)?;
6946 }
6947 Ok(())
6948 }
6949
6950 #[allow(clippy::too_many_arguments)]
6954 pub fn moe_gate_up_silu8_dev_q8_csr(
6955 &self,
6956 table: &CudaSlice<u64>,
6957 sel: &CudaSlice<i32>,
6958 aq: &CudaSlice<i8>,
6959 ad: &CudaSlice<f32>,
6960 n_pairs: usize,
6961 in_f: usize,
6962 n_ff: usize,
6963 n_used: usize,
6964 n_expert: usize,
6965 qt_g: i32,
6966 qt_u: i32,
6967 rb_g: usize,
6968 rb_u: usize,
6969 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6970 let f = if qt_g == crate::QT_NVFP4 {
6973 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6974 } else {
6975 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6976 };
6977 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6978 let cfg = LaunchConfig {
6979 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6980 block_dim: (32, 1, 1),
6981 shared_mem_bytes: 0,
6982 };
6983 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6984 in_f as i32,
6985 n_ff as i32,
6986 n_expert as i32,
6987 n_used as i32,
6988 n_pairs as i32,
6989 rb_g as i64,
6990 rb_u as i64,
6991 );
6992 let __s_b = self.gpu.stream();
6993 let mut b = __s_b.launch_builder(&f);
6994 b.arg(table)
6995 .arg(sel)
6996 .arg(aq)
6997 .arg(ad)
6998 .arg(&mut act)
6999 .arg(&inf)
7000 .arg(&nff)
7001 .arg(&ne)
7002 .arg(&qt_g)
7003 .arg(&qt_u)
7004 .arg(&rbg)
7005 .arg(&rbu)
7006 .arg(&nu)
7007 .arg(&npi);
7008 unsafe {
7009 b.launch(cfg)?;
7010 }
7011 Ok(act)
7012 }
7013
7014 #[allow(clippy::too_many_arguments)]
7018 pub fn moe_down8_fma_dev_q8_variant(
7019 &self,
7020 variant: &str,
7021 table: &CudaSlice<u64>,
7022 sel: &cudarc::driver::CudaView<i32>,
7023 w: &cudarc::driver::CudaView<f32>,
7024 aq2: &CudaSlice<i8>,
7025 ad2: &CudaSlice<f32>,
7026 dst: &mut cudarc::driver::CudaViewMut<f32>,
7027 in_f: usize,
7028 out_f: usize,
7029 n_used: usize,
7030 n_expert: usize,
7031 qt: i32,
7032 rb: usize,
7033 ) -> Result<(), Box<dyn std::error::Error>> {
7034 let (inf, outf, nu, ne, rbi) = (
7035 in_f as i32,
7036 out_f as i32,
7037 n_used as i32,
7038 n_expert as i32,
7039 rb as i64,
7040 );
7041 let (f, cfg) = match variant {
7042 "w8h2" | "w8h2v" => (
7043 self.func(if variant == "w8h2" {
7044 "moe_down8_fma_dev_q8_w8h2"
7045 } else {
7046 "moe_down8_fma_dev_q8_w8h2v"
7047 }),
7048 LaunchConfig {
7049 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7050 block_dim: (32, n_used as u32, 1),
7051 shared_mem_bytes: 0,
7052 },
7053 ),
7054 "w8h2r2" | "w8h2r2v" => (
7055 self.func(if variant == "w8h2r2" {
7056 "moe_down8_fma_dev_q8_w8h2r2"
7057 } else {
7058 "moe_down8_fma_dev_q8_w8h2r2v"
7059 }),
7060 LaunchConfig {
7061 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7062 block_dim: (32, n_used as u32, 1),
7063 shared_mem_bytes: 0,
7064 },
7065 ),
7066 _ => (
7067 self.func("moe_down8_fma_dev_q8"),
7068 LaunchConfig {
7069 grid_dim: (out_f as u32, 1, 1),
7070 block_dim: (32, 1, 1),
7071 shared_mem_bytes: 0,
7072 },
7073 ),
7074 };
7075 let __s_b = self.gpu.stream();
7076 let mut b = __s_b.launch_builder(&f);
7077 b.arg(table)
7078 .arg(sel)
7079 .arg(w)
7080 .arg(aq2)
7081 .arg(ad2)
7082 .arg(dst)
7083 .arg(&inf)
7084 .arg(&outf)
7085 .arg(&nu)
7086 .arg(&ne)
7087 .arg(&qt)
7088 .arg(&rbi);
7089 unsafe {
7090 b.launch(cfg)?;
7091 }
7092 Ok(())
7093 }
7094
7095 #[allow(clippy::too_many_arguments)]
7097 pub fn moe_gate_up_silu8_dev_q8_variant(
7098 &self,
7099 variant: &str,
7100 table: &CudaSlice<u64>,
7101 sel: &cudarc::driver::CudaView<i32>,
7102 aq: &CudaSlice<i8>,
7103 ad: &CudaSlice<f32>,
7104 in_f: usize,
7105 n_ff: usize,
7106 n_used: usize,
7107 n_expert: usize,
7108 qt_g: i32,
7109 qt_u: i32,
7110 rb_g: usize,
7111 rb_u: usize,
7112 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7113 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7114 let (inf, nff, ne, rbg, rbu) = (
7115 in_f as i32,
7116 n_ff as i32,
7117 n_expert as i32,
7118 rb_g as i64,
7119 rb_u as i64,
7120 );
7121 let f = self.func(if variant == "v" {
7122 "moe_gate_up_silu8_dev_q8_v"
7123 } else {
7124 "moe_gate_up_silu8_dev_q8"
7125 });
7126 let cfg = LaunchConfig {
7127 grid_dim: (n_ff as u32, n_used as u32, 1),
7128 block_dim: (32, 1, 1),
7129 shared_mem_bytes: 0,
7130 };
7131 let __s_b = self.gpu.stream();
7132 let mut b = __s_b.launch_builder(&f);
7133 b.arg(table)
7134 .arg(sel)
7135 .arg(aq)
7136 .arg(ad)
7137 .arg(&mut act)
7138 .arg(&inf)
7139 .arg(&nff)
7140 .arg(&ne)
7141 .arg(&qt_g)
7142 .arg(&qt_u)
7143 .arg(&rbg)
7144 .arg(&rbu);
7145 unsafe {
7146 b.launch(cfg)?;
7147 }
7148 Ok(act)
7149 }
7150
7151 pub fn moe_gate_up_silu8_dev(
7152 &self,
7153 table: &CudaSlice<u64>,
7154 sel: &cudarc::driver::CudaView<i32>,
7155 x: &cudarc::driver::CudaView<f32>,
7156 in_f: usize,
7157 n_ff: usize,
7158 n_used: usize,
7159 n_expert: usize,
7160 qt_g: i32,
7161 qt_u: i32,
7162 rb_g: usize,
7163 rb_u: usize,
7164 macros: &CudaSlice<f32>,
7165 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7166 let f = self.func("moe_gate_up_silu8_dev");
7167 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
7169 grid_dim: (n_ff as u32, n_used as u32, 1),
7170 block_dim: (256, 1, 1),
7171 shared_mem_bytes: 0,
7172 };
7173 let (inf, nff, ne, rbg, rbu) = (
7174 in_f as i32,
7175 n_ff as i32,
7176 n_expert as i32,
7177 rb_g as i64,
7178 rb_u as i64,
7179 );
7180 let __s_b = self.gpu.stream();
7181 let mut b = __s_b.launch_builder(&f);
7182 b.arg(table)
7183 .arg(sel)
7184 .arg(x)
7185 .arg(&mut act)
7186 .arg(&inf)
7187 .arg(&nff)
7188 .arg(&ne)
7189 .arg(&qt_g)
7190 .arg(&qt_u)
7191 .arg(&rbg)
7192 .arg(&rbu)
7193 .arg(macros);
7194 unsafe {
7195 b.launch(cfg)?;
7196 }
7197 Ok(act)
7198 }
7199
7200 #[allow(clippy::too_many_arguments)]
7203 pub fn moe_down8_fma_dev(
7204 &self,
7205 table: &CudaSlice<u64>,
7206 sel: &cudarc::driver::CudaView<i32>,
7207 w: &cudarc::driver::CudaView<f32>,
7208 act: &CudaSlice<f32>,
7209 dst: &mut cudarc::driver::CudaViewMut<f32>,
7210 in_f: usize,
7211 out_f: usize,
7212 n_used: usize,
7213 n_expert: usize,
7214 qt: i32,
7215 rb: usize,
7216 ) -> Result<(), Box<dyn std::error::Error>> {
7217 let f = self.func("moe_down8_fma_dev");
7218 let cfg = LaunchConfig {
7219 grid_dim: (out_f as u32, 1, 1),
7220 block_dim: (256, 1, 1),
7221 shared_mem_bytes: 0,
7222 };
7223 let (inf, outf, nu, ne, rbv) = (
7224 in_f as i32,
7225 out_f as i32,
7226 n_used as i32,
7227 n_expert as i32,
7228 rb as i64,
7229 );
7230 let __s_b = self.gpu.stream();
7231 let mut b = __s_b.launch_builder(&f);
7232 b.arg(table)
7233 .arg(sel)
7234 .arg(w)
7235 .arg(act)
7236 .arg(dst)
7237 .arg(&inf)
7238 .arg(&outf)
7239 .arg(&nu)
7240 .arg(&ne)
7241 .arg(&qt)
7242 .arg(&rbv);
7243 unsafe {
7244 b.launch(cfg)?;
7245 }
7246 Ok(())
7247 }
7248
7249 pub fn axpy_into(
7251 &self,
7252 src: &CudaSlice<f32>,
7253 alpha: f32,
7254 dst: &mut cudarc::driver::CudaViewMut<f32>,
7255 n: usize,
7256 ) -> Result<(), Box<dyn std::error::Error>> {
7257 let f = self.func("axpy_f32");
7258 let cfg = LaunchConfig::for_num_elems(n as u32);
7259 let (a, ni) = (alpha, n as i32);
7260 let __s_b = self.gpu.stream();
7261 let mut b = __s_b.launch_builder(&f);
7262 b.arg(src).arg(dst).arg(&a).arg(&ni);
7263 unsafe {
7264 b.launch(cfg)?;
7265 }
7266 Ok(())
7267 }
7268
7269 pub fn axpy_host_into(
7271 &self,
7272 src: &cudarc::driver::CudaView<'_, f32>,
7273 alpha: f32,
7274 dst: &mut cudarc::driver::CudaViewMut<f32>,
7275 n: usize,
7276 ) -> Result<(), Box<dyn std::error::Error>> {
7277 let f = self.func("axpy_host_f32");
7278 let cfg = LaunchConfig::for_num_elems(n as u32);
7279 let (a, ni) = (alpha, n as i32);
7280 let __s_b = self.gpu.stream();
7281 let mut b = __s_b.launch_builder(&f);
7282 b.arg(src).arg(dst).arg(&a).arg(&ni);
7283 unsafe {
7284 b.launch(cfg)?;
7285 }
7286 Ok(())
7287 }
7288
7289 pub fn add_scaled_rows(
7291 &self,
7292 src: &CudaSlice<f32>,
7293 scale: &CudaSlice<f32>,
7294 dst: &mut CudaSlice<f32>,
7295 ncols: usize,
7296 nrows: usize,
7297 ) -> Result<(), Box<dyn std::error::Error>> {
7298 let f = self.func("add_scaled_rows_f32");
7299 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7300 let (nc, nr) = (ncols as i32, nrows as i32);
7301 let __s_b = self.gpu.stream();
7302 let mut b = __s_b.launch_builder(&f);
7303 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7304 unsafe {
7305 b.launch(cfg)?;
7306 }
7307 Ok(())
7308 }
7309
7310 pub fn gather_rows(
7314 &self,
7315 src: &CudaSlice<f32>,
7316 idx: &CudaSlice<i32>,
7317 dst: &mut CudaSlice<f32>,
7318 ncols: usize,
7319 m_e: usize,
7320 ) -> Result<(), Box<dyn std::error::Error>> {
7321 let f = self.func("gather_rows_f32");
7322 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7323 let (nc, me) = (ncols as i32, m_e as i32);
7324 let __s_b = self.gpu.stream();
7325 let mut b = __s_b.launch_builder(&f);
7326 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7327 unsafe {
7328 b.launch(cfg)?;
7329 }
7330 Ok(())
7331 }
7332
7333 pub fn scatter_slot(
7338 &self,
7339 src: &CudaSlice<f32>,
7340 tok_idx: &CudaSlice<i32>,
7341 slot_idx: &CudaSlice<i32>,
7342 weight: &CudaSlice<f32>,
7343 dst: &mut CudaSlice<f32>,
7344 wbuf: &mut CudaSlice<f32>,
7345 ncols: usize,
7346 n_used: usize,
7347 m_e: usize,
7348 ) -> Result<(), Box<dyn std::error::Error>> {
7349 let f = self.func("scatter_add_slot_f32");
7350 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7351 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7352 let __s_b = self.gpu.stream();
7353 let mut b = __s_b.launch_builder(&f);
7354 b.arg(src)
7355 .arg(tok_idx)
7356 .arg(slot_idx)
7357 .arg(weight)
7358 .arg(dst)
7359 .arg(wbuf)
7360 .arg(&nc)
7361 .arg(&nu)
7362 .arg(&me);
7363 unsafe {
7364 b.launch(cfg)?;
7365 }
7366 Ok(())
7367 }
7368
7369 pub fn reduce_slots(
7373 &self,
7374 slots: &CudaSlice<f32>,
7375 wbuf: &CudaSlice<f32>,
7376 dst: &mut CudaSlice<f32>,
7377 ncols: usize,
7378 n_used: usize,
7379 t: usize,
7380 ) -> Result<(), Box<dyn std::error::Error>> {
7381 let f = self.func("reduce_slots_f32");
7382 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7383 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7384 let __s_b = self.gpu.stream();
7385 let mut b = __s_b.launch_builder(&f);
7386 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7387 unsafe {
7388 b.launch(cfg)?;
7389 }
7390 Ok(())
7391 }
7392
7393 pub fn reduce_slots_host(
7398 &self,
7399 slots: &CudaSlice<f32>,
7400 wbuf: &CudaSlice<f32>,
7401 dst: &mut CudaSlice<f32>,
7402 ncols: usize,
7403 n_used: usize,
7404 t: usize,
7405 ) -> Result<(), Box<dyn std::error::Error>> {
7406 let f = self.func("reduce_slots_host_f32");
7407 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7408 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7409 let __s_b = self.gpu.stream();
7410 let mut b = __s_b.launch_builder(&f);
7411 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7412 unsafe {
7413 b.launch(cfg)?;
7414 }
7415 Ok(())
7416 }
7417
7418 pub fn quantize_q8_1_view(
7425 &self,
7426 x: &cudarc::driver::CudaView<f32>,
7427 m: usize,
7428 in_f: usize,
7429 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7430 let f = self.func("quantize_q8_1");
7431 let nblk = in_f / 32;
7432 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7433 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7434 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7435 let (inf, mi) = (in_f as i32, m as i32);
7436 let __s_b = self.gpu.stream();
7437 let mut b = __s_b.launch_builder(&f);
7438 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7439 unsafe {
7440 b.launch(cfg)?;
7441 }
7442 Ok((q, d))
7443 }
7444
7445 pub fn quantize_q8_1(
7446 &self,
7447 x: &CudaSlice<f32>,
7448 m: usize,
7449 in_f: usize,
7450 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7451 let nblk = in_f / 32;
7452 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);
7456 let (inf, mi) = (in_f as i32, m as i32);
7457 if Self::pdl_on() && Self::pdl_wb_on() {
7458 {
7459 use cudarc::driver::{DevicePtr, DevicePtrMut};
7460 let s = &self.gpu.stream();
7461 let (px, _g0) = x.device_ptr(s);
7462 let (pq, _g1) = q.device_ptr_mut(s);
7463 let (pd, _g2) = d.device_ptr_mut(s);
7464 let mut ps = [
7465 &px as *const _ as *mut std::ffi::c_void,
7466 &pq as *const _ as *mut _,
7467 &pd as *const _ as *mut _,
7468 &inf as *const _ as *mut _,
7469 &mi as *const _ as *mut _,
7470 ];
7471 unsafe {
7472 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7473 }
7474 }
7475 return Ok((q, d));
7476 }
7477 let f = self.func("quantize_q8_1");
7478 let __s_b = self.gpu.stream();
7479 let mut b = __s_b.launch_builder(&f);
7480 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7481 unsafe {
7482 b.launch(cfg)?;
7483 }
7484 Ok((q, d))
7485 }
7486
7487 pub fn quantize_fp4_act(
7491 &self,
7492 x: &CudaSlice<f32>,
7493 m: usize,
7494 in_f: usize,
7495 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7496 let f = self.func("quantize_fp4_act");
7497 let nb16 = in_f / 16;
7498 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);
7501 let (inf, mi) = (in_f as i32, m as i32);
7502 let __s_b = self.gpu.stream();
7503 let mut b = __s_b.launch_builder(&f);
7504 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7505 unsafe {
7506 b.launch(cfg)?;
7507 }
7508 Ok((aq4, ad4))
7509 }
7510
7511 pub fn qmatvec_gemm_nvfp4_fp4(
7516 &self,
7517 bytes: &CudaSlice<u8>,
7518 x: &CudaSlice<f32>,
7519 m: usize,
7520 in_f: usize,
7521 out_f: usize,
7522 row_bytes: usize,
7523 scale: f32,
7524 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7525 assert!(
7526 in_f % 64 == 0,
7527 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7528 );
7529 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7530 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7531 if scale != 1.0 {
7532 self.scale_inplace(&mut y, scale, m * out_f)?;
7533 }
7534 Ok(y)
7535 }
7536
7537 fn fp4_gemm_launch(
7540 &self,
7541 bytes: &CudaSlice<u8>,
7542 aq4: &CudaSlice<u32>,
7543 ad4: &CudaSlice<u8>,
7544 m: usize,
7545 in_f: usize,
7546 out_f: usize,
7547 row_bytes: usize,
7548 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7549 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7550 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7552 const BN: u32 = 256;
7553 let cfg = LaunchConfig {
7554 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7555 block_dim: (32, 4, 1),
7556 shared_mem_bytes: 0,
7557 };
7558 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7559 let __s_b = self.gpu.stream();
7560 let mut b = __s_b.launch_builder(&f);
7561 b.arg(bytes)
7562 .arg(aq4)
7563 .arg(ad4)
7564 .arg(&mut y)
7565 .arg(&inf)
7566 .arg(&outf)
7567 .arg(&mi)
7568 .arg(&rb);
7569 unsafe {
7570 b.launch(cfg)?;
7571 }
7572 Ok(y)
7573 }
7574
7575 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7577 &self,
7578 bytes: &CudaSlice<u8>,
7579 x: &CudaSlice<f32>,
7580 m: usize,
7581 in_f: usize,
7582 out_f: usize,
7583 row_bytes: usize,
7584 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7585 assert!(
7586 in_f % 64 == 0,
7587 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7588 );
7589 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7590 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7591 }
7592
7593 pub fn qmatvec_q8_0_fast(
7595 &self,
7596 w: &CudaSlice<u8>,
7597 x: &CudaSlice<f32>,
7598 m: usize,
7599 in_f: usize,
7600 out_f: usize,
7601 row_bytes: usize,
7602 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7603 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7604 let f = self.func("qmatvec_q8_0_dp4a");
7605 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7607 grid_dim: (out_f as u32, m as u32, 1),
7608 block_dim: (128, 1, 1),
7609 shared_mem_bytes: 0,
7610 };
7611 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7612 let __s_b = self.gpu.stream();
7613 let mut b = __s_b.launch_builder(&f);
7614 b.arg(w)
7615 .arg(&aq)
7616 .arg(&ad)
7617 .arg(&mut y)
7618 .arg(&inf)
7619 .arg(&outf)
7620 .arg(&mi)
7621 .arg(&rb);
7622 unsafe {
7623 b.launch(cfg)?;
7624 }
7625 Ok(y)
7626 }
7627
7628 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7631 &self,
7632 w: &CudaSlice<u8>,
7633 x: &CudaSlice<f32>,
7634 m: usize,
7635 in_f: usize,
7636 out_f: usize,
7637 row_bytes: usize,
7638 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7639 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7640 let f = self.func("qmatvec_q4_K_dp4a");
7641 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7643 grid_dim: (out_f as u32, m as u32, 1),
7644 block_dim: (128, 1, 1),
7645 shared_mem_bytes: 0,
7646 };
7647 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7648 let __s_b = self.gpu.stream();
7649 let mut b = __s_b.launch_builder(&f);
7650 b.arg(w)
7651 .arg(&aq)
7652 .arg(&ad)
7653 .arg(&mut y)
7654 .arg(&inf)
7655 .arg(&outf)
7656 .arg(&mi)
7657 .arg(&rb);
7658 unsafe {
7659 b.launch(cfg)?;
7660 }
7661 Ok(y)
7662 }
7663
7664 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7667 &self,
7668 w: &CudaSlice<u8>,
7669 x: &CudaSlice<f32>,
7670 m: usize,
7671 in_f: usize,
7672 out_f: usize,
7673 row_bytes: usize,
7674 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7675 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7676 let f = self.func("qmatvec_q6_K_dp4a");
7677 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7679 grid_dim: (out_f as u32, m as u32, 1),
7680 block_dim: (128, 1, 1),
7681 shared_mem_bytes: 0,
7682 };
7683 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7684 let __s_b = self.gpu.stream();
7685 let mut b = __s_b.launch_builder(&f);
7686 b.arg(w)
7687 .arg(&aq)
7688 .arg(&ad)
7689 .arg(&mut y)
7690 .arg(&inf)
7691 .arg(&outf)
7692 .arg(&mi)
7693 .arg(&rb);
7694 unsafe {
7695 b.launch(cfg)?;
7696 }
7697 Ok(y)
7698 }
7699
7700 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7703 &self,
7704 w: &CudaSlice<u8>,
7705 x: &CudaSlice<f32>,
7706 m: usize,
7707 in_f: usize,
7708 out_f: usize,
7709 row_bytes: usize,
7710 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7711 self.qmatvec_dp4a_named(
7712 "qmatvec_q5_K_dp4a",
7713 &w.slice(0..w.len()),
7714 x,
7715 m,
7716 in_f,
7717 out_f,
7718 row_bytes,
7719 )
7720 }
7721 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7724 &self,
7725 w: &CudaSlice<u8>,
7726 x: &CudaSlice<f32>,
7727 m: usize,
7728 in_f: usize,
7729 out_f: usize,
7730 row_bytes: usize,
7731 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7732 self.qmatvec_dp4a_named(
7733 "qmatvec_q3_K_dp4a",
7734 &w.slice(0..w.len()),
7735 x,
7736 m,
7737 in_f,
7738 out_f,
7739 row_bytes,
7740 )
7741 }
7742 pub fn qmatvec_nvfp4_fast_rp(
7744 &self,
7745 w: &CudaSlice<u8>,
7746 x: &CudaSlice<f32>,
7747 m: usize,
7748 in_f: usize,
7749 out_f: usize,
7750 row_bytes: usize,
7751 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7752 assert!(
7753 in_f % 64 == 0,
7754 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7755 );
7756 self.qmatvec_dp4a_named(
7757 "qmatvec_nvfp4_dp4a_rp",
7758 &w.slice(0..w.len()),
7759 x,
7760 m,
7761 in_f,
7762 out_f,
7763 row_bytes,
7764 )
7765 }
7766 pub fn qmatvec_nvfp4_fast(
7768 &self,
7769 w: &cudarc::driver::CudaView<'_, u8>,
7770 x: &CudaSlice<f32>,
7771 m: usize,
7772 in_f: usize,
7773 out_f: usize,
7774 row_bytes: usize,
7775 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7776 assert!(
7779 in_f % 64 == 0,
7780 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7781 );
7782 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7783 }
7784 pub fn qmatvec_nvfp4_fast_v2(
7787 &self,
7788 w: &cudarc::driver::CudaView<'_, u8>,
7789 x: &CudaSlice<f32>,
7790 m: usize,
7791 in_f: usize,
7792 out_f: usize,
7793 row_bytes: usize,
7794 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7795 assert!(
7796 in_f % 64 == 0,
7797 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7798 );
7799 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7800 }
7801 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7804 &self,
7805 w: &CudaSlice<u8>,
7806 x: &CudaSlice<f32>,
7807 m: usize,
7808 in_f: usize,
7809 out_f: usize,
7810 row_bytes: usize,
7811 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7812 self.qmatvec_dp4a_named(
7813 "qmatvec_iq4_XS_dp4a",
7814 &w.slice(0..w.len()),
7815 x,
7816 m,
7817 in_f,
7818 out_f,
7819 row_bytes,
7820 )
7821 }
7822
7823 fn qmatvec_dp4a_named(
7825 &self,
7826 name: &str,
7827 w: &cudarc::driver::CudaView<'_, u8>,
7828 x: &CudaSlice<f32>,
7829 m: usize,
7830 in_f: usize,
7831 out_f: usize,
7832 row_bytes: usize,
7833 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7834 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7835 let f = self.func(name);
7836 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7838 grid_dim: (out_f as u32, m as u32, 1),
7839 block_dim: (128, 1, 1),
7840 shared_mem_bytes: 0,
7841 };
7842 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7843 let __s_b = self.gpu.stream();
7844 let mut b = __s_b.launch_builder(&f);
7845 b.arg(w)
7846 .arg(&aq)
7847 .arg(&ad)
7848 .arg(&mut y)
7849 .arg(&inf)
7850 .arg(&outf)
7851 .arg(&mi)
7852 .arg(&rb);
7853 unsafe {
7854 b.launch(cfg)?;
7855 }
7856 Ok(y)
7857 }
7858
7859 #[allow(clippy::too_many_arguments)]
7865 pub fn qmatvec_nvfp4_fast_prequant_into(
7866 &self,
7867 w: &CudaSlice<u8>,
7868 aq: &CudaSlice<i8>,
7869 ad: &CudaSlice<f32>,
7870 y: &mut CudaSlice<f32>,
7871 m: usize,
7872 in_f: usize,
7873 out_f: usize,
7874 row_bytes: usize,
7875 ) -> Result<(), Box<dyn std::error::Error>> {
7876 assert!(
7877 in_f % 64 == 0,
7878 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7879 );
7880 if y.len() < m * out_f {
7881 return Err(format!(
7882 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7883 y.len()
7884 )
7885 .into());
7886 }
7887 let f = self.func("qmatvec_nvfp4_dp4a");
7888 let cfg = LaunchConfig {
7889 grid_dim: (out_f as u32, m as u32, 1),
7890 block_dim: (128, 1, 1),
7891 shared_mem_bytes: 0,
7892 };
7893 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7894 let __s_b = self.gpu.stream();
7895 let mut b = __s_b.launch_builder(&f);
7896 b.arg(w)
7897 .arg(aq)
7898 .arg(ad)
7899 .arg(y)
7900 .arg(&inf)
7901 .arg(&outf)
7902 .arg(&mi)
7903 .arg(&rb);
7904 unsafe {
7905 b.launch(cfg)?;
7906 }
7907 Ok(())
7908 }
7909
7910 #[allow(clippy::too_many_arguments)]
7913 pub fn matvec_f32_qkv_into(
7914 &self,
7915 wq: &CudaSlice<f32>,
7916 wk: &CudaSlice<f32>,
7917 wv: &CudaSlice<f32>,
7918 wg: &CudaSlice<f32>,
7919 x: &CudaSlice<f32>,
7920 yq: &mut CudaSlice<f32>,
7921 yk: &mut CudaSlice<f32>,
7922 yv: &mut CudaSlice<f32>,
7923 yg: &mut CudaSlice<f32>,
7924 in_f: usize,
7925 out_q: usize,
7926 out_kv: usize,
7927 out_g: usize,
7928 ) -> Result<(), Box<dyn std::error::Error>> {
7929 if in_f % 4 != 0
7930 || wq.len() != out_q * in_f
7931 || wk.len() != out_kv * in_f
7932 || wv.len() != out_kv * in_f
7933 || wg.len() < out_g * in_f
7934 || x.len() < in_f
7935 || yq.len() < out_q
7936 || yk.len() < out_kv
7937 || yv.len() < out_kv
7938 || (out_g > 0 && yg.len() < out_g)
7939 {
7940 return Err(format!(
7941 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7942 wq={} wk={} wv={} wg={}",
7943 wq.len(),
7944 wk.len(),
7945 wv.len(),
7946 wg.len()
7947 )
7948 .into());
7949 }
7950 let f = self.func("matvec_f32_qkv");
7951 let cfg = LaunchConfig {
7952 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7953 block_dim: (128, 1, 1),
7954 shared_mem_bytes: 0,
7955 };
7956 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7957 let __s_b = self.gpu.stream();
7958 let mut b = __s_b.launch_builder(&f);
7959 b.arg(wq)
7960 .arg(wk)
7961 .arg(wv)
7962 .arg(wg)
7963 .arg(x)
7964 .arg(yq)
7965 .arg(yk)
7966 .arg(yv)
7967 .arg(yg)
7968 .arg(&inf)
7969 .arg(&oq)
7970 .arg(&okv)
7971 .arg(&og);
7972 unsafe {
7973 b.launch(cfg)?;
7974 }
7975 Ok(())
7976 }
7977
7978 #[allow(clippy::too_many_arguments)]
7984 #[allow(clippy::too_many_arguments)]
7987 pub fn qmatvec_nvfp4_sel_gu_into(
7988 &self,
7989 gate_bank: &CudaSlice<u8>,
7990 up_bank: &CudaSlice<u8>,
7991 sel: &CudaSlice<i32>,
7992 aq: &CudaSlice<i8>,
7993 ad: &CudaSlice<f32>,
7994 yg: &mut CudaSlice<f32>,
7995 yu: &mut CudaSlice<f32>,
7996 n_sel: usize,
7997 in_f: usize,
7998 out_f: usize,
7999 row_bytes: usize,
8000 expert_stride: usize,
8001 ) -> Result<(), Box<dyn std::error::Error>> {
8002 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8003 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8004 return Err("NVFP4 gu sel geometry".into());
8005 }
8006 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
8009 let rpw = *RPW.get_or_init(|| {
8010 std::env::var("MEMRA_SEL_GU_RPW")
8011 .ok()
8012 .and_then(|v| v.parse().ok())
8013 .filter(|r| *r == 2 || *r == 4)
8014 .unwrap_or(1)
8015 });
8016 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
8017 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8020 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
8021 let f = self.func(match (wpr, rpw) {
8022 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
8023 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
8024 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
8025 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
8026 });
8027 let cfg = LaunchConfig {
8028 grid_dim: if wpr {
8029 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
8030 } else if rpw == 1 {
8031 ((2 * out_f) as u32, n_sel as u32, 1)
8032 } else {
8033 ((out_f / rpw) as u32, n_sel as u32, 1)
8034 },
8035 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
8036 shared_mem_bytes: 0,
8037 };
8038 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8039 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8040 let (ars, adrs) = (0i64, 0i64);
8041 let __s_b = self.gpu.stream();
8042 let mut b = __s_b.launch_builder(&f);
8043 b.arg(gate_bank)
8044 .arg(up_bank)
8045 .arg(sel)
8046 .arg(aq)
8047 .arg(ad)
8048 .arg(yg)
8049 .arg(yu)
8050 .arg(&inf)
8051 .arg(&outf)
8052 .arg(&ns)
8053 .arg(&rb)
8054 .arg(&es)
8055 .arg(&ars)
8056 .arg(&adrs);
8057 unsafe {
8058 b.launch(cfg)?;
8059 }
8060 Ok(())
8061 }
8062
8063 #[allow(clippy::too_many_arguments)]
8071 pub fn qmatvec_nvfp4_sel_down8_into(
8072 &self,
8073 bank: &CudaSlice<u8>,
8074 sel: &CudaSlice<i32>,
8075 aq: &CudaSlice<i8>,
8076 ad: &CudaSlice<f32>,
8077 route_w: &CudaSlice<f32>,
8078 md: &CudaSlice<f32>,
8079 dst: &mut CudaSlice<f32>,
8080 n_sel: usize,
8081 in_f: usize,
8082 out_f: usize,
8083 row_bytes: usize,
8084 expert_stride: usize,
8085 act_row_stride: usize,
8086 ad_row_stride: usize,
8087 ) -> Result<(), Box<dyn std::error::Error>> {
8088 if in_f % 64 != 0
8089 || n_sel == 0
8090 || n_sel > 8
8091 || (in_f >> 5) > 32
8092 || dst.len() < out_f
8093 || sel.len() < n_sel
8094 || route_w.len() < n_sel
8095 {
8096 return Err(format!(
8097 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
8098 dst.len()
8099 )
8100 .into());
8101 }
8102 if !crate::tp::nvfp4_bank_v2_on() {
8103 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
8104 }
8105 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
8106 let cfg = LaunchConfig {
8107 grid_dim: (out_f as u32, 1, 1),
8108 block_dim: (32, n_sel as u32, 1),
8109 shared_mem_bytes: 0,
8110 };
8111 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8112 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8113 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8114 let __s_b = self.gpu.stream();
8115 let mut b = __s_b.launch_builder(&f);
8116 b.arg(bank)
8117 .arg(sel)
8118 .arg(aq)
8119 .arg(ad)
8120 .arg(route_w)
8121 .arg(md)
8122 .arg(dst)
8123 .arg(&inf)
8124 .arg(&outf)
8125 .arg(&ns)
8126 .arg(&rb)
8127 .arg(&es)
8128 .arg(&ars)
8129 .arg(&adrs);
8130 unsafe {
8131 b.launch(cfg)?;
8132 }
8133 Ok(())
8134 }
8135
8136 #[allow(clippy::too_many_arguments)]
8140 pub fn qmatvec_nvfp4_sel_down8_rows_into(
8141 &self,
8142 bank: &CudaSlice<u8>,
8143 sel: &CudaSlice<i32>,
8144 aq: &CudaSlice<i8>,
8145 ad: &CudaSlice<f32>,
8146 route_w: &CudaSlice<f32>,
8147 md: &CudaSlice<f32>,
8148 dst: &mut CudaSlice<f32>,
8149 t: usize,
8150 n_sel_col: 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 ) -> Result<(), Box<dyn std::error::Error>> {
8158 let n_sel = t * n_sel_col;
8159 if in_f % 64 != 0
8160 || n_sel_col == 0
8161 || n_sel_col > 8
8162 || t == 0
8163 || t > 64
8164 || (in_f >> 5) > 32
8165 || dst.len() < t * out_f
8166 || sel.len() < n_sel
8167 || route_w.len() < n_sel
8168 {
8169 return Err(format!(
8170 "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
8171 dst.len()
8172 )
8173 .into());
8174 }
8175 if !crate::tp::nvfp4_bank_v2_on() {
8176 return Err(
8177 "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8178 );
8179 }
8180 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8181 let cfg = LaunchConfig {
8182 grid_dim: (out_f as u32, t as u32, 1),
8183 block_dim: (32, n_sel_col as u32, 1),
8184 shared_mem_bytes: 0,
8185 };
8186 let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8187 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8188 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8189 let __s_b = self.gpu.stream();
8190 let mut b = __s_b.launch_builder(&f);
8191 b.arg(bank)
8192 .arg(sel)
8193 .arg(aq)
8194 .arg(ad)
8195 .arg(route_w)
8196 .arg(md)
8197 .arg(dst)
8198 .arg(&inf)
8199 .arg(&outf)
8200 .arg(&nsc)
8201 .arg(&rb)
8202 .arg(&es)
8203 .arg(&ars)
8204 .arg(&adrs);
8205 unsafe {
8206 b.launch(cfg)?;
8207 }
8208 Ok(())
8209 }
8210
8211 #[allow(clippy::too_many_arguments)]
8214 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8215 &self,
8216 gate_bank: &CudaSlice<u8>,
8217 up_bank: &CudaSlice<u8>,
8218 sel: &CudaSlice<i32>,
8219 aq: &CudaSlice<i8>,
8220 ad: &CudaSlice<f32>,
8221 yg: &mut CudaSlice<f32>,
8222 yu: &mut CudaSlice<f32>,
8223 n_sel: usize,
8224 in_f: usize,
8225 out_f: usize,
8226 row_bytes: usize,
8227 expert_stride: usize,
8228 owner: usize,
8229 ) -> Result<(), Box<dyn std::error::Error>> {
8230 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8231 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8232 return Err("NVFP4 gu ep geometry".into());
8233 }
8234 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8235 let cfg = LaunchConfig {
8236 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8237 block_dim: (128, 1, 1),
8238 shared_mem_bytes: 0,
8239 };
8240 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8241 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8242 let (ars, adrs) = (0i64, 0i64);
8243 let __s_b = self.gpu.stream();
8244 let mut b = __s_b.launch_builder(&f);
8245 b.arg(gate_bank)
8246 .arg(up_bank)
8247 .arg(sel)
8248 .arg(aq)
8249 .arg(ad)
8250 .arg(yg)
8251 .arg(yu)
8252 .arg(&inf)
8253 .arg(&outf)
8254 .arg(&ns)
8255 .arg(&rb)
8256 .arg(&es)
8257 .arg(&ars)
8258 .arg(&adrs)
8259 .arg(&own);
8260 unsafe {
8261 b.launch(cfg)?;
8262 }
8263 Ok(())
8264 }
8265
8266 #[allow(clippy::too_many_arguments)]
8268 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8269 &self,
8270 gate: &CudaSlice<f32>,
8271 up: &CudaSlice<f32>,
8272 gmac: &CudaSlice<f32>,
8273 umac: &CudaSlice<f32>,
8274 sel: &CudaSlice<i32>,
8275 limit: Option<f32>,
8276 out_q: &mut CudaSlice<i8>,
8277 out_d: &mut CudaSlice<f32>,
8278 n_per: usize,
8279 n_sel: usize,
8280 owner: usize,
8281 ) -> Result<(), Box<dyn std::error::Error>> {
8282 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8283 return Err("NVFP4 silu ep geometry".into());
8284 }
8285 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8286 let warps = n_sel * n_per / 32;
8287 let cfg = LaunchConfig {
8288 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8289 block_dim: (128, 1, 1),
8290 shared_mem_bytes: 0,
8291 };
8292 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8293 let (lim, has) = match limit {
8294 Some(l) => (l, 1i32),
8295 None => (0.0f32, 0i32),
8296 };
8297 let __s_b = self.gpu.stream();
8298 let mut b = __s_b.launch_builder(&f);
8299 b.arg(gate)
8300 .arg(up)
8301 .arg(gmac)
8302 .arg(umac)
8303 .arg(sel)
8304 .arg(&lim)
8305 .arg(&has)
8306 .arg(out_q)
8307 .arg(out_d)
8308 .arg(&np)
8309 .arg(&ns)
8310 .arg(&own);
8311 unsafe {
8312 b.launch(cfg)?;
8313 }
8314 Ok(())
8315 }
8316
8317 #[allow(clippy::too_many_arguments)]
8319 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8320 &self,
8321 bank: &CudaSlice<u8>,
8322 sel: &CudaSlice<i32>,
8323 aq: &CudaSlice<i8>,
8324 ad: &CudaSlice<f32>,
8325 route_w: &CudaSlice<f32>,
8326 md: &CudaSlice<f32>,
8327 dst: &mut CudaSlice<f32>,
8328 n_sel: usize,
8329 in_f: usize,
8330 out_f: usize,
8331 row_bytes: usize,
8332 expert_stride: usize,
8333 act_row_stride: usize,
8334 ad_row_stride: usize,
8335 owner: usize,
8336 ) -> Result<(), Box<dyn std::error::Error>> {
8337 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8338 return Err("NVFP4 down8 ep geometry".into());
8339 }
8340 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8341 let cfg = LaunchConfig {
8342 grid_dim: (out_f as u32, 1, 1),
8343 block_dim: (32, n_sel as u32, 1),
8344 shared_mem_bytes: 0,
8345 };
8346 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8347 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8348 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8349 let __s_b = self.gpu.stream();
8350 let mut b = __s_b.launch_builder(&f);
8351 b.arg(bank)
8352 .arg(sel)
8353 .arg(aq)
8354 .arg(ad)
8355 .arg(route_w)
8356 .arg(md)
8357 .arg(dst)
8358 .arg(&inf)
8359 .arg(&outf)
8360 .arg(&ns)
8361 .arg(&rb)
8362 .arg(&es)
8363 .arg(&ars)
8364 .arg(&adrs)
8365 .arg(&own);
8366 unsafe {
8367 b.launch(cfg)?;
8368 }
8369 Ok(())
8370 }
8371
8372 pub fn qmatvec_nvfp4_sel_into(
8373 &self,
8374 bank: &CudaSlice<u8>,
8375 sel: &CudaSlice<i32>,
8376 aq: &CudaSlice<i8>,
8377 ad: &CudaSlice<f32>,
8378 y: &mut CudaSlice<f32>,
8379 n_sel: usize,
8380 in_f: usize,
8381 out_f: usize,
8382 row_bytes: usize,
8383 expert_stride: usize,
8384 act_row_stride: usize,
8385 ad_row_stride: usize,
8386 ) -> Result<(), Box<dyn std::error::Error>> {
8387 assert!(
8388 in_f % 64 == 0,
8389 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8390 );
8391 if y.len() < n_sel * out_f || sel.len() < n_sel {
8392 return Err(format!(
8393 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8394 y.len(),
8395 sel.len()
8396 )
8397 .into());
8398 }
8399 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8406 let mode = *MR.get_or_init(|| {
8407 if crate::tp::nvfp4_bank_v2_on() {
8408 3
8409 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8410 2
8411 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8412 1
8413 } else {
8414 0
8415 }
8416 });
8417 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8418 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8422 let v2s = mode == 3
8423 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8424 && row_bytes % 16 == 0
8425 && in_f <= 4096;
8426 let f = match (mode, v2s) {
8427 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8428 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8429 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8430 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8431 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8432 };
8433 let nsb = in_f >> 5;
8438 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8439 32
8440 } else if mode == 1 {
8441 512
8442 } else {
8443 128
8444 };
8445 let cfg = LaunchConfig {
8446 grid_dim: (
8447 if v2s {
8448 (out_f as u32).div_ceil(8)
8449 } else {
8450 match mode {
8451 2 => (out_f as u32).div_ceil(16),
8452 1 => (out_f as u32).div_ceil(4),
8453 _ => out_f as u32,
8454 }
8455 },
8456 n_sel as u32,
8457 1,
8458 ),
8459 block_dim: (fit_block, 1, 1),
8460 shared_mem_bytes: 0,
8461 };
8462 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8463 let (rb, es, ars, adrs) = (
8464 row_bytes as i64,
8465 expert_stride as i64,
8466 act_row_stride as i64,
8467 ad_row_stride as i64,
8468 );
8469 let __s_b = self.gpu.stream();
8470 let mut b = __s_b.launch_builder(&f);
8471 b.arg(bank)
8472 .arg(sel)
8473 .arg(aq)
8474 .arg(ad)
8475 .arg(y)
8476 .arg(&inf)
8477 .arg(&outf)
8478 .arg(&ns)
8479 .arg(&rb)
8480 .arg(&es)
8481 .arg(&ars)
8482 .arg(&adrs);
8483 unsafe {
8484 b.launch(cfg)?;
8485 }
8486 Ok(())
8487 }
8488
8489 #[allow(clippy::too_many_arguments)]
8494 pub fn silu_mul_scaled_q8_1_sel_into(
8495 &self,
8496 gate: &CudaSlice<f32>,
8497 up: &CudaSlice<f32>,
8498 gmac: &CudaSlice<f32>,
8499 umac: &CudaSlice<f32>,
8500 sel: &CudaSlice<i32>,
8501 limit: Option<f32>,
8502 out_q: &mut CudaSlice<i8>,
8503 out_d: &mut CudaSlice<f32>,
8504 n_per: usize,
8505 n_sel: usize,
8506 ) -> Result<(), Box<dyn std::error::Error>> {
8507 let n = n_per * n_sel;
8508 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8509 return Err(format!(
8510 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8511 out_q.len(),
8512 out_d.len()
8513 )
8514 .into());
8515 }
8516 if let Some(limit) = limit {
8517 if limit <= 1e-6 {
8518 return Err(format!(
8519 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8520 )
8521 .into());
8522 }
8523 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8524 let cfg = LaunchConfig::for_num_elems(n as u32);
8525 let (np, ns) = (n_per as i32, n_sel as i32);
8526 let __s_b = self.gpu.stream();
8527 let mut b = __s_b.launch_builder(&f);
8528 b.arg(gate)
8529 .arg(up)
8530 .arg(gmac)
8531 .arg(umac)
8532 .arg(sel)
8533 .arg(&limit)
8534 .arg(out_q)
8535 .arg(out_d)
8536 .arg(&np)
8537 .arg(&ns);
8538 unsafe {
8539 b.launch(cfg)?;
8540 }
8541 return Ok(());
8542 }
8543 let f = self.func("silu_mul_scaled_q8_1_sel");
8544 let cfg = LaunchConfig::for_num_elems(n as u32);
8545 let (np, ns) = (n_per as i32, n_sel as i32);
8546 let __s_b = self.gpu.stream();
8547 let mut b = __s_b.launch_builder(&f);
8548 b.arg(gate)
8549 .arg(up)
8550 .arg(gmac)
8551 .arg(umac)
8552 .arg(sel)
8553 .arg(out_q)
8554 .arg(out_d)
8555 .arg(&np)
8556 .arg(&ns);
8557 unsafe {
8558 b.launch(cfg)?;
8559 }
8560 Ok(())
8561 }
8562
8563 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8564 Ok(self.gpu.stream().clone_htod(v)?)
8565 }
8566 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8567 Ok(self.gpu.stream().clone_htod(v)?)
8568 }
8569 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8571 Ok(self.gpu.stream().clone_htod(v)?)
8572 }
8573 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8574 Ok(self.gpu.stream().clone_htod(v)?)
8575 }
8576 pub fn dtoh_view(
8578 &self,
8579 d: &cudarc::driver::CudaView<f32>,
8580 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8581 let v = self.gpu.stream().clone_dtoh(d)?;
8582 self.gpu.stream().synchronize()?;
8583 Ok(v)
8584 }
8585 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8586 let v = self.gpu.stream().clone_dtoh(d)?;
8587 self.gpu.stream().synchronize()?;
8588 Ok(v)
8589 }
8590 pub fn dtoh_pair(
8594 &self,
8595 a: &CudaSlice<f32>,
8596 b: &CudaSlice<f32>,
8597 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8598 let av = self.gpu.stream().clone_dtoh(a)?;
8599 let bv = self.gpu.stream().clone_dtoh(b)?;
8600 self.gpu.stream().synchronize()?;
8601 Ok((av, bv))
8602 }
8603 pub fn dtoh_pair_views(
8606 &self,
8607 a: &cudarc::driver::CudaView<f32>,
8608 b: &cudarc::driver::CudaView<f32>,
8609 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8610 let av = self.gpu.stream().clone_dtoh(a)?;
8611 let bv = self.gpu.stream().clone_dtoh(b)?;
8612 self.gpu.stream().synchronize()?;
8613 Ok((av, bv))
8614 }
8615 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8617 let v = self.gpu.stream().clone_dtoh(d)?;
8618 self.gpu.stream().synchronize()?;
8619 Ok(v)
8620 }
8621 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8623 let v = self.gpu.stream().clone_dtoh(d)?;
8624 self.gpu.stream().synchronize()?;
8625 Ok(v)
8626 }
8627 pub fn dtoh_u8_view(
8628 &self,
8629 d: &cudarc::driver::CudaView<u8>,
8630 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8631 let v = self.gpu.stream().clone_dtoh(d)?;
8632 self.gpu.stream().synchronize()?;
8633 Ok(v)
8634 }
8635 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8636 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8637 self.keep_if_capturing(&s);
8638 Ok(s)
8639 }
8640
8641 pub fn prob_of_token_device(
8650 &self,
8651 logits: &CudaSlice<f32>,
8652 tok: &CudaSlice<u32>,
8653 n_vocab: usize,
8654 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8655 let nb = ARGMAX_NB;
8656 let mut part = self.alloc_uninit::<f32>(nb)?;
8657 let mut p = self.alloc_uninit::<f32>(1)?;
8658 let f1 = self.func("prob_of_token_partial_f32");
8659 let cfg1 = LaunchConfig {
8660 grid_dim: (nb as u32, 1, 1),
8661 block_dim: (256, 1, 1),
8662 shared_mem_bytes: 0,
8663 };
8664 let nv = n_vocab as i32;
8665 let __s_b1 = self.gpu.stream();
8666 let mut b1 = __s_b1.launch_builder(&f1);
8667 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8668 unsafe {
8669 b1.launch(cfg1)?;
8670 }
8671 let f2 = self.func("prob_of_token_final_f32");
8672 let cfg2 = LaunchConfig {
8673 grid_dim: (1, 1, 1),
8674 block_dim: (256, 1, 1),
8675 shared_mem_bytes: 0,
8676 };
8677 let nbi = nb as i32;
8678 let __s_b2 = self.gpu.stream();
8679 let mut b2 = __s_b2.launch_builder(&f2);
8680 b2.arg(&part).arg(&mut p).arg(&nbi);
8681 unsafe {
8682 b2.launch(cfg2)?;
8683 }
8684 Ok(p)
8685 }
8686
8687 pub fn prob_of_token_device_col(
8694 &self,
8695 logits: &CudaSlice<f32>,
8696 tok_all: &CudaSlice<u32>,
8697 tok_idx: usize,
8698 p_out: &mut CudaSlice<f32>,
8699 p_idx: usize,
8700 n_vocab: usize,
8701 ) -> Result<(), Box<dyn std::error::Error>> {
8702 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8703 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8704 let nb = ARGMAX_NB;
8705 let mut part = self.alloc_uninit::<f32>(nb)?;
8706 let f1 = self.func("prob_of_token_partial_f32");
8707 let cfg1 = LaunchConfig {
8708 grid_dim: (nb as u32, 1, 1),
8709 block_dim: (256, 1, 1),
8710 shared_mem_bytes: 0,
8711 };
8712 let nv = n_vocab as i32;
8713 let __s_b1 = self.gpu.stream();
8714 let mut b1 = __s_b1.launch_builder(&f1);
8715 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8716 unsafe {
8717 b1.launch(cfg1)?;
8718 }
8719 let f2 = self.func("prob_of_token_final_f32");
8720 let cfg2 = LaunchConfig {
8721 grid_dim: (1, 1, 1),
8722 block_dim: (256, 1, 1),
8723 shared_mem_bytes: 0,
8724 };
8725 let nbi = nb as i32;
8726 let __s_b2 = self.gpu.stream();
8727 let mut b2 = __s_b2.launch_builder(&f2);
8728 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8729 unsafe {
8730 b2.launch(cfg2)?;
8731 }
8732 Ok(())
8733 }
8734
8735 pub fn prob_of_token_device_into(
8736 &self,
8737 logits: &CudaSlice<f32>,
8738 tok: &CudaSlice<u32>,
8739 p_out: &mut CudaSlice<f32>,
8740 n_vocab: usize,
8741 ) -> Result<(), Box<dyn std::error::Error>> {
8742 let nb = ARGMAX_NB;
8743 let mut part = self.alloc_uninit::<f32>(nb)?;
8744 let f1 = self.func("prob_of_token_partial_f32");
8745 let cfg1 = LaunchConfig {
8746 grid_dim: (nb as u32, 1, 1),
8747 block_dim: (256, 1, 1),
8748 shared_mem_bytes: 0,
8749 };
8750 let nv = n_vocab as i32;
8751 let __s_b1 = self.gpu.stream();
8752 let mut b1 = __s_b1.launch_builder(&f1);
8753 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8754 unsafe {
8755 b1.launch(cfg1)?;
8756 }
8757 let f2 = self.func("prob_of_token_final_f32");
8758 let cfg2 = LaunchConfig {
8759 grid_dim: (1, 1, 1),
8760 block_dim: (256, 1, 1),
8761 shared_mem_bytes: 0,
8762 };
8763 let nbi = nb as i32;
8764 let __s_b2 = self.gpu.stream();
8765 let mut b2 = __s_b2.launch_builder(&f2);
8766 b2.arg(&part).arg(p_out).arg(&nbi);
8767 unsafe {
8768 b2.launch(cfg2)?;
8769 }
8770 Ok(())
8771 }
8772
8773 pub fn u32_hist_append(
8776 &self,
8777 tok: &CudaSlice<u32>,
8778 hist: &mut CudaSlice<u32>,
8779 idx: &mut CudaSlice<i32>,
8780 ) -> Result<(), Box<dyn std::error::Error>> {
8781 let f = self.func("u32_hist_append");
8782 let cfg = LaunchConfig {
8783 grid_dim: (1, 1, 1),
8784 block_dim: (32, 1, 1),
8785 shared_mem_bytes: 0,
8786 };
8787 let __s_b = self.gpu.stream();
8788 let mut b = __s_b.launch_builder(&f);
8789 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8790 unsafe {
8791 b.launch(cfg)?;
8792 }
8793 Ok(())
8794 }
8795
8796 pub fn argmax_token_device(
8797 &self,
8798 logits: &CudaSlice<f32>,
8799 n_vocab: usize,
8800 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8801 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8802 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8803 Ok(tok)
8804 }
8805 pub fn argmax_token_device_into(
8812 &self,
8813 logits: &CudaSlice<f32>,
8814 tok: &mut CudaSlice<u32>,
8815 n_vocab: usize,
8816 ) -> Result<(), Box<dyn std::error::Error>> {
8817 let nb = ARGMAX_NB;
8818 let f1 = self.func("argmax_partial_f32");
8819 let f2 = self.func("argmax_final_f32");
8820 let mut guard = self.argmax_partials.lock().unwrap();
8821 if guard.is_none() {
8822 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8825 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8826 *guard = Some((pv, pi));
8827 }
8828 let (part_v, part_i) = guard.as_mut().unwrap();
8829 let nv = n_vocab as i32;
8830 let nbi = nb as i32;
8831 let cfg1 = LaunchConfig {
8833 grid_dim: (nb as u32, 1, 1),
8834 block_dim: (256, 1, 1),
8835 shared_mem_bytes: 0,
8836 };
8837 let __s_b1 = self.gpu.stream();
8838 let mut b1 = __s_b1.launch_builder(&f1);
8839 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8840 unsafe {
8841 b1.launch(cfg1)?;
8842 }
8843 let cfg2 = LaunchConfig {
8845 grid_dim: (1, 1, 1),
8846 block_dim: (256, 1, 1),
8847 shared_mem_bytes: 0,
8848 };
8849 let __s_b2 = self.gpu.stream();
8850 let mut b2 = __s_b2.launch_builder(&f2);
8851 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8852 unsafe {
8853 b2.launch(cfg2)?;
8854 }
8855 Ok(())
8856 }
8857 pub fn argmax_token_device_col(
8863 &self,
8864 logits: &CudaSlice<f32>,
8865 col: usize,
8866 n_vocab: usize,
8867 toks: &mut CudaSlice<u32>,
8868 out_idx: usize,
8869 ) -> Result<(), Box<dyn std::error::Error>> {
8870 let nb = ARGMAX_NB;
8871 let f1 = self.func("argmax_partial_f32");
8872 let f2 = self.func("argmax_final_f32");
8873 let mut guard = self.argmax_partials.lock().unwrap();
8874 if guard.is_none() {
8875 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8876 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8877 *guard = Some((pv, pi));
8878 }
8879 let (part_v, part_i) = guard.as_mut().unwrap();
8880 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8881 let nv = n_vocab as i32;
8882 let nbi = nb as i32;
8883 let cfg1 = LaunchConfig {
8884 grid_dim: (nb as u32, 1, 1),
8885 block_dim: (256, 1, 1),
8886 shared_mem_bytes: 0,
8887 };
8888 let __s_b1 = self.gpu.stream();
8889 let mut b1 = __s_b1.launch_builder(&f1);
8890 b1.arg(&col_view)
8891 .arg(&mut *part_v)
8892 .arg(&mut *part_i)
8893 .arg(&nv);
8894 unsafe {
8895 b1.launch(cfg1)?;
8896 }
8897 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8898 let cfg2 = LaunchConfig {
8899 grid_dim: (1, 1, 1),
8900 block_dim: (256, 1, 1),
8901 shared_mem_bytes: 0,
8902 };
8903 let __s_b2 = self.gpu.stream();
8904 let mut b2 = __s_b2.launch_builder(&f2);
8905 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8906 unsafe {
8907 b2.launch(cfg2)?;
8908 }
8909 Ok(())
8910 }
8911 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8913 Ok(self.gpu.stream().clone_htod(v)?)
8914 }
8915 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8916 let v = self.gpu.stream().clone_dtoh(d)?;
8917 self.gpu.stream().synchronize()?;
8918 Ok(v)
8919 }
8920
8921 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8922 let v = self.gpu.stream().clone_dtoh(d)?;
8923 self.gpu.stream().synchronize()?;
8924 Ok(v)
8925 }
8926 pub fn htod_u32_into(
8930 &self,
8931 dst: &mut CudaSlice<u32>,
8932 src: &[u32],
8933 ) -> Result<(), Box<dyn std::error::Error>> {
8934 let mut view = dst.slice_mut(0..src.len());
8935 self.gpu.stream().memcpy_htod(src, &mut view)?;
8936 Ok(())
8937 }
8938
8939 pub fn htod_i32_into(
8942 &self,
8943 dst: &mut CudaSlice<i32>,
8944 src: &[i32],
8945 ) -> Result<(), Box<dyn std::error::Error>> {
8946 let mut view = dst.slice_mut(0..src.len());
8947 self.gpu.stream().memcpy_htod(src, &mut view)?;
8948 Ok(())
8949 }
8950
8951 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8952 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8953 self.keep_if_capturing(&s);
8954 Ok(s)
8955 }
8956 pub fn embed_gather_device_into(
8959 &self,
8960 embd: &CudaSlice<u8>,
8961 token_d: &CudaSlice<u32>,
8962 x_out: &mut CudaSlice<f32>,
8963 n_embd: usize,
8964 qtype: i32,
8965 row_bytes: usize,
8966 ) -> Result<(), Box<dyn std::error::Error>> {
8967 let f = self.func("embed_gather_u32");
8968 let cfg = LaunchConfig {
8969 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8970 block_dim: (256, 1, 1),
8971 shared_mem_bytes: 0,
8972 };
8973 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8974 let __s_b = self.gpu.stream();
8975 let mut b = __s_b.launch_builder(&f);
8976 b.arg(embd)
8977 .arg(token_d)
8978 .arg(x_out)
8979 .arg(&ne)
8980 .arg(&qt)
8981 .arg(&rb);
8982 unsafe {
8983 b.launch(cfg)?;
8984 }
8985 Ok(())
8986 }
8987 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8989 let v = self.gpu.stream().clone_dtoh(d)?;
8990 self.gpu.stream().synchronize()?;
8991 Ok(v[0])
8992 }
8993 pub fn i32_set_k(
9000 &self,
9001 dst: &mut CudaSlice<i32>,
9002 v: i32,
9003 ) -> Result<(), Box<dyn std::error::Error>> {
9004 let f = self.func("i32_set_k");
9005 let cfg = LaunchConfig {
9006 grid_dim: (1, 1, 1),
9007 block_dim: (1, 1, 1),
9008 shared_mem_bytes: 0,
9009 };
9010 let idx = 0i32;
9011 let __s_b = self.gpu.stream();
9012 let mut b = __s_b.launch_builder(&f);
9013 b.arg(dst).arg(&v).arg(&idx);
9014 unsafe {
9015 b.launch(cfg)?;
9016 }
9017 Ok(())
9018 }
9019
9020 pub fn set_i32_one(
9021 &self,
9022 d: &mut CudaSlice<i32>,
9023 v: i32,
9024 ) -> Result<(), Box<dyn std::error::Error>> {
9025 self.gpu.stream().memcpy_htod(&[v], d)?;
9026 Ok(())
9027 }
9028 pub fn set_u32_one(
9031 &self,
9032 d: &mut CudaSlice<u32>,
9033 v: u32,
9034 ) -> Result<(), Box<dyn std::error::Error>> {
9035 self.gpu.stream().memcpy_htod(&[v], d)?;
9036 Ok(())
9037 }
9038 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9040 let v = self.gpu.stream().clone_dtoh(d)?;
9041 self.gpu.stream().synchronize()?;
9042 Ok(v[0])
9043 }
9044 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9046 Ok(self.gpu.stream().clone_htod(bytes)?)
9047 }
9048 pub fn embed_gather_device(
9052 &self,
9053 embd: &CudaSlice<u8>,
9054 token_d: &CudaSlice<u32>,
9055 n_embd: usize,
9056 qtype: i32,
9057 row_bytes: usize,
9058 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9059 let f = self.func("embed_gather_u32");
9060 let mut x = self.alloc_uninit::<f32>(n_embd)?;
9061 let cfg = LaunchConfig {
9062 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9063 block_dim: (256, 1, 1),
9064 shared_mem_bytes: 0,
9065 };
9066 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9067 let __s_b = self.gpu.stream();
9068 let mut b = __s_b.launch_builder(&f);
9069 b.arg(embd)
9070 .arg(token_d)
9071 .arg(&mut x)
9072 .arg(&ne)
9073 .arg(&qt)
9074 .arg(&rb);
9075 unsafe {
9076 b.launch(cfg)?;
9077 }
9078 Ok(x)
9079 }
9080
9081 pub fn embed_gather_device_t(
9085 &self,
9086 embd: &CudaSlice<u8>,
9087 tokens: &[u32],
9088 n_embd: usize,
9089 qtype: i32,
9090 row_bytes: usize,
9091 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9092 let t = tokens.len();
9093 let tok_d = self.gpu.stream().clone_htod(tokens)?;
9094 let f = self.func("embed_gather_u32_t");
9095 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9096 let cfg = LaunchConfig {
9097 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9098 block_dim: (256, 1, 1),
9099 shared_mem_bytes: 0,
9100 };
9101 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9102 let __s_b = self.gpu.stream();
9103 let mut b = __s_b.launch_builder(&f);
9104 b.arg(embd)
9105 .arg(&tok_d)
9106 .arg(&mut x)
9107 .arg(&ne)
9108 .arg(&qt)
9109 .arg(&rb)
9110 .arg(&ti);
9111 unsafe {
9112 b.launch(cfg)?;
9113 }
9114 Ok(x)
9115 }
9116
9117 pub fn embed_gather_device_tv(
9122 &self,
9123 embd: &CudaSlice<u8>,
9124 tok_v: &cudarc::driver::CudaView<u32>,
9125 t: usize,
9126 n_embd: usize,
9127 qtype: i32,
9128 row_bytes: usize,
9129 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9130 let f = self.func("embed_gather_u32_t");
9131 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9132 let cfg = LaunchConfig {
9133 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9134 block_dim: (256, 1, 1),
9135 shared_mem_bytes: 0,
9136 };
9137 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9138 let __s_b = self.gpu.stream();
9139 let mut b = __s_b.launch_builder(&f);
9140 b.arg(embd)
9141 .arg(tok_v)
9142 .arg(&mut x)
9143 .arg(&ne)
9144 .arg(&qt)
9145 .arg(&rb)
9146 .arg(&ti);
9147 unsafe {
9148 b.launch(cfg)?;
9149 }
9150 Ok(x)
9151 }
9152
9153 pub fn embed_gather_device_td(
9154 &self,
9155 embd: &CudaSlice<u8>,
9156 tok_d: &CudaSlice<u32>,
9157 t: usize,
9158 n_embd: usize,
9159 qtype: i32,
9160 row_bytes: usize,
9161 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9162 let f = self.func("embed_gather_u32_t");
9163 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9164 let cfg = LaunchConfig {
9165 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9166 block_dim: (256, 1, 1),
9167 shared_mem_bytes: 0,
9168 };
9169 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9170 let __s_b = self.gpu.stream();
9171 let mut b = __s_b.launch_builder(&f);
9172 b.arg(embd)
9173 .arg(tok_d)
9174 .arg(&mut x)
9175 .arg(&ne)
9176 .arg(&qt)
9177 .arg(&rb)
9178 .arg(&ti);
9179 unsafe {
9180 b.launch(cfg)?;
9181 }
9182 Ok(x)
9183 }
9184
9185 #[inline]
9191 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9193 if self
9194 .capture_keep_on
9195 .load(std::sync::atomic::Ordering::Relaxed)
9196 {
9197 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9198 }
9199 }
9200
9201 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9202 &self,
9203 n: usize,
9204 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9205 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9206 {
9210 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9211 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9212 use cudarc::driver::DevicePtrMut;
9214 let n_bytes = s.len() * std::mem::size_of::<T>();
9215 let stream = self.gpu.stream();
9216 let (p_, _g) = s.device_ptr_mut(&stream);
9217 unsafe {
9218 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9219 .result()?;
9220 }
9221 }
9222 }
9223 self.keep_if_capturing(&s);
9224 Ok(s)
9225 }
9226
9227 pub fn uninit_q8_pair(
9232 &self,
9233 n: usize,
9234 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9235 Ok((
9236 self.alloc_uninit::<i8>(n)?,
9237 self.alloc_uninit::<f32>(n / 32)?,
9238 ))
9239 }
9240
9241 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9242 self.alloc_uninit::<f32>(n)
9243 }
9244
9245 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9247 self.alloc_uninit::<i8>(n)
9248 }
9249
9250 #[allow(clippy::too_many_arguments)]
9254 pub fn rms_norm3(
9255 &self,
9256 x: &CudaSlice<f32>,
9257 w0: &CudaSlice<f32>,
9258 w1: &CudaSlice<f32>,
9259 w2: &CudaSlice<f32>,
9260 d0: &mut CudaSlice<f32>,
9261 d1: &mut CudaSlice<f32>,
9262 d2: &mut CudaSlice<f32>,
9263 ncols: usize,
9264 nrows: usize,
9265 eps: f32,
9266 ) -> Result<(), Box<dyn std::error::Error>> {
9267 let f = self.func("rms_norm3_f32");
9268 let cfg = LaunchConfig {
9269 grid_dim: (nrows as u32, 1, 1),
9270 block_dim: (rms_block(), 1, 1),
9271 shared_mem_bytes: 0,
9272 };
9273 let (nc, e) = (ncols as i32, eps);
9274 let __s_b = self.gpu.stream();
9275 let mut b = __s_b.launch_builder(&f);
9276 b.arg(x)
9277 .arg(w0)
9278 .arg(w1)
9279 .arg(w2)
9280 .arg(d0)
9281 .arg(d1)
9282 .arg(d2)
9283 .arg(&nc)
9284 .arg(&e);
9285 unsafe {
9286 b.launch(cfg)?;
9287 }
9288 Ok(())
9289 }
9290
9291 #[allow(clippy::too_many_arguments)]
9293 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9296 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9297 *WARP_ON.get_or_init(|| {
9298 std::env::var("MEMRA_QKVNORM_W")
9299 .map(|v| v != "0")
9300 .unwrap_or(true)
9301 }) && ncols % 4 == 0
9302 && rows >= 64
9303 }
9304
9305 #[allow(clippy::too_many_arguments)]
9308 pub fn rms_norm_qkv_w4b(
9309 &self,
9310 q: &CudaSlice<f32>,
9311 k: &CudaSlice<f32>,
9312 v: &CudaSlice<f32>,
9313 wq: &CudaSlice<f32>,
9314 wk: &CudaSlice<f32>,
9315 wv: &CudaSlice<f32>,
9316 dq: &mut CudaSlice<f32>,
9317 dk: &mut CudaSlice<f32>,
9318 dv: &mut CudaSlice<f32>,
9319 dvb: &mut CudaSlice<u8>,
9320 ncols: usize,
9321 rq: usize,
9322 rk: usize,
9323 eps: f32,
9324 vf16: bool,
9325 ) -> Result<(), Box<dyn std::error::Error>> {
9326 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9327 let f = self.func("rms_norm_qkv_w4b_f32");
9328 let rows = (rq + 2 * rk) as u32;
9329 let cfg = LaunchConfig {
9330 grid_dim: (rows.div_ceil(8), 1, 1),
9331 block_dim: (256, 1, 1),
9332 shared_mem_bytes: 0,
9333 };
9334 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9335 let vf = vf16 as i32;
9336 let __s_b = self.gpu.stream();
9337 let mut b = __s_b.launch_builder(&f);
9338 b.arg(q)
9339 .arg(k)
9340 .arg(v)
9341 .arg(wq)
9342 .arg(wk)
9343 .arg(wv)
9344 .arg(dq)
9345 .arg(dk)
9346 .arg(dv)
9347 .arg(&mut *dvb)
9348 .arg(&nc)
9349 .arg(&rqi)
9350 .arg(&rki)
9351 .arg(&rvi)
9352 .arg(&e)
9353 .arg(&vf);
9354 unsafe {
9355 b.launch(cfg)?;
9356 }
9357 Ok(())
9358 }
9359
9360 pub fn rms_norm_qkv(
9361 &self,
9362 q: &CudaSlice<f32>,
9363 k: &CudaSlice<f32>,
9364 v: &CudaSlice<f32>,
9365 wq: &CudaSlice<f32>,
9366 wk: &CudaSlice<f32>,
9367 wv: &CudaSlice<f32>,
9368 dq: &mut CudaSlice<f32>,
9369 dk: &mut CudaSlice<f32>,
9370 dv: &mut CudaSlice<f32>,
9371 ncols: usize,
9372 rq: usize,
9373 rk: usize,
9374 eps: f32,
9375 ) -> Result<(), Box<dyn std::error::Error>> {
9376 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9380 let warp_on = *WARP_ON.get_or_init(|| {
9381 std::env::var("MEMRA_QKVNORM_W")
9382 .map(|v| v != "0")
9383 .unwrap_or(true)
9384 });
9385 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9388 let f = self.func("rms_norm_qkv_w4_f32");
9389 let rows = (rq + 2 * rk) as u32;
9390 let cfg = LaunchConfig {
9391 grid_dim: (rows.div_ceil(8), 1, 1),
9392 block_dim: (256, 1, 1),
9393 shared_mem_bytes: 0,
9394 };
9395 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9396 let __s_b = self.gpu.stream();
9397 let mut b = __s_b.launch_builder(&f);
9398 b.arg(q)
9399 .arg(k)
9400 .arg(v)
9401 .arg(wq)
9402 .arg(wk)
9403 .arg(wv)
9404 .arg(dq)
9405 .arg(dk)
9406 .arg(dv)
9407 .arg(&nc)
9408 .arg(&rqi)
9409 .arg(&rki)
9410 .arg(&rvi)
9411 .arg(&e);
9412 unsafe {
9413 b.launch(cfg)?;
9414 }
9415 return Ok(());
9416 }
9417 let f = self.func("rms_norm_qkv_f32");
9418 let grid = (rq + 2 * rk) as u32;
9419 let cfg = LaunchConfig {
9420 grid_dim: (grid, 1, 1),
9421 block_dim: (rms_block(), 1, 1),
9422 shared_mem_bytes: 0,
9423 };
9424 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9425 let __s_b = self.gpu.stream();
9426 let mut b = __s_b.launch_builder(&f);
9427 b.arg(q)
9428 .arg(k)
9429 .arg(v)
9430 .arg(wq)
9431 .arg(wk)
9432 .arg(wv)
9433 .arg(dq)
9434 .arg(dk)
9435 .arg(dv)
9436 .arg(&nc)
9437 .arg(&rqi)
9438 .arg(&rki)
9439 .arg(&e);
9440 unsafe {
9441 b.launch(cfg)?;
9442 }
9443 Ok(())
9444 }
9445
9446 #[allow(clippy::too_many_arguments)]
9448 pub fn rms_norm2x(
9449 &self,
9450 a: &CudaSlice<f32>,
9451 bb: &CudaSlice<f32>,
9452 wa: &CudaSlice<f32>,
9453 wb: &CudaSlice<f32>,
9454 da: &mut CudaSlice<f32>,
9455 db: &mut CudaSlice<f32>,
9456 ncols: usize,
9457 nrows: usize,
9458 eps: f32,
9459 ) -> Result<(), Box<dyn std::error::Error>> {
9460 let f = self.func("rms_norm2x_f32");
9461 let cfg = LaunchConfig {
9462 grid_dim: (2 * nrows as u32, 1, 1),
9463 block_dim: (rms_block(), 1, 1),
9464 shared_mem_bytes: 0,
9465 };
9466 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9467 let __s_b = self.gpu.stream();
9468 let mut b = __s_b.launch_builder(&f);
9469 b.arg(a)
9470 .arg(bb)
9471 .arg(wa)
9472 .arg(wb)
9473 .arg(da)
9474 .arg(db)
9475 .arg(&nc)
9476 .arg(&nr)
9477 .arg(&e);
9478 unsafe {
9479 b.launch(cfg)?;
9480 }
9481 Ok(())
9482 }
9483
9484 pub fn softcap(
9486 &self,
9487 y: &mut CudaSlice<f32>,
9488 cap: f32,
9489 n: usize,
9490 ) -> Result<(), Box<dyn std::error::Error>> {
9491 let f = self.func("softcap_f32");
9492 let cfg = LaunchConfig::for_num_elems(n as u32);
9493 let ni = n as i32;
9494 let __s_b = self.gpu.stream();
9495 let mut b = __s_b.launch_builder(&f);
9496 b.arg(y).arg(&cap).arg(&ni);
9497 unsafe {
9498 b.launch(cfg)?;
9499 }
9500 Ok(())
9501 }
9502
9503 pub fn mask_ids_rows(
9506 &self,
9507 y: &mut CudaSlice<f32>,
9508 ids: &CudaSlice<i32>,
9509 n_ids: usize,
9510 n_vocab: usize,
9511 t: usize,
9512 ) -> Result<(), Box<dyn std::error::Error>> {
9513 let f = self.func("mask_ids_rows_f32");
9514 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9515 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9516 let __s_b = self.gpu.stream();
9517 let mut b = __s_b.launch_builder(&f);
9518 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9519 unsafe {
9520 b.launch(cfg)?;
9521 }
9522 Ok(())
9523 }
9524
9525 #[allow(clippy::too_many_arguments)]
9527 pub fn add_scale_rms_norm(
9528 &self,
9529 a: &CudaSlice<f32>,
9530 b_in: &CudaSlice<f32>,
9531 c: f32,
9532 w: &CudaSlice<f32>,
9533 res: &mut CudaSlice<f32>,
9534 dst: &mut CudaSlice<f32>,
9535 ncols: usize,
9536 nrows: usize,
9537 eps: f32,
9538 ) -> Result<(), Box<dyn std::error::Error>> {
9539 let f = self.func("add_scale_rms_norm_f32");
9540 let cfg = LaunchConfig {
9541 grid_dim: (nrows as u32, 1, 1),
9542 block_dim: (rms_block(), 1, 1),
9543 shared_mem_bytes: 0,
9544 };
9545 let (nc, e2) = (ncols as i32, eps);
9546 let __s_b = self.gpu.stream();
9547 let mut b = __s_b.launch_builder(&f);
9548 b.arg(a)
9549 .arg(b_in)
9550 .arg(&c)
9551 .arg(w)
9552 .arg(res)
9553 .arg(dst)
9554 .arg(&nc)
9555 .arg(&e2);
9556 unsafe {
9557 b.launch(cfg)?;
9558 }
9559 Ok(())
9560 }
9561
9562 #[allow(clippy::too_many_arguments)]
9565 pub fn add_scale_rms_norm_q8_1(
9566 &self,
9567 a: &CudaSlice<f32>,
9568 b_in: &CudaSlice<f32>,
9569 c: f32,
9570 w: &CudaSlice<f32>,
9571 res: &mut CudaSlice<f32>,
9572 ncols: usize,
9573 nrows: usize,
9574 eps: f32,
9575 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9576 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9577 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9578 let (nc, e2) = (ncols as i32, eps);
9579 if Self::pdl_on() && Self::pdl_wb_on() {
9580 {
9581 use cudarc::driver::{DevicePtr, DevicePtrMut};
9582 let s = &self.gpu.stream();
9583 let (pa, _g0) = a.device_ptr(s);
9584 let (pb, _g1) = b_in.device_ptr(s);
9585 let (pw, _g2) = w.device_ptr(s);
9586 let (pr, _g3) = res.device_ptr_mut(s);
9587 let (pq, _g4) = out_q.device_ptr_mut(s);
9588 let (pd, _g5) = out_d.device_ptr_mut(s);
9589 let mut ps = [
9590 &pa as *const _ as *mut std::ffi::c_void,
9591 &pb as *const _ as *mut _,
9592 &c as *const _ as *mut _,
9593 &pw as *const _ as *mut _,
9594 &pr as *const _ as *mut _,
9595 &pq as *const _ as *mut _,
9596 &pd as *const _ as *mut _,
9597 &nc as *const _ as *mut _,
9598 &e2 as *const _ as *mut _,
9599 ];
9600 unsafe {
9601 self.launch_pdl(
9602 "add_scale_rms_norm_q8_1",
9603 (nrows as u32, 1, 1),
9604 (rms_block(), 1, 1),
9605 &mut ps,
9606 )?;
9607 }
9608 }
9609 return Ok((out_q, out_d));
9610 }
9611 let f = self.func("add_scale_rms_norm_q8_1");
9612 let cfg = LaunchConfig {
9613 grid_dim: (nrows as u32, 1, 1),
9614 block_dim: (rms_block(), 1, 1),
9615 shared_mem_bytes: 0,
9616 };
9617 let __s_b = self.gpu.stream();
9618 let mut b = __s_b.launch_builder(&f);
9619 b.arg(a)
9620 .arg(b_in)
9621 .arg(&c)
9622 .arg(w)
9623 .arg(res)
9624 .arg(&mut out_q)
9625 .arg(&mut out_d)
9626 .arg(&nc)
9627 .arg(&e2);
9628 unsafe {
9629 b.launch(cfg)?;
9630 }
9631 Ok((out_q, out_d))
9632 }
9633
9634 #[allow(clippy::too_many_arguments)]
9636 pub fn add_scale_rms_norm_q8_1_into(
9637 &self,
9638 a: &CudaSlice<f32>,
9639 b_in: &CudaSlice<f32>,
9640 c: f32,
9641 w: &CudaSlice<f32>,
9642 res: &mut CudaSlice<f32>,
9643 ncols: usize,
9644 nrows: usize,
9645 eps: f32,
9646 out_q: &mut CudaSlice<i8>,
9647 out_d: &mut CudaSlice<f32>,
9648 ) -> Result<(), Box<dyn std::error::Error>> {
9649 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9650 let (nc, e2) = (ncols as i32, eps);
9651 if Self::pdl_on() && Self::pdl_wb_on() {
9652 use cudarc::driver::{DevicePtr, DevicePtrMut};
9653 let s = &self.gpu.stream();
9654 let (pa, _g0) = a.device_ptr(s);
9655 let (pb, _g1) = b_in.device_ptr(s);
9656 let (pw, _g2) = w.device_ptr(s);
9657 let (pr, _g3) = res.device_ptr_mut(s);
9658 let (pq, _g4) = out_q.device_ptr_mut(s);
9659 let (pd, _g5) = out_d.device_ptr_mut(s);
9660 let mut ps = [
9661 &pa as *const _ as *mut std::ffi::c_void,
9662 &pb as *const _ as *mut _,
9663 &c as *const _ as *mut _,
9664 &pw as *const _ as *mut _,
9665 &pr as *const _ as *mut _,
9666 &pq as *const _ as *mut _,
9667 &pd as *const _ as *mut _,
9668 &nc as *const _ as *mut _,
9669 &e2 as *const _ as *mut _,
9670 ];
9671 unsafe {
9672 self.launch_pdl(
9673 "add_scale_rms_norm_q8_1",
9674 (nrows as u32, 1, 1),
9675 (rms_block(), 1, 1),
9676 &mut ps,
9677 )?;
9678 }
9679 return Ok(());
9680 }
9681 let f = self.func("add_scale_rms_norm_q8_1");
9682 let cfg = LaunchConfig {
9683 grid_dim: (nrows as u32, 1, 1),
9684 block_dim: (rms_block(), 1, 1),
9685 shared_mem_bytes: 0,
9686 };
9687 let __s_b = self.gpu.stream();
9688 let mut b = __s_b.launch_builder(&f);
9689 b.arg(a)
9690 .arg(b_in)
9691 .arg(&c)
9692 .arg(w)
9693 .arg(res)
9694 .arg(&mut *out_q)
9695 .arg(&mut *out_d)
9696 .arg(&nc)
9697 .arg(&e2);
9698 unsafe {
9699 b.launch(cfg)?;
9700 }
9701 Ok(())
9702 }
9703
9704 #[allow(clippy::too_many_arguments)]
9707 pub fn rms_pre_add_scale_rms_norm_q8_1(
9708 &self,
9709 a: &CudaSlice<f32>,
9710 wa: &CudaSlice<f32>,
9711 b_in: &CudaSlice<f32>,
9712 c: f32,
9713 w: &CudaSlice<f32>,
9714 res: &mut CudaSlice<f32>,
9715 ncols: usize,
9716 nrows: usize,
9717 eps: f32,
9718 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9719 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9720 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9721 let (nc, e2) = (ncols as i32, eps);
9722 if Self::pdl_on() {
9723 {
9724 use cudarc::driver::{DevicePtr, DevicePtrMut};
9725 let s = &self.gpu.stream();
9726 let (pa, _g0) = a.device_ptr(s);
9727 let (pwa, _g1) = wa.device_ptr(s);
9728 let (pb, _g2) = b_in.device_ptr(s);
9729 let (pw, _g3) = w.device_ptr(s);
9730 let (pr, _g4) = res.device_ptr_mut(s);
9731 let (pq, _g5) = out_q.device_ptr_mut(s);
9732 let (pd, _g6) = out_d.device_ptr_mut(s);
9733 let mut ps = [
9734 &pa as *const _ as *mut std::ffi::c_void,
9735 &pwa as *const _ as *mut _,
9736 &pb as *const _ as *mut _,
9737 &c as *const _ as *mut _,
9738 &pw as *const _ as *mut _,
9739 &pr as *const _ as *mut _,
9740 &pq as *const _ as *mut _,
9741 &pd as *const _ as *mut _,
9742 &nc as *const _ as *mut _,
9743 &e2 as *const _ as *mut _,
9744 ];
9745 unsafe {
9746 self.launch_pdl(
9747 "rms_pre_add_scale_rms_norm_q8_1",
9748 (nrows as u32, 1, 1),
9749 (rms_block(), 1, 1),
9750 &mut ps,
9751 )?;
9752 }
9753 }
9754 return Ok((out_q, out_d));
9755 }
9756 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9757 let cfg = LaunchConfig {
9758 grid_dim: (nrows as u32, 1, 1),
9759 block_dim: (rms_block(), 1, 1),
9760 shared_mem_bytes: 0,
9761 };
9762 let __s_b = self.gpu.stream();
9763 let mut b = __s_b.launch_builder(&f);
9764 b.arg(a)
9765 .arg(wa)
9766 .arg(b_in)
9767 .arg(&c)
9768 .arg(w)
9769 .arg(res)
9770 .arg(&mut out_q)
9771 .arg(&mut out_d)
9772 .arg(&nc)
9773 .arg(&e2);
9774 unsafe {
9775 b.launch(cfg)?;
9776 }
9777 Ok((out_q, out_d))
9778 }
9779
9780 pub fn gelu_tanh_mul_q8_1(
9783 &self,
9784 gate: &CudaSlice<f32>,
9785 up: &cudarc::driver::CudaView<f32>,
9786 act: &mut CudaSlice<f32>,
9787 ncols: usize,
9788 nrows: usize,
9789 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9790 debug_assert!(ncols % 128 == 0);
9791 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9792 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9793 let nc = ncols as i32;
9794 if Self::pdl_on() {
9795 {
9796 use cudarc::driver::{DevicePtr, DevicePtrMut};
9797 let s = &self.gpu.stream();
9798 let (pg, _g0) = gate.device_ptr(s);
9799 let (pu, _g1) = up.device_ptr(s);
9800 let (pact, _g2) = act.device_ptr_mut(s);
9801 let (pq, _g3) = out_q.device_ptr_mut(s);
9802 let (pd, _g4) = out_d.device_ptr_mut(s);
9803 let mut ps = [
9804 &pg as *const _ as *mut std::ffi::c_void,
9805 &pu as *const _ as *mut _,
9806 &pact as *const _ as *mut _,
9807 &pq as *const _ as *mut _,
9808 &pd as *const _ as *mut _,
9809 &nc as *const _ as *mut _,
9810 ];
9811 unsafe {
9812 self.launch_pdl(
9813 "gelu_tanh_mul_q8_1",
9814 (nrows as u32, 1, 1),
9815 (rms_block(), 1, 1),
9816 &mut ps,
9817 )?;
9818 }
9819 }
9820 return Ok((out_q, out_d));
9821 }
9822 let f = self.func("gelu_tanh_mul_q8_1");
9823 let cfg = LaunchConfig {
9824 grid_dim: (nrows as u32, 1, 1),
9825 block_dim: (rms_block(), 1, 1),
9826 shared_mem_bytes: 0,
9827 };
9828 let __s_b = self.gpu.stream();
9829 let mut b = __s_b.launch_builder(&f);
9830 b.arg(gate)
9831 .arg(up)
9832 .arg(act)
9833 .arg(&mut out_q)
9834 .arg(&mut out_d)
9835 .arg(&nc);
9836 unsafe {
9837 b.launch(cfg)?;
9838 }
9839 Ok((out_q, out_d))
9840 }
9841
9842 #[allow(clippy::too_many_arguments)]
9844 pub fn gelu_tanh_mul_q8_1_into(
9845 &self,
9846 gate: &CudaSlice<f32>,
9847 up: &cudarc::driver::CudaView<f32>,
9848 act: &mut CudaSlice<f32>,
9849 ncols: usize,
9850 nrows: usize,
9851 out_q: &mut CudaSlice<i8>,
9852 out_d: &mut CudaSlice<f32>,
9853 ) -> Result<(), Box<dyn std::error::Error>> {
9854 debug_assert!(ncols % 128 == 0);
9855 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9856 let nc = ncols as i32;
9857 if Self::pdl_on() {
9858 use cudarc::driver::{DevicePtr, DevicePtrMut};
9859 let s = &self.gpu.stream();
9860 let (pg, _g0) = gate.device_ptr(s);
9861 let (pu, _g1) = up.device_ptr(s);
9862 let (pact, _g2) = act.device_ptr_mut(s);
9863 let (pq, _g3) = out_q.device_ptr_mut(s);
9864 let (pd, _g4) = out_d.device_ptr_mut(s);
9865 let mut ps = [
9866 &pg as *const _ as *mut std::ffi::c_void,
9867 &pu as *const _ as *mut _,
9868 &pact as *const _ as *mut _,
9869 &pq as *const _ as *mut _,
9870 &pd as *const _ as *mut _,
9871 &nc as *const _ as *mut _,
9872 ];
9873 unsafe {
9874 self.launch_pdl(
9875 "gelu_tanh_mul_q8_1",
9876 (nrows as u32, 1, 1),
9877 (rms_block(), 1, 1),
9878 &mut ps,
9879 )?;
9880 }
9881 return Ok(());
9882 }
9883 let f = self.func("gelu_tanh_mul_q8_1");
9884 let cfg = LaunchConfig {
9885 grid_dim: (nrows as u32, 1, 1),
9886 block_dim: (rms_block(), 1, 1),
9887 shared_mem_bytes: 0,
9888 };
9889 let __s_b = self.gpu.stream();
9890 let mut b = __s_b.launch_builder(&f);
9891 b.arg(gate)
9892 .arg(up)
9893 .arg(&mut *act)
9894 .arg(&mut *out_q)
9895 .arg(&mut *out_d)
9896 .arg(&nc);
9897 unsafe {
9898 b.launch(cfg)?;
9899 }
9900 Ok(())
9901 }
9902
9903 #[allow(clippy::too_many_arguments)]
9905 pub fn add_rms_norm3_q8z(
9906 &self,
9907 a: &CudaSlice<f32>,
9908 b_in: &CudaSlice<f32>,
9909 w0: &CudaSlice<f32>,
9910 w1: &CudaSlice<f32>,
9911 w2: &CudaSlice<f32>,
9912 res: &mut CudaSlice<f32>,
9913 out1: &mut CudaSlice<f32>,
9914 ncols: usize,
9915 nrows: usize,
9916 eps: f32,
9917 ) -> Result<
9918 (
9919 (CudaSlice<i8>, CudaSlice<f32>),
9920 (CudaSlice<i8>, CudaSlice<f32>),
9921 ),
9922 Box<dyn std::error::Error>,
9923 > {
9924 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9925 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9926 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9927 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9928 let f = self.func("add_rms_norm3_q8z_f32");
9929 let cfg = LaunchConfig {
9930 grid_dim: (nrows as u32, 1, 1),
9931 block_dim: (rms_block(), 1, 1),
9932 shared_mem_bytes: 0,
9933 };
9934 let (nc, e2) = (ncols as i32, eps);
9935 let __s_b = self.gpu.stream();
9936 let mut b = __s_b.launch_builder(&f);
9937 b.arg(a)
9938 .arg(b_in)
9939 .arg(w0)
9940 .arg(w1)
9941 .arg(w2)
9942 .arg(res)
9943 .arg(&mut q0)
9944 .arg(&mut d0)
9945 .arg(out1)
9946 .arg(&mut q2)
9947 .arg(&mut d2)
9948 .arg(&nc)
9949 .arg(&e2);
9950 unsafe {
9951 b.launch(cfg)?;
9952 }
9953 Ok(((q0, d0), (q2, d2)))
9954 }
9955
9956 #[allow(clippy::too_many_arguments)]
9958 pub fn add_rms_norm3(
9959 &self,
9960 a: &CudaSlice<f32>,
9961 b_in: &CudaSlice<f32>,
9962 w0: &CudaSlice<f32>,
9963 w1: &CudaSlice<f32>,
9964 w2: &CudaSlice<f32>,
9965 res: &mut CudaSlice<f32>,
9966 d0: &mut CudaSlice<f32>,
9967 d1: &mut CudaSlice<f32>,
9968 d2: &mut CudaSlice<f32>,
9969 ncols: usize,
9970 nrows: usize,
9971 eps: f32,
9972 ) -> Result<(), Box<dyn std::error::Error>> {
9973 let f = self.func("add_rms_norm3_f32");
9974 let cfg = LaunchConfig {
9975 grid_dim: (nrows as u32, 1, 1),
9976 block_dim: (rms_block(), 1, 1),
9977 shared_mem_bytes: 0,
9978 };
9979 let (nc, e2) = (ncols as i32, eps);
9980 let __s_b = self.gpu.stream();
9981 let mut b = __s_b.launch_builder(&f);
9982 b.arg(a)
9983 .arg(b_in)
9984 .arg(w0)
9985 .arg(w1)
9986 .arg(w2)
9987 .arg(res)
9988 .arg(d0)
9989 .arg(d1)
9990 .arg(d2)
9991 .arg(&nc)
9992 .arg(&e2);
9993 unsafe {
9994 b.launch(cfg)?;
9995 }
9996 Ok(())
9997 }
9998
9999 pub fn add_scale(
10001 &self,
10002 a: &CudaSlice<f32>,
10003 b_in: &CudaSlice<f32>,
10004 c: f32,
10005 dst: &mut CudaSlice<f32>,
10006 n: usize,
10007 ) -> Result<(), Box<dyn std::error::Error>> {
10008 let f = self.func("add_scale_f32");
10009 let cfg = LaunchConfig::for_num_elems(n as u32);
10010 let ni = n as i32;
10011 let __s_b = self.gpu.stream();
10012 let mut b = __s_b.launch_builder(&f);
10013 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10014 unsafe {
10015 b.launch(cfg)?;
10016 }
10017 Ok(())
10018 }
10019
10020 pub fn layer_norm_bias(
10022 &self,
10023 x: &CudaSlice<f32>,
10024 w: &CudaSlice<f32>,
10025 b: &CudaSlice<f32>,
10026 dst: &mut CudaSlice<f32>,
10027 ncols: usize,
10028 nrows: usize,
10029 eps: f32,
10030 ) -> Result<(), Box<dyn std::error::Error>> {
10031 let f = self.func("layer_norm_bias_f32");
10032 let (nc, e) = (ncols as i32, eps);
10033 let cfg = LaunchConfig {
10034 grid_dim: (nrows as u32, 1, 1),
10035 block_dim: (256, 1, 1),
10036 shared_mem_bytes: 0,
10037 };
10038 let __s_b = self.gpu.stream();
10039 let mut lb = __s_b.launch_builder(&f);
10040 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10041 unsafe {
10042 lb.launch(cfg)?;
10043 }
10044 Ok(())
10045 }
10046
10047 pub fn gelu_tanh(
10049 &self,
10050 x: &CudaSlice<f32>,
10051 dst: &mut CudaSlice<f32>,
10052 n: usize,
10053 ) -> Result<(), Box<dyn std::error::Error>> {
10054 let f = self.func("gelu_tanh_f32");
10055 let ni = n as i64;
10056 let cfg = LaunchConfig {
10057 grid_dim: (n.div_ceil(256) as u32, 1, 1),
10058 block_dim: (256, 1, 1),
10059 shared_mem_bytes: 0,
10060 };
10061 let __s_b = self.gpu.stream();
10062 let mut lb = __s_b.launch_builder(&f);
10063 lb.arg(x).arg(&mut *dst).arg(&ni);
10064 unsafe {
10065 lb.launch(cfg)?;
10066 }
10067 Ok(())
10068 }
10069
10070 pub fn row_softmax(
10072 &self,
10073 x: &mut CudaSlice<f32>,
10074 ncols: usize,
10075 nrows: usize,
10076 ) -> Result<(), Box<dyn std::error::Error>> {
10077 let f = self.func("row_softmax_f32");
10078 let nc = ncols as i32;
10079 let cfg = LaunchConfig {
10080 grid_dim: (nrows as u32, 1, 1),
10081 block_dim: (256, 1, 1),
10082 shared_mem_bytes: 0,
10083 };
10084 let __s_b = self.gpu.stream();
10085 let mut lb = __s_b.launch_builder(&f);
10086 lb.arg(&mut *x).arg(&nc);
10087 unsafe {
10088 lb.launch(cfg)?;
10089 }
10090 Ok(())
10091 }
10092
10093 pub fn rms_norm(
10094 &self,
10095 x: &CudaSlice<f32>,
10096 w: &CudaSlice<f32>,
10097 dst: &mut CudaSlice<f32>,
10098 ncols: usize,
10099 nrows: usize,
10100 eps: f32,
10101 ) -> Result<(), Box<dyn std::error::Error>> {
10102 let (nc, e) = (ncols as i32, eps);
10103 let kname = if Self::norm_ilp_on() {
10104 "rms_norm_f32_v2"
10105 } else {
10106 "rms_norm_f32"
10107 };
10108 if Self::pdl_on() && Self::pdl_wb_on() {
10109 use cudarc::driver::{DevicePtr, DevicePtrMut};
10110 let s = &self.gpu.stream();
10111 let (px, _g0) = x.device_ptr(s);
10112 let (pw, _g1) = w.device_ptr(s);
10113 let (pd, _g2) = dst.device_ptr_mut(s);
10114 let mut ps = [
10115 &px as *const _ as *mut std::ffi::c_void,
10116 &pw as *const _ as *mut _,
10117 &pd as *const _ as *mut _,
10118 &nc as *const _ as *mut _,
10119 &e as *const _ as *mut _,
10120 ];
10121 unsafe {
10122 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10123 }
10124 return Ok(());
10125 }
10126 let f = self.func(kname);
10127 let cfg = LaunchConfig {
10128 grid_dim: (nrows as u32, 1, 1),
10129 block_dim: (rms_block(), 1, 1),
10130 shared_mem_bytes: 0,
10131 };
10132 let __s_b = self.gpu.stream();
10133 let mut b = __s_b.launch_builder(&f);
10134 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10135 unsafe {
10136 b.launch(cfg)?;
10137 }
10138 Ok(())
10139 }
10140
10141 pub fn rms_norm_decode(
10149 &self,
10150 x: &CudaSlice<f32>,
10151 w: &CudaSlice<f32>,
10152 dst: &mut CudaSlice<f32>,
10153 ncols: usize,
10154 nrows: usize,
10155 eps: f32,
10156 ) -> Result<(), Box<dyn std::error::Error>> {
10157 let f = self.func(if Self::norm_ilp_on() {
10158 "rms_norm_f32_v2"
10159 } else {
10160 "rms_norm_f32"
10161 });
10162 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_b = self.gpu.stream();
10169 let mut b = __s_b.launch_builder(&f);
10170 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10171 unsafe {
10172 b.launch(cfg)?;
10173 }
10174 Ok(())
10175 }
10176
10177 pub fn rms_norm_q8_1(
10181 &self,
10182 x: &CudaSlice<f32>,
10183 w: &CudaSlice<f32>,
10184 ncols: usize,
10185 nrows: usize,
10186 eps: f32,
10187 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10188 let nblk = ncols / 32;
10189 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10190 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10191 let (nc, e) = (ncols as i32, eps);
10192 if Self::pdl_on() {
10193 {
10194 use cudarc::driver::{DevicePtr, DevicePtrMut};
10195 let s = &self.gpu.stream();
10196 let (px, _g0) = x.device_ptr(s);
10197 let (pw, _g1) = w.device_ptr(s);
10198 let (pq, _g2) = q.device_ptr_mut(s);
10199 let (pd, _g3) = d.device_ptr_mut(s);
10200 let mut ps = [
10201 &px as *const _ as *mut std::ffi::c_void,
10202 &pw as *const _ as *mut _,
10203 &pq as *const _ as *mut _,
10204 &pd as *const _ as *mut _,
10205 &nc as *const _ as *mut _,
10206 &e as *const _ as *mut _,
10207 ];
10208 unsafe {
10209 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10210 }
10211 }
10212 return Ok((q, d));
10213 }
10214 let f = self.func("rms_norm_q8_1");
10215 let cfg = LaunchConfig {
10218 grid_dim: (nrows as u32, 1, 1),
10219 block_dim: (1024, 1, 1),
10220 shared_mem_bytes: 0,
10221 };
10222 let __s_b = self.gpu.stream();
10223 let mut b = __s_b.launch_builder(&f);
10224 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10225 unsafe {
10226 b.launch(cfg)?;
10227 }
10228 Ok((q, d))
10229 }
10230
10231 pub fn rms_norm_q8_1_into(
10234 &self,
10235 x: &CudaSlice<f32>,
10236 w: &CudaSlice<f32>,
10237 ncols: usize,
10238 nrows: usize,
10239 eps: f32,
10240 q: &mut CudaSlice<i8>,
10241 d: &mut CudaSlice<f32>,
10242 ) -> Result<(), Box<dyn std::error::Error>> {
10243 let nblk = ncols / 32;
10244 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10245 let (nc, e) = (ncols as i32, eps);
10246 if Self::pdl_on() {
10247 use cudarc::driver::{DevicePtr, DevicePtrMut};
10248 let s = &self.gpu.stream();
10249 let (px, _g0) = x.device_ptr(s);
10250 let (pw, _g1) = w.device_ptr(s);
10251 let (pq, _g2) = q.device_ptr_mut(s);
10252 let (pd, _g3) = d.device_ptr_mut(s);
10253 let mut ps = [
10254 &px as *const _ as *mut std::ffi::c_void,
10255 &pw as *const _ as *mut _,
10256 &pq 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("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10263 }
10264 return Ok(());
10265 }
10266 let f = self.func("rms_norm_q8_1");
10267 let cfg = LaunchConfig {
10268 grid_dim: (nrows as u32, 1, 1),
10269 block_dim: (1024, 1, 1),
10270 shared_mem_bytes: 0,
10271 };
10272 let __s_b = self.gpu.stream();
10273 let mut b = __s_b.launch_builder(&f);
10274 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10275 unsafe {
10276 b.launch(cfg)?;
10277 }
10278 Ok(())
10279 }
10280
10281 pub fn quantize_q8_1_into(
10283 &self,
10284 x: &CudaSlice<f32>,
10285 m: usize,
10286 in_f: usize,
10287 q: &mut CudaSlice<i8>,
10288 d: &mut CudaSlice<f32>,
10289 ) -> Result<(), Box<dyn std::error::Error>> {
10290 let nblk = in_f / 32;
10291 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10292 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10293 let (inf, mi) = (in_f as i32, m as i32);
10294 if Self::pdl_on() && Self::pdl_wb_on() {
10295 use cudarc::driver::{DevicePtr, DevicePtrMut};
10296 let s = &self.gpu.stream();
10297 let (px, _g0) = x.device_ptr(s);
10298 let (pq, _g1) = q.device_ptr_mut(s);
10299 let (pd, _g2) = d.device_ptr_mut(s);
10300 let mut ps = [
10301 &px as *const _ as *mut std::ffi::c_void,
10302 &pq as *const _ as *mut _,
10303 &pd as *const _ as *mut _,
10304 &inf as *const _ as *mut _,
10305 &mi as *const _ as *mut _,
10306 ];
10307 unsafe {
10308 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10309 }
10310 return Ok(());
10311 }
10312 let f = self.func("quantize_q8_1");
10313 let __s_b = self.gpu.stream();
10314 let mut b = __s_b.launch_builder(&f);
10315 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10316 unsafe {
10317 b.launch(cfg)?;
10318 }
10319 Ok(())
10320 }
10321
10322 pub fn add_rms_norm_q8_1(
10326 &self,
10327 a: &CudaSlice<f32>,
10328 b_in: &CudaSlice<f32>,
10329 w: &CudaSlice<f32>,
10330 res: &mut CudaSlice<f32>,
10331 ncols: usize,
10332 nrows: usize,
10333 eps: f32,
10334 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10335 let nblk = ncols / 32;
10336 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10337 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10338 let f = self.func("add_rms_norm_q8_1");
10339 let cfg = LaunchConfig {
10341 grid_dim: (nrows as u32, 1, 1),
10342 block_dim: (1024, 1, 1),
10343 shared_mem_bytes: 0,
10344 };
10345 let (nc, e) = (ncols as i32, eps);
10346 let __s_bld = self.gpu.stream();
10347 let mut bld = __s_bld.launch_builder(&f);
10348 bld.arg(a)
10349 .arg(b_in)
10350 .arg(w)
10351 .arg(res)
10352 .arg(&mut q)
10353 .arg(&mut d)
10354 .arg(&nc)
10355 .arg(&e);
10356 unsafe {
10357 bld.launch(cfg)?;
10358 }
10359 Ok((q, d))
10360 }
10361
10362 #[allow(clippy::too_many_arguments)]
10368 pub fn join_add_rms_norm_raw(
10369 &self,
10370 a0_raw: u64,
10371 a1_raw: u64,
10372 x: &CudaSlice<f32>,
10373 w: &CudaSlice<f32>,
10374 res: &mut CudaSlice<f32>,
10375 dst: &mut CudaSlice<f32>,
10376 ncols: usize,
10377 eps: f32,
10378 ) -> Result<(), Box<dyn std::error::Error>> {
10379 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10380 return Err("join_add_rms_norm geometry".into());
10381 }
10382 let f = self.func("join_add_rms_norm_f32");
10383 let cfg = LaunchConfig {
10384 grid_dim: (1, 1, 1),
10385 block_dim: (rms_block(), 1, 1),
10386 shared_mem_bytes: 0,
10387 };
10388 let (nc, e) = (ncols as i32, eps);
10389 let __s_b = self.gpu.stream();
10390 let mut b = __s_b.launch_builder(&f);
10391 b.arg(&a0_raw)
10392 .arg(&a1_raw)
10393 .arg(x)
10394 .arg(w)
10395 .arg(&mut *res)
10396 .arg(&mut *dst)
10397 .arg(&nc)
10398 .arg(&e);
10399 unsafe {
10400 b.launch(cfg)?;
10401 }
10402 Ok(())
10403 }
10404
10405 pub fn add_rms_norm(
10406 &self,
10407 a: &CudaSlice<f32>,
10408 b: &CudaSlice<f32>,
10409 w: &CudaSlice<f32>,
10410 res: &mut CudaSlice<f32>,
10411 dst: &mut CudaSlice<f32>,
10412 ncols: usize,
10413 nrows: usize,
10414 eps: f32,
10415 ) -> Result<(), Box<dyn std::error::Error>> {
10416 let (nc, e) = (ncols as i32, eps);
10417 let kname = if Self::norm_ilp_on() {
10418 "add_rms_norm_f32_v2"
10419 } else {
10420 "add_rms_norm_f32"
10421 };
10422 if Self::pdl_on() && Self::pdl_wb_on() {
10423 use cudarc::driver::{DevicePtr, DevicePtrMut};
10424 let s = &self.gpu.stream();
10425 let (pa, _g0) = a.device_ptr(s);
10426 let (pb, _g1) = b.device_ptr(s);
10427 let (pw, _g2) = w.device_ptr(s);
10428 let (pr, _g3) = res.device_ptr_mut(s);
10429 let (pd, _g4) = dst.device_ptr_mut(s);
10430 let mut ps = [
10431 &pa as *const _ as *mut std::ffi::c_void,
10432 &pb as *const _ as *mut _,
10433 &pw as *const _ as *mut _,
10434 &pr as *const _ as *mut _,
10435 &pd as *const _ as *mut _,
10436 &nc as *const _ as *mut _,
10437 &e as *const _ as *mut _,
10438 ];
10439 unsafe {
10440 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10441 }
10442 return Ok(());
10443 }
10444 let f = self.func(kname);
10445 let cfg = LaunchConfig {
10446 grid_dim: (nrows as u32, 1, 1),
10447 block_dim: (rms_block(), 1, 1),
10448 shared_mem_bytes: 0,
10449 };
10450 let __s_b2 = self.gpu.stream();
10451 let mut b2 = __s_b2.launch_builder(&f);
10452 b2.arg(a)
10453 .arg(b)
10454 .arg(w)
10455 .arg(&mut *res)
10456 .arg(&mut *dst)
10457 .arg(&nc)
10458 .arg(&e);
10459 unsafe {
10460 b2.launch(cfg)?;
10461 }
10462 Ok(())
10463 }
10464
10465 #[allow(clippy::too_many_arguments)]
10468 pub fn rms_pre_add_rms_norm(
10469 &self,
10470 a: &CudaSlice<f32>,
10471 wa: &CudaSlice<f32>,
10472 b: &CudaSlice<f32>,
10473 w: &CudaSlice<f32>,
10474 res: &mut CudaSlice<f32>,
10475 dst: &mut CudaSlice<f32>,
10476 ncols: usize,
10477 nrows: usize,
10478 eps: f32,
10479 ) -> Result<(), Box<dyn std::error::Error>> {
10480 let f = self.func("rms_pre_add_rms_norm_f32");
10481 let cfg = LaunchConfig {
10482 grid_dim: (nrows as u32, 1, 1),
10483 block_dim: (rms_block(), 1, 1),
10484 shared_mem_bytes: 0,
10485 };
10486 let (nc, e) = (ncols as i32, eps);
10487 let __s_b2 = self.gpu.stream();
10488 let mut b2 = __s_b2.launch_builder(&f);
10489 b2.arg(a)
10490 .arg(wa)
10491 .arg(b)
10492 .arg(w)
10493 .arg(&mut *res)
10494 .arg(&mut *dst)
10495 .arg(&nc)
10496 .arg(&e);
10497 unsafe {
10498 b2.launch(cfg)?;
10499 }
10500 Ok(())
10501 }
10502
10503 #[allow(clippy::too_many_arguments)]
10505 pub fn rms_pre_add_rms_norm_q8z(
10506 &self,
10507 a: &CudaSlice<f32>,
10508 wa: &CudaSlice<f32>,
10509 b: &CudaSlice<f32>,
10510 w: &CudaSlice<f32>,
10511 res: &mut CudaSlice<f32>,
10512 dst: &mut CudaSlice<f32>,
10513 ncols: usize,
10514 nrows: usize,
10515 eps: f32,
10516 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10517 debug_assert!(ncols % 128 == 0);
10518 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10519 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10520 let (nc, e) = (ncols as i32, eps);
10521 if Self::pdl_on() {
10522 {
10523 use cudarc::driver::{DevicePtr, DevicePtrMut};
10524 let s = &self.gpu.stream();
10525 let (pa, _g0) = a.device_ptr(s);
10526 let (pwa, _g1) = wa.device_ptr(s);
10527 let (pb, _g2) = b.device_ptr(s);
10528 let (pw, _g3) = w.device_ptr(s);
10529 let (pr, _g4) = res.device_ptr_mut(s);
10530 let (pdst, _g5) = dst.device_ptr_mut(s);
10531 let (pq, _g6) = out_q.device_ptr_mut(s);
10532 let (pd, _g7) = out_d.device_ptr_mut(s);
10533 let mut ps = [
10534 &pa as *const _ as *mut std::ffi::c_void,
10535 &pwa as *const _ as *mut _,
10536 &pb as *const _ as *mut _,
10537 &pw as *const _ as *mut _,
10538 &pr as *const _ as *mut _,
10539 &pdst as *const _ as *mut _,
10540 &pq as *const _ as *mut _,
10541 &pd as *const _ as *mut _,
10542 &nc as *const _ as *mut _,
10543 &e as *const _ as *mut _,
10544 ];
10545 unsafe {
10546 self.launch_pdl(
10547 "rms_pre_add_rms_norm_q8z_f32",
10548 (nrows as u32, 1, 1),
10549 (rms_block(), 1, 1),
10550 &mut ps,
10551 )?;
10552 }
10553 }
10554 return Ok((out_q, out_d));
10555 }
10556 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10557 let cfg = LaunchConfig {
10558 grid_dim: (nrows as u32, 1, 1),
10559 block_dim: (rms_block(), 1, 1),
10560 shared_mem_bytes: 0,
10561 };
10562 let __s_b2 = self.gpu.stream();
10563 let mut b2 = __s_b2.launch_builder(&f);
10564 b2.arg(a)
10565 .arg(wa)
10566 .arg(b)
10567 .arg(w)
10568 .arg(&mut *res)
10569 .arg(&mut *dst)
10570 .arg(&mut out_q)
10571 .arg(&mut out_d)
10572 .arg(&nc)
10573 .arg(&e);
10574 unsafe {
10575 b2.launch(cfg)?;
10576 }
10577 Ok((out_q, out_d))
10578 }
10579
10580 #[allow(clippy::too_many_arguments)]
10584 pub fn rms_pre_add_rms_norm_q8z_into(
10585 &self,
10586 a: &CudaSlice<f32>,
10587 wa: &CudaSlice<f32>,
10588 b: &CudaSlice<f32>,
10589 w: &CudaSlice<f32>,
10590 res: &mut CudaSlice<f32>,
10591 dst: &mut CudaSlice<f32>,
10592 ncols: usize,
10593 nrows: usize,
10594 eps: f32,
10595 out_q: &mut CudaSlice<i8>,
10596 out_d: &mut CudaSlice<f32>,
10597 ) -> Result<(), Box<dyn std::error::Error>> {
10598 debug_assert!(ncols % 128 == 0);
10599 let (nc, e) = (ncols as i32, eps);
10600 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10601 let cfg = LaunchConfig {
10602 grid_dim: (nrows as u32, 1, 1),
10603 block_dim: (rms_block(), 1, 1),
10604 shared_mem_bytes: 0,
10605 };
10606 let __s_b = self.gpu.stream();
10607 let mut b2 = __s_b.launch_builder(&f);
10608 b2.arg(a)
10609 .arg(wa)
10610 .arg(b)
10611 .arg(w)
10612 .arg(&mut *res)
10613 .arg(&mut *dst)
10614 .arg(&mut *out_q)
10615 .arg(&mut *out_d)
10616 .arg(&nc)
10617 .arg(&e);
10618 unsafe {
10619 b2.launch(cfg)?;
10620 }
10621 Ok(())
10622 }
10623
10624 #[allow(clippy::too_many_arguments)]
10627 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10628 &self,
10629 a: &CudaSlice<f32>,
10630 wa: &CudaSlice<f32>,
10631 b_in: &CudaSlice<f32>,
10632 c: f32,
10633 w: &CudaSlice<f32>,
10634 res: &mut CudaSlice<f32>,
10635 ncols: usize,
10636 nrows: usize,
10637 eps: f32,
10638 out_q: &mut CudaSlice<i8>,
10639 out_d: &mut CudaSlice<f32>,
10640 ) -> Result<(), Box<dyn std::error::Error>> {
10641 debug_assert!(ncols % 128 == 0);
10642 let (nc, e2) = (ncols as i32, eps);
10643 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10644 let cfg = LaunchConfig {
10645 grid_dim: (nrows as u32, 1, 1),
10646 block_dim: (rms_block(), 1, 1),
10647 shared_mem_bytes: 0,
10648 };
10649 let __s_b = self.gpu.stream();
10650 let mut b2 = __s_b.launch_builder(&f);
10651 b2.arg(a)
10652 .arg(wa)
10653 .arg(b_in)
10654 .arg(&c)
10655 .arg(w)
10656 .arg(&mut *res)
10657 .arg(&mut *out_q)
10658 .arg(&mut *out_d)
10659 .arg(&nc)
10660 .arg(&e2);
10661 unsafe {
10662 b2.launch(cfg)?;
10663 }
10664 Ok(())
10665 }
10666
10667 pub fn g4_pnfold_on() -> bool {
10675 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10676 *ON.get_or_init(|| {
10677 std::env::var("MEMRA_G4_PNFOLD")
10678 .map(|v| v != "0")
10679 .unwrap_or(true)
10680 })
10681 }
10682
10683 pub fn build_q4_out_concat3(
10687 &self,
10688 w0: &crate::model::GpuTensor,
10689 w1: &crate::model::GpuTensor,
10690 w2: &crate::model::GpuTensor,
10691 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10692 use crate::model::GpuTensor;
10693 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10694 match w {
10695 GpuTensor::Quant {
10696 qtype,
10697 row_bytes,
10698 rp,
10699 ..
10700 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10701 _ => None,
10702 }
10703 };
10704 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10705 else {
10706 return Ok(None);
10707 };
10708 if rb0 != rb1
10709 || rb0 != rb2
10710 || w0.in_features() != w1.in_features()
10711 || w0.in_features() != w2.in_features()
10712 {
10713 return Ok(None);
10714 }
10715 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10716 match w {
10717 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10718 _ => unreachable!(),
10719 }
10720 }
10721 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10722 let total = rb0 * (o0 + o1 + o2);
10723 let mut cat = self.alloc_u8(total)?;
10724 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10725 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10726 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10727 Ok(Some(GpuTensor::Quant {
10728 bytes: cat,
10729 qtype: QT_Q4_0,
10730 row_bytes: rb0,
10731 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10732 scale: 1.0,
10733 rp: false,
10734 #[cfg(memra_cutlass)]
10735 cutlass: None,
10736 fp8: None,
10737 blk: None,
10738 rp4: None,
10739 f16: None,
10740 }))
10741 }
10742
10743 fn full_width_rope_only(
10761 kernel: &str,
10762 n_rot: usize,
10763 head_dim: usize,
10764 ) -> Result<(), Box<dyn std::error::Error>> {
10765 if n_rot == head_dim {
10766 return Ok(());
10767 }
10768 Err(format!(
10769 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10770 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10771 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10772 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10773 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10774 )
10775 .into())
10776 }
10777
10778 #[allow(clippy::too_many_arguments)]
10782 pub fn rms_norm_qkv_rope_cat(
10783 &self,
10784 qkv: &CudaSlice<f32>,
10785 wq: &CudaSlice<f32>,
10786 wk: &CudaSlice<f32>,
10787 wv: &CudaSlice<f32>,
10788 q: &mut CudaSlice<f32>,
10789 k: &mut CudaSlice<f32>,
10790 v: &mut CudaSlice<f32>,
10791 head_dim: usize,
10792 n_rot: usize,
10793 rq: usize,
10794 rk: usize,
10795 pos: &CudaSlice<i32>,
10796 nh_q: usize,
10797 nh_k: usize,
10798 base: f32,
10799 freq_scale: f32,
10800 ff: Option<&CudaSlice<f32>>,
10801 eps: f32,
10802 ) -> Result<(), Box<dyn std::error::Error>> {
10803 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10804 let rows = rq + rk + rk;
10805 let theta_scale = base.powf(-2.0 / head_dim as f32);
10806 let (nc, rqi, rki, nhq, nhk) = (
10807 head_dim as i32,
10808 rq as i32,
10809 rk as i32,
10810 nh_q as i32,
10811 nh_k as i32,
10812 );
10813 if Self::pdl_on() {
10814 use cudarc::driver::{DevicePtr, DevicePtrMut};
10815 let s = &self.gpu.stream();
10816 let (pqkv, _g0) = qkv.device_ptr(s);
10817 let (pwq, _g1) = wq.device_ptr(s);
10818 let (pwk, _g2) = wk.device_ptr(s);
10819 let (pwv, _g3) = wv.device_ptr(s);
10820 let (pq, _g4) = q.device_ptr_mut(s);
10821 let (pk, _g5) = k.device_ptr_mut(s);
10822 let (pv, _g6) = v.device_ptr_mut(s);
10823 let (ppos, _g7) = pos.device_ptr(s);
10824 let (pff, _g8) = match ff {
10825 Some(t) => {
10826 let (p, g) = t.device_ptr(s);
10827 (p, Some(g))
10828 }
10829 None => (0, None),
10830 };
10831 let mut ps = [
10832 &pqkv as *const _ as *mut std::ffi::c_void,
10833 &pwq as *const _ as *mut _,
10834 &pwk as *const _ as *mut _,
10835 &pwv as *const _ as *mut _,
10836 &pq as *const _ as *mut _,
10837 &pk as *const _ as *mut _,
10838 &pv as *const _ as *mut _,
10839 &nc as *const _ as *mut _,
10840 &rqi as *const _ as *mut _,
10841 &rki as *const _ as *mut _,
10842 &ppos as *const _ as *mut _,
10843 &nhq as *const _ as *mut _,
10844 &nhk as *const _ as *mut _,
10845 &theta_scale as *const _ as *mut _,
10846 &freq_scale as *const _ as *mut _,
10847 &pff as *const _ as *mut _,
10848 &eps as *const _ as *mut _,
10849 ];
10850 unsafe {
10851 self.launch_pdl(
10852 "rms_norm_qkv_rope_cat_f32",
10853 (rows as u32, 1, 1),
10854 (rms_block(), 1, 1),
10855 &mut ps,
10856 )?;
10857 }
10858 return Ok(());
10859 }
10860 let f = self.func("rms_norm_qkv_rope_cat_f32");
10861 let cfg = LaunchConfig {
10862 grid_dim: (rows as u32, 1, 1),
10863 block_dim: (rms_block(), 1, 1),
10864 shared_mem_bytes: 0,
10865 };
10866 let __s_b = self.gpu.stream();
10867 let mut b = __s_b.launch_builder(&f);
10868 match ff {
10869 Some(t) => {
10870 b.arg(qkv)
10871 .arg(wq)
10872 .arg(wk)
10873 .arg(wv)
10874 .arg(&mut *q)
10875 .arg(&mut *k)
10876 .arg(&mut *v)
10877 .arg(&nc)
10878 .arg(&rqi)
10879 .arg(&rki)
10880 .arg(pos)
10881 .arg(&nhq)
10882 .arg(&nhk)
10883 .arg(&theta_scale)
10884 .arg(&freq_scale)
10885 .arg(t)
10886 .arg(&eps);
10887 unsafe {
10888 b.launch(cfg)?;
10889 }
10890 }
10891 None => {
10892 let null: u64 = 0;
10893 b.arg(qkv)
10894 .arg(wq)
10895 .arg(wk)
10896 .arg(wv)
10897 .arg(&mut *q)
10898 .arg(&mut *k)
10899 .arg(&mut *v)
10900 .arg(&nc)
10901 .arg(&rqi)
10902 .arg(&rki)
10903 .arg(pos)
10904 .arg(&nhq)
10905 .arg(&nhk)
10906 .arg(&theta_scale)
10907 .arg(&freq_scale)
10908 .arg(&null)
10909 .arg(&eps);
10910 unsafe {
10911 b.launch(cfg)?;
10912 }
10913 }
10914 }
10915 Ok(())
10916 }
10917
10918 #[allow(clippy::too_many_arguments)]
10922 pub fn rms_norm_qkv_rope(
10923 &self,
10924 q0: &CudaSlice<f32>,
10925 k0: &CudaSlice<f32>,
10926 v0: &CudaSlice<f32>,
10927 wq: &CudaSlice<f32>,
10928 wk: &CudaSlice<f32>,
10929 wv: &CudaSlice<f32>,
10930 q: &mut CudaSlice<f32>,
10931 k: &mut CudaSlice<f32>,
10932 v: &mut CudaSlice<f32>,
10933 head_dim: usize,
10934 n_rot: usize,
10935 rq: usize,
10936 rk: usize,
10937 pos: &CudaSlice<i32>,
10938 nh_q: usize,
10939 nh_k: usize,
10940 base: f32,
10941 freq_scale: f32,
10942 ff: Option<&CudaSlice<f32>>,
10943 eps: f32,
10944 ) -> Result<(), Box<dyn std::error::Error>> {
10945 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10946 let f = self.func("rms_norm_qkv_rope_f32");
10947 let rows = rq + rk + rk; let cfg = LaunchConfig {
10949 grid_dim: (rows as u32, 1, 1),
10950 block_dim: (rms_block(), 1, 1),
10951 shared_mem_bytes: 0,
10952 };
10953 let theta_scale = base.powf(-2.0 / head_dim as f32);
10954 let (nc, rqi, rki, nhq, nhk) = (
10955 head_dim as i32,
10956 rq as i32,
10957 rk as i32,
10958 nh_q as i32,
10959 nh_k as i32,
10960 );
10961 let __s_b = self.gpu.stream();
10962 let mut b = __s_b.launch_builder(&f);
10963 match ff {
10964 Some(t) => {
10965 b.arg(q0)
10966 .arg(k0)
10967 .arg(v0)
10968 .arg(wq)
10969 .arg(wk)
10970 .arg(wv)
10971 .arg(&mut *q)
10972 .arg(&mut *k)
10973 .arg(&mut *v)
10974 .arg(&nc)
10975 .arg(&rqi)
10976 .arg(&rki)
10977 .arg(pos)
10978 .arg(&nhq)
10979 .arg(&nhk)
10980 .arg(&theta_scale)
10981 .arg(&freq_scale)
10982 .arg(t)
10983 .arg(&eps);
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 unsafe {
11010 b.launch(cfg)?;
11011 }
11012 }
11013 }
11014 Ok(())
11015 }
11016
11017 #[allow(clippy::too_many_arguments)]
11023 pub fn rms_norm_qkv_rope_append_dc(
11024 &self,
11025 q0: &CudaSlice<f32>,
11026 k0: &CudaSlice<f32>,
11027 v0: &CudaSlice<f32>,
11028 wq: &CudaSlice<f32>,
11029 wk: &CudaSlice<f32>,
11030 wv: &CudaSlice<f32>,
11031 q: &mut CudaSlice<f32>,
11032 k: &mut CudaSlice<f32>,
11033 v: &mut CudaSlice<f32>,
11034 head_dim: usize,
11035 n_rot: usize,
11036 rq: usize,
11037 rk: usize,
11038 pos: &CudaSlice<i32>,
11039 nh_q: usize,
11040 nh_k: usize,
11041 base: f32,
11042 freq_scale: f32,
11043 ff: Option<&CudaSlice<f32>>,
11044 eps: f32,
11045 kc: &mut CudaSlice<u8>,
11046 vc: &mut CudaSlice<u8>,
11047 t_dev: &CudaSlice<i32>,
11048 k_tok_bytes: usize,
11049 v_tok_bytes: usize,
11050 g: bool,
11051 ) -> Result<(), Box<dyn std::error::Error>> {
11052 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11053 let rows = rq + rk + rk;
11054 let theta_scale = base.powf(-2.0 / head_dim as f32);
11055 let (nc, rqi, rki, nhq, nhk) = (
11056 head_dim as i32,
11057 rq as i32,
11058 rk as i32,
11059 nh_q as i32,
11060 nh_k as i32,
11061 );
11062 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11063 if Self::pdl_on() && Self::pdl_wb_on() {
11064 use cudarc::driver::{DevicePtr, DevicePtrMut};
11065 let s = &self.gpu.stream();
11066 let (p0, _a0) = q0.device_ptr(s);
11067 let (p1, _a1) = k0.device_ptr(s);
11068 let (p2, _a2) = v0.device_ptr(s);
11069 let (pwq, _a3) = wq.device_ptr(s);
11070 let (pwk, _a4) = wk.device_ptr(s);
11071 let (pwv, _a5) = wv.device_ptr(s);
11072 let (pq, _a6) = q.device_ptr_mut(s);
11073 let (pk, _a7) = k.device_ptr_mut(s);
11074 let (pv, _a8) = v.device_ptr_mut(s);
11075 let (pp, _a9) = pos.device_ptr(s);
11076 let pff: u64 = match ff {
11077 Some(t) => {
11078 let (p, _gg) = t.device_ptr(s);
11079 p as u64
11080 }
11081 None => 0,
11082 };
11083 let (pkc, _a10) = kc.device_ptr_mut(s);
11084 let (pvc, _a11) = vc.device_ptr_mut(s);
11085 let (pt, _a12) = t_dev.device_ptr(s);
11086 let mut ps = [
11087 &p0 as *const _ as *mut std::ffi::c_void,
11088 &p1 as *const _ as *mut _,
11089 &p2 as *const _ as *mut _,
11090 &pwq as *const _ as *mut _,
11091 &pwk as *const _ as *mut _,
11092 &pwv as *const _ as *mut _,
11093 &pq as *const _ as *mut _,
11094 &pk as *const _ as *mut _,
11095 &pv as *const _ as *mut _,
11096 &nc as *const _ as *mut _,
11097 &rqi as *const _ as *mut _,
11098 &rki as *const _ as *mut _,
11099 &pp as *const _ as *mut _,
11100 &nhq as *const _ as *mut _,
11101 &nhk as *const _ as *mut _,
11102 &theta_scale as *const _ as *mut _,
11103 &freq_scale as *const _ as *mut _,
11104 &pff as *const _ as *mut _,
11105 &eps as *const _ as *mut _,
11106 &pkc as *const _ as *mut _,
11107 &pvc as *const _ as *mut _,
11108 &pt as *const _ as *mut _,
11109 &ktb as *const _ as *mut _,
11110 &vtb as *const _ as *mut _,
11111 ];
11112 unsafe {
11113 self.launch_pdl_flash(
11114 g,
11115 "rms_norm_qkv_rope_append_dc_f32",
11116 (rows as u32, 1, 1),
11117 (rms_block(), 1, 1),
11118 0,
11119 &mut ps,
11120 )?;
11121 }
11122 return Ok(());
11123 }
11124 let f = if g {
11125 self.func_g("rms_norm_qkv_rope_append_dc_f32")
11126 } else {
11127 self.func("rms_norm_qkv_rope_append_dc_f32")
11128 };
11129 let cfg = LaunchConfig {
11130 grid_dim: (rows as u32, 1, 1),
11131 block_dim: (rms_block(), 1, 1),
11132 shared_mem_bytes: 0,
11133 };
11134 let __s_b = self.gpu.stream();
11135 let mut b = __s_b.launch_builder(&f);
11136 match ff {
11137 Some(t) => {
11138 b.arg(q0)
11139 .arg(k0)
11140 .arg(v0)
11141 .arg(wq)
11142 .arg(wk)
11143 .arg(wv)
11144 .arg(&mut *q)
11145 .arg(&mut *k)
11146 .arg(&mut *v)
11147 .arg(&nc)
11148 .arg(&rqi)
11149 .arg(&rki)
11150 .arg(pos)
11151 .arg(&nhq)
11152 .arg(&nhk)
11153 .arg(&theta_scale)
11154 .arg(&freq_scale)
11155 .arg(t)
11156 .arg(&eps)
11157 .arg(&mut *kc)
11158 .arg(&mut *vc)
11159 .arg(t_dev)
11160 .arg(&ktb)
11161 .arg(&vtb);
11162 unsafe {
11163 b.launch(cfg)?;
11164 }
11165 }
11166 None => {
11167 let null: u64 = 0;
11168 b.arg(q0)
11169 .arg(k0)
11170 .arg(v0)
11171 .arg(wq)
11172 .arg(wk)
11173 .arg(wv)
11174 .arg(&mut *q)
11175 .arg(&mut *k)
11176 .arg(&mut *v)
11177 .arg(&nc)
11178 .arg(&rqi)
11179 .arg(&rki)
11180 .arg(pos)
11181 .arg(&nhq)
11182 .arg(&nhk)
11183 .arg(&theta_scale)
11184 .arg(&freq_scale)
11185 .arg(&null)
11186 .arg(&eps)
11187 .arg(&mut *kc)
11188 .arg(&mut *vc)
11189 .arg(t_dev)
11190 .arg(&ktb)
11191 .arg(&vtb);
11192 unsafe {
11193 b.launch(cfg)?;
11194 }
11195 }
11196 }
11197 Ok(())
11198 }
11199
11200 #[allow(clippy::too_many_arguments)]
11208 pub fn rms_norm_qkv_rope_append(
11209 &self,
11210 q0: &CudaSlice<f32>,
11211 k0: &CudaSlice<f32>,
11212 v0: &CudaSlice<f32>,
11213 wq: &CudaSlice<f32>,
11214 wk: &CudaSlice<f32>,
11215 wv: &CudaSlice<f32>,
11216 q: &mut CudaSlice<f32>,
11217 k: &mut CudaSlice<f32>,
11218 v: &mut CudaSlice<f32>,
11219 head_dim: usize,
11220 n_rot: usize,
11221 rq: usize,
11222 rk: usize,
11223 pos: &CudaSlice<i32>,
11224 nh_q: usize,
11225 nh_k: usize,
11226 base: f32,
11227 freq_scale: f32,
11228 ff: Option<&CudaSlice<f32>>,
11229 eps: f32,
11230 kc: &mut CudaSlice<u8>,
11231 vc: &mut CudaSlice<u8>,
11232 t: usize,
11233 k_tok_bytes: usize,
11234 v_tok_bytes: usize,
11235 g: bool,
11236 ) -> Result<(), Box<dyn std::error::Error>> {
11237 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11238 let rows = rq + rk + rk;
11239 let theta_scale = base.powf(-2.0 / head_dim as f32);
11240 let (nc, rqi, rki, nhq, nhk) = (
11241 head_dim as i32,
11242 rq as i32,
11243 rk as i32,
11244 nh_q as i32,
11245 nh_k as i32,
11246 );
11247 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11248 let ti = t as i32;
11249 if Self::pdl_on() && Self::pdl_wb_on() {
11250 use cudarc::driver::{DevicePtr, DevicePtrMut};
11251 let s = &self.gpu.stream();
11252 let (p0, _a0) = q0.device_ptr(s);
11253 let (p1, _a1) = k0.device_ptr(s);
11254 let (p2, _a2) = v0.device_ptr(s);
11255 let (pwq, _a3) = wq.device_ptr(s);
11256 let (pwk, _a4) = wk.device_ptr(s);
11257 let (pwv, _a5) = wv.device_ptr(s);
11258 let (pq, _a6) = q.device_ptr_mut(s);
11259 let (pk, _a7) = k.device_ptr_mut(s);
11260 let (pv, _a8) = v.device_ptr_mut(s);
11261 let (pp, _a9) = pos.device_ptr(s);
11262 let pff: u64 = match ff {
11263 Some(t) => {
11264 let (p, _gg) = t.device_ptr(s);
11265 p as u64
11266 }
11267 None => 0,
11268 };
11269 let (pkc, _a10) = kc.device_ptr_mut(s);
11270 let (pvc, _a11) = vc.device_ptr_mut(s);
11271 let mut ps = [
11272 &p0 as *const _ as *mut std::ffi::c_void,
11273 &p1 as *const _ as *mut _,
11274 &p2 as *const _ as *mut _,
11275 &pwq as *const _ as *mut _,
11276 &pwk as *const _ as *mut _,
11277 &pwv as *const _ as *mut _,
11278 &pq as *const _ as *mut _,
11279 &pk as *const _ as *mut _,
11280 &pv as *const _ as *mut _,
11281 &nc as *const _ as *mut _,
11282 &rqi as *const _ as *mut _,
11283 &rki as *const _ as *mut _,
11284 &pp as *const _ as *mut _,
11285 &nhq as *const _ as *mut _,
11286 &nhk as *const _ as *mut _,
11287 &theta_scale as *const _ as *mut _,
11288 &freq_scale as *const _ as *mut _,
11289 &pff as *const _ as *mut _,
11290 &eps as *const _ as *mut _,
11291 &pkc as *const _ as *mut _,
11292 &pvc as *const _ as *mut _,
11293 &ti as *const _ as *mut _,
11294 &ktb as *const _ as *mut _,
11295 &vtb as *const _ as *mut _,
11296 ];
11297 unsafe {
11298 self.launch_pdl_flash(
11299 g,
11300 "rms_norm_qkv_rope_append_f32",
11301 (rows as u32, 1, 1),
11302 (rms_block(), 1, 1),
11303 0,
11304 &mut ps,
11305 )?;
11306 }
11307 return Ok(());
11308 }
11309 let f = if g {
11310 self.func_g("rms_norm_qkv_rope_append_f32")
11311 } else {
11312 self.func("rms_norm_qkv_rope_append_f32")
11313 };
11314 let cfg = LaunchConfig {
11315 grid_dim: (rows as u32, 1, 1),
11316 block_dim: (rms_block(), 1, 1),
11317 shared_mem_bytes: 0,
11318 };
11319 let __s_b = self.gpu.stream();
11320 let mut b = __s_b.launch_builder(&f);
11321 let null: u64 = 0;
11322 b.arg(q0)
11323 .arg(k0)
11324 .arg(v0)
11325 .arg(wq)
11326 .arg(wk)
11327 .arg(wv)
11328 .arg(&mut *q)
11329 .arg(&mut *k)
11330 .arg(&mut *v)
11331 .arg(&nc)
11332 .arg(&rqi)
11333 .arg(&rki)
11334 .arg(pos)
11335 .arg(&nhq)
11336 .arg(&nhk)
11337 .arg(&theta_scale)
11338 .arg(&freq_scale);
11339 match ff {
11340 Some(t) => {
11341 b.arg(t);
11342 }
11343 None => {
11344 b.arg(&null);
11345 }
11346 }
11347 b.arg(&eps)
11348 .arg(&mut *kc)
11349 .arg(&mut *vc)
11350 .arg(&ti)
11351 .arg(&ktb)
11352 .arg(&vtb);
11353 unsafe {
11354 b.launch(cfg)?;
11355 }
11356 Ok(())
11357 }
11358
11359 pub fn add_q8_1(
11360 &self,
11361 a: &CudaSlice<f32>,
11362 b: &CudaSlice<f32>,
11363 res: &mut CudaSlice<f32>,
11364 ncols: usize,
11365 nrows: usize,
11366 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11367 debug_assert!(ncols % 128 == 0);
11368 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11369 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11370 let f = self.func("add_q8_1_f32");
11371 let cfg = LaunchConfig {
11372 grid_dim: (nrows as u32, 1, 1),
11373 block_dim: (rms_block(), 1, 1),
11374 shared_mem_bytes: 0,
11375 };
11376 let nc = ncols as i32;
11377 let __s_b2 = self.gpu.stream();
11378 let mut b2 = __s_b2.launch_builder(&f);
11379 b2.arg(a)
11380 .arg(b)
11381 .arg(&mut *res)
11382 .arg(&mut out_q)
11383 .arg(&mut out_d)
11384 .arg(&nc);
11385 unsafe {
11386 b2.launch(cfg)?;
11387 }
11388 Ok((out_q, out_d))
11389 }
11390
11391 pub fn rms_pre_add_q8_1(
11395 &self,
11396 a: &CudaSlice<f32>,
11397 wa: &CudaSlice<f32>,
11398 b: &CudaSlice<f32>,
11399 res: &mut CudaSlice<f32>,
11400 ncols: usize,
11401 nrows: usize,
11402 eps: f32,
11403 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11404 debug_assert!(ncols % 128 == 0);
11405 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11406 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11407 let f = self.func("rms_pre_add_q8_1_f32");
11408 let cfg = LaunchConfig {
11409 grid_dim: (nrows as u32, 1, 1),
11410 block_dim: (rms_block(), 1, 1),
11411 shared_mem_bytes: 0,
11412 };
11413 let (nc, ep) = (ncols as i32, eps);
11414 let __s_b2 = self.gpu.stream();
11415 let mut b2 = __s_b2.launch_builder(&f);
11416 b2.arg(a)
11417 .arg(wa)
11418 .arg(b)
11419 .arg(&mut *res)
11420 .arg(&mut out_q)
11421 .arg(&mut out_d)
11422 .arg(&nc)
11423 .arg(&ep);
11424 unsafe {
11425 b2.launch(cfg)?;
11426 }
11427 Ok((out_q, out_d))
11428 }
11429
11430 pub fn l2_v2_on(ncols: usize) -> bool {
11434 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11435 }
11436
11437 pub fn l2_norm_pp(
11438 &self,
11439 x: &CudaSlice<f32>,
11440 dst: &mut CudaSlice<f32>,
11441 dst16: Option<&mut CudaSlice<u8>>,
11442 ncols: usize,
11443 nrows: usize,
11444 eps: f32,
11445 ) -> Result<(), Box<dyn std::error::Error>> {
11446 if Self::l2_v2_on(ncols) {
11447 let f = self.func("l2_norm_pp_v2_f32");
11448 let rows_per_block = 8u32; let cfg = LaunchConfig {
11450 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11451 block_dim: (256, 1, 1),
11452 shared_mem_bytes: 0,
11453 };
11454 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11455 let d16: u64 = match dst16 {
11457 Some(d) => self.addr_u8(d),
11458 None => 0,
11459 };
11460 let __s_b = self.gpu.stream();
11461 let mut b = __s_b.launch_builder(&f);
11462 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11463 unsafe {
11464 b.launch(cfg)?;
11465 }
11466 return Ok(());
11467 }
11468 self.l2_norm(x, dst, ncols, nrows, eps)
11469 }
11470
11471 pub fn l2_norm(
11472 &self,
11473 x: &CudaSlice<f32>,
11474 dst: &mut CudaSlice<f32>,
11475 ncols: usize,
11476 nrows: usize,
11477 eps: f32,
11478 ) -> Result<(), Box<dyn std::error::Error>> {
11479 let f = self.func("l2_norm_f32");
11480 let cfg = LaunchConfig {
11481 grid_dim: (nrows as u32, 1, 1),
11482 block_dim: (256, 1, 1),
11483 shared_mem_bytes: 0,
11484 };
11485 let (nc, e) = (ncols as i32, eps);
11486 let __s_b = self.gpu.stream();
11487 let mut b = __s_b.launch_builder(&f);
11488 b.arg(x).arg(dst).arg(&nc).arg(&e);
11489 unsafe {
11490 b.launch(cfg)?;
11491 }
11492 Ok(())
11493 }
11494
11495 pub fn l2_norm_decode(
11501 &self,
11502 x: &CudaSlice<f32>,
11503 dst: &mut CudaSlice<f32>,
11504 ncols: usize,
11505 nrows: usize,
11506 eps: f32,
11507 ) -> Result<(), Box<dyn std::error::Error>> {
11508 let f = self.func("l2_norm_f32");
11509 let cfg = LaunchConfig {
11510 grid_dim: (nrows as u32, 1, 1),
11511 block_dim: (32, 1, 1),
11512 shared_mem_bytes: 0,
11513 };
11514 let (nc, e) = (ncols as i32, eps);
11515 let __s_b = self.gpu.stream();
11516 let mut b = __s_b.launch_builder(&f);
11517 b.arg(x).arg(dst).arg(&nc).arg(&e);
11518 unsafe {
11519 b.launch(cfg)?;
11520 }
11521 Ok(())
11522 }
11523
11524 pub fn rope_neox(
11526 &self,
11527 x: &mut CudaSlice<f32>,
11528 pos: &CudaSlice<i32>,
11529 head_dim: usize,
11530 n_dims: usize,
11531 n_heads: usize,
11532 n_tokens: usize,
11533 freq_base: f32,
11534 freq_scale: f32,
11535 ) -> Result<(), Box<dyn std::error::Error>> {
11536 let f = self.func("rope_neox_f32");
11537 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11538 let grid = (n_heads * n_tokens) as u32;
11539 let cfg = LaunchConfig {
11540 grid_dim: (grid, 1, 1),
11541 block_dim: ((head_dim / 2) as u32, 1, 1),
11542 shared_mem_bytes: 0,
11543 };
11544 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11545 let __s_b = self.gpu.stream();
11546 let mut b = __s_b.launch_builder(&f);
11547 b.arg(x)
11548 .arg(pos)
11549 .arg(&hd)
11550 .arg(&nd)
11551 .arg(&nh)
11552 .arg(&theta_scale)
11553 .arg(&freq_scale);
11554 unsafe {
11555 b.launch(cfg)?;
11556 }
11557 Ok(())
11558 }
11559
11560 pub fn rope_neox_ff(
11562 &self,
11563 x: &mut CudaSlice<f32>,
11564 pos: &CudaSlice<i32>,
11565 head_dim: usize,
11566 n_dims: usize,
11567 n_heads: usize,
11568 n_tokens: usize,
11569 freq_base: f32,
11570 freq_scale: f32,
11571 ff: &CudaSlice<f32>,
11572 ) -> Result<(), Box<dyn std::error::Error>> {
11573 let f = self.func("rope_neox_ff_f32");
11574 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11575 let grid = (n_heads * n_tokens) as u32;
11576 let cfg = LaunchConfig {
11577 grid_dim: (grid, 1, 1),
11578 block_dim: ((head_dim / 2) as u32, 1, 1),
11579 shared_mem_bytes: 0,
11580 };
11581 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11582 let __s_b = self.gpu.stream();
11583 let mut b = __s_b.launch_builder(&f);
11584 b.arg(x)
11585 .arg(pos)
11586 .arg(&hd)
11587 .arg(&nd)
11588 .arg(&nh)
11589 .arg(&theta_scale)
11590 .arg(&freq_scale)
11591 .arg(ff);
11592 unsafe {
11593 b.launch(cfg)?;
11594 }
11595 Ok(())
11596 }
11597
11598 #[allow(clippy::too_many_arguments)]
11600 pub fn rope_neox2(
11601 &self,
11602 q: &mut CudaSlice<f32>,
11603 k: &mut CudaSlice<f32>,
11604 pos: &CudaSlice<i32>,
11605 head_dim: usize,
11606 n_dims: usize,
11607 nh_q: usize,
11608 nh_k: usize,
11609 n_tokens: usize,
11610 freq_base: f32,
11611 freq_scale: f32,
11612 ff: Option<&CudaSlice<f32>>,
11613 ) -> Result<(), Box<dyn std::error::Error>> {
11614 let f = self.func("rope_neox2_f32");
11615 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11616 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11617 let cfg = LaunchConfig {
11618 grid_dim: (grid, 1, 1),
11619 block_dim: ((head_dim / 2) as u32, 1, 1),
11620 shared_mem_bytes: 0,
11621 };
11622 let (hd, nd, nq, nk, nt) = (
11623 head_dim as i32,
11624 n_dims as i32,
11625 nh_q as i32,
11626 nh_k as i32,
11627 n_tokens as i32,
11628 );
11629 let __s_b = self.gpu.stream();
11630 let mut b = __s_b.launch_builder(&f);
11631 b.arg(q)
11632 .arg(k)
11633 .arg(pos)
11634 .arg(&hd)
11635 .arg(&nd)
11636 .arg(&nq)
11637 .arg(&nk)
11638 .arg(&nt)
11639 .arg(&theta_scale)
11640 .arg(&freq_scale);
11641 match ff {
11642 Some(ffv) => {
11643 b.arg(ffv);
11644 unsafe {
11645 b.launch(cfg)?;
11646 }
11647 }
11648 None => {
11649 let null: u64 = 0;
11650 b.arg(&null);
11651 unsafe {
11652 b.launch(cfg)?;
11653 }
11654 }
11655 }
11656 Ok(())
11657 }
11658
11659 pub fn gelu_tanh_mul(
11661 &self,
11662 gate: &CudaSlice<f32>,
11663 up: &CudaSlice<f32>,
11664 dst: &mut CudaSlice<f32>,
11665 n: usize,
11666 ) -> Result<(), Box<dyn std::error::Error>> {
11667 let f = self.func("gelu_tanh_mul_f32");
11668 let cfg = LaunchConfig::for_num_elems(n as u32);
11669 let ni = n as i32;
11670 let __s_b = self.gpu.stream();
11671 let mut b = __s_b.launch_builder(&f);
11672 b.arg(gate).arg(up).arg(dst).arg(&ni);
11673 unsafe {
11674 b.launch(cfg)?;
11675 }
11676 Ok(())
11677 }
11678
11679 pub fn silu_mul(
11680 &self,
11681 gate: &CudaSlice<f32>,
11682 up: &CudaSlice<f32>,
11683 dst: &mut CudaSlice<f32>,
11684 n: usize,
11685 ) -> Result<(), Box<dyn std::error::Error>> {
11686 let f = self.func("silu_mul_f32");
11687 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11689 let ni = n as i32;
11690 let __s_b = self.gpu.stream();
11691 let mut b = __s_b.launch_builder(&f);
11692 b.arg(gate).arg(up).arg(dst).arg(&ni);
11693 unsafe {
11694 b.launch(cfg)?;
11695 }
11696 Ok(())
11697 }
11698
11699 pub fn silu_mul_host_expf(
11701 &self,
11702 gate: &CudaSlice<f32>,
11703 up: &CudaSlice<f32>,
11704 dst: &mut CudaSlice<f32>,
11705 n: usize,
11706 ) -> Result<(), Box<dyn std::error::Error>> {
11707 let f = self.func("silu_mul_host_expf_f32");
11708 let cfg = LaunchConfig::for_num_elems(n as u32);
11709 let ni = n as i32;
11710 let __s_b = self.gpu.stream();
11711 let mut b = __s_b.launch_builder(&f);
11712 b.arg(gate).arg(up).arg(dst).arg(&ni);
11713 unsafe {
11714 b.launch(cfg)?;
11715 }
11716 Ok(())
11717 }
11718
11719 pub fn silu_clamped_mul_host_expf(
11721 &self,
11722 gate: &CudaSlice<f32>,
11723 up: &CudaSlice<f32>,
11724 limit: f32,
11725 dst: &mut CudaSlice<f32>,
11726 n: usize,
11727 ) -> Result<(), Box<dyn std::error::Error>> {
11728 if !limit.is_finite() || limit <= 0.0 {
11729 return Err(
11730 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11731 );
11732 }
11733 let f = self.func("silu_clamped_mul_host_expf_f32");
11734 let cfg = LaunchConfig::for_num_elems(n as u32);
11735 let ni = n as i32;
11736 let __s_b = self.gpu.stream();
11737 let mut b = __s_b.launch_builder(&f);
11738 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11739 unsafe {
11740 b.launch(cfg)?;
11741 }
11742 Ok(())
11743 }
11744
11745 pub fn silu_mul_f16out(
11748 &self,
11749 gate: &CudaSlice<f32>,
11750 up: &CudaSlice<f32>,
11751 dst: &mut CudaSlice<f32>,
11752 dst16: &mut CudaSlice<u8>,
11753 n: usize,
11754 ) -> Result<(), Box<dyn std::error::Error>> {
11755 let f = self.func("silu_mul_f16out_f32");
11756 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11757 let ni = n as i32;
11758 let __s_b = self.gpu.stream();
11759 let mut b = __s_b.launch_builder(&f);
11760 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11761 unsafe {
11762 b.launch(cfg)?;
11763 }
11764 Ok(())
11765 }
11766
11767 pub fn silu_mul_scaled(
11774 &self,
11775 gate: &CudaSlice<f32>,
11776 up: &CudaSlice<f32>,
11777 gs: f32,
11778 us: f32,
11779 dst: &mut CudaSlice<f32>,
11780 n: usize,
11781 ) -> Result<(), Box<dyn std::error::Error>> {
11782 let f = self.func("silu_mul_scaled_f32");
11783 let cfg = LaunchConfig::for_num_elems(n as u32);
11784 let ni = n as i32;
11785 let (gsf, usf) = (gs, us);
11786 let __s_b = self.gpu.stream();
11787 let mut b = __s_b.launch_builder(&f);
11788 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11789 unsafe {
11790 b.launch(cfg)?;
11791 }
11792 Ok(())
11793 }
11794
11795 #[allow(clippy::too_many_arguments)]
11799 pub fn swigluoai_mul_scaled(
11800 &self,
11801 gate: &CudaSlice<f32>,
11802 up: &CudaSlice<f32>,
11803 gs: f32,
11804 us: f32,
11805 alpha: f32,
11806 limit: f32,
11807 dst: &mut CudaSlice<f32>,
11808 n: usize,
11809 ) -> Result<(), Box<dyn std::error::Error>> {
11810 let f = self.func("swigluoai_mul_scaled_f32");
11811 let cfg = LaunchConfig::for_num_elems(n as u32);
11812 let ni = n as i32;
11813 let __s_b = self.gpu.stream();
11814 let mut b = __s_b.launch_builder(&f);
11815 b.arg(gate)
11816 .arg(up)
11817 .arg(&gs)
11818 .arg(&us)
11819 .arg(&alpha)
11820 .arg(&limit)
11821 .arg(dst)
11822 .arg(&ni);
11823 unsafe {
11824 b.launch(cfg)?;
11825 }
11826 Ok(())
11827 }
11828
11829 pub fn silu_mul_scaled_q8_1(
11837 &self,
11838 gate: &CudaSlice<f32>,
11839 up: &CudaSlice<f32>,
11840 gs: f32,
11841 us: f32,
11842 n: usize,
11843 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11844 let f = self.func("silu_mul_scaled_q8_1");
11845 let nblk = n / 32;
11846 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);
11850 let (gsf, usf, ni) = (gs, us, n as i32);
11851 let __s_b = self.gpu.stream();
11852 let mut b = __s_b.launch_builder(&f);
11853 b.arg(gate)
11854 .arg(up)
11855 .arg(&gsf)
11856 .arg(&usf)
11857 .arg(&mut aq)
11858 .arg(&mut ad)
11859 .arg(&ni);
11860 unsafe {
11861 b.launch(cfg)?;
11862 }
11863 Ok((aq, ad))
11864 }
11865
11866 pub fn add(
11867 &self,
11868 a: &CudaSlice<f32>,
11869 b_in: &CudaSlice<f32>,
11870 dst: &mut CudaSlice<f32>,
11871 n: usize,
11872 ) -> Result<(), Box<dyn std::error::Error>> {
11873 let f = self.func("add_f32");
11874 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11876 let ni = n as i32;
11877 let __s_bld = self.gpu.stream();
11878 let mut bld = __s_bld.launch_builder(&f);
11879 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11880 unsafe {
11881 bld.launch(cfg)?;
11882 }
11883 Ok(())
11884 }
11885
11886 pub fn mul(
11887 &self,
11888 a: &CudaSlice<f32>,
11889 b_in: &CudaSlice<f32>,
11890 dst: &mut CudaSlice<f32>,
11891 n: usize,
11892 ) -> Result<(), Box<dyn std::error::Error>> {
11893 let f = self.func("mul_f32");
11894 let cfg = LaunchConfig::for_num_elems(n as u32);
11895 let ni = n as i32;
11896 let __s_bld = self.gpu.stream();
11897 let mut bld = __s_bld.launch_builder(&f);
11898 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11899 unsafe {
11900 bld.launch(cfg)?;
11901 }
11902 Ok(())
11903 }
11904
11905 pub fn matmul(
11908 &self,
11909 w: &crate::model::GpuTensor,
11910 x: &CudaSlice<f32>,
11911 m: usize,
11912 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11913 use crate::model::GpuTensor;
11914 let in_f = w.in_features();
11915 let out_f = w.out_features();
11916 #[allow(non_snake_case)]
11924 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11927 usize::MAX
11928 } else {
11929 16usize
11930 };
11931
11932 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11957 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11958 return Ok(y);
11959 }
11960 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11967 return Ok(y);
11968 }
11969 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11972 return Ok(y);
11973 }
11974 }
11975 if let GpuTensor::Quant { qtype, .. } = w {
11990 if *qtype == QT_F8_E4M3_BLK {
11991 if m >= GEMM_M_THRESHOLD {
11992 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11993 return Ok(y);
11994 }
11995 }
11996 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11997 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11998 return Ok(y);
11999 }
12000 }
12001 }
12002 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12003 return self.qmatvec_mmq(w, x, m);
12004 }
12005 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12006 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12007 return self.qmatvec_gemm(w, &aq, &ad, m);
12008 }
12009 if m >= GEMM_M_THRESHOLD {
12012 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12013 return Ok(y);
12014 }
12015 }
12016 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12020 if m == 1 && fast {
12025 if let GpuTensor::Quant {
12026 bytes,
12027 qtype,
12028 row_bytes,
12029 rp,
12030 rp4,
12031 scale,
12032 ..
12033 } = w
12034 {
12035 if self.mmvq_supports(*qtype) {
12036 let (bytes, rp) = match rp4 {
12040 Some(m4) => (m4, true),
12041 None => (bytes, *rp),
12042 };
12043 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12044 return self.qmatvec_mmvq(
12045 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12046 );
12047 }
12048 }
12049 }
12050 if (2..=16).contains(&m)
12066 && fast
12067 && std::env::var("MEMRA_NO_BATCHED").is_err()
12068 && (m <= 4 || Self::b8_enabled())
12069 {
12070 let m_ok = m <= 8
12080 || matches!(w, GpuTensor::Quant { qtype, .. }
12081 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12082 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12083 if m_ok {
12084 if let GpuTensor::Quant {
12085 bytes,
12086 qtype,
12087 row_bytes,
12088 rp,
12089 rp4,
12090 ..
12091 } = w
12092 {
12093 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12094 let (bytes, rp) = match rp4 {
12095 Some(m4) => (m4, true),
12096 None => (bytes, *rp),
12097 };
12098 let mcols = Self::batched_mcols(m);
12099 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12100 let mut y = self.qmatvec_mmvq_batched(
12101 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12102 )?;
12103 if let GpuTensor::Quant { scale, .. } = w {
12104 if *scale != 1.0 {
12105 self.scale_inplace(&mut y, *scale, m * out_f)?;
12106 }
12107 }
12108 return Ok(y);
12109 }
12110 }
12111 }
12112 }
12113 if fast {
12119 if let GpuTensor::Quant {
12120 bytes,
12121 qtype,
12122 row_bytes,
12123 scale,
12124 ..
12125 } = w
12126 {
12127 if *qtype == QT_F8_E4M3 {
12128 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12129 return self.qmatvec_mmvq(
12130 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12131 );
12132 }
12133 }
12134 }
12135 let mut y = match w {
12136 GpuTensor::Quant {
12137 bytes,
12138 qtype,
12139 row_bytes,
12140 ..
12141 } if fast && *qtype == QT_Q8_0 => {
12142 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12143 }
12144 GpuTensor::Quant {
12145 bytes,
12146 qtype,
12147 row_bytes,
12148 ..
12149 } if fast && *qtype == QT_Q4_K => {
12150 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12151 }
12152 GpuTensor::Quant {
12153 bytes,
12154 qtype,
12155 row_bytes,
12156 ..
12157 } if fast && *qtype == QT_Q6_K => {
12158 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12159 }
12160 GpuTensor::Quant {
12161 bytes,
12162 qtype,
12163 row_bytes,
12164 ..
12165 } if fast && *qtype == QT_Q5_K => {
12166 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12167 }
12168 GpuTensor::Quant {
12169 bytes,
12170 qtype,
12171 row_bytes,
12172 ..
12173 } if fast && *qtype == QT_Q3_K => {
12174 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12175 }
12176 GpuTensor::Quant {
12177 bytes,
12178 qtype,
12179 row_bytes,
12180 rp,
12181 ..
12182 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12183 if *rp {
12184 "qmatvec_nvfp4_dp4a_rp"
12185 } else {
12186 "qmatvec_nvfp4_dp4a"
12187 },
12188 &bytes.slice(0..bytes.len()),
12189 x,
12190 m,
12191 in_f,
12192 out_f,
12193 *row_bytes,
12194 )?,
12195 GpuTensor::Quant {
12199 bytes,
12200 qtype,
12201 row_bytes,
12202 ..
12203 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12204 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12205 }
12206 GpuTensor::Quant {
12211 bytes,
12212 qtype,
12213 row_bytes,
12214 rp,
12215 ..
12216 } =>
12217 {
12220 self.qmatvec(
12221 bytes,
12222 x,
12223 m,
12224 in_f,
12225 out_f,
12226 if *rp && *qtype == QT_NVFP4 {
12227 QT_NVFP4_RP
12228 } else {
12229 *qtype
12230 },
12231 *row_bytes,
12232 )?
12233 }
12234 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12235 GpuTensor::FloatBf16 { data, .. } => {
12238 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12245 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12246 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12247 y
12248 } else {
12249 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12250 }
12251 }
12252 };
12253 if let GpuTensor::Quant { scale, .. } = w {
12255 if *scale != 1.0 {
12256 self.scale_inplace(&mut y, *scale, m * out_f)?;
12257 }
12258 }
12259 Ok(y)
12260 }
12261
12262 pub fn stage_a_raw_needed() -> bool {
12272 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12273 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12274 }
12275
12276 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12279 use crate::model::GpuTensor;
12280 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12281 return false;
12282 }
12283 match w {
12284 GpuTensor::Quant { qtype, .. } => {
12291 matches!(
12292 *qtype,
12293 QT_Q8_0
12294 | QT_Q4_K
12295 | QT_Q6_K
12296 | QT_Q5_K
12297 | QT_Q3_K
12298 | QT_NVFP4
12299 | QT_F8_E4M3
12300 | QT_F8_E4M3_BLK
12301 | QT_Q4_0
12302 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12303 }
12304 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12305 }
12306 }
12307
12308 pub fn matmul_pre(
12313 &self,
12314 w: &crate::model::GpuTensor,
12315 aq: &CudaSlice<i8>,
12316 ad: &CudaSlice<f32>,
12317 x_fallback: &CudaSlice<f32>,
12318 m: usize,
12319 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12320 use crate::model::GpuTensor;
12321 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12327 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12330 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12331 return Ok(y);
12332 }
12333 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12336 return Ok(y);
12337 }
12338 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12340 return Ok(y);
12341 }
12342 }
12343 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12349 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12350 return Ok(y);
12351 }
12352 }
12353 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12354 return Ok(y);
12355 }
12356 if m >= 16
12361 && w.out_features() >= 128
12362 && self.mmq_supports(w)
12363 && !self.verify_exact_on()
12364 && x_raw_ok
12365 {
12366 return self.qmatvec_mmq(w, x_fallback, m);
12367 }
12368 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12371 if let Some(y) =
12372 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12373 {
12374 return Ok(y);
12375 }
12376 }
12377 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12380 return self.qmatvec_gemm(w, aq, ad, m);
12381 }
12382 if !self.uses_q8_1_fast(w) {
12401 if !x_raw_ok {
12402 return Err(format!(
12403 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12404 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12405 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12406 activation (see Engine::rms_norm_decode, which is bit-identical to \
12407 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12408 x_fallback.len(),
12409 m,
12410 w.in_features(),
12411 m * w.in_features()
12412 )
12413 .into());
12414 }
12415 return self.matmul(w, x_fallback, m);
12416 }
12417 let in_f = w.in_features();
12418 let out_f = w.out_features();
12419 let (bytes, qtype, row_bytes, scale, rp) = match w {
12420 GpuTensor::Quant {
12421 bytes,
12422 qtype,
12423 row_bytes,
12424 scale,
12425 rp,
12426 ..
12427 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12428 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12429 };
12430 let (mbytes, mrp) = match w {
12433 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12434 _ => (bytes, rp),
12435 };
12436 if m == 1 && self.mmvq_supports(qtype) {
12440 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12441 }
12442 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12455 && std::env::var("MEMRA_NO_BATCHED").is_err()
12456 && (m <= 4 || Self::b8_enabled())
12457 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12461 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12462 {
12463 let mcols = Self::batched_mcols(m);
12464 return self.qmatvec_mmvq_batched(
12465 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12466 );
12467 }
12468 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12474 let (b2, r2) = if qtype == QT_Q4_0 {
12475 (mbytes, mrp)
12476 } else {
12477 (bytes, rp)
12478 };
12479 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12480 }
12481 let name = match qtype {
12482 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12483 QT_Q4_K => "qmatvec_q4_K_dp4a",
12484 QT_Q6_K => "qmatvec_q6_K_dp4a",
12485 QT_Q5_K => "qmatvec_q5_K_dp4a",
12486 QT_Q3_K => "qmatvec_q3_K_dp4a",
12487 QT_NVFP4 => {
12488 if rp {
12489 "qmatvec_nvfp4_dp4a_rp"
12490 } else {
12491 "qmatvec_nvfp4_dp4a"
12492 }
12493 }
12494 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12495 _ => unreachable!(),
12496 };
12497 let f = self.func(name);
12498 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12500 grid_dim: (out_f as u32, m as u32, 1),
12501 block_dim: (128, 1, 1),
12502 shared_mem_bytes: 0,
12503 };
12504 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12505 let __s_b = self.gpu.stream();
12506 let mut b = __s_b.launch_builder(&f);
12507 b.arg(bytes)
12508 .arg(aq)
12509 .arg(ad)
12510 .arg(&mut y)
12511 .arg(&inf)
12512 .arg(&outf)
12513 .arg(&mi)
12514 .arg(&rb);
12515 unsafe {
12516 b.launch(cfg)?;
12517 }
12518 if scale != 1.0 {
12519 self.scale_inplace(&mut y, scale, m * out_f)?;
12520 }
12521 Ok(y)
12522 }
12523
12524 pub fn matmul_decode_exact(
12532 &self,
12533 w: &crate::model::GpuTensor,
12534 x: &CudaSlice<f32>,
12535 m: usize,
12536 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12537 use crate::model::GpuTensor;
12538 if let GpuTensor::Float { data, .. } = w {
12546 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12547 }
12548 if let GpuTensor::FloatBf16 { data, .. } = w {
12551 let (in_f, out_f) = (w.in_features(), w.out_features());
12552 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12555 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12556 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12557 return Ok(y);
12558 }
12559 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12560 }
12561 if !self.uses_q8_1_fast(w) {
12562 return self.matmul(w, x, m);
12563 }
12564 let in_f = w.in_features();
12565 let out_f = w.out_features();
12566 let (bytes, qtype, row_bytes, scale, rp) = match w {
12567 GpuTensor::Quant {
12568 bytes,
12569 qtype,
12570 row_bytes,
12571 scale,
12572 rp,
12573 ..
12574 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12575 _ => return self.matmul(w, x, m),
12576 };
12577 let (bytes, rp) = match w {
12580 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12581 _ => (bytes, rp),
12582 };
12583 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12584 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12588 return Ok(y);
12589 }
12590 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12599 && std::env::var("MEMRA_NO_BATCHED").is_err()
12600 && (m <= 4 || Self::b8_enabled())
12601 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12604 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12605 {
12606 let mcols = Self::batched_mcols(m);
12607 return self.qmatvec_mmvq_batched(
12608 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12609 );
12610 }
12611 if self.mmvq_supports(qtype) {
12612 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12615 }
12616 self.matmul_pre(w, &aq, &ad, x, m)
12619 }
12620
12621 pub fn matmul_decode_exact_pre(
12631 &self,
12632 w: &crate::model::GpuTensor,
12633 aq: &CudaSlice<i8>,
12634 ad: &CudaSlice<f32>,
12635 m: usize,
12636 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12637 use crate::model::GpuTensor;
12638 debug_assert!(
12639 self.uses_q8_1_fast(w),
12640 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12641 );
12642 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12644 return Ok(y);
12645 }
12646 let in_f = w.in_features();
12647 let out_f = w.out_features();
12648 let (bytes, qtype, row_bytes, scale, rp) = match w {
12649 GpuTensor::Quant {
12650 bytes,
12651 qtype,
12652 row_bytes,
12653 scale,
12654 rp,
12655 ..
12656 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12657 _ => {
12658 return Err(
12659 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12660 );
12661 }
12662 };
12663 let (bytes, rp) = match w {
12665 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12666 _ => (bytes, rp),
12667 };
12668 if (2..=16).contains(&m)
12670 && self.batched_supports(qtype)
12671 && self.mmvq_supports(qtype)
12672 && std::env::var("MEMRA_NO_BATCHED").is_err()
12673 && (m <= 4 || Self::b8_enabled())
12674 && (m <= 8
12675 || qtype == QT_Q4_0
12676 || qtype == QT_Q6_K
12677 || qtype == QT_F8_E4M3
12678 || qtype == QT_NVFP4
12679 || qtype == QT_Q4_K
12680 || qtype == QT_Q5_K
12681 || qtype == QT_Q8_0)
12682 {
12683 let mcols = Self::batched_mcols(m);
12684 return self.qmatvec_mmvq_batched(
12685 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12686 );
12687 }
12688 if self.mmvq_supports(qtype) {
12689 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12690 }
12691 let x0 = self.zeros(0)?;
12694 self.matmul_pre(w, aq, ad, &x0, m)
12695 }
12696
12697 pub fn matmul_decode_exact_dual_pre(
12706 &self,
12707 w0: &crate::model::GpuTensor,
12708 w1: &crate::model::GpuTensor,
12709 aq: &CudaSlice<i8>,
12710 ad: &CudaSlice<f32>,
12711 m: usize,
12712 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12713 {
12714 use crate::model::GpuTensor;
12715 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12716 let on = *ON.get_or_init(|| {
12717 std::env::var("MEMRA_SPEC_DUAL_T")
12718 .map(|v| v != "0")
12719 .unwrap_or(true)
12720 });
12721 if !on
12722 || !(2..=7).contains(&m)
12723 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12724 || !self.uses_q8_1_fast(w0)
12725 || !self.uses_q8_1_fast(w1)
12726 {
12727 return Ok(None);
12728 }
12729 if !self.mmvq_supports(QT_NVFP4) {
12734 return Ok(None);
12735 }
12736 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12737 if w1.in_features() != in_f || w1.out_features() != out_f {
12738 return Ok(None);
12739 }
12740 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12741 (
12742 GpuTensor::Quant {
12743 bytes: b0,
12744 qtype: q0,
12745 row_bytes: rb0,
12746 scale: s0,
12747 rp: rp0,
12748 rp4: None,
12749 ..
12750 },
12751 GpuTensor::Quant {
12752 bytes: b1,
12753 qtype: q1,
12754 row_bytes: rb1,
12755 scale: s1,
12756 rp: rp1,
12757 rp4: None,
12758 ..
12759 },
12760 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12761 (b0, b1, *rb0, *s0, *s1, *rp0)
12762 }
12763 _ => return Ok(None),
12764 };
12765 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12768 {
12769 return Ok(None);
12770 }
12771 let (y0, y1) =
12772 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12773 Ok(Some(((y0, s0), (y1, s1))))
12774 }
12775
12776 pub fn matmul_decode_exact_group4_pre(
12788 &self,
12789 ws: [&crate::model::GpuTensor; 4],
12790 aq: &CudaSlice<i8>,
12791 ad: &CudaSlice<f32>,
12792 m: usize,
12793 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12794 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12795 let on = *ON.get_or_init(|| {
12796 std::env::var("MEMRA_TK_GDN_GROUP")
12797 .map(|v| v != "0")
12798 .unwrap_or(true)
12799 });
12800 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12801 }
12802
12803 pub fn matmul_decode_exact_group3_pre(
12808 &self,
12809 ws: [&crate::model::GpuTensor; 3],
12810 aq: &CudaSlice<i8>,
12811 ad: &CudaSlice<f32>,
12812 m: usize,
12813 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12814 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12815 let on = *ON.get_or_init(|| {
12816 std::env::var("MEMRA_TK_FA_GROUP")
12817 .map(|v| v != "0")
12818 .unwrap_or(true)
12819 });
12820 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12821 }
12822
12823 fn matmul_decode_exact_group_pre(
12827 &self,
12828 ws: &[&crate::model::GpuTensor],
12829 aq: &CudaSlice<i8>,
12830 ad: &CudaSlice<f32>,
12831 m: usize,
12832 on: bool,
12833 tag: &'static str,
12834 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12835 use crate::model::GpuTensor;
12836 if !on
12837 || !(2..=16).contains(&m)
12838 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12839 || (m > 4 && !Self::b8_enabled())
12840 || !self.mmvq_supports(QT_NVFP4)
12841 || !self.batched_supports(QT_NVFP4)
12842 {
12843 return Ok(None);
12844 }
12845 let in_f = ws[0].in_features();
12846 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12847 for w in ws {
12848 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12849 return Ok(None);
12850 }
12851 match w {
12852 GpuTensor::Quant {
12853 bytes,
12854 qtype,
12855 scale,
12856 rp: true,
12857 rp4: None,
12858 ..
12859 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12860 parts.push((bytes, w.out_features(), *scale));
12861 }
12862 _ => return Ok(None),
12863 }
12864 }
12865 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12867 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12868 let mcols = if (5..=7).contains(&m) && b567 {
12869 m
12870 } else {
12871 Self::batched_mcols(m)
12872 };
12873 let kname: &'static str = match mcols {
12874 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12875 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12876 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12877 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12878 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12879 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12880 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12881 _ => return Ok(None),
12882 };
12883 if std::env::var("MEMRA_DEBUG").is_ok() {
12886 use std::sync::Mutex;
12887 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12888 let mut seen = SEEN.lock().unwrap();
12889 if !seen.contains(&tag) {
12890 seen.push(tag);
12891 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12892 }
12893 }
12894 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12897 let three = parts.len() == 3;
12898 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12899 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12900 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12901 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12904 let cfg = LaunchConfig {
12905 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12906 block_dim: (32, ROWS_PER_BLOCK, 1),
12907 shared_mem_bytes: 0,
12908 };
12909 let (inf, mi) = (in_f as i32, m as i32);
12910 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12911 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12912 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12913 let s3 = if three { 1.0f32 } else { parts[3].2 };
12914 let w3 = if three { parts[0].0 } else { parts[3].0 };
12915 let f = self.func(kname);
12916 let __s_b = self.gpu.stream();
12917 let mut b = __s_b.launch_builder(&f);
12918 b.arg(parts[0].0)
12919 .arg(parts[1].0)
12920 .arg(parts[2].0)
12921 .arg(w3)
12922 .arg(aq)
12923 .arg(ad)
12924 .arg(&mut y0)
12925 .arg(&mut y1)
12926 .arg(&mut y2)
12927 .arg(&mut y3)
12928 .arg(&inf)
12929 .arg(&n0)
12930 .arg(&n1)
12931 .arg(&n2)
12932 .arg(&n3)
12933 .arg(&mi)
12934 .arg(&s0)
12935 .arg(&s1)
12936 .arg(&s2)
12937 .arg(&s3);
12938 unsafe {
12939 b.launch(cfg)?;
12940 }
12941 Ok(Some(if three {
12942 vec![y0, y1, y2]
12943 } else {
12944 vec![y0, y1, y2, y3]
12945 }))
12946 }
12947
12948 pub fn matmul_decode_exact_dual(
12964 &self,
12965 w0: &crate::model::GpuTensor,
12966 w1: &crate::model::GpuTensor,
12967 x: &CudaSlice<f32>,
12968 m: usize,
12969 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12970 use crate::model::GpuTensor;
12971 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12972 let on = *ON.get_or_init(|| {
12973 std::env::var("MEMRA_SPEC_DUAL_T")
12974 .map(|v| v != "0")
12975 .unwrap_or(true)
12976 });
12977 if !on
12978 || !(2..=4).contains(&m)
12979 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12980 || !self.uses_q8_1_fast(w0)
12981 || !self.uses_q8_1_fast(w1)
12982 {
12983 return Ok(None);
12984 }
12985 if !self.mmvq_supports(QT_NVFP4) {
12990 return Ok(None);
12991 }
12992 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12993 if w1.in_features() != in_f || w1.out_features() != out_f {
12994 return Ok(None);
12995 }
12996 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12997 (
12998 GpuTensor::Quant {
12999 bytes: b0,
13000 qtype: q0,
13001 row_bytes: rb0,
13002 scale: s0,
13003 rp: rp0,
13004 rp4: None,
13005 ..
13006 },
13007 GpuTensor::Quant {
13008 bytes: b1,
13009 qtype: q1,
13010 row_bytes: rb1,
13011 scale: s1,
13012 rp: rp1,
13013 rp4: None,
13014 ..
13015 },
13016 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13017 (b0, b1, *rb0, *s0, *s1, *rp0)
13018 }
13019 _ => return Ok(None),
13020 };
13021 if std::env::var("MEMRA_DEBUG").is_ok() {
13024 static ONCE: std::sync::Once = std::sync::Once::new();
13025 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13026 }
13027 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13028 let (y0, y1) =
13029 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13030 let mut y0 = y0;
13031 let mut y1 = y1;
13032 if s0 != 1.0 {
13033 self.scale_inplace(&mut y0, s0, m * out_f)?;
13034 }
13035 if s1 != 1.0 {
13036 self.scale_inplace(&mut y1, s1, m * out_f)?;
13037 }
13038 Ok(Some((y0, y1)))
13039 }
13040
13041 #[allow(clippy::too_many_arguments)]
13046 pub fn qmatvec_batched_dual_raw(
13047 &self,
13048 b0: &CudaSlice<u8>,
13049 b1: &CudaSlice<u8>,
13050 aq: &CudaSlice<i8>,
13051 ad: &CudaSlice<f32>,
13052 m: usize,
13053 in_f: usize,
13054 out_f: usize,
13055 row_bytes: usize,
13056 rp: bool,
13057 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13058 const ROWS_PER_BLOCK: u32 = 4;
13059 let mcols = Self::batched_mcols(m);
13060 let tiny_rp1 = rp
13063 && mcols == 4
13064 && out_f <= 128
13065 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13066 let (name, rows_per_block) = if tiny_rp1 {
13067 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13068 } else {
13069 match (mcols, rp, m) {
13070 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13071 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13072 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13073 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13074 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13075 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13076 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13077 _ => {
13078 return Err(
13079 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13080 );
13081 }
13082 }
13083 };
13084 let f = self.func(name);
13085 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13086 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13087 let cfg = LaunchConfig {
13088 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13089 block_dim: (32, ROWS_PER_BLOCK, 1),
13090 shared_mem_bytes: 0,
13091 };
13092 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13093 let __s_b = self.gpu.stream();
13094 let mut b = __s_b.launch_builder(&f);
13095 b.arg(b0)
13096 .arg(b1)
13097 .arg(aq)
13098 .arg(ad)
13099 .arg(&mut y0)
13100 .arg(&mut y1)
13101 .arg(&inf)
13102 .arg(&outf)
13103 .arg(&mi)
13104 .arg(&rb);
13105 unsafe {
13106 b.launch(cfg)?;
13107 }
13108 Ok((y0, y1))
13109 }
13110
13111 pub fn matmul_pre_dual_noscale(
13123 &self,
13124 w0: &crate::model::GpuTensor,
13125 w1: &crate::model::GpuTensor,
13126 aq: &CudaSlice<i8>,
13127 ad: &CudaSlice<f32>,
13128 m: usize,
13129 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13130 {
13131 use crate::model::GpuTensor;
13132 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13133 return Ok(None);
13134 }
13135 if !self.mmvq_supports(QT_NVFP4) {
13145 return Ok(None);
13146 }
13147 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13148 if w1.in_features() != in_f || w1.out_features() != out_f {
13149 return Ok(None);
13150 }
13151 let no_mirror =
13164 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13165 if self.q8_ffn_fuse2_on()
13166 && no_mirror(w0)
13167 && no_mirror(w1)
13168 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13169 {
13170 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13171 return Ok(Some(((y0, 1.0), (y1, 1.0))));
13172 }
13173 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13183 let (y0, y1) =
13184 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13185 return Ok(Some(((y0, p0.3), (y1, p1.3))));
13186 }
13187 let (b0, q0, rb0, s0, rp0) = match w0 {
13188 GpuTensor::Quant {
13189 bytes,
13190 qtype,
13191 row_bytes,
13192 scale,
13193 rp,
13194 ..
13195 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13196 _ => return Ok(None),
13197 };
13198 let (b1, q1, rb1, s1, rp1) = match w1 {
13199 GpuTensor::Quant {
13200 bytes,
13201 qtype,
13202 row_bytes,
13203 scale,
13204 rp,
13205 ..
13206 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13207 _ => return Ok(None),
13208 };
13209 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13210 return Ok(None);
13211 }
13212 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13214 let rows_per_block = ROWS_PER_BLOCK * RPW;
13215 let f = self.func(if rp0 {
13216 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13217 } else {
13218 "qmatvec_nvfp4_mmvq_dual_mr2"
13219 });
13220 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13221 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13222 let cfg = LaunchConfig {
13223 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13224 block_dim: (32, ROWS_PER_BLOCK, 1),
13225 shared_mem_bytes: 0,
13226 };
13227 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13228 let one = 1.0f32;
13231 let __s_b = self.gpu.stream();
13232 let mut b = __s_b.launch_builder(&f);
13233 b.arg(b0)
13234 .arg(b1)
13235 .arg(aq)
13236 .arg(ad)
13237 .arg(&mut y0)
13238 .arg(&mut y1)
13239 .arg(&inf)
13240 .arg(&outf)
13241 .arg(&mi)
13242 .arg(&rb)
13243 .arg(&one)
13244 .arg(&one);
13245 unsafe {
13246 b.launch(cfg)?;
13247 }
13248 Ok(Some(((y0, s0), (y1, s1))))
13249 }
13250
13251 #[allow(clippy::too_many_arguments)]
13259 pub fn matmul_nvfp4_fused3(
13260 &self,
13261 w0: &crate::model::GpuTensor,
13262 w1: &crate::model::GpuTensor,
13263 w2: &crate::model::GpuTensor,
13264 aq: &CudaSlice<i8>,
13265 ad: &CudaSlice<f32>,
13266 m: usize,
13267 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13268 {
13269 use crate::model::GpuTensor;
13270 if !self.mmvq_supports(QT_NVFP4)
13277 || !self.uses_q8_1_fast(w0)
13278 || !self.uses_q8_1_fast(w1)
13279 || !self.uses_q8_1_fast(w2)
13280 {
13281 return Ok(None);
13282 }
13283 if (9..=16).contains(&m) {
13286 return Ok(
13287 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13288 Some(mut ys) => {
13289 let y2 = ys.pop().unwrap();
13290 let y1 = ys.pop().unwrap();
13291 let y0 = ys.pop().unwrap();
13292 Some((y0, y1, y2))
13293 }
13294 None => None,
13295 },
13296 );
13297 }
13298 if !(1..=8).contains(&m) {
13299 return Ok(None);
13300 }
13301 if m > 1 {
13302 let in_f = w0.in_features();
13303 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13304 || !self.batched_supports(QT_NVFP4)
13305 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13306 || (m > 4 && !Self::b8_enabled())
13307 || in_f % 512 != 0
13308 || in_f / 64 > 272
13309 {
13310 return Ok(None);
13311 }
13312 }
13313 let unpack = |w: &crate::model::GpuTensor| match w {
13314 GpuTensor::Quant {
13315 bytes,
13316 qtype,
13317 scale,
13318 rp,
13319 ..
13320 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13321 _ => None,
13322 };
13323 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13324 return Ok(None);
13325 };
13326 let in_f = w0.in_features();
13327 if w1.in_features() != in_f || w2.in_features() != in_f {
13328 return Ok(None);
13329 }
13330 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13331 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13333 let rows_pb = ROWS_PER_BLOCK * RPW;
13334 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13335 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13336 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13337 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13338 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13339 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13342 if m > 1 {
13343 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13345 return Ok(None);
13346 }
13347 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13348 let cfg = LaunchConfig {
13349 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13350 block_dim: (32, ROWS_PER_BLOCK, 1),
13351 shared_mem_bytes: 0,
13352 };
13353 let __s_b = self.gpu.stream();
13354 let mut b = __s_b.launch_builder(&f);
13355 b.arg(b0)
13356 .arg(b1)
13357 .arg(b2)
13358 .arg(aq)
13359 .arg(ad)
13360 .arg(&mut y0)
13361 .arg(&mut y1)
13362 .arg(&mut y2)
13363 .arg(&inf)
13364 .arg(&oi0)
13365 .arg(&oi1)
13366 .arg(&oi2)
13367 .arg(&mi);
13368 unsafe {
13369 b.launch(cfg)?;
13370 }
13371 return Ok(Some((y0, y1, y2)));
13372 }
13373 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13374 let cfg = LaunchConfig {
13375 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13376 block_dim: (32, ROWS_PER_BLOCK, 1),
13377 shared_mem_bytes: 0,
13378 };
13379 let __s_b = self.gpu.stream();
13380 let mut b = __s_b.launch_builder(&f);
13381 b.arg(b0)
13382 .arg(b1)
13383 .arg(b2)
13384 .arg(aq)
13385 .arg(ad)
13386 .arg(&mut y0)
13387 .arg(&mut y1)
13388 .arg(&mut y2)
13389 .arg(&inf)
13390 .arg(&oi0)
13391 .arg(&oi1)
13392 .arg(&oi2)
13393 .arg(&mi)
13394 .arg(&p0.1)
13395 .arg(&p1.1)
13396 .arg(&p2.1);
13397 unsafe {
13398 b.launch(cfg)?;
13399 }
13400 Ok(Some((y0, y1, y2)))
13401 }
13402
13403 pub fn matmul_nvfp4_fused2(
13412 &self,
13413 w0: &crate::model::GpuTensor,
13414 w1: &crate::model::GpuTensor,
13415 aq: &CudaSlice<i8>,
13416 ad: &CudaSlice<f32>,
13417 m: usize,
13418 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13419 use crate::model::GpuTensor;
13420 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13421 let off =
13422 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13423 if off
13426 || m != 1
13427 || !self.mmvq_supports(QT_NVFP4)
13428 || !self.uses_q8_1_fast(w0)
13429 || !self.uses_q8_1_fast(w1)
13430 {
13431 return Ok(None);
13432 }
13433 let unpack = |w: &crate::model::GpuTensor| match w {
13434 GpuTensor::Quant {
13435 bytes,
13436 qtype,
13437 scale,
13438 rp,
13439 ..
13440 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13441 _ => None,
13442 };
13443 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13444 return Ok(None);
13445 };
13446 let in_f = w0.in_features();
13447 if w1.in_features() != in_f {
13448 return Ok(None);
13449 }
13450 let (o0, o1) = (w0.out_features(), w1.out_features());
13451 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13453 let rows_pb = ROWS_PER_BLOCK * RPW;
13454 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13455 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13456 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13457 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13458 let cfg = LaunchConfig {
13459 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13460 block_dim: (32, ROWS_PER_BLOCK, 1),
13461 shared_mem_bytes: 0,
13462 };
13463 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13464 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13467 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13470 {
13471 use cudarc::driver::{DevicePtr, DevicePtrMut};
13472 let s = &self.gpu.stream();
13473 let (pw0, _g0) = b0.device_ptr(s);
13474 let (pw1, _g1) = b1.device_ptr(s);
13475 let (paq, _g2) = aq.device_ptr(s);
13476 let (pad, _g3) = ad.device_ptr(s);
13477 let (py0, _g4) = y0.device_ptr_mut(s);
13478 let (py1, _g5) = y1.device_ptr_mut(s);
13479 let (s0, s1) = (p0.1, p1.1);
13480 let mut ps = [
13481 &pw0 as *const _ as *mut std::ffi::c_void,
13482 &pw1 as *const _ as *mut _,
13483 &paq as *const _ as *mut _,
13484 &pad as *const _ as *mut _,
13485 &py0 as *const _ as *mut _,
13486 &py1 as *const _ as *mut _,
13487 &inf as *const _ as *mut _,
13488 &oi0 as *const _ as *mut _,
13489 &oi1 as *const _ as *mut _,
13490 &mi as *const _ as *mut _,
13491 &s0 as *const _ as *mut _,
13492 &s1 as *const _ as *mut _,
13493 ];
13494 unsafe {
13495 self.launch_pdl(
13496 "qmatvec_nvfp4_mmvq_fused2_rp",
13497 cfg.grid_dim,
13498 cfg.block_dim,
13499 &mut ps,
13500 )?;
13501 }
13502 }
13503 return Ok(Some((y0, y1)));
13504 }
13505 let __s_b = self.gpu.stream();
13506 let mut b = __s_b.launch_builder(&f);
13507 b.arg(b0)
13508 .arg(b1)
13509 .arg(aq)
13510 .arg(ad)
13511 .arg(&mut y0)
13512 .arg(&mut y1)
13513 .arg(&inf)
13514 .arg(&oi0)
13515 .arg(&oi1)
13516 .arg(&mi)
13517 .arg(&p0.1)
13518 .arg(&p1.1);
13519 unsafe {
13520 b.launch(cfg)?;
13521 }
13522 Ok(Some((y0, y1)))
13523 }
13524
13525 pub fn matmul_nvfp4_fused2_into(
13530 &self,
13531 w0: &crate::model::GpuTensor,
13532 w1: &crate::model::GpuTensor,
13533 aq: &CudaSlice<i8>,
13534 ad: &CudaSlice<f32>,
13535 y0: &mut CudaSlice<f32>,
13536 y1: &mut CudaSlice<f32>,
13537 ) -> Result<bool, Box<dyn std::error::Error>> {
13538 use crate::model::GpuTensor;
13539 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13540 let off =
13541 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13542 if off
13543 || !self.mmvq_supports(QT_NVFP4)
13544 || !self.uses_q8_1_fast(w0)
13545 || !self.uses_q8_1_fast(w1)
13546 {
13547 return Ok(false);
13548 }
13549 let unpack = |w: &crate::model::GpuTensor| match w {
13550 GpuTensor::Quant {
13551 bytes,
13552 qtype,
13553 scale,
13554 rp,
13555 ..
13556 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13557 _ => None,
13558 };
13559 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13560 return Ok(false);
13561 };
13562 let in_f = w0.in_features();
13563 if w1.in_features() != in_f {
13564 return Ok(false);
13565 }
13566 let (o0, o1) = (w0.out_features(), w1.out_features());
13567 if y0.len() < o0 || y1.len() < o1 {
13568 return Ok(false);
13569 }
13570 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13572 let rows_pb = ROWS_PER_BLOCK * RPW;
13573 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13574 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13575 let cfg = LaunchConfig {
13576 grid_dim: (nb(o0) + nb(o1), 1, 1),
13577 block_dim: (32, ROWS_PER_BLOCK, 1),
13578 shared_mem_bytes: 0,
13579 };
13580 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13581 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
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(aq)
13589 .arg(ad)
13590 .arg(&mut *y0)
13591 .arg(&mut *y1)
13592 .arg(&inf)
13593 .arg(&oi0)
13594 .arg(&oi1)
13595 .arg(&mi)
13596 .arg(&p0.1)
13597 .arg(&p1.1);
13598 unsafe {
13599 b.launch(cfg)?;
13600 }
13601 Ok(true)
13602 }
13603
13604 #[allow(clippy::type_complexity)]
13609 pub fn matmul_nvfp4_fused4(
13610 &self,
13611 w0: &crate::model::GpuTensor,
13612 w1: &crate::model::GpuTensor,
13613 w2: &crate::model::GpuTensor,
13614 w3: &crate::model::GpuTensor,
13615 aq: &CudaSlice<i8>,
13616 ad: &CudaSlice<f32>,
13617 m: usize,
13618 ) -> Result<
13619 Option<(
13620 CudaSlice<f32>,
13621 CudaSlice<f32>,
13622 CudaSlice<f32>,
13623 CudaSlice<f32>,
13624 )>,
13625 Box<dyn std::error::Error>,
13626 > {
13627 use crate::model::GpuTensor;
13628 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13635 || !self.mmvq_supports(QT_NVFP4)
13636 || !self.uses_q8_1_fast(w0)
13637 || !self.uses_q8_1_fast(w1)
13638 || !self.uses_q8_1_fast(w2)
13639 || !self.uses_q8_1_fast(w3)
13640 {
13641 return Ok(None);
13642 }
13643 if (9..=16).contains(&m) {
13648 return Ok(
13649 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13650 Some(mut ys) => {
13651 let y3 = ys.pop().unwrap();
13652 let y2 = ys.pop().unwrap();
13653 let y1 = ys.pop().unwrap();
13654 let y0 = ys.pop().unwrap();
13655 Some((y0, y1, y2, y3))
13656 }
13657 None => None,
13658 },
13659 );
13660 }
13661 if !(1..=8).contains(&m) {
13662 return Ok(None);
13663 }
13664 if m > 1 {
13665 let in_f = w0.in_features();
13668 if !self.batched_supports(QT_NVFP4)
13669 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13670 || (m > 4 && !Self::b8_enabled())
13671 || in_f % 512 != 0
13672 || in_f / 64 > 272
13673 {
13674 return Ok(None);
13675 }
13676 }
13677 let unpack = |w: &crate::model::GpuTensor| match w {
13678 GpuTensor::Quant {
13679 bytes,
13680 qtype,
13681 scale,
13682 rp,
13683 ..
13684 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13685 _ => None,
13686 };
13687 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13688 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13689 else {
13690 return Ok(None);
13691 };
13692 let in_f = w0.in_features();
13693 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13694 return Ok(None);
13695 }
13696 let (o0, o1, o2, o3) = (
13697 w0.out_features(),
13698 w1.out_features(),
13699 w2.out_features(),
13700 w3.out_features(),
13701 );
13702 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13704 let rows_pb = ROWS_PER_BLOCK * RPW;
13705 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13706 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13707 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13708 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13709 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13710 let (inf, oi0, oi1, oi2, oi3, mi) = (
13711 in_f as i32,
13712 o0 as i32,
13713 o1 as i32,
13714 o2 as i32,
13715 o3 as i32,
13716 m as i32,
13717 );
13718 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13721 if m > 1 {
13722 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13725 return Ok(None);
13726 }
13727 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13728 let cfg = LaunchConfig {
13729 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13730 block_dim: (32, ROWS_PER_BLOCK, 1),
13731 shared_mem_bytes: 0,
13732 };
13733 let __s_b = self.gpu.stream();
13734 let mut b = __s_b.launch_builder(&f);
13735 b.arg(b0)
13736 .arg(b1)
13737 .arg(b2)
13738 .arg(b3)
13739 .arg(aq)
13740 .arg(ad)
13741 .arg(&mut y0)
13742 .arg(&mut y1)
13743 .arg(&mut y2)
13744 .arg(&mut y3)
13745 .arg(&inf)
13746 .arg(&oi0)
13747 .arg(&oi1)
13748 .arg(&oi2)
13749 .arg(&oi3)
13750 .arg(&mi);
13751 unsafe {
13752 b.launch(cfg)?;
13753 }
13754 return Ok(Some((y0, y1, y2, y3)));
13755 }
13756 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13757 let cfg = LaunchConfig {
13758 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13759 block_dim: (32, ROWS_PER_BLOCK, 1),
13760 shared_mem_bytes: 0,
13761 };
13762 let __s_b = self.gpu.stream();
13763 let mut b = __s_b.launch_builder(&f);
13764 b.arg(b0)
13765 .arg(b1)
13766 .arg(b2)
13767 .arg(b3)
13768 .arg(aq)
13769 .arg(ad)
13770 .arg(&mut y0)
13771 .arg(&mut y1)
13772 .arg(&mut y2)
13773 .arg(&mut y3)
13774 .arg(&inf)
13775 .arg(&oi0)
13776 .arg(&oi1)
13777 .arg(&oi2)
13778 .arg(&oi3)
13779 .arg(&mi)
13780 .arg(&p0.1)
13781 .arg(&p1.1)
13782 .arg(&p2.1)
13783 .arg(&p3.1);
13784 unsafe {
13785 b.launch(cfg)?;
13786 }
13787 Ok(Some((y0, y1, y2, y3)))
13788 }
13789
13790 pub fn matmul_q8_fused2(
13798 &self,
13799 w0: &crate::model::GpuTensor,
13800 w1: &crate::model::GpuTensor,
13801 aq: &CudaSlice<i8>,
13802 ad: &CudaSlice<f32>,
13803 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13804 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13810 return Ok(Some(self.e4m3_fused2_core(
13811 p0.0,
13812 p1.0,
13813 aq,
13814 ad,
13815 w0.in_features(),
13816 p0.1,
13817 p1.1,
13818 p0.2,
13819 p0.3,
13820 p1.3,
13821 )?));
13822 }
13823 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13824 return Ok(None);
13825 };
13826 Ok(Some(self.q8_fused2_core(
13827 p0.0,
13828 p1.0,
13829 aq,
13830 ad,
13831 w0.in_features(),
13832 p0.1,
13833 p1.1,
13834 p0.2,
13835 )?))
13836 }
13837
13838 #[allow(clippy::too_many_arguments)]
13839 fn q8_fused2_core(
13840 &self,
13841 b0: &CudaSlice<u8>,
13842 b1: &CudaSlice<u8>,
13843 aq: &CudaSlice<i8>,
13844 ad: &CudaSlice<f32>,
13845 in_f: usize,
13846 out0: usize,
13847 out1: usize,
13848 row_bytes: usize,
13849 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13850 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13852 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13853 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13854 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13855 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13856 let cfg = LaunchConfig {
13857 grid_dim: (nb0 + nb1, 1, 1),
13858 block_dim: (32, ROWS_PER_BLOCK, 1),
13859 shared_mem_bytes: 0,
13860 };
13861 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13862 let __s_b = self.gpu.stream();
13863 let mut b = __s_b.launch_builder(&f);
13864 b.arg(b0)
13865 .arg(b1)
13866 .arg(aq)
13867 .arg(ad)
13868 .arg(&mut y0)
13869 .arg(&mut y1)
13870 .arg(&inf)
13871 .arg(&o0)
13872 .arg(&o1)
13873 .arg(&rbl);
13874 unsafe {
13875 b.launch(cfg)?;
13876 }
13877 Ok((y0, y1))
13878 }
13879
13880 pub fn matmul_q8_fused2_x(
13886 &self,
13887 w0: &crate::model::GpuTensor,
13888 w1: &crate::model::GpuTensor,
13889 x: &CudaSlice<f32>,
13890 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13891 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13892 return Ok(None);
13893 }
13894 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13895 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13896 return Ok(Some(self.e4m3_fused2_core(
13897 p0.0,
13898 p1.0,
13899 &aq,
13900 &ad,
13901 w0.in_features(),
13902 p0.1,
13903 p1.1,
13904 p0.2,
13905 p0.3,
13906 p1.3,
13907 )?));
13908 }
13909 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13910 return Ok(None);
13911 };
13912 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13913 Ok(Some(self.q8_fused2_core(
13914 p0.0,
13915 p1.0,
13916 &aq,
13917 &ad,
13918 w0.in_features(),
13919 p0.1,
13920 p1.1,
13921 p0.2,
13922 )?))
13923 }
13924
13925 #[allow(clippy::too_many_arguments)]
13928 pub fn qmatvec_q8_fused2_raw(
13929 &self,
13930 b0: &CudaSlice<u8>,
13931 b1: &CudaSlice<u8>,
13932 x: &CudaSlice<f32>,
13933 in_f: usize,
13934 out0: usize,
13935 out1: usize,
13936 row_bytes: usize,
13937 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13938 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13939 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13940 }
13941
13942 pub fn matmul_q4_fused3(
13948 &self,
13949 w0: &crate::model::GpuTensor,
13950 w1: &crate::model::GpuTensor,
13951 w2: &crate::model::GpuTensor,
13952 aq: &CudaSlice<i8>,
13953 ad: &CudaSlice<f32>,
13954 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13955 {
13956 use crate::model::GpuTensor;
13957 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13958 match w {
13959 GpuTensor::Quant {
13960 qtype, row_bytes, ..
13961 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13962 _ => None,
13963 }
13964 };
13965 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13966 return Ok(None);
13967 };
13968 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13969 return Ok(None);
13970 }
13971 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13975 match w {
13976 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13977 Some(m) => (m, true),
13978 None => (bytes, *rp),
13979 },
13980 _ => unreachable!(),
13981 }
13982 }
13983 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13984 if rp0 != rp1 || rp1 != rp2 {
13985 return Ok(None);
13986 }
13987 let rp = rp0;
13988 let rpb: u32 = 4;
13989 let mr1 = rp && Self::q40_mr1_on();
13993 let nb = |o: usize| {
13994 if mr1 {
13995 (o as u32).div_ceil(rpb)
13996 } else {
13997 (o as u32).div_ceil(2).div_ceil(rpb)
13998 }
13999 };
14000 let grid = nb(o0) + nb(o1) + nb(o2);
14001 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14002 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14003 let mut y2 = self.alloc_uninit::<f32>(o2)?;
14004 let f = self.func(if mr1 {
14005 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14006 } else if rp {
14007 "qmatvec_q4_0_mmvq_fused3_rp"
14008 } else {
14009 "qmatvec_q4_0_mmvq_fused3"
14010 });
14011 let cfg = LaunchConfig {
14012 grid_dim: (grid, 1, 1),
14013 block_dim: (32, rpb, 1),
14014 shared_mem_bytes: 0,
14015 };
14016 let inf = w0.in_features() as i32;
14017 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14018 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14019 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14022 {
14023 use cudarc::driver::{DevicePtr, DevicePtrMut};
14024 let s = &self.gpu.stream();
14025 let (p0, _g0) = b0.device_ptr(s);
14026 let (p1, _g1) = b1.device_ptr(s);
14027 let (p2, _g2) = b2.device_ptr(s);
14028 let (paq, _g3) = aq.device_ptr(s);
14029 let (pad, _g4) = ad.device_ptr(s);
14030 let (py0, _g5) = y0.device_ptr_mut(s);
14031 let (py1, _g6) = y1.device_ptr_mut(s);
14032 let (py2, _g7) = y2.device_ptr_mut(s);
14033 let mut ps = [
14034 &p0 as *const _ as *mut std::ffi::c_void,
14035 &p1 as *const _ as *mut _,
14036 &p2 as *const _ as *mut _,
14037 &paq as *const _ as *mut _,
14038 &pad as *const _ as *mut _,
14039 &py0 as *const _ as *mut _,
14040 &py1 as *const _ as *mut _,
14041 &py2 as *const _ as *mut _,
14042 &inf as *const _ as *mut _,
14043 &oo0 as *const _ as *mut _,
14044 &oo1 as *const _ as *mut _,
14045 &oo2 as *const _ as *mut _,
14046 &r0 as *const _ as *mut _,
14047 &r1 as *const _ as *mut _,
14048 &r2 as *const _ as *mut _,
14049 ];
14050 unsafe {
14051 self.launch_pdl(
14052 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14053 (grid, 1, 1),
14054 (32, rpb, 1),
14055 &mut ps,
14056 )?;
14057 }
14058 }
14059 return Ok(Some((y0, y1, y2)));
14060 }
14061 let __s_b = self.gpu.stream();
14062 let mut b = __s_b.launch_builder(&f);
14063 b.arg(b0)
14064 .arg(b1)
14065 .arg(b2)
14066 .arg(aq)
14067 .arg(ad)
14068 .arg(&mut y0)
14069 .arg(&mut y1)
14070 .arg(&mut y2)
14071 .arg(&inf)
14072 .arg(&oo0)
14073 .arg(&oo1)
14074 .arg(&oo2)
14075 .arg(&r0)
14076 .arg(&r1)
14077 .arg(&r2);
14078 unsafe {
14079 b.launch(cfg)?;
14080 }
14081 Ok(Some((y0, y1, y2)))
14082 }
14083
14084 #[allow(clippy::too_many_arguments)]
14087 pub fn matmul_q4_fused3_into(
14088 &self,
14089 w0: &crate::model::GpuTensor,
14090 w1: &crate::model::GpuTensor,
14091 w2: &crate::model::GpuTensor,
14092 aq: &CudaSlice<i8>,
14093 ad: &CudaSlice<f32>,
14094 y0: &mut CudaSlice<f32>,
14095 y1: &mut CudaSlice<f32>,
14096 y2: &mut CudaSlice<f32>,
14097 ) -> Result<bool, Box<dyn std::error::Error>> {
14098 use crate::model::GpuTensor;
14099 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14100 match w {
14101 GpuTensor::Quant {
14102 qtype, row_bytes, ..
14103 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14104 _ => None,
14105 }
14106 };
14107 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14108 return Ok(false);
14109 };
14110 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14111 return Ok(false);
14112 }
14113 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14114 match w {
14115 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14116 Some(m) => (m, true),
14117 None => (bytes, *rp),
14118 },
14119 _ => unreachable!(),
14120 }
14121 }
14122 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14123 if rp0 != rp1 || rp1 != rp2 {
14124 return Ok(false);
14125 }
14126 let rp = rp0;
14127 let rpb: u32 = 4;
14128 let mr1 = rp && Self::q40_mr1_on();
14129 let nb = |o: usize| {
14130 if mr1 {
14131 (o as u32).div_ceil(rpb)
14132 } else {
14133 (o as u32).div_ceil(2).div_ceil(rpb)
14134 }
14135 };
14136 let grid = nb(o0) + nb(o1) + nb(o2);
14137 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14138 let f = self.func(if mr1 {
14139 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14140 } else if rp {
14141 "qmatvec_q4_0_mmvq_fused3_rp"
14142 } else {
14143 "qmatvec_q4_0_mmvq_fused3"
14144 });
14145 let cfg = LaunchConfig {
14146 grid_dim: (grid, 1, 1),
14147 block_dim: (32, rpb, 1),
14148 shared_mem_bytes: 0,
14149 };
14150 let inf = w0.in_features() as i32;
14151 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14152 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14153 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14155 use cudarc::driver::{DevicePtr, DevicePtrMut};
14156 let s = &self.gpu.stream();
14157 let (p0, _g0) = b0.device_ptr(s);
14158 let (p1, _g1) = b1.device_ptr(s);
14159 let (p2, _g2) = b2.device_ptr(s);
14160 let (paq, _g3) = aq.device_ptr(s);
14161 let (pad, _g4) = ad.device_ptr(s);
14162 let (py0, _g5) = y0.device_ptr_mut(s);
14163 let (py1, _g6) = y1.device_ptr_mut(s);
14164 let (py2, _g7) = y2.device_ptr_mut(s);
14165 let mut ps = [
14166 &p0 as *const _ as *mut std::ffi::c_void,
14167 &p1 as *const _ as *mut _,
14168 &p2 as *const _ as *mut _,
14169 &paq as *const _ as *mut _,
14170 &pad as *const _ as *mut _,
14171 &py0 as *const _ as *mut _,
14172 &py1 as *const _ as *mut _,
14173 &py2 as *const _ as *mut _,
14174 &inf as *const _ as *mut _,
14175 &oo0 as *const _ as *mut _,
14176 &oo1 as *const _ as *mut _,
14177 &oo2 as *const _ as *mut _,
14178 &r0 as *const _ as *mut _,
14179 &r1 as *const _ as *mut _,
14180 &r2 as *const _ as *mut _,
14181 ];
14182 unsafe {
14183 self.launch_pdl(
14184 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14185 (grid, 1, 1),
14186 (32, rpb, 1),
14187 &mut ps,
14188 )?;
14189 }
14190 return Ok(true);
14191 }
14192 let __s_b = self.gpu.stream();
14193 let mut b = __s_b.launch_builder(&f);
14194 b.arg(b0)
14195 .arg(b1)
14196 .arg(b2)
14197 .arg(aq)
14198 .arg(ad)
14199 .arg(&mut *y0)
14200 .arg(&mut *y1)
14201 .arg(&mut *y2)
14202 .arg(&inf)
14203 .arg(&oo0)
14204 .arg(&oo1)
14205 .arg(&oo2)
14206 .arg(&r0)
14207 .arg(&r1)
14208 .arg(&r2);
14209 unsafe {
14210 b.launch(cfg)?;
14211 }
14212 Ok(true)
14213 }
14214
14215 pub fn matmul_q4_fused2(
14217 &self,
14218 w0: &crate::model::GpuTensor,
14219 w1: &crate::model::GpuTensor,
14220 aq: &CudaSlice<i8>,
14221 ad: &CudaSlice<f32>,
14222 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14223 use crate::model::GpuTensor;
14224 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14225 match w {
14226 GpuTensor::Quant {
14227 qtype, row_bytes, ..
14228 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14229 _ => None,
14230 }
14231 };
14232 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14233 return Ok(None);
14234 };
14235 if w0.in_features() != w1.in_features() {
14236 return Ok(None);
14237 }
14238 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14240 match w {
14241 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14242 Some(m) => (m, true),
14243 None => (bytes, *rp),
14244 },
14245 _ => unreachable!(),
14246 }
14247 }
14248 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14249 if rp0 != rp1 {
14250 return Ok(None);
14251 }
14252 let rp = rp0;
14253 let rpb: u32 = 4;
14254 let mr1 = rp && Self::q40_mr1_on();
14256 let nb = |o: usize| {
14257 if mr1 {
14258 (o as u32).div_ceil(rpb)
14259 } else {
14260 (o as u32).div_ceil(2).div_ceil(rpb)
14261 }
14262 };
14263 let grid = nb(o0) + nb(o1);
14264 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14265 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14266 let f = self.func(if mr1 {
14267 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14268 } else if rp {
14269 "qmatvec_q4_0_mmvq_fused2_rp"
14270 } else {
14271 "qmatvec_q4_0_mmvq_fused2"
14272 });
14273 let cfg = LaunchConfig {
14274 grid_dim: (grid, 1, 1),
14275 block_dim: (32, rpb, 1),
14276 shared_mem_bytes: 0,
14277 };
14278 let inf = w0.in_features() as i32;
14279 let (oo0, oo1) = (o0 as i32, o1 as i32);
14280 let (r0, r1) = (rb0 as i64, rb1 as i64);
14281 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14283 {
14284 use cudarc::driver::{DevicePtr, DevicePtrMut};
14285 let s = &self.gpu.stream();
14286 let (p0, _g0) = b0.device_ptr(s);
14287 let (p1, _g1) = b1.device_ptr(s);
14288 let (paq, _g2) = aq.device_ptr(s);
14289 let (pad, _g3) = ad.device_ptr(s);
14290 let (py0, _g4) = y0.device_ptr_mut(s);
14291 let (py1, _g5) = y1.device_ptr_mut(s);
14292 let mut ps = [
14293 &p0 as *const _ as *mut std::ffi::c_void,
14294 &p1 as *const _ as *mut _,
14295 &paq as *const _ as *mut _,
14296 &pad as *const _ as *mut _,
14297 &py0 as *const _ as *mut _,
14298 &py1 as *const _ as *mut _,
14299 &inf as *const _ as *mut _,
14300 &oo0 as *const _ as *mut _,
14301 &oo1 as *const _ as *mut _,
14302 &r0 as *const _ as *mut _,
14303 &r1 as *const _ as *mut _,
14304 ];
14305 unsafe {
14306 self.launch_pdl(
14307 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14308 (grid, 1, 1),
14309 (32, rpb, 1),
14310 &mut ps,
14311 )?;
14312 }
14313 }
14314 return Ok(Some((y0, y1)));
14315 }
14316 let __s_b = self.gpu.stream();
14317 let mut b = __s_b.launch_builder(&f);
14318 b.arg(b0)
14319 .arg(b1)
14320 .arg(aq)
14321 .arg(ad)
14322 .arg(&mut y0)
14323 .arg(&mut y1)
14324 .arg(&inf)
14325 .arg(&oo0)
14326 .arg(&oo1)
14327 .arg(&r0)
14328 .arg(&r1);
14329 unsafe {
14330 b.launch(cfg)?;
14331 }
14332 Ok(Some((y0, y1)))
14333 }
14334
14335 pub fn matmul_q4_fused2_into(
14337 &self,
14338 w0: &crate::model::GpuTensor,
14339 w1: &crate::model::GpuTensor,
14340 aq: &CudaSlice<i8>,
14341 ad: &CudaSlice<f32>,
14342 y0: &mut CudaSlice<f32>,
14343 y1: &mut CudaSlice<f32>,
14344 ) -> Result<bool, Box<dyn std::error::Error>> {
14345 use crate::model::GpuTensor;
14346 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14347 match w {
14348 GpuTensor::Quant {
14349 qtype, row_bytes, ..
14350 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14351 _ => None,
14352 }
14353 };
14354 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14355 return Ok(false);
14356 };
14357 if w0.in_features() != w1.in_features() {
14358 return Ok(false);
14359 }
14360 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14361 match w {
14362 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14363 Some(m) => (m, true),
14364 None => (bytes, *rp),
14365 },
14366 _ => unreachable!(),
14367 }
14368 }
14369 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14370 if rp0 != rp1 {
14371 return Ok(false);
14372 }
14373 let rp = rp0;
14374 let rpb: u32 = 4;
14375 let mr1 = rp && Self::q40_mr1_on();
14376 let nb = |o: usize| {
14377 if mr1 {
14378 (o as u32).div_ceil(rpb)
14379 } else {
14380 (o as u32).div_ceil(2).div_ceil(rpb)
14381 }
14382 };
14383 let grid = nb(o0) + nb(o1);
14384 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14385 let f = self.func(if mr1 {
14386 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14387 } else if rp {
14388 "qmatvec_q4_0_mmvq_fused2_rp"
14389 } else {
14390 "qmatvec_q4_0_mmvq_fused2"
14391 });
14392 let cfg = LaunchConfig {
14393 grid_dim: (grid, 1, 1),
14394 block_dim: (32, rpb, 1),
14395 shared_mem_bytes: 0,
14396 };
14397 let inf = w0.in_features() as i32;
14398 let (oo0, oo1) = (o0 as i32, o1 as i32);
14399 let (r0, r1) = (rb0 as i64, rb1 as i64);
14400 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14402 use cudarc::driver::{DevicePtr, DevicePtrMut};
14403 let s = &self.gpu.stream();
14404 let (p0, _g0) = b0.device_ptr(s);
14405 let (p1, _g1) = b1.device_ptr(s);
14406 let (paq, _g2) = aq.device_ptr(s);
14407 let (pad, _g3) = ad.device_ptr(s);
14408 let (py0, _g4) = y0.device_ptr_mut(s);
14409 let (py1, _g5) = y1.device_ptr_mut(s);
14410 let mut ps = [
14411 &p0 as *const _ as *mut std::ffi::c_void,
14412 &p1 as *const _ as *mut _,
14413 &paq as *const _ as *mut _,
14414 &pad as *const _ as *mut _,
14415 &py0 as *const _ as *mut _,
14416 &py1 as *const _ as *mut _,
14417 &inf as *const _ as *mut _,
14418 &oo0 as *const _ as *mut _,
14419 &oo1 as *const _ as *mut _,
14420 &r0 as *const _ as *mut _,
14421 &r1 as *const _ as *mut _,
14422 ];
14423 unsafe {
14424 self.launch_pdl(
14425 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14426 (grid, 1, 1),
14427 (32, rpb, 1),
14428 &mut ps,
14429 )?;
14430 }
14431 return Ok(true);
14432 }
14433 let __s_b = self.gpu.stream();
14434 let mut b = __s_b.launch_builder(&f);
14435 b.arg(b0)
14436 .arg(b1)
14437 .arg(aq)
14438 .arg(ad)
14439 .arg(&mut *y0)
14440 .arg(&mut *y1)
14441 .arg(&inf)
14442 .arg(&oo0)
14443 .arg(&oo1)
14444 .arg(&r0)
14445 .arg(&r1);
14446 unsafe {
14447 b.launch(cfg)?;
14448 }
14449 Ok(true)
14450 }
14451
14452 pub fn matmul_q4_fused2_batched(
14457 &self,
14458 w0: &crate::model::GpuTensor,
14459 w1: &crate::model::GpuTensor,
14460 aq: &CudaSlice<i8>,
14461 ad: &CudaSlice<f32>,
14462 m: usize,
14463 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14464 use crate::model::GpuTensor;
14465 if m < 2 || m > 8 {
14466 return Ok(None);
14467 }
14468 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14469 match w {
14470 GpuTensor::Quant {
14471 qtype, row_bytes, ..
14472 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14473 _ => None,
14474 }
14475 };
14476 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14477 return Ok(None);
14478 };
14479 if w0.in_features() != w1.in_features() {
14480 return Ok(None);
14481 }
14482 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14483 match w {
14484 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14485 Some(mr) => (mr, true),
14486 None => (bytes, *rp),
14487 },
14488 _ => unreachable!(),
14489 }
14490 }
14491 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14492 if !rp0 || !rp1 {
14493 return Ok(None);
14494 }
14495 let mcols = Self::batched_mcols(m);
14496 let rpb: u32 = 4;
14497 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14498 let grid = nb(o0) + nb(o1);
14499 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14500 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14501 let f = self.func(match mcols {
14502 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14503 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14504 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14505 });
14506 let cfg = LaunchConfig {
14507 grid_dim: (grid, 1, 1),
14508 block_dim: (32, rpb, 1),
14509 shared_mem_bytes: 0,
14510 };
14511 let inf = w0.in_features() as i32;
14512 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14513 let rb = rb0 as i64;
14514 let __s_b = self.gpu.stream();
14515 let mut b = __s_b.launch_builder(&f);
14516 b.arg(b0)
14517 .arg(b1)
14518 .arg(aq)
14519 .arg(ad)
14520 .arg(&mut y0)
14521 .arg(&mut y1)
14522 .arg(&inf)
14523 .arg(&oo0)
14524 .arg(&oo1)
14525 .arg(&mi)
14526 .arg(&rb);
14527 unsafe {
14528 b.launch(cfg)?;
14529 }
14530 Ok(Some((y0, y1)))
14531 }
14532
14533 #[allow(clippy::too_many_arguments)]
14536 pub fn matmul_q4_fused3_batched(
14537 &self,
14538 w0: &crate::model::GpuTensor,
14539 w1: &crate::model::GpuTensor,
14540 w2: &crate::model::GpuTensor,
14541 aq: &CudaSlice<i8>,
14542 ad: &CudaSlice<f32>,
14543 m: usize,
14544 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14545 {
14546 use crate::model::GpuTensor;
14547 if m < 2 || m > 8 {
14548 return Ok(None);
14549 }
14550 let q4 = |w: &GpuTensor| -> Option<usize> {
14551 match w {
14552 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14553 _ => None,
14554 }
14555 };
14556 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14557 return Ok(None);
14558 };
14559 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14560 return Ok(None);
14561 }
14562 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14563 match w {
14564 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14565 Some(mr) => (mr, true),
14566 None => (bytes, *rp),
14567 },
14568 _ => unreachable!(),
14569 }
14570 }
14571 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14572 if !rp0 || !rp1 || !rp2 {
14573 return Ok(None);
14574 }
14575 let mcols = Self::batched_mcols(m);
14576 let rpb: u32 = 4;
14577 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14578 let grid = nb(o0) + nb(o1) + nb(o2);
14579 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14580 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14581 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14582 let f = self.func(match mcols {
14583 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14584 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14585 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14586 });
14587 let cfg = LaunchConfig {
14588 grid_dim: (grid, 1, 1),
14589 block_dim: (32, rpb, 1),
14590 shared_mem_bytes: 0,
14591 };
14592 let inf = w0.in_features() as i32;
14593 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14594 let rb = 0i64;
14595 let __s_b = self.gpu.stream();
14596 let mut b = __s_b.launch_builder(&f);
14597 b.arg(b0)
14598 .arg(b1)
14599 .arg(b2)
14600 .arg(aq)
14601 .arg(ad)
14602 .arg(&mut y0)
14603 .arg(&mut y1)
14604 .arg(&mut y2)
14605 .arg(&inf)
14606 .arg(&oo0)
14607 .arg(&oo1)
14608 .arg(&oo2)
14609 .arg(&mi)
14610 .arg(&rb);
14611 unsafe {
14612 b.launch(cfg)?;
14613 }
14614 Ok(Some((y0, y1, y2)))
14615 }
14616
14617 pub fn matmul_q8_fused3(
14618 &self,
14619 w0: &crate::model::GpuTensor,
14620 w1: &crate::model::GpuTensor,
14621 w2: &crate::model::GpuTensor,
14622 aq: &CudaSlice<i8>,
14623 ad: &CudaSlice<f32>,
14624 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14625 {
14626 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14629 return Ok(Some(self.e4m3_fused3_core(
14630 p0.0,
14631 p1.0,
14632 p2.0,
14633 aq,
14634 ad,
14635 w0.in_features(),
14636 p0.1,
14637 p1.1,
14638 p2.1,
14639 p0.2,
14640 p0.3,
14641 p1.3,
14642 p2.3,
14643 )?));
14644 }
14645 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14646 return Ok(None);
14647 };
14648 Ok(Some(self.q8_fused3_core(
14649 p0.0,
14650 p1.0,
14651 p2.0,
14652 aq,
14653 ad,
14654 w0.in_features(),
14655 p0.1,
14656 p1.1,
14657 p2.1,
14658 p0.2,
14659 )?))
14660 }
14661
14662 #[allow(clippy::too_many_arguments)]
14663 fn q8_fused3_core(
14664 &self,
14665 b0: &CudaSlice<u8>,
14666 b1: &CudaSlice<u8>,
14667 b2: &CudaSlice<u8>,
14668 aq: &CudaSlice<i8>,
14669 ad: &CudaSlice<f32>,
14670 in_f: usize,
14671 out0: usize,
14672 out1: usize,
14673 out2: usize,
14674 row_bytes: usize,
14675 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14676 const ROWS_PER_BLOCK: u32 = 4;
14677 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14678 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14679 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14680 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14681 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14682 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14683 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14684 let cfg = LaunchConfig {
14685 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14686 block_dim: (32, ROWS_PER_BLOCK, 1),
14687 shared_mem_bytes: 0,
14688 };
14689 let (inf, o0, o1, o2, rbl) = (
14690 in_f as i32,
14691 out0 as i32,
14692 out1 as i32,
14693 out2 as i32,
14694 row_bytes as i64,
14695 );
14696 let __s_b = self.gpu.stream();
14697 let mut b = __s_b.launch_builder(&f);
14698 b.arg(b0)
14699 .arg(b1)
14700 .arg(b2)
14701 .arg(aq)
14702 .arg(ad)
14703 .arg(&mut y0)
14704 .arg(&mut y1)
14705 .arg(&mut y2)
14706 .arg(&inf)
14707 .arg(&o0)
14708 .arg(&o1)
14709 .arg(&o2)
14710 .arg(&rbl);
14711 unsafe {
14712 b.launch(cfg)?;
14713 }
14714 Ok((y0, y1, y2))
14715 }
14716
14717 #[allow(clippy::too_many_arguments)]
14719 pub fn qmatvec_q8_fused3_raw(
14720 &self,
14721 b0: &CudaSlice<u8>,
14722 b1: &CudaSlice<u8>,
14723 b2: &CudaSlice<u8>,
14724 x: &CudaSlice<f32>,
14725 in_f: usize,
14726 out0: usize,
14727 out1: usize,
14728 out2: usize,
14729 row_bytes: usize,
14730 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14731 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14732 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14733 }
14734
14735 pub fn matmul_q8_fused2_t(
14746 &self,
14747 w0: &crate::model::GpuTensor,
14748 w1: &crate::model::GpuTensor,
14749 aq: &CudaSlice<i8>,
14750 ad: &CudaSlice<f32>,
14751 m: usize,
14752 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14753 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14757 return Ok(None);
14758 }
14759 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14762 if m > 4 && !Self::b8_enabled() {
14763 return Ok(None);
14764 }
14765 return Ok(Some(self.e4m3_fused2_t_core(
14766 p0.0,
14767 p1.0,
14768 aq,
14769 ad,
14770 m,
14771 w0.in_features(),
14772 p0.1,
14773 p1.1,
14774 p0.2,
14775 p0.3,
14776 p1.3,
14777 )?));
14778 }
14779 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14780 return Ok(None);
14781 };
14782 Ok(Some(self.q8_fused2_t_core(
14783 p0.0,
14784 p1.0,
14785 aq,
14786 ad,
14787 m,
14788 w0.in_features(),
14789 p0.1,
14790 p1.1,
14791 p0.2,
14792 )?))
14793 }
14794
14795 #[allow(clippy::too_many_arguments)]
14796 fn q8_fused2_t_core(
14797 &self,
14798 b0: &CudaSlice<u8>,
14799 b1: &CudaSlice<u8>,
14800 aq: &CudaSlice<i8>,
14801 ad: &CudaSlice<f32>,
14802 m: usize,
14803 in_f: usize,
14804 out0: usize,
14805 out1: usize,
14806 row_bytes: usize,
14807 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14808 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14810 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14811 let f = self.func(match Self::batched_mcols(m) {
14812 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14813 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14814 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14816 });
14817 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14818 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14819 let cfg = LaunchConfig {
14820 grid_dim: (nb0 + nb1, 1, 1),
14821 block_dim: (32, ROWS_PER_BLOCK, 1),
14822 shared_mem_bytes: 0,
14823 };
14824 let (inf, o0, o1, mi, rbl) = (
14825 in_f as i32,
14826 out0 as i32,
14827 out1 as i32,
14828 m as i32,
14829 row_bytes as i64,
14830 );
14831 let __s_b = self.gpu.stream();
14832 let mut b = __s_b.launch_builder(&f);
14833 b.arg(b0)
14834 .arg(b1)
14835 .arg(aq)
14836 .arg(ad)
14837 .arg(&mut y0)
14838 .arg(&mut y1)
14839 .arg(&inf)
14840 .arg(&o0)
14841 .arg(&o1)
14842 .arg(&mi)
14843 .arg(&rbl);
14844 unsafe {
14845 b.launch(cfg)?;
14846 }
14847 Ok((y0, y1))
14848 }
14849
14850 #[allow(clippy::too_many_arguments)]
14853 pub fn qmatvec_q8_fused2_t_raw(
14854 &self,
14855 b0: &CudaSlice<u8>,
14856 b1: &CudaSlice<u8>,
14857 x: &CudaSlice<f32>,
14858 m: usize,
14859 in_f: usize,
14860 out0: usize,
14861 out1: usize,
14862 row_bytes: usize,
14863 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14864 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14865 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14866 }
14867
14868 #[allow(clippy::too_many_arguments)]
14871 pub fn matmul_q8_fused3_t(
14872 &self,
14873 w0: &crate::model::GpuTensor,
14874 w1: &crate::model::GpuTensor,
14875 w2: &crate::model::GpuTensor,
14876 aq: &CudaSlice<i8>,
14877 ad: &CudaSlice<f32>,
14878 m: usize,
14879 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14880 {
14881 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14882 return Ok(None);
14883 }
14884 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14885 return Ok(Some(self.e4m3_fused3_t_core(
14886 p0.0,
14887 p1.0,
14888 p2.0,
14889 aq,
14890 ad,
14891 m,
14892 w0.in_features(),
14893 p0.1,
14894 p1.1,
14895 p2.1,
14896 p0.2,
14897 p0.3,
14898 p1.3,
14899 p2.3,
14900 )?));
14901 }
14902 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14903 return Ok(None);
14904 };
14905 Ok(Some(self.q8_fused3_t_core(
14906 p0.0,
14907 p1.0,
14908 p2.0,
14909 aq,
14910 ad,
14911 m,
14912 w0.in_features(),
14913 p0.1,
14914 p1.1,
14915 p2.1,
14916 p0.2,
14917 )?))
14918 }
14919
14920 #[allow(clippy::too_many_arguments)]
14921 fn q8_fused3_t_core(
14922 &self,
14923 b0: &CudaSlice<u8>,
14924 b1: &CudaSlice<u8>,
14925 b2: &CudaSlice<u8>,
14926 aq: &CudaSlice<i8>,
14927 ad: &CudaSlice<f32>,
14928 m: usize,
14929 in_f: usize,
14930 out0: usize,
14931 out1: usize,
14932 out2: usize,
14933 row_bytes: usize,
14934 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14935 const ROWS_PER_BLOCK: u32 = 4;
14936 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14937 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14938 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14939 let f = self.func(if Self::batched_mcols(m) == 2 {
14940 "qmatvec_q8_0_mmvq_fused3_b2"
14941 } else {
14942 "qmatvec_q8_0_mmvq_fused3_b4"
14943 });
14944 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14945 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14946 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14947 let cfg = LaunchConfig {
14948 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14949 block_dim: (32, ROWS_PER_BLOCK, 1),
14950 shared_mem_bytes: 0,
14951 };
14952 let (inf, o0, o1, o2, mi, rbl) = (
14953 in_f as i32,
14954 out0 as i32,
14955 out1 as i32,
14956 out2 as i32,
14957 m as i32,
14958 row_bytes as i64,
14959 );
14960 let __s_b = self.gpu.stream();
14961 let mut b = __s_b.launch_builder(&f);
14962 b.arg(b0)
14963 .arg(b1)
14964 .arg(b2)
14965 .arg(aq)
14966 .arg(ad)
14967 .arg(&mut y0)
14968 .arg(&mut y1)
14969 .arg(&mut y2)
14970 .arg(&inf)
14971 .arg(&o0)
14972 .arg(&o1)
14973 .arg(&o2)
14974 .arg(&mi)
14975 .arg(&rbl);
14976 unsafe {
14977 b.launch(cfg)?;
14978 }
14979 Ok((y0, y1, y2))
14980 }
14981
14982 #[allow(clippy::too_many_arguments)]
14984 pub fn qmatvec_q8_fused3_t_raw(
14985 &self,
14986 b0: &CudaSlice<u8>,
14987 b1: &CudaSlice<u8>,
14988 b2: &CudaSlice<u8>,
14989 x: &CudaSlice<f32>,
14990 m: usize,
14991 in_f: usize,
14992 out0: usize,
14993 out1: usize,
14994 out2: usize,
14995 row_bytes: usize,
14996 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14997 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14998 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14999 }
15000
15001 pub fn q8_ffn_fuse2_on(&self) -> bool {
15005 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15006 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15007 }
15008
15009 #[allow(clippy::type_complexity)]
15015 fn q8_fused_params<'w, const N: usize>(
15016 &self,
15017 ws: &[&'w crate::model::GpuTensor; N],
15018 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15019 use crate::model::GpuTensor;
15020 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15021 return None;
15022 }
15023 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15024 return None;
15025 }
15026 let in_f = ws[0].in_features();
15027 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15028 for (i, w) in ws.iter().enumerate() {
15029 match w {
15030 GpuTensor::Quant {
15031 bytes,
15032 qtype,
15033 row_bytes,
15034 scale,
15035 ..
15036 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15037 out[i] = Some((bytes, w.out_features(), *row_bytes))
15038 }
15039 _ => return None,
15040 }
15041 }
15042 Some(out.map(|o| o.unwrap()))
15043 }
15044
15045 pub fn e4m3_dual_on(&self) -> bool {
15048 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15049 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15050 }
15051
15052 #[allow(clippy::type_complexity)]
15064 fn e4m3_fused_params<'w, const N: usize>(
15065 &self,
15066 ws: &[&'w crate::model::GpuTensor; N],
15067 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15068 use crate::model::GpuTensor;
15069 if !self.e4m3_dual_on() {
15070 return None;
15071 }
15072 let in_f = ws[0].in_features();
15073 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15074 for (i, w) in ws.iter().enumerate() {
15075 match w {
15076 GpuTensor::Quant {
15077 bytes,
15078 qtype,
15079 row_bytes,
15080 scale,
15081 rp,
15082 rp4,
15083 ..
15084 } if *qtype == QT_F8_E4M3
15085 && w.in_features() == in_f
15086 && *row_bytes == in_f
15087 && !*rp
15088 && rp4.is_none() =>
15089 {
15090 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15091 }
15092 _ => return None,
15093 }
15094 }
15095 Some(out.map(|o| o.unwrap()))
15096 }
15097
15098 #[allow(clippy::too_many_arguments)]
15102 fn e4m3_fused2_core(
15103 &self,
15104 b0: &CudaSlice<u8>,
15105 b1: &CudaSlice<u8>,
15106 aq: &CudaSlice<i8>,
15107 ad: &CudaSlice<f32>,
15108 in_f: usize,
15109 out0: usize,
15110 out1: usize,
15111 row_bytes: usize,
15112 ws0: f32,
15113 ws1: f32,
15114 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15115 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15117 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15118 let f = self.func("qmatvec_e4m3_mmvq_fused2");
15119 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15120 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15121 let cfg = LaunchConfig {
15122 grid_dim: (nb0 + nb1, 1, 1),
15123 block_dim: (32, ROWS_PER_BLOCK, 1),
15124 shared_mem_bytes: 0,
15125 };
15126 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15127 let __s_b = self.gpu.stream();
15128 let mut b = __s_b.launch_builder(&f);
15129 b.arg(b0)
15130 .arg(b1)
15131 .arg(aq)
15132 .arg(ad)
15133 .arg(&mut y0)
15134 .arg(&mut y1)
15135 .arg(&inf)
15136 .arg(&o0)
15137 .arg(&o1)
15138 .arg(&rbl)
15139 .arg(&ws0)
15140 .arg(&ws1);
15141 unsafe {
15142 b.launch(cfg)?;
15143 }
15144 Ok((y0, y1))
15145 }
15146
15147 #[allow(clippy::too_many_arguments)]
15149 fn e4m3_fused3_core(
15150 &self,
15151 b0: &CudaSlice<u8>,
15152 b1: &CudaSlice<u8>,
15153 b2: &CudaSlice<u8>,
15154 aq: &CudaSlice<i8>,
15155 ad: &CudaSlice<f32>,
15156 in_f: usize,
15157 out0: usize,
15158 out1: usize,
15159 out2: usize,
15160 row_bytes: usize,
15161 ws0: f32,
15162 ws1: f32,
15163 ws2: f32,
15164 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15165 const ROWS_PER_BLOCK: u32 = 4;
15166 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15167 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15168 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15169 let f = self.func("qmatvec_e4m3_mmvq_fused3");
15170 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15171 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15172 let mut y2 = self.alloc_uninit::<f32>(out2)?;
15173 let cfg = LaunchConfig {
15174 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15175 block_dim: (32, ROWS_PER_BLOCK, 1),
15176 shared_mem_bytes: 0,
15177 };
15178 let (inf, o0, o1, o2, rbl) = (
15179 in_f as i32,
15180 out0 as i32,
15181 out1 as i32,
15182 out2 as i32,
15183 row_bytes as i64,
15184 );
15185 let __s_b = self.gpu.stream();
15186 let mut b = __s_b.launch_builder(&f);
15187 b.arg(b0)
15188 .arg(b1)
15189 .arg(b2)
15190 .arg(aq)
15191 .arg(ad)
15192 .arg(&mut y0)
15193 .arg(&mut y1)
15194 .arg(&mut y2)
15195 .arg(&inf)
15196 .arg(&o0)
15197 .arg(&o1)
15198 .arg(&o2)
15199 .arg(&rbl)
15200 .arg(&ws0)
15201 .arg(&ws1)
15202 .arg(&ws2);
15203 unsafe {
15204 b.launch(cfg)?;
15205 }
15206 Ok((y0, y1, y2))
15207 }
15208
15209 #[allow(clippy::too_many_arguments)]
15213 fn e4m3_fused2_t_core(
15214 &self,
15215 b0: &CudaSlice<u8>,
15216 b1: &CudaSlice<u8>,
15217 aq: &CudaSlice<i8>,
15218 ad: &CudaSlice<f32>,
15219 m: usize,
15220 in_f: usize,
15221 out0: usize,
15222 out1: usize,
15223 row_bytes: usize,
15224 ws0: f32,
15225 ws1: f32,
15226 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15227 const ROWS_PER_BLOCK: u32 = 4;
15228 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15229 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15230 let f = self.func(match Self::batched_mcols(m) {
15231 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15232 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15233 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15234 });
15235 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15236 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15237 let cfg = LaunchConfig {
15238 grid_dim: (nb0 + nb1, 1, 1),
15239 block_dim: (32, ROWS_PER_BLOCK, 1),
15240 shared_mem_bytes: 0,
15241 };
15242 let (inf, o0, o1, mi, rbl) = (
15243 in_f as i32,
15244 out0 as i32,
15245 out1 as i32,
15246 m as i32,
15247 row_bytes as i64,
15248 );
15249 let __s_b = self.gpu.stream();
15250 let mut b = __s_b.launch_builder(&f);
15251 b.arg(b0)
15252 .arg(b1)
15253 .arg(aq)
15254 .arg(ad)
15255 .arg(&mut y0)
15256 .arg(&mut y1)
15257 .arg(&inf)
15258 .arg(&o0)
15259 .arg(&o1)
15260 .arg(&mi)
15261 .arg(&rbl);
15262 unsafe {
15263 b.launch(cfg)?;
15264 }
15265 if ws0 != 1.0 {
15266 self.scale_inplace(&mut y0, ws0, m * out0)?;
15267 }
15268 if ws1 != 1.0 {
15269 self.scale_inplace(&mut y1, ws1, m * out1)?;
15270 }
15271 Ok((y0, y1))
15272 }
15273
15274 #[allow(clippy::too_many_arguments)]
15276 fn e4m3_fused3_t_core(
15277 &self,
15278 b0: &CudaSlice<u8>,
15279 b1: &CudaSlice<u8>,
15280 b2: &CudaSlice<u8>,
15281 aq: &CudaSlice<i8>,
15282 ad: &CudaSlice<f32>,
15283 m: usize,
15284 in_f: usize,
15285 out0: usize,
15286 out1: usize,
15287 out2: usize,
15288 row_bytes: usize,
15289 ws0: f32,
15290 ws1: f32,
15291 ws2: f32,
15292 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15293 const ROWS_PER_BLOCK: u32 = 4;
15294 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15295 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15296 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15297 let f = self.func(if Self::batched_mcols(m) == 2 {
15298 "qmatvec_e4m3_mmvq_fused3_b2"
15299 } else {
15300 "qmatvec_e4m3_mmvq_fused3_b4"
15301 });
15302 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15303 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15304 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15305 let cfg = LaunchConfig {
15306 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15307 block_dim: (32, ROWS_PER_BLOCK, 1),
15308 shared_mem_bytes: 0,
15309 };
15310 let (inf, o0, o1, o2, mi, rbl) = (
15311 in_f as i32,
15312 out0 as i32,
15313 out1 as i32,
15314 out2 as i32,
15315 m as i32,
15316 row_bytes as i64,
15317 );
15318 let __s_b = self.gpu.stream();
15319 let mut b = __s_b.launch_builder(&f);
15320 b.arg(b0)
15321 .arg(b1)
15322 .arg(b2)
15323 .arg(aq)
15324 .arg(ad)
15325 .arg(&mut y0)
15326 .arg(&mut y1)
15327 .arg(&mut y2)
15328 .arg(&inf)
15329 .arg(&o0)
15330 .arg(&o1)
15331 .arg(&o2)
15332 .arg(&mi)
15333 .arg(&rbl);
15334 unsafe {
15335 b.launch(cfg)?;
15336 }
15337 if ws0 != 1.0 {
15338 self.scale_inplace(&mut y0, ws0, m * out0)?;
15339 }
15340 if ws1 != 1.0 {
15341 self.scale_inplace(&mut y1, ws1, m * out1)?;
15342 }
15343 if ws2 != 1.0 {
15344 self.scale_inplace(&mut y2, ws2, m * out2)?;
15345 }
15346 Ok((y0, y1, y2))
15347 }
15348
15349 pub fn qmatvec_e4m3_blk_mmvq(
15359 &self,
15360 bytes: &CudaSlice<u8>,
15361 aq: &CudaSlice<i8>,
15362 ad: &CudaSlice<f32>,
15363 scales: &CudaSlice<f32>,
15364 m: usize,
15365 in_f: usize,
15366 out_f: usize,
15367 row_bytes: usize,
15368 scale_cols: usize,
15369 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15370 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15372 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15373 )?;
15374 Ok(y)
15375 }
15376
15377 #[allow(clippy::too_many_arguments)]
15379 pub fn qmatvec_e4m3_blk_mmvq_into(
15380 &self,
15381 bytes: &CudaSlice<u8>,
15382 aq: &CudaSlice<i8>,
15383 ad: &CudaSlice<f32>,
15384 scales: &CudaSlice<f32>,
15385 m: usize,
15386 in_f: usize,
15387 out_f: usize,
15388 row_bytes: usize,
15389 scale_cols: usize,
15390 y: &mut CudaSlice<f32>,
15391 ) -> Result<(), Box<dyn std::error::Error>> {
15392 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15394 let cfg = LaunchConfig {
15395 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15396 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15399 let (inf, outf, mi, rb, sc) = (
15400 in_f as i32,
15401 out_f as i32,
15402 m as i32,
15403 row_bytes as i64,
15404 scale_cols as i32,
15405 );
15406 let __s_b = self.gpu.stream();
15407 let mut b = __s_b.launch_builder(&f);
15408 b.arg(bytes)
15409 .arg(aq)
15410 .arg(ad)
15411 .arg(scales)
15412 .arg(&mut *y)
15413 .arg(&inf)
15414 .arg(&outf)
15415 .arg(&mi)
15416 .arg(&rb)
15417 .arg(&sc);
15418 unsafe {
15419 b.launch(cfg)?;
15420 }
15421 Ok(())
15422 }
15423
15424 #[allow(clippy::too_many_arguments)]
15430 pub fn qmatvec_e4m3_blk_mmvq_batched(
15431 &self,
15432 bytes: &CudaSlice<u8>,
15433 aq: &CudaSlice<i8>,
15434 ad: &CudaSlice<f32>,
15435 scales: &CudaSlice<f32>,
15436 m: usize,
15437 in_f: usize,
15438 out_f: usize,
15439 row_bytes: usize,
15440 scale_cols: usize,
15441 mcols: usize,
15442 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15443 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15445 let name = match mcols {
15446 2 => "qmatvec_e4m3_blk_mmvq_b2",
15447 4 => "qmatvec_e4m3_blk_mmvq_b4",
15448 8 => "qmatvec_e4m3_blk_mmvq_b8",
15449 16 => "qmatvec_e4m3_blk_mmvq_b16",
15450 _ => {
15451 return Err(
15452 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15453 );
15454 }
15455 };
15456 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15457 let f = self.func(name);
15458 let cfg = LaunchConfig {
15459 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15460 block_dim: (32, ROWS_PER_BLOCK, 1),
15461 shared_mem_bytes: 0,
15462 };
15463 let (inf, outf, mi, rb, sc) = (
15464 in_f as i32,
15465 out_f as i32,
15466 m as i32,
15467 row_bytes as i64,
15468 scale_cols as i32,
15469 );
15470 let __s_b = self.gpu.stream();
15471 let mut b = __s_b.launch_builder(&f);
15472 b.arg(bytes)
15473 .arg(aq)
15474 .arg(ad)
15475 .arg(scales)
15476 .arg(&mut y)
15477 .arg(&inf)
15478 .arg(&outf)
15479 .arg(&mi)
15480 .arg(&rb)
15481 .arg(&sc);
15482 unsafe {
15483 b.launch(cfg)?;
15484 }
15485 Ok(y)
15486 }
15487
15488 #[allow(clippy::too_many_arguments)]
15491 pub fn qmatvec_e4m3_blk_batched_raw(
15492 &self,
15493 bytes: &CudaSlice<u8>,
15494 x: &CudaSlice<f32>,
15495 scales: &CudaSlice<f32>,
15496 m: usize,
15497 in_f: usize,
15498 out_f: usize,
15499 row_bytes: usize,
15500 scale_cols: usize,
15501 mcols: usize,
15502 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15503 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15504 self.qmatvec_e4m3_blk_mmvq_batched(
15505 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15506 )
15507 }
15508
15509 #[allow(clippy::too_many_arguments)]
15512 pub fn qmatvec_e4m3_blk_mmvq_raw(
15513 &self,
15514 bytes: &CudaSlice<u8>,
15515 x: &CudaSlice<f32>,
15516 scales: &CudaSlice<f32>,
15517 m: usize,
15518 in_f: usize,
15519 out_f: usize,
15520 row_bytes: usize,
15521 scale_cols: usize,
15522 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15523 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15524 self.qmatvec_e4m3_blk_mmvq(
15525 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15526 )
15527 }
15528
15529 #[allow(clippy::too_many_arguments)]
15532 pub fn qmatvec_e4m3_fused2_raw(
15533 &self,
15534 b0: &CudaSlice<u8>,
15535 b1: &CudaSlice<u8>,
15536 x: &CudaSlice<f32>,
15537 in_f: usize,
15538 out0: usize,
15539 out1: usize,
15540 row_bytes: usize,
15541 ws0: f32,
15542 ws1: f32,
15543 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15544 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15545 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15546 }
15547
15548 #[allow(clippy::too_many_arguments)]
15549 pub fn qmatvec_e4m3_fused3_raw(
15550 &self,
15551 b0: &CudaSlice<u8>,
15552 b1: &CudaSlice<u8>,
15553 b2: &CudaSlice<u8>,
15554 x: &CudaSlice<f32>,
15555 in_f: usize,
15556 out0: usize,
15557 out1: usize,
15558 out2: usize,
15559 row_bytes: usize,
15560 ws0: f32,
15561 ws1: f32,
15562 ws2: f32,
15563 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15564 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15565 self.e4m3_fused3_core(
15566 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15567 )
15568 }
15569
15570 #[allow(clippy::too_many_arguments)]
15571 pub fn qmatvec_e4m3_fused2_t_raw(
15572 &self,
15573 b0: &CudaSlice<u8>,
15574 b1: &CudaSlice<u8>,
15575 x: &CudaSlice<f32>,
15576 m: usize,
15577 in_f: usize,
15578 out0: usize,
15579 out1: usize,
15580 row_bytes: usize,
15581 ws0: f32,
15582 ws1: f32,
15583 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15584 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15585 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15586 }
15587
15588 #[allow(clippy::too_many_arguments)]
15589 pub fn qmatvec_e4m3_fused3_t_raw(
15590 &self,
15591 b0: &CudaSlice<u8>,
15592 b1: &CudaSlice<u8>,
15593 b2: &CudaSlice<u8>,
15594 x: &CudaSlice<f32>,
15595 m: usize,
15596 in_f: usize,
15597 out0: usize,
15598 out1: usize,
15599 out2: usize,
15600 row_bytes: usize,
15601 ws0: f32,
15602 ws1: f32,
15603 ws2: f32,
15604 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15605 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15606 self.e4m3_fused3_t_core(
15607 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15608 )
15609 }
15610
15611 fn try_e4m3_blk_pre(
15622 &self,
15623 w: &crate::model::GpuTensor,
15624 aq: &CudaSlice<i8>,
15625 ad: &CudaSlice<f32>,
15626 m: usize,
15627 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15628 use crate::model::GpuTensor;
15629 if let GpuTensor::Quant {
15630 bytes,
15631 qtype,
15632 row_bytes,
15633 blk: Some(g),
15634 ..
15635 } = w
15636 {
15637 if *qtype == QT_F8_E4M3_BLK {
15638 if (2..=16).contains(&m)
15644 && std::env::var("MEMRA_NO_BATCHED").is_err()
15645 && (m <= 4 || Self::b8_enabled())
15646 {
15647 let mcols = Self::batched_mcols(m);
15648 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15649 bytes,
15650 aq,
15651 ad,
15652 &g.scales,
15653 m,
15654 w.in_features(),
15655 w.out_features(),
15656 *row_bytes,
15657 g.cols,
15658 mcols,
15659 )?));
15660 }
15661 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15662 bytes,
15663 aq,
15664 ad,
15665 &g.scales,
15666 m,
15667 w.in_features(),
15668 w.out_features(),
15669 *row_bytes,
15670 g.cols,
15671 )?));
15672 }
15673 }
15674 Ok(None)
15675 }
15676
15677 fn try_e4m3_blk_prefill(
15724 &self,
15725 w: &crate::model::GpuTensor,
15726 x: &CudaSlice<f32>,
15727 m: usize,
15728 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15729 use crate::model::GpuTensor;
15730 let GpuTensor::Quant {
15731 bytes,
15732 qtype,
15733 blk: Some(g),
15734 ..
15735 } = w
15736 else {
15737 return Ok(None);
15738 };
15739 if *qtype != QT_F8_E4M3_BLK {
15740 return Ok(None);
15741 }
15742 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15747 return Ok(Some(y));
15748 }
15749 let (in_f, out_f) = (w.in_features(), w.out_features());
15750 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15751 let tmp = GpuTensor::Quant {
15752 bytes: slab,
15753 qtype: QT_Q8_0,
15754 row_bytes: in_f / 32 * 34,
15755 ne: vec![in_f as u64, out_f as u64],
15756 scale: 1.0,
15757 rp: false,
15758 #[cfg(memra_cutlass)]
15759 cutlass: None,
15760 fp8: None,
15761 blk: None,
15762 f16: None,
15763 rp4: None,
15764 };
15765 Ok(Some(self.matmul(&tmp, x, m)?))
15767 }
15768
15769 pub fn matmul_pre_noscale(
15770 &self,
15771 w: &crate::model::GpuTensor,
15772 aq: &CudaSlice<i8>,
15773 ad: &CudaSlice<f32>,
15774 m: usize,
15775 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15776 use crate::model::GpuTensor;
15777 if m == 1 {
15781 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15782 return Ok(Some((y, 1.0)));
15783 }
15784 }
15785 if m != 1 || !self.uses_q8_1_fast(w) {
15787 return Ok(None);
15788 }
15789 let in_f = w.in_features();
15790 let out_f = w.out_features();
15791 let (bytes, qtype, row_bytes, scale, rp) = match w {
15792 GpuTensor::Quant {
15793 bytes,
15794 qtype,
15795 row_bytes,
15796 scale,
15797 rp,
15798 ..
15799 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15800 _ => return Ok(None),
15801 };
15802 if self.mmvq_supports(qtype) {
15804 let (mbytes, mrp) = match w {
15806 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15807 _ => (bytes, rp),
15808 };
15809 let y = self.qmatvec_mmvq(
15810 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15811 )?;
15812 return Ok(Some((y, scale)));
15813 }
15814 let name = match qtype {
15816 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15817 QT_Q4_K => "qmatvec_q4_K_dp4a",
15818 QT_Q6_K => "qmatvec_q6_K_dp4a",
15819 QT_Q5_K => "qmatvec_q5_K_dp4a",
15820 QT_Q3_K => "qmatvec_q3_K_dp4a",
15821 QT_NVFP4 => {
15822 if rp {
15823 "qmatvec_nvfp4_dp4a_rp"
15824 } else {
15825 "qmatvec_nvfp4_dp4a"
15826 }
15827 }
15828 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15829 _ => return Ok(None),
15830 };
15831 let f = self.func(name);
15832 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15833 let cfg = LaunchConfig {
15834 grid_dim: (out_f as u32, m as u32, 1),
15835 block_dim: (128, 1, 1),
15836 shared_mem_bytes: 0,
15837 };
15838 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15839 let __s_b = self.gpu.stream();
15840 let mut b = __s_b.launch_builder(&f);
15841 b.arg(bytes)
15842 .arg(aq)
15843 .arg(ad)
15844 .arg(&mut y)
15845 .arg(&inf)
15846 .arg(&outf)
15847 .arg(&mi)
15848 .arg(&rb);
15849 unsafe {
15850 b.launch(cfg)?;
15851 }
15852 Ok(Some((y, scale)))
15853 }
15854
15855 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15858 if qtype == QT_F8_E4M3 {
15863 return true;
15864 }
15865 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15866 return false;
15867 }
15868 matches!(
15869 qtype,
15870 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15871 )
15872 }
15873
15874 pub fn qmatvec_mmvq(
15879 &self,
15880 bytes: &CudaSlice<u8>,
15881 aq: &CudaSlice<i8>,
15882 ad: &CudaSlice<f32>,
15883 m: usize,
15884 in_f: usize,
15885 out_f: usize,
15886 qtype: i32,
15887 row_bytes: usize,
15888 scale: f32,
15889 rp: bool,
15890 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15891 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15893 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15894 )?;
15895 Ok(y)
15896 }
15897
15898 #[allow(clippy::too_many_arguments)]
15900 pub fn qmatvec_mmvq_into(
15901 &self,
15902 bytes: &CudaSlice<u8>,
15903 aq: &CudaSlice<i8>,
15904 ad: &CudaSlice<f32>,
15905 m: usize,
15906 in_f: usize,
15907 out_f: usize,
15908 qtype: i32,
15909 row_bytes: usize,
15910 scale: f32,
15911 rp: bool,
15912 y: &mut CudaSlice<f32>,
15913 ) -> Result<(), Box<dyn std::error::Error>> {
15914 debug_assert!(y.len() >= m * out_f);
15915 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15921 && rp
15922 && m == 1
15923 && out_f >= 64
15924 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15925 && {
15926 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15927 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15928 }
15929 {
15930 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15931 let cfg = LaunchConfig {
15932 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15933 block_dim: (32, 2, 1),
15934 shared_mem_bytes: 0,
15935 };
15936 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15937 let __s_b = self.gpu.stream();
15938 let mut b = __s_b.launch_builder(&f);
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 unsafe {
15948 b.launch(cfg)?;
15949 }
15950 if scale != 1.0 {
15951 self.scale_inplace(y, scale, out_f)?;
15952 }
15953 return Ok(());
15954 }
15955 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15964 2
15965 } else {
15966 1
15967 };
15968 if m == 1 && qtype == QT_Q4_0 {
15973 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15974 mr = *Q40MR.get_or_init(|| {
15977 std::env::var("MEMRA_Q40_MR")
15978 .ok()
15979 .and_then(|v| v.parse().ok())
15980 .unwrap_or(1)
15981 });
15982 }
15983 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15994 let q5_force = q5_mode.as_deref() == Some("2");
15995 let q5_il = qtype == QT_Q5_K
15998 && m == 1
15999 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16000 if q5_il && !q5_force && out_f > 65536 {
16001 mr = 1;
16002 }
16003 if qtype == QT_Q4_0 && rp && mr != 1 {
16006 mr = 2;
16007 }
16008 if qtype == QT_Q8_0 && rp {
16012 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16013 mr = *Q80MR.get_or_init(|| {
16014 std::env::var("MEMRA_Q80_MR")
16015 .ok()
16016 .and_then(|v| v.parse().ok())
16017 .unwrap_or(1)
16018 });
16019 }
16020 let name = match (qtype, mr, rp) {
16021 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16022 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16023 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16024 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16025 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16026 (QT_Q5_K, 2, _) => {
16027 if q5_il {
16028 "qmatvec_q5_K_mmvq_mr2_il"
16029 } else {
16030 "qmatvec_q5_K_mmvq_mr2"
16031 }
16032 }
16033 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16034 (QT_Q8_0, _, true)
16039 if in_f % 1024 == 0 && {
16040 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16041 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16042 } =>
16043 {
16044 "qmatvec_q8_0_mmvq_rpca"
16045 }
16046 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16047 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16048 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16052 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16053 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16054 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16055 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16056 (QT_Q5_K, _, _) => {
16057 if q5_il {
16058 "qmatvec_q5_K_mmvq_il"
16059 } else {
16060 "qmatvec_q5_K_mmvq"
16061 }
16062 }
16063 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16064 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16065 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16066 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16067 };
16068 let f = self.func(name);
16069 let rows_per_block = ROWS_PER_BLOCK * mr;
16071 let cfg = LaunchConfig {
16072 grid_dim: (
16073 (out_f as u32 + rows_per_block - 1) / rows_per_block,
16074 m as u32,
16075 1,
16076 ),
16077 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
16080 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16081 let __s_b = self.gpu.stream();
16082 let mut b = __s_b.launch_builder(&f);
16083 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16088 if Self::pdl_on()
16091 && Self::pdl_mmvq_on()
16092 && Self::pdl_nvfp4q8_on()
16093 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16094 {
16095 use cudarc::driver::{DevicePtr, DevicePtrMut};
16096 let s = &self.gpu.stream();
16097 let (pw, _g0) = bytes.device_ptr(s);
16098 let (paq, _g1) = aq.device_ptr(s);
16099 let (pad, _g2) = ad.device_ptr(s);
16100 let (py, _g3) = y.device_ptr_mut(s);
16101 let mut ps = [
16102 &pw as *const _ as *mut std::ffi::c_void,
16103 &paq as *const _ as *mut _,
16104 &pad as *const _ as *mut _,
16105 &py as *const _ as *mut _,
16106 &inf as *const _ as *mut _,
16107 &outf as *const _ as *mut _,
16108 &mi as *const _ as *mut _,
16109 &rb as *const _ as *mut _,
16110 &scale as *const _ as *mut _,
16111 ];
16112 unsafe {
16113 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16114 }
16115 return Ok(());
16116 }
16117 b.arg(bytes)
16118 .arg(aq)
16119 .arg(ad)
16120 .arg(&mut *y)
16121 .arg(&inf)
16122 .arg(&outf)
16123 .arg(&mi)
16124 .arg(&rb)
16125 .arg(&scale);
16126 unsafe {
16127 b.launch(cfg)?;
16128 }
16129 } else if Self::pdl_on()
16130 && Self::pdl_mmvq_on()
16131 && (matches!(
16132 name,
16133 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16134 ) || (Self::pdl_nvfp4q8_on()
16135 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16136 {
16137 {
16141 use cudarc::driver::{DevicePtr, DevicePtrMut};
16142 let s = &self.gpu.stream();
16143 let (pw, _g0) = bytes.device_ptr(s);
16144 let (paq, _g1) = aq.device_ptr(s);
16145 let (pad, _g2) = ad.device_ptr(s);
16146 let (py, _g3) = y.device_ptr_mut(s);
16147 let mut ps = [
16148 &pw as *const _ as *mut std::ffi::c_void,
16149 &paq as *const _ as *mut _,
16150 &pad as *const _ as *mut _,
16151 &py as *const _ as *mut _,
16152 &inf as *const _ as *mut _,
16153 &outf as *const _ as *mut _,
16154 &mi as *const _ as *mut _,
16155 &rb as *const _ as *mut _,
16156 ];
16157 unsafe {
16158 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16159 }
16160 }
16161 if scale != 1.0 {
16162 self.scale_inplace(y, scale, m * out_f)?;
16163 }
16164 } else {
16165 b.arg(bytes)
16166 .arg(aq)
16167 .arg(ad)
16168 .arg(&mut *y)
16169 .arg(&inf)
16170 .arg(&outf)
16171 .arg(&mi)
16172 .arg(&rb);
16173 unsafe {
16174 b.launch(cfg)?;
16175 }
16176 if scale != 1.0 {
16177 self.scale_inplace(y, scale, m * out_f)?;
16178 }
16179 }
16180 Ok(())
16181 }
16182
16183 pub fn qmatvec_mmvq_raw(
16187 &self,
16188 bytes: &CudaSlice<u8>,
16189 x: &CudaSlice<f32>,
16190 m: usize,
16191 in_f: usize,
16192 out_f: usize,
16193 qtype: i32,
16194 row_bytes: usize,
16195 rp: bool,
16196 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16197 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16198 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16199 }
16200
16201 pub fn batched_supports(&self, qtype: i32) -> bool {
16205 matches!(
16206 qtype,
16207 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16208 )
16209 }
16210
16211 pub fn iq_fast_enabled() -> bool {
16219 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16220 *ON.get_or_init(|| {
16221 std::env::var("MEMRA_IQ_FAST")
16222 .map(|v| v != "0")
16223 .unwrap_or(true)
16224 })
16225 }
16226
16227 pub fn b8_enabled() -> bool {
16230 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16231 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16232 }
16233
16234 pub fn batched_mcols(m: usize) -> usize {
16236 if m == 2 {
16237 2
16238 } else if m <= 4 {
16239 4
16240 } else if m <= 8 {
16241 8
16242 } else {
16243 16
16244 }
16245 }
16246
16247 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16252 Some(match (qtype, mcols) {
16253 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16254 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16255 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16256 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16262 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16263 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16264 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16265 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16268 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16269 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16270 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16271 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16274 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16275 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16276 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16277 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16278 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16279 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16280 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16281 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16285 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16286 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16287 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16288 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16292 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16293 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16294 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16295 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16296 _ => return None,
16297 })
16298 }
16299
16300 pub fn sm_count(&self) -> i32 {
16335 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16336 *SMS.get_or_init(|| {
16337 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16338 self.gpu
16339 .ctx
16340 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16341 .unwrap_or(82)
16342 })
16343 }
16344
16345 pub fn batched_variant(
16346 &self,
16347 _m: usize,
16348 in_f: usize,
16349 out_f: usize,
16350 qtype: i32,
16351 row_bytes: usize,
16352 mcols: usize,
16353 rp: bool,
16354 ) -> &'static str {
16355 if qtype == QT_Q8_0 {
16360 return if rp { "rp" } else { "base" };
16361 }
16362 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16363 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16364 Ok("base") => "base",
16365 Ok("pf") => "pf",
16366 Ok("r2") => "r2",
16367 Ok("r2w8") => "r2w8",
16368 Ok("pfr2") => "pfr2",
16369 Ok("ca") => "ca",
16370 Ok("car2") => "car2",
16371 Ok("rp") => "rp",
16374 Ok("rpr2") => "rpr2",
16375 Ok("rpr2w8") => "rpr2w8",
16376 Ok("rpca") => "rpca",
16379 Ok("rpcar2") => "rpcar2",
16380 Ok("rpsc") => "rpsc",
16387 Ok("rpms") => "rpms",
16388 Ok("rpmsc") => "rpmsc",
16389 Ok("rpks") => "rpks",
16390 Ok("rpksc") => "rpksc",
16391 _ => "auto",
16392 });
16393 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16397 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16402 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16403 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16404 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16405 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16406 let sms = *SMS.get_or_init(|| {
16407 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16408 self.gpu
16409 .ctx
16410 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16411 .unwrap_or(82)
16412 });
16413 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16433 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16436 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16437 Ok("base") => "base",
16438 Ok("r2") => "r2",
16439 Ok("r2w8") => "r2w8",
16440 _ => "auto",
16441 });
16442 let variant: &'static str = if qtype == QT_Q4_0 {
16443 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16447 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16448 Ok("base") => "base",
16454 Ok("r2") => "r2",
16455 Ok("ms") => "ms",
16456 Ok("sm") => "sm",
16457 Ok("la") => "la",
16458 _ => "auto",
16459 });
16460 let v = if q40 != "auto" {
16461 q40
16462 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16463 "r2"
16464 } else {
16465 "base"
16466 };
16467 if rp {
16472 match v {
16473 "ms" => "r2ms_rp",
16474 "sm" => "r2sm_rp",
16475 "la" => "r2la_rp",
16476 "r2" => "r2_rp",
16477 _ => "rp",
16478 }
16479 } else if matches!(v, "ms" | "sm" | "la") {
16480 "r2"
16481 } else {
16482 v
16483 }
16484 } else if qtype != QT_NVFP4 && !kq_r2 {
16485 "base"
16486 } else if kq_r2 && rp {
16487 "rp"
16491 } else if kq_r2 {
16492 if kq_bv != "auto" {
16495 if kq_bv == "r2w8" && mcols != 4 {
16496 "r2"
16497 } else {
16498 kq_bv
16499 }
16500 } else if bv != "auto" {
16501 match bv {
16502 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16503 "r2w8" | "rpr2w8" => {
16504 if mcols != 4 {
16505 "r2"
16506 } else {
16507 "r2w8"
16508 }
16509 }
16510 _ => "base", }
16512 } else {
16513 let blocks = (out_f + 7) / 8;
16514 let waves = blocks as f64 / (7 * sms as usize) as f64;
16515 let filled = blocks >= 4 * sms as usize;
16516 let use_r2 = if qtype == QT_Q4_K {
16517 filled
16518 } else {
16519 waves >= 2.0
16520 };
16521 if use_r2 { "r2" } else { "base" }
16522 }
16523 } else if bv != "auto" {
16524 let v = if bv == "r2w8" && mcols == 2 {
16529 "r2"
16530 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16531 "pf"
16532 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16533 "r2"
16534 } else if bv == "pfr2" && mcols == 8 {
16535 "r2"
16536 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16537 "rpr2"
16538 }
16539 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16541 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16542 } else if bv == "rpcar2" && mcols == 2 {
16543 "rpca"
16544 }
16545 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16548 "rpr2"
16549 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16550 "rpr2"
16551 } else {
16552 bv
16553 };
16554 if rp {
16555 match v {
16556 "base" | "pf" | "ca" | "rp" => "rp",
16557 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16558 "r2w8" | "rpr2w8" => {
16559 if mcols == 2 {
16560 "rpr2"
16561 } else {
16562 "rpr2w8"
16563 }
16564 }
16565 other => other, }
16567 } else {
16568 v
16569 }
16570 } else if mcols == 8 {
16571 if rp {
16582 if sc_ok { "rpsc" } else { "rpr2w8" }
16583 } else {
16584 "r2w8"
16585 }
16586 } else if mcols >= 4 {
16587 let blocks = (out_f + 7) / 8;
16591 let r7 = 7 * sms as usize;
16592 let r8 = 8 * sms as usize;
16593 let waves = blocks as f64 / r7 as f64;
16594 let filled = blocks >= 4 * sms as usize;
16595 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16599 if rp { "rpr2w8" } else { "r2w8" }
16603 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16604 if rp { "rpr2" } else { "r2" }
16607 } else {
16608 if rp { "rp" } else { "pf" }
16612 }
16613 } else if in_f >= 6144 {
16614 if rp { "rpr2" } else { "r2" }
16618 } else if rp {
16619 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16624 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16625 "rpsc"
16626 } else {
16627 "rp"
16628 }
16629 } else {
16630 "base"
16631 };
16632 variant
16633 }
16634
16635 pub fn qmatvec_mmvq_batched(
16636 &self,
16637 bytes: &CudaSlice<u8>,
16638 aq: &CudaSlice<i8>,
16639 ad: &CudaSlice<f32>,
16640 m: usize,
16641 in_f: usize,
16642 out_f: usize,
16643 qtype: i32,
16644 row_bytes: usize,
16645 mcols: usize,
16646 scale: f32,
16647 rp: bool,
16648 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16649 const ROWS_PER_BLOCK: u32 = 4;
16650 let forced: Option<&'static str> = {
16655 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16656 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16657 .as_deref()
16658 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16659 };
16660 let variant = match forced {
16661 Some(v) if !rp || v.contains("rp") => v,
16662 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16663 };
16664 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16665 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16666 })?;
16667 let variant = if mcols == 16 {
16671 if rp { "rp" } else { "base" }
16672 } else {
16673 variant
16674 };
16675 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16682 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16683 if b567
16684 && qtype == QT_NVFP4
16685 && rp
16686 && mcols == 8
16687 && (5..=7).contains(&m)
16688 && matches!(variant, "rpsc" | "rpr2w8")
16689 {
16690 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16691 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16693 let cfg = LaunchConfig {
16694 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16695 block_dim: (32, ROWS_PER_BLOCK, 1),
16696 shared_mem_bytes: 0,
16697 };
16698 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16699 let __s_b = self.gpu.stream();
16700 let mut b = __s_b.launch_builder(&f);
16701 b.arg(bytes)
16702 .arg(aq)
16703 .arg(ad)
16704 .arg(&mut y)
16705 .arg(&inf)
16706 .arg(&outf)
16707 .arg(&mi)
16708 .arg(&rb);
16709 unsafe {
16710 b.launch(cfg)?;
16711 }
16712 if scale != 1.0 {
16713 self.scale_inplace(&mut y, scale, m * out_f)?;
16714 }
16715 return Ok(y);
16716 }
16717 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16718 "base" => (base_name.into(), ROWS_PER_BLOCK),
16719 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16720 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16721 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16722 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16726 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16727 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16728 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16729 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16730 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16731 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16732 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16734 debug_assert!(
16735 !rp || name.contains("_rp"),
16736 "rp weight dispatched to a GGUF-layout kernel"
16737 );
16738 let f = self.func(&name);
16739 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16740 let smem = if name.contains("_r2sm_rp") {
16742 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16743 } else {
16744 0
16745 };
16746 let cfg = LaunchConfig {
16747 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16748 block_dim: (32, ROWS_PER_BLOCK, 1),
16749 shared_mem_bytes: smem,
16750 };
16751 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16752 let __s_b = self.gpu.stream();
16753 let mut b = __s_b.launch_builder(&f);
16754 b.arg(bytes)
16755 .arg(aq)
16756 .arg(ad)
16757 .arg(&mut y)
16758 .arg(&inf)
16759 .arg(&outf)
16760 .arg(&mi)
16761 .arg(&rb);
16762 unsafe {
16763 b.launch(cfg)?;
16764 }
16765 if scale != 1.0 {
16766 self.scale_inplace(&mut y, scale, m * out_f)?;
16767 }
16768 Ok(y)
16769 }
16770
16771 pub fn qmatvec_batched_raw(
16775 &self,
16776 bytes: &CudaSlice<u8>,
16777 x: &CudaSlice<f32>,
16778 m: usize,
16779 in_f: usize,
16780 out_f: usize,
16781 qtype: i32,
16782 row_bytes: usize,
16783 mcols: usize,
16784 rp: bool,
16785 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16786 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16787 self.qmatvec_mmvq_batched(
16788 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16789 )
16790 }
16791
16792 pub fn qmatvec_nvfp4_batched_raw(
16794 &self,
16795 bytes: &CudaSlice<u8>,
16796 x: &CudaSlice<f32>,
16797 m: usize,
16798 in_f: usize,
16799 out_f: usize,
16800 row_bytes: usize,
16801 mcols: usize,
16802 rp: bool,
16803 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16804 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16805 }
16806
16807 fn try_fp4_gemm(
16811 &self,
16812 w: &crate::model::GpuTensor,
16813 x: &CudaSlice<f32>,
16814 m: usize,
16815 in_f: usize,
16816 out_f: usize,
16817 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16818 use crate::model::GpuTensor;
16819 if cfg!(memra_portable_cuda) {
16820 return Ok(None);
16821 }
16822 if std::env::var("MEMRA_FP4").is_ok() {
16825 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16826 }
16827 if std::env::var("MEMRA_FP4").is_err() {
16828 return Ok(None);
16829 }
16830 #[cfg(memra_cutlass)]
16839 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16840 if let GpuTensor::Quant {
16841 bytes,
16842 qtype,
16843 scale,
16844 row_bytes,
16845 cutlass,
16846 ..
16847 } = w
16848 {
16849 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16850 if let Some(cw) = cutlass {
16851 let y = self.cutlass_fp4_gemm(
16853 &cw.b_packed,
16854 &cw.sfb_swizzled,
16855 x,
16856 *scale,
16857 m,
16858 out_f,
16859 in_f,
16860 )?;
16861 return Ok(Some(y));
16862 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16863 let (b_packed, sfb_sw) =
16868 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16869 let y =
16870 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16871 return Ok(Some(y));
16872 }
16873 }
16874 }
16875 }
16876 if let GpuTensor::Quant {
16877 bytes,
16878 qtype,
16879 row_bytes,
16880 scale,
16881 rp,
16882 ..
16883 } = w
16884 {
16885 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16888 let y =
16889 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16890 return Ok(Some(y));
16891 }
16892 }
16893 Ok(None)
16894 }
16895
16896 pub fn rms_norm_f16out(
16900 &self,
16901 x: &CudaSlice<f32>,
16902 w: &CudaSlice<f32>,
16903 dst: &mut CudaSlice<f32>,
16904 dst16: &mut CudaSlice<u8>,
16905 ncols: usize,
16906 nrows: usize,
16907 eps: f32,
16908 ) -> Result<(), Box<dyn std::error::Error>> {
16909 let f = self.func("rms_norm_f16out_f32");
16910 let cfg = LaunchConfig {
16911 grid_dim: (nrows as u32, 1, 1),
16912 block_dim: (rms_block(), 1, 1),
16913 shared_mem_bytes: 0,
16914 };
16915 let (nc, e) = (ncols as i32, eps);
16916 let __s_b = self.gpu.stream();
16917 let mut b = __s_b.launch_builder(&f);
16918 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16919 unsafe {
16920 b.launch(cfg)?;
16921 }
16922 Ok(())
16923 }
16924
16925 #[allow(clippy::too_many_arguments)]
16928 pub fn add_rms_norm_f16out(
16929 &self,
16930 a: &CudaSlice<f32>,
16931 b: &CudaSlice<f32>,
16932 w: &CudaSlice<f32>,
16933 res: &mut CudaSlice<f32>,
16934 dst: &mut CudaSlice<f32>,
16935 dst16: &mut CudaSlice<u8>,
16936 ncols: usize,
16937 nrows: usize,
16938 eps: f32,
16939 ) -> Result<(), Box<dyn std::error::Error>> {
16940 let f = self.func("add_rms_norm_f16out_f32");
16941 let cfg = LaunchConfig {
16942 grid_dim: (nrows as u32, 1, 1),
16943 block_dim: (rms_block(), 1, 1),
16944 shared_mem_bytes: 0,
16945 };
16946 let (nc, e) = (ncols as i32, eps);
16947 let __s_lb = self.gpu.stream();
16948 let mut lb = __s_lb.launch_builder(&f);
16949 lb.arg(a)
16950 .arg(b)
16951 .arg(w)
16952 .arg(res)
16953 .arg(dst)
16954 .arg(dst16)
16955 .arg(&nc)
16956 .arg(&e);
16957 unsafe {
16958 lb.launch(cfg)?;
16959 }
16960 Ok(())
16961 }
16962
16963 pub fn matmul_group_xh(
16966 &self,
16967 ws: &[&crate::model::GpuTensor],
16968 x: &CudaSlice<f32>,
16969 xh: &CudaSlice<u8>,
16970 m: usize,
16971 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16972 let mut out = Vec::with_capacity(ws.len());
16973 let in_f = ws[0].in_features();
16974 for w in ws {
16975 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16976 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16977 out.push(y);
16978 continue;
16979 }
16980 }
16981 out.push(self.matmul(w, x, m)?);
16982 }
16983 Ok(out)
16984 }
16985
16986 pub fn gdn_pad_mask(
16989 &self,
16990 beta: &mut CudaSlice<f32>,
16991 g_log: &mut CudaSlice<f32>,
16992 len_d: &CudaSlice<i32>,
16993 h: usize,
16994 t: usize,
16995 ) -> Result<(), Box<dyn std::error::Error>> {
16996 let f = self.func("gdn_pad_mask_f32");
16997 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16998 let (hi, ti) = (h as i32, t as i32);
16999 let __s_b = self.gpu.stream();
17000 let mut b = __s_b.launch_builder(&f);
17001 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17002 unsafe {
17003 b.launch(cfg)?;
17004 }
17005 Ok(())
17006 }
17007
17008 pub fn row_gather_dev(
17011 &self,
17012 src: &CudaSlice<f32>,
17013 dst: &mut CudaSlice<f32>,
17014 len_d: &CudaSlice<i32>,
17015 ncols: usize,
17016 ) -> Result<(), Box<dyn std::error::Error>> {
17017 let f = self.func("row_gather_dev_f32");
17018 let cfg = LaunchConfig::for_num_elems(ncols as u32);
17019 let nc = ncols as i32;
17020 let __s_b = self.gpu.stream();
17021 let mut b = __s_b.launch_builder(&f);
17022 b.arg(src).arg(dst).arg(len_d).arg(&nc);
17023 unsafe {
17024 b.launch(cfg)?;
17025 }
17026 Ok(())
17027 }
17028
17029 pub fn matmul_group(
17036 &self,
17037 ws: &[&crate::model::GpuTensor],
17038 x: &CudaSlice<f32>,
17039 m: usize,
17040 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17041 use crate::model::GpuTensor;
17042 let mut out = Vec::with_capacity(ws.len());
17043 let any_mirror = ws
17044 .iter()
17045 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17046 if m >= 16 && any_mirror && !self.verify_exact_on() {
17047 let in_f = ws[0].in_features();
17048 let xh = self.f16_act(x, m * in_f, in_f)?;
17049 for w in ws {
17050 if w.in_features() == in_f {
17051 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17052 out.push(y);
17053 continue;
17054 }
17055 }
17056 out.push(self.matmul(w, x, m)?);
17057 }
17058 return Ok(out);
17059 }
17060 for w in ws {
17061 out.push(self.matmul(w, x, m)?);
17062 }
17063 Ok(out)
17064 }
17065
17066 pub fn matmul_group_multi(
17073 &self,
17074 ws: &[&crate::model::GpuTensor],
17075 xs: &[&CudaSlice<f32>],
17076 ms: &[usize],
17077 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17078 assert_eq!(xs.len(), ms.len());
17079 let in_f = ws[0].in_features();
17080 let total: usize = ms.iter().sum();
17081 let mut xcat = self.uninit(total * in_f)?;
17082 let mut off = 0usize;
17083 for (x, &m) in xs.iter().zip(ms) {
17084 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17085 off += m;
17086 }
17087 let ys = self.matmul_group(ws, &xcat, total)?;
17088 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17089 for (w, y) in ws.iter().zip(ys) {
17090 let out_f = w.out_features();
17091 let mut off = 0usize;
17092 for (s, &m) in ms.iter().enumerate() {
17093 let mut ys_s = self.uninit(m * out_f)?;
17094 let src = y.slice(off * out_f..(off + m) * out_f);
17095 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17096 out[s].push(ys_s);
17097 off += m;
17098 }
17099 }
17100 Ok(out)
17101 }
17102
17103 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17113 use crate::model::GpuTensor;
17114 if !legacy_quant_gemm_allowed(
17115 cfg!(memra_portable_cuda),
17116 cfg!(memra_hopper_mma),
17117 std::env::var_os("MEMRA_NO_GEMM").is_some(),
17118 ) {
17119 return false;
17120 }
17121 match w {
17122 GpuTensor::Quant { qtype, .. } => {
17123 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17124 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17125 }
17126 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17127 }
17128 }
17129
17130 pub fn qmatvec_gemm(
17137 &self,
17138 w: &crate::model::GpuTensor,
17139 aq: &CudaSlice<i8>,
17140 ad: &CudaSlice<f32>,
17141 m: usize,
17142 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17143 use crate::model::GpuTensor;
17144 let in_f = w.in_features();
17145 let out_f = w.out_features();
17146 let (bytes, qtype, row_bytes, scale, rp) = match w {
17147 GpuTensor::Quant {
17148 bytes,
17149 qtype,
17150 row_bytes,
17151 scale,
17152 rp,
17153 ..
17154 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17155 _ => unreachable!("gemm_supports guaranteed Quant"),
17156 };
17157 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17163 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17164 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17165 if scale != 1.0 {
17166 self.scale_inplace(&mut y, scale, m * out_f)?;
17167 }
17168 return Ok(y);
17169 }
17170 }
17171 let name = match qtype {
17172 QT_Q8_0 => "qmatvec_gemm_q8_0",
17173 QT_Q4_K => "qmatvec_gemm_q4_K",
17174 QT_Q4_0 => {
17175 if rp {
17176 "qmatvec_gemm_q4_0_rp"
17177 } else {
17178 "qmatvec_gemm_q4_0"
17179 }
17180 }
17181 QT_Q5_K => "qmatvec_gemm_q5_K",
17182 QT_Q6_K => "qmatvec_gemm_q6_K",
17183 QT_NVFP4 => {
17184 if rp {
17185 "qmatvec_gemm_nvfp4_rp"
17186 } else {
17187 "qmatvec_gemm_nvfp4"
17188 }
17189 }
17190 _ => unreachable!(),
17191 };
17192 let f = self.func(name);
17193 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);
17198 let k1_tile = if is_k1 {
17200 k1_launch_override().unwrap_or((128, 128, 8))
17201 } else {
17202 (128, 128, 8)
17203 };
17204 let (bm, bn): (u32, u32) = if is_k1 {
17205 (k1_tile.0, k1_tile.1)
17206 } else {
17207 (64, 256)
17208 };
17209 let warps: u32 = if is_k1 {
17210 k1_tile.2
17211 } else {
17212 match qtype {
17213 QT_NVFP4 => 8,
17214 _ => 4,
17215 }
17216 };
17217 let cfg = LaunchConfig {
17218 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17219 block_dim: (32, warps, 1),
17220 shared_mem_bytes: 0,
17221 };
17222 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17223 let __s_b = self.gpu.stream();
17224 let mut b = __s_b.launch_builder(&f);
17225 b.arg(bytes)
17226 .arg(aq)
17227 .arg(ad)
17228 .arg(&mut y)
17229 .arg(&inf)
17230 .arg(&outf)
17231 .arg(&mi)
17232 .arg(&rb);
17233 unsafe {
17234 b.launch(cfg)?;
17235 }
17236 if scale != 1.0 {
17237 self.scale_inplace(&mut y, scale, m * out_f)?;
17238 }
17239 Ok(y)
17240 }
17241
17242 pub fn qmatvec_gemm_raw(
17247 &self,
17248 bytes: &CudaSlice<u8>,
17249 x: &CudaSlice<f32>,
17250 m: usize,
17251 in_f: usize,
17252 out_f: usize,
17253 qtype: i32,
17254 row_bytes: usize,
17255 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17256 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17257 let name = match qtype {
17258 QT_Q8_0 => "qmatvec_gemm_q8_0",
17259 QT_Q4_K => "qmatvec_gemm_q4_K",
17260 QT_Q4_0 => "qmatvec_gemm_q4_0",
17261 QT_Q5_K => "qmatvec_gemm_q5_K",
17262 QT_Q6_K => "qmatvec_gemm_q6_K",
17263 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17264 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17265 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17266 };
17267 let f = self.func(name);
17268 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);
17272 let k1_tile = if is_k1 {
17274 k1_launch_override().unwrap_or((128, 128, 8))
17275 } else {
17276 (128, 128, 8)
17277 };
17278 let (bm, bn): (u32, u32) = if is_k1 {
17279 (k1_tile.0, k1_tile.1)
17280 } else {
17281 (64, 256)
17282 };
17283 let warps: u32 = if is_k1 {
17284 k1_tile.2
17285 } else {
17286 match qtype {
17287 QT_NVFP4 | QT_NVFP4_RP => 8,
17288 _ => 4,
17289 }
17290 };
17291 let cfg = LaunchConfig {
17292 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17293 block_dim: (32, warps, 1),
17294 shared_mem_bytes: 0,
17295 };
17296 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17297 let __s_b = self.gpu.stream();
17298 let mut b = __s_b.launch_builder(&f);
17299 b.arg(bytes)
17300 .arg(&aq)
17301 .arg(&ad)
17302 .arg(&mut y)
17303 .arg(&inf)
17304 .arg(&outf)
17305 .arg(&mi)
17306 .arg(&rb);
17307 unsafe {
17308 b.launch(cfg)?;
17309 }
17310 Ok(y)
17311 }
17312
17313 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17320 &self,
17321 rp4: &CudaSlice<u8>,
17322 aq: &CudaSlice<i8>,
17323 ad: &CudaSlice<f32>,
17324 m: usize,
17325 in_f: usize,
17326 out_f: usize,
17327 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17328 assert!(
17329 out_f % 64 == 0 && in_f % 32 == 0,
17330 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17331 );
17332 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17333 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17335 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17336 block_dim: (128, 1, 1),
17337 shared_mem_bytes: 0,
17338 };
17339 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17340 let __s_b = self.gpu.stream();
17341 let mut b = __s_b.launch_builder(&f);
17342 b.arg(rp4)
17343 .arg(aq)
17344 .arg(ad)
17345 .arg(&mut y)
17346 .arg(&inf)
17347 .arg(&outf)
17348 .arg(&mi);
17349 unsafe {
17350 b.launch(cfg)?;
17351 }
17352 Ok(y)
17353 }
17354
17355 pub fn scale_inplace(
17357 &self,
17358 y: &mut CudaSlice<f32>,
17359 s: f32,
17360 n: usize,
17361 ) -> Result<(), Box<dyn std::error::Error>> {
17362 let f = self.func("scale_f32");
17363 let cfg = LaunchConfig::for_num_elems(n as u32);
17364 let (sf, ni) = (s, n as i32);
17365 let __s_b = self.gpu.stream();
17366 let mut b = __s_b.launch_builder(&f);
17367 b.arg(y).arg(&sf).arg(&ni);
17368 unsafe {
17369 b.launch(cfg)?;
17370 }
17371 Ok(())
17372 }
17373
17374 pub fn bf16_to_f32(
17379 &self,
17380 data: &cudarc::driver::CudaView<'_, u8>,
17381 n: usize,
17382 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17383 let mut out = self.alloc_uninit::<f32>(n)?;
17384 let f = self.func("bf16_to_f32");
17385 let cfg = LaunchConfig::for_num_elems(n as u32);
17386 let ni = n as i32;
17387 let __s_b = self.gpu.stream();
17388 let mut b = __s_b.launch_builder(&f);
17389 b.arg(data).arg(&mut out).arg(&ni);
17390 unsafe {
17391 b.launch(cfg)?;
17392 }
17393 Ok(out)
17394 }
17395
17396 fn linear_bf16_chunked(
17403 &self,
17404 x: &CudaSlice<f32>,
17405 data: &CudaSlice<u8>,
17406 m: usize,
17407 in_f: usize,
17408 out_f: usize,
17409 exact: bool,
17410 canonical_chunk_rows: Option<usize>,
17411 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17412 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17416 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17417 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17418 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17419 let started = timing.then(std::time::Instant::now);
17420 let result =
17421 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17422 if let Some(started) = started {
17423 use std::sync::atomic::Ordering;
17424 self.stream().synchronize()?;
17425 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17426 + started.elapsed().as_nanos() as u64;
17427 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17428 + (in_f * out_f * 2) as u64;
17429 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17430 if calls % 1024 == 0 {
17431 eprintln!(
17432 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17433 weight_gb={:.2}",
17434 ns as f64 / 1.0e6,
17435 ns as f64 / calls as f64 / 1.0e3,
17436 wb as f64 / 1.0e9,
17437 );
17438 }
17439 }
17440 result
17441 }
17442
17443 pub(crate) fn bf16_mmv_on() -> bool {
17448 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17449 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17450 }
17451
17452 fn matvec_bf16(
17455 &self,
17456 data: &CudaSlice<u8>,
17457 x: &CudaSlice<f32>,
17458 in_f: usize,
17459 out_f: usize,
17460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17461 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17462 return Err(format!(
17463 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17464 data.len(),
17465 x.len()
17466 )
17467 .into());
17468 }
17469 let mut y = self.alloc_uninit::<f32>(out_f)?;
17470 let f = self.func("matvec_bf16_f32acc");
17471 let cfg = LaunchConfig {
17472 grid_dim: (out_f as u32, 1, 1),
17473 block_dim: (mmv_block(), 1, 1),
17474 shared_mem_bytes: 0,
17475 };
17476 let ini = in_f as i32;
17477 let __s_bld = self.gpu.stream();
17478 let mut bld = __s_bld.launch_builder(&f);
17479 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17480 unsafe {
17481 bld.launch(cfg)?;
17482 }
17483 Ok(y)
17484 }
17485
17486 #[allow(clippy::too_many_arguments)]
17490 #[allow(clippy::too_many_arguments)]
17495 #[allow(clippy::too_many_arguments)]
17500 pub fn qk_norm_rope_append_inc_dcw_rows(
17501 &self,
17502 q_raw_t: &CudaSlice<f32>,
17503 k_raw_t: &CudaSlice<f32>,
17504 v_raw_t: &CudaSlice<f32>,
17505 qw: &CudaSlice<f32>,
17506 kw: &CudaSlice<f32>,
17507 q_out_t: &mut CudaSlice<f32>,
17508 k_out_t: &mut CudaSlice<f32>,
17509 tab: &CudaSlice<u64>,
17510 pos_t: &CudaSlice<i32>,
17511 same_session: bool,
17512 t: usize,
17513 kv_dim_k: usize,
17514 kv_dim_v: usize,
17515 k_tok_bytes: usize,
17516 v_tok_bytes: usize,
17517 head_dim: usize,
17518 n_dims: usize,
17519 nh_q: usize,
17520 nh_k: usize,
17521 eps: f32,
17522 freq_base: f32,
17523 freq_scale: f32,
17524 ff: Option<&CudaSlice<f32>>,
17525 ) -> Result<(), Box<dyn std::error::Error>> {
17526 if head_dim != 128
17527 || kv_dim_v != kv_dim_k
17528 || kv_dim_k != nh_k * head_dim
17529 || t == 0
17530 || t > 32
17531 || tab.len() < t * 6
17532 || pos_t.len() < t
17533 || q_raw_t.len() < t * nh_q * head_dim
17534 || k_raw_t.len() < t * nh_k * head_dim
17535 || v_raw_t.len() < t * kv_dim_v
17536 || q_out_t.len() < t * nh_q * head_dim
17537 || k_out_t.len() < t * nh_k * head_dim
17538 {
17539 return Err(format!(
17540 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17541 nh_q={nh_q} nh_k={nh_k}"
17542 )
17543 .into());
17544 }
17545 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17546 let same_t: i32 = if same_session { t as i32 } else { 0 };
17547 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17548 let cfg = LaunchConfig {
17549 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17550 block_dim: (128, 1, 1),
17551 shared_mem_bytes: 0,
17552 };
17553 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17554 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17555 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17556 let null: u64 = 0;
17557 let __s_b = self.gpu.stream();
17558 let mut b = __s_b.launch_builder(&f);
17559 b.arg(q_raw_t)
17560 .arg(k_raw_t)
17561 .arg(v_raw_t)
17562 .arg(qw)
17563 .arg(kw)
17564 .arg(q_out_t)
17565 .arg(k_out_t)
17566 .arg(tab)
17567 .arg(pos_t)
17568 .arg(&same_t)
17569 .arg(&kvk)
17570 .arg(&kvv)
17571 .arg(&ktb)
17572 .arg(&vtb)
17573 .arg(&hd)
17574 .arg(&nd)
17575 .arg(&nq)
17576 .arg(&nk)
17577 .arg(&eps)
17578 .arg(&theta_scale)
17579 .arg(&freq_scale);
17580 match ff {
17581 Some(freqs) => {
17582 b.arg(freqs);
17583 }
17584 None => {
17585 b.arg(&null);
17586 }
17587 }
17588 unsafe {
17589 b.launch(cfg)?;
17590 }
17591 Ok(())
17592 }
17593
17594 pub fn qk_norm_rope_append_inc_dcw(
17595 &self,
17596 q_raw: &CudaSlice<f32>,
17597 k_raw: &CudaSlice<f32>,
17598 v_raw: &CudaSlice<f32>,
17599 qw: &CudaSlice<f32>,
17600 kw: &CudaSlice<f32>,
17601 q_out: &mut CudaSlice<f32>,
17602 k_out: &mut CudaSlice<f32>,
17603 pos: &CudaSlice<i32>,
17604 k_plane: &mut CudaSlice<u8>,
17605 v_plane: &mut CudaSlice<u8>,
17606 len_dev: &CudaSlice<i32>,
17609 base_dev: Option<&CudaSlice<i32>>,
17610 done_ctr: &mut CudaSlice<u32>,
17611 kv_dim_k: usize,
17612 kv_dim_v: usize,
17613 k_tok_bytes: usize,
17614 v_tok_bytes: usize,
17615 head_dim: usize,
17616 n_dims: usize,
17617 nh_q: usize,
17618 nh_k: usize,
17619 eps: f32,
17620 freq_base: f32,
17621 freq_scale: f32,
17622 ff: Option<&CudaSlice<f32>>,
17623 ) -> Result<(), Box<dyn std::error::Error>> {
17624 if head_dim != 128
17625 || kv_dim_v != kv_dim_k
17626 || kv_dim_k != nh_k * head_dim
17627 || q_raw.len() < nh_q * head_dim
17628 || k_raw.len() < nh_k * head_dim
17629 || v_raw.len() < kv_dim_v
17630 || q_out.len() < nh_q * head_dim
17631 || k_out.len() < nh_k * head_dim
17632 || pos.is_empty()
17633 || done_ctr.is_empty()
17634 {
17635 return Err(format!(
17636 "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}"
17637 )
17638 .into());
17639 }
17640 let f = self.func("qk_norm_rope_append_inc_dcw");
17641 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17642 let cfg = LaunchConfig {
17643 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17644 block_dim: (128, 1, 1),
17645 shared_mem_bytes: 0,
17646 };
17647 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17648 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17649 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17650 let null: u64 = 0;
17651 let __s_b = self.gpu.stream();
17652 let mut b = __s_b.launch_builder(&f);
17653 b.arg(q_raw)
17654 .arg(k_raw)
17655 .arg(v_raw)
17656 .arg(qw)
17657 .arg(kw)
17658 .arg(q_out)
17659 .arg(k_out)
17660 .arg(pos)
17661 .arg(&mut *k_plane)
17662 .arg(&mut *v_plane)
17663 .arg(len_dev);
17664 match base_dev {
17665 Some(base) => {
17666 b.arg(base);
17667 }
17668 None => {
17669 b.arg(&null);
17670 }
17671 }
17672 b.arg(&mut *done_ctr)
17673 .arg(&kvk)
17674 .arg(&kvv)
17675 .arg(&ktb)
17676 .arg(&vtb)
17677 .arg(&hd)
17678 .arg(&nd)
17679 .arg(&nq)
17680 .arg(&eps)
17681 .arg(&theta_scale)
17682 .arg(&freq_scale);
17683 match ff {
17684 Some(freqs) => {
17685 b.arg(freqs);
17686 }
17687 None => {
17688 b.arg(&null);
17689 }
17690 }
17691 unsafe {
17692 b.launch(cfg)?;
17693 }
17694 Ok(())
17695 }
17696
17697 pub fn qk_norm_rope_into(
17698 &self,
17699 q_raw: &CudaSlice<f32>,
17700 k_raw: &CudaSlice<f32>,
17701 qw: &CudaSlice<f32>,
17702 kw: &CudaSlice<f32>,
17703 q_out: &mut CudaSlice<f32>,
17704 k_out: &mut CudaSlice<f32>,
17705 pos: &CudaSlice<i32>,
17706 head_dim: usize,
17707 n_dims: usize,
17708 nh_q: usize,
17709 nh_k: usize,
17710 eps: f32,
17711 freq_base: f32,
17712 freq_scale: f32,
17713 ff: Option<&CudaSlice<f32>>,
17714 ) -> Result<(), Box<dyn std::error::Error>> {
17715 if head_dim > 512
17716 || q_raw.len() < nh_q * head_dim
17717 || k_raw.len() < nh_k * head_dim
17718 || q_out.len() < nh_q * head_dim
17719 || k_out.len() < nh_k * head_dim
17720 || qw.len() < head_dim
17721 || kw.len() < head_dim
17722 || pos.is_empty()
17723 {
17724 return Err(format!(
17725 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17726 )
17727 .into());
17728 }
17729 let f = self.func("qk_norm_rope_f32");
17730 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17731 let cfg = LaunchConfig {
17732 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17733 block_dim: (128, 1, 1),
17734 shared_mem_bytes: 0,
17735 };
17736 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17737 let __s_b = self.gpu.stream();
17738 let mut b = __s_b.launch_builder(&f);
17739 b.arg(q_raw)
17740 .arg(k_raw)
17741 .arg(qw)
17742 .arg(kw)
17743 .arg(q_out)
17744 .arg(k_out)
17745 .arg(pos)
17746 .arg(&hd)
17747 .arg(&nd)
17748 .arg(&nq)
17749 .arg(&eps)
17750 .arg(&theta_scale)
17751 .arg(&freq_scale);
17752 match ff {
17753 Some(ffv) => {
17754 b.arg(ffv);
17755 unsafe {
17756 b.launch(cfg)?;
17757 }
17758 }
17759 None => {
17760 let null: u64 = 0;
17761 b.arg(&null);
17762 unsafe {
17763 b.launch(cfg)?;
17764 }
17765 }
17766 }
17767 Ok(())
17768 }
17769
17770 #[allow(clippy::too_many_arguments)]
17773 pub fn matvec_f32_b4_into(
17774 &self,
17775 w: [&CudaSlice<f32>; 4],
17776 x: &CudaSlice<f32>,
17777 y: &mut CudaSlice<f32>,
17778 block_cols: usize,
17779 out_f: usize,
17780 ) -> Result<(), Box<dyn std::error::Error>> {
17781 if block_cols % 4 != 0
17782 || x.len() < 4 * block_cols
17783 || y.len() < out_f
17784 || w.iter().any(|w| w.len() != out_f * block_cols)
17785 {
17786 return Err(format!(
17787 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17788 x.len()
17789 )
17790 .into());
17791 }
17792 let f = self.func("matvec_f32_b4");
17793 let cfg = LaunchConfig {
17794 grid_dim: (out_f as u32, 1, 1),
17795 block_dim: (128, 1, 1),
17796 shared_mem_bytes: 0,
17797 };
17798 let (bc, of) = (block_cols as i32, out_f as i32);
17799 let __s_b = self.gpu.stream();
17800 let mut b = __s_b.launch_builder(&f);
17801 b.arg(w[0])
17802 .arg(w[1])
17803 .arg(w[2])
17804 .arg(w[3])
17805 .arg(x)
17806 .arg(y)
17807 .arg(&bc)
17808 .arg(&of);
17809 unsafe {
17810 b.launch(cfg)?;
17811 }
17812 Ok(())
17813 }
17814
17815 pub fn axpy_rows_seq_into(
17818 &self,
17819 x: &CudaSlice<f32>,
17820 w: &CudaSlice<f32>,
17821 y: &mut CudaSlice<f32>,
17822 width: usize,
17823 n_rows: usize,
17824 ) -> Result<(), Box<dyn std::error::Error>> {
17825 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17826 return Err(format!(
17827 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17828 x.len(),
17829 w.len(),
17830 y.len()
17831 )
17832 .into());
17833 }
17834 let f = self.func("axpy_rows_seq_f32");
17835 let cfg = LaunchConfig::for_num_elems(width as u32);
17836 let (wi, nr) = (width as i32, n_rows as i32);
17837 let __s_b = self.gpu.stream();
17838 let mut b = __s_b.launch_builder(&f);
17839 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17840 unsafe {
17841 b.launch(cfg)?;
17842 }
17843 Ok(())
17844 }
17845
17846 #[allow(clippy::too_many_arguments)]
17850 pub fn axpy_rows_seq_md_off_into(
17851 &self,
17852 x: &CudaSlice<f32>,
17853 w_route: &CudaSlice<f32>,
17854 md: &CudaSlice<f32>,
17855 sel: &CudaSlice<i32>,
17856 y: &mut CudaSlice<f32>,
17857 width: usize,
17858 n_rows: usize,
17859 row0: usize,
17860 ) -> Result<(), Box<dyn std::error::Error>> {
17861 if x.len() < (row0 + n_rows) * width
17862 || w_route.len() < row0 + n_rows
17863 || sel.len() < row0 + n_rows
17864 || y.len() < width
17865 {
17866 return Err(format!(
17867 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17868 rows={n_rows} row0={row0}",
17869 x.len(),
17870 w_route.len(),
17871 sel.len(),
17872 y.len()
17873 )
17874 .into());
17875 }
17876 let f = self.func("axpy_rows_seq_md_off_f32");
17877 let cfg = LaunchConfig::for_num_elems(width as u32);
17878 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17879 let __s_b = self.gpu.stream();
17880 let mut b = __s_b.launch_builder(&f);
17881 b.arg(x)
17882 .arg(w_route)
17883 .arg(md)
17884 .arg(sel)
17885 .arg(y)
17886 .arg(&wi)
17887 .arg(&nr)
17888 .arg(&r0);
17889 unsafe {
17890 b.launch(cfg)?;
17891 }
17892 Ok(())
17893 }
17894
17895 #[allow(clippy::too_many_arguments)]
17900 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17901 &self,
17902 gate_bank: &CudaSlice<u8>,
17903 up_bank: &CudaSlice<u8>,
17904 sel: &CudaSlice<i32>,
17905 aq: &CudaSlice<i8>,
17906 ad: &CudaSlice<f32>,
17907 yg: &mut CudaSlice<f32>,
17908 yu: &mut CudaSlice<f32>,
17909 n_sel: usize,
17910 n_sel_col: usize,
17911 in_f: usize,
17912 out_f: usize,
17913 row_bytes: usize,
17914 expert_stride: usize,
17915 act_row_stride: usize,
17916 ad_row_stride: usize,
17917 ) -> Result<(), Box<dyn std::error::Error>> {
17918 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17919 if yg.len() < n_sel * out_f
17920 || yu.len() < n_sel * out_f
17921 || sel.len() < n_sel
17922 || n_sel_col == 0
17923 || n_sel % n_sel_col != 0
17924 {
17925 return Err("NVFP4 gu tcol geometry".into());
17926 }
17927 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17928 let cfg = LaunchConfig {
17929 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17930 block_dim: (128, 1, 1),
17931 shared_mem_bytes: 0,
17932 };
17933 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17934 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17935 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17936 let __s_b = self.gpu.stream();
17937 let mut b = __s_b.launch_builder(&f);
17938 b.arg(gate_bank)
17939 .arg(up_bank)
17940 .arg(sel)
17941 .arg(aq)
17942 .arg(ad)
17943 .arg(yg)
17944 .arg(yu)
17945 .arg(&inf)
17946 .arg(&outf)
17947 .arg(&ns)
17948 .arg(&rb)
17949 .arg(&es)
17950 .arg(&ars)
17951 .arg(&adrs)
17952 .arg(&nsc);
17953 unsafe {
17954 b.launch(cfg)?;
17955 }
17956 Ok(())
17957 }
17958
17959 #[allow(clippy::too_many_arguments)]
17962 pub fn axpy_rows_seq_md_into(
17963 &self,
17964 x: &CudaSlice<f32>,
17965 w_route: &CudaSlice<f32>,
17966 md: &CudaSlice<f32>,
17967 sel: &CudaSlice<i32>,
17968 y: &mut CudaSlice<f32>,
17969 width: usize,
17970 n_rows: usize,
17971 ) -> Result<(), Box<dyn std::error::Error>> {
17972 if x.len() < n_rows * width
17973 || w_route.len() < n_rows
17974 || sel.len() < n_rows
17975 || y.len() < width
17976 {
17977 return Err(format!(
17978 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17979 x.len(),
17980 w_route.len(),
17981 sel.len(),
17982 y.len()
17983 )
17984 .into());
17985 }
17986 let f = self.func("axpy_rows_seq_md_f32");
17987 let cfg = LaunchConfig::for_num_elems(width as u32);
17988 let (wi, nr) = (width as i32, n_rows as i32);
17989 let __s_b = self.gpu.stream();
17990 let mut b = __s_b.launch_builder(&f);
17991 b.arg(x)
17992 .arg(w_route)
17993 .arg(md)
17994 .arg(sel)
17995 .arg(y)
17996 .arg(&wi)
17997 .arg(&nr);
17998 unsafe {
17999 b.launch(cfg)?;
18000 }
18001 Ok(())
18002 }
18003
18004 #[allow(clippy::too_many_arguments)]
18006 #[allow(clippy::too_many_arguments)]
18010 pub fn matvec_bf16_qkvg_tcol_into(
18011 &self,
18012 wq: &CudaSlice<u8>,
18013 wk: &CudaSlice<u8>,
18014 wv: &CudaSlice<u8>,
18015 wg: &CudaSlice<u8>,
18016 x_t: &CudaSlice<f32>,
18017 yq: &mut CudaSlice<f32>,
18018 yk: &mut CudaSlice<f32>,
18019 yv: &mut CudaSlice<f32>,
18020 yg: &mut CudaSlice<f32>,
18021 in_f: usize,
18022 out_q: usize,
18023 out_kv: usize,
18024 out_g: usize,
18025 t: usize,
18026 ) -> Result<(), Box<dyn std::error::Error>> {
18027 if t == 0
18028 || t > 8
18029 || in_f % 8 != 0
18030 || x_t.len() < t * in_f
18031 || yq.len() < t * out_q
18032 || yk.len() < t * out_kv
18033 || yv.len() < t * out_kv
18034 || (out_g > 0 && yg.len() < t * out_g)
18035 {
18036 return Err("matvec_bf16_qkvg_tcol geometry".into());
18037 }
18038 let grid = out_q + 2 * out_kv + out_g;
18039 let cfg = LaunchConfig {
18040 grid_dim: (grid as u32, 1, 1),
18041 block_dim: (mmv_block(), 1, 1),
18042 shared_mem_bytes: 0,
18043 };
18044 let (ini, oq, okv, og, ti) = (
18045 in_f as i32,
18046 out_q as i32,
18047 out_kv as i32,
18048 out_g as i32,
18049 t as i32,
18050 );
18051 let __s_b = self.gpu.stream();
18052 let f = self.func("matvec_bf16_qkvg_tcol");
18058 let mut b = __s_b.launch_builder(&f);
18059 b.arg(wq)
18060 .arg(wk)
18061 .arg(wv)
18062 .arg(wg)
18063 .arg(x_t)
18064 .arg(yq)
18065 .arg(yk)
18066 .arg(yv)
18067 .arg(yg)
18068 .arg(&ini)
18069 .arg(&oq)
18070 .arg(&okv)
18071 .arg(&og)
18072 .arg(&ti);
18073 unsafe {
18074 b.launch(cfg)?;
18075 }
18076 Ok(())
18077 }
18078
18079 pub fn matvec_bf16_qkvg_into(
18080 &self,
18081 wq: &CudaSlice<u8>,
18082 wk: &CudaSlice<u8>,
18083 wv: &CudaSlice<u8>,
18084 wg: &CudaSlice<u8>,
18085 x: &CudaSlice<f32>,
18086 yq: &mut CudaSlice<f32>,
18087 yk: &mut CudaSlice<f32>,
18088 yv: &mut CudaSlice<f32>,
18089 yg: &mut CudaSlice<f32>,
18090 in_f: usize,
18091 out_q: usize,
18092 out_kv: usize,
18093 out_g: usize,
18094 ) -> Result<(), Box<dyn std::error::Error>> {
18095 if in_f % 8 != 0
18096 || wq.len() != out_q * in_f * 2
18097 || wk.len() != out_kv * in_f * 2
18098 || wv.len() != out_kv * in_f * 2
18099 || wg.len() < out_g * in_f * 2
18100 || x.len() < in_f
18101 || yq.len() < out_q
18102 || yk.len() < out_kv
18103 || yv.len() < out_kv
18104 || (out_g > 0 && yg.len() < out_g)
18105 {
18106 return Err(format!(
18107 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18108 )
18109 .into());
18110 }
18111 let f = self.func("matvec_bf16_qkvg");
18112 let cfg = LaunchConfig {
18113 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18114 block_dim: (mmv_block(), 1, 1),
18115 shared_mem_bytes: 0,
18116 };
18117 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18118 let __s_b = self.gpu.stream();
18119 let mut b = __s_b.launch_builder(&f);
18120 b.arg(wq)
18121 .arg(wk)
18122 .arg(wv)
18123 .arg(wg)
18124 .arg(x)
18125 .arg(yq)
18126 .arg(yk)
18127 .arg(yv)
18128 .arg(yg)
18129 .arg(&inf)
18130 .arg(&oq)
18131 .arg(&okv)
18132 .arg(&og);
18133 unsafe {
18134 b.launch(cfg)?;
18135 }
18136 Ok(())
18137 }
18138
18139 pub fn matvec_bf16_b4_into(
18141 &self,
18142 w: [&CudaSlice<u8>; 4],
18143 x: &CudaSlice<f32>,
18144 y: &mut CudaSlice<f32>,
18145 block_cols: usize,
18146 out_f: usize,
18147 ) -> Result<(), Box<dyn std::error::Error>> {
18148 if block_cols % 8 != 0
18149 || x.len() < 4 * block_cols
18150 || y.len() < out_f
18151 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18152 {
18153 return Err(format!(
18154 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18155 x.len()
18156 )
18157 .into());
18158 }
18159 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18162 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18163 let f = self.func(if x2 {
18164 "matvec_bf16_b4_x2"
18165 } else {
18166 "matvec_bf16_b4"
18167 });
18168 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18169 let cfg = LaunchConfig {
18170 grid_dim: (grid as u32, 1, 1),
18171 block_dim: (mmv_block(), 1, 1),
18172 shared_mem_bytes: 0,
18173 };
18174 let (bc, of) = (block_cols as i32, out_f as i32);
18175 let __s_b = self.gpu.stream();
18176 let mut b = __s_b.launch_builder(&f);
18177 b.arg(w[0])
18178 .arg(w[1])
18179 .arg(w[2])
18180 .arg(w[3])
18181 .arg(x)
18182 .arg(y)
18183 .arg(&bc)
18184 .arg(&of);
18185 unsafe {
18186 b.launch(cfg)?;
18187 }
18188 Ok(())
18189 }
18190
18191 pub fn matvec_bf16_b4_tcol_into(
18197 &self,
18198 w: [&CudaSlice<u8>; 4],
18199 x_t: &CudaSlice<f32>,
18200 y_t: &mut CudaSlice<f32>,
18201 block_cols: usize,
18202 out_f: usize,
18203 t: usize,
18204 ) -> Result<(), Box<dyn std::error::Error>> {
18205 if block_cols % 8 != 0
18206 || t == 0
18207 || t > 8
18208 || x_t.len() < t * 4 * block_cols
18209 || y_t.len() < t * out_f
18210 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18211 {
18212 return Err(format!(
18213 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18214 x_t.len()
18215 )
18216 .into());
18217 }
18218 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18219 return Err(
18220 "b4 tcol verify is qualified against the plain b4 kernel only \
18221 (MEMRA_B4_X2=1 is a different t=1 program)"
18222 .into(),
18223 );
18224 }
18225 let cfg = LaunchConfig {
18229 grid_dim: (out_f as u32, 1, 1),
18230 block_dim: (mmv_block(), 1, 1),
18231 shared_mem_bytes: 0,
18232 };
18233 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18234 let __s_b = self.gpu.stream();
18235 let f = self.func("matvec_bf16_b4_tcol");
18236 let mut b = __s_b.launch_builder(&f);
18237 b.arg(w[0])
18238 .arg(w[1])
18239 .arg(w[2])
18240 .arg(w[3])
18241 .arg(x_t)
18242 .arg(y_t)
18243 .arg(&bc)
18244 .arg(&of)
18245 .arg(&ti);
18246 unsafe {
18247 b.launch(cfg)?;
18248 }
18249 Ok(())
18250 }
18251
18252 pub fn q8_0_row_bytes(in_f: usize) -> usize {
18255 in_f / 32 * 34
18256 }
18257
18258 pub fn encode_q8_0_from_bf16(
18262 &self,
18263 w_bf16: &CudaSlice<u8>,
18264 out: &mut CudaSlice<u8>,
18265 in_f: usize,
18266 out_f: usize,
18267 ) -> Result<(), Box<dyn std::error::Error>> {
18268 if in_f % 32 != 0
18269 || w_bf16.len() < in_f * out_f * 2
18270 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18271 {
18272 return Err(format!(
18273 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18274 w_bf16.len(),
18275 out.len()
18276 )
18277 .into());
18278 }
18279 let f = self.func("encode_q8_0_rows_from_bf16");
18280 const PAIRS_PER_BLOCK: u32 = 4;
18283 let pairs = (out_f * (in_f / 32)) as u64;
18284 let cfg = LaunchConfig {
18285 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18286 block_dim: (32, PAIRS_PER_BLOCK, 1),
18287 shared_mem_bytes: 0,
18288 };
18289 let (ini, outi) = (in_f as i32, out_f as i32);
18290 let __s_b = self.gpu.stream();
18291 let mut b = __s_b.launch_builder(&f);
18292 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18293 unsafe {
18294 b.launch(cfg)?;
18295 }
18296 Ok(())
18297 }
18298
18299 #[allow(clippy::too_many_arguments)]
18304 pub fn qmatvec_q8_0_qkv_rp_into(
18305 &self,
18306 wq: &CudaSlice<u8>,
18307 wk: &CudaSlice<u8>,
18308 wv: &CudaSlice<u8>,
18309 aq: &CudaSlice<i8>,
18310 ad: &CudaSlice<f32>,
18311 yq: &mut CudaSlice<f32>,
18312 yk: &mut CudaSlice<f32>,
18313 yv: &mut CudaSlice<f32>,
18314 in_f: usize,
18315 out_q: usize,
18316 out_kv: usize,
18317 ) -> Result<(), Box<dyn std::error::Error>> {
18318 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
18320 let nblk = in_f / 32;
18321 if in_f % 32 != 0
18322 || aq.len() < in_f
18323 || ad.len() < nblk
18324 || yq.len() < out_q
18325 || yk.len() < out_kv
18326 || yv.len() < out_kv
18327 || wq.len() < out_q * nblk * 34
18328 || wk.len() < out_kv * nblk * 34
18329 || wv.len() < out_kv * nblk * 34
18330 {
18331 return Err(
18332 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18333 );
18334 }
18335 let f = self.func("qmatvec_q8_0_qkv_rp");
18336 let cfg = LaunchConfig {
18337 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18338 block_dim: (32, ROWS_PER_BLOCK, 1),
18339 shared_mem_bytes: 0,
18340 };
18341 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18342 let __s_b = self.gpu.stream();
18343 let mut b = __s_b.launch_builder(&f);
18344 b.arg(wq)
18345 .arg(wk)
18346 .arg(wv)
18347 .arg(aq)
18348 .arg(ad)
18349 .arg(yq)
18350 .arg(yk)
18351 .arg(yv)
18352 .arg(&ini)
18353 .arg(&oq)
18354 .arg(&okv);
18355 unsafe {
18356 b.launch(cfg)?;
18357 }
18358 Ok(())
18359 }
18360
18361 #[allow(clippy::too_many_arguments)]
18365 pub fn qmatvec_q8_0_b4_rp_into(
18366 &self,
18367 w: [&CudaSlice<u8>; 4],
18368 aq: &CudaSlice<i8>,
18369 ad: &CudaSlice<f32>,
18370 y: &mut CudaSlice<f32>,
18371 block_cols: usize,
18372 out_f: usize,
18373 ) -> Result<(), Box<dyn std::error::Error>> {
18374 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
18376 if block_cols % 32 != 0
18377 || aq.len() < 4 * block_cols
18378 || ad.len() < 4 * nblk
18379 || y.len() < out_f
18380 || w.iter().any(|p| p.len() < out_f * nblk * 34)
18381 {
18382 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18383 }
18384 let f = self.func("qmatvec_q8_0_b4_rp");
18385 let cfg = LaunchConfig {
18386 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18387 block_dim: (32, ROWS_PER_BLOCK, 1),
18388 shared_mem_bytes: 0,
18389 };
18390 let (bc, of) = (block_cols as i32, out_f as i32);
18391 let __s_b = self.gpu.stream();
18392 let mut b = __s_b.launch_builder(&f);
18393 b.arg(w[0])
18394 .arg(w[1])
18395 .arg(w[2])
18396 .arg(w[3])
18397 .arg(aq)
18398 .arg(ad)
18399 .arg(y)
18400 .arg(&bc)
18401 .arg(&of);
18402 unsafe {
18403 b.launch(cfg)?;
18404 }
18405 Ok(())
18406 }
18407
18408 fn matvec_bf16_via_q8_mirror(
18411 &self,
18412 data: &CudaSlice<u8>,
18413 x: &CudaSlice<f32>,
18414 y: &mut CudaSlice<f32>,
18415 in_f: usize,
18416 out_f: usize,
18417 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18418 use cudarc::driver::DevicePtr;
18419 let key = {
18420 let s = self.gpu.stream();
18421 let (p, _g) = data.device_ptr(&s);
18422 p as u64
18423 };
18424 {
18425 let mut mirrors = self
18426 .w8_mirrors
18427 .lock()
18428 .map_err(|_| "w8 mirror map is poisoned")?;
18429 if !mirrors.contains_key(&key) {
18430 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18431 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18432 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18433 mirrors.insert(key, planar);
18434 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18440 eprintln!(
18441 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18442 mirrors.len()
18443 );
18444 }
18445 }
18446 }
18447 let nblk = in_f / 32;
18448 {
18449 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18450 if !act.contains_key(&in_f) {
18451 let aq = self.alloc_uninit::<i8>(in_f)?;
18452 let ad = self.alloc_uninit::<f32>(nblk)?;
18453 act.insert(in_f, (aq, ad));
18454 }
18455 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18456 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18457 }
18458 let mirrors = self
18459 .w8_mirrors
18460 .lock()
18461 .map_err(|_| "w8 mirror map is poisoned")?;
18462 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18463 let mirror = mirrors.get(&key).expect("built above");
18464 let (aq, ad) = act.get(&in_f).expect("built above");
18465 self.qmatvec_mmvq_into(
18466 mirror,
18467 aq,
18468 ad,
18469 1,
18470 in_f,
18471 out_f,
18472 QT_Q8_0,
18473 Self::q8_0_row_bytes(in_f),
18474 1.0,
18475 true,
18476 y,
18477 )?;
18478 Ok(Some(()))
18479 }
18480
18481 pub fn matvec_bf16_into(
18482 &self,
18483 data: &CudaSlice<u8>,
18484 x: &CudaSlice<f32>,
18485 y: &mut CudaSlice<f32>,
18486 in_f: usize,
18487 out_f: usize,
18488 ) -> Result<(), Box<dyn std::error::Error>> {
18489 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18490 return Err(format!(
18491 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18492 data.len(),
18493 x.len(),
18494 y.len()
18495 )
18496 .into());
18497 }
18498 if step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18505 if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
18506 return Ok(());
18507 }
18508 }
18509 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18513 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18514 && in_f <= 2048;
18515 if x4 {
18516 let f = self.func("matvec_bf16_f32acc_x4");
18517 let cfg = LaunchConfig {
18518 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18519 block_dim: (mmv_block(), 1, 1),
18520 shared_mem_bytes: 0,
18521 };
18522 let (ini, outi) = (in_f as i32, out_f as i32);
18523 let __s_b = self.gpu.stream();
18524 let mut b = __s_b.launch_builder(&f);
18525 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18526 unsafe {
18527 b.launch(cfg)?;
18528 }
18529 return Ok(());
18530 }
18531 let f = self.func("matvec_bf16_f32acc");
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(data).arg(x).arg(y).arg(&ini);
18541 unsafe {
18542 b.launch(cfg)?;
18543 }
18544 Ok(())
18545 }
18546
18547 pub fn matvec_bf16_view_into(
18550 &self,
18551 data: &cudarc::driver::CudaView<'_, u8>,
18552 x: &CudaSlice<f32>,
18553 y: &mut CudaSlice<f32>,
18554 in_f: usize,
18555 out_f: usize,
18556 ) -> Result<(), Box<dyn std::error::Error>> {
18557 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18558 return Err(format!(
18559 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18560 data.len(),
18561 x.len(),
18562 y.len()
18563 )
18564 .into());
18565 }
18566 let f = self.func("matvec_bf16_f32acc");
18567 let cfg = LaunchConfig {
18568 grid_dim: (out_f as u32, 1, 1),
18569 block_dim: (mmv_block(), 1, 1),
18570 shared_mem_bytes: 0,
18571 };
18572 let ini = in_f as i32;
18573 let __s_b = self.gpu.stream();
18574 let mut b = __s_b.launch_builder(&f);
18575 b.arg(data).arg(x).arg(y).arg(&ini);
18576 unsafe {
18577 b.launch(cfg)?;
18578 }
18579 Ok(())
18580 }
18581
18582 pub fn matvec_bf16_raw_out(
18585 &self,
18586 w: &CudaSlice<u8>,
18587 x: &CudaSlice<f32>,
18588 y_raw: u64,
18589 in_f: usize,
18590 out_f: usize,
18591 ) -> Result<(), Box<dyn std::error::Error>> {
18592 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18593 return Err("matvec_bf16_raw_out geometry".into());
18594 }
18595 let f = self.func("matvec_bf16_f32acc");
18596 let cfg = LaunchConfig {
18597 grid_dim: (out_f as u32, 1, 1),
18598 block_dim: (mmv_block(), 1, 1),
18599 shared_mem_bytes: 0,
18600 };
18601 let ini = in_f as i32;
18602 let __s_b = self.gpu.stream();
18603 let mut b = __s_b.launch_builder(&f);
18604 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18605 unsafe {
18606 b.launch(cfg)?;
18607 }
18608 Ok(())
18609 }
18610
18611 pub fn add3_raw(
18615 &self,
18616 a: &CudaSlice<f32>,
18617 b: &CudaSlice<f32>,
18618 sh_raw: u64,
18619 scale_raw: u64,
18620 dst: &mut CudaSlice<f32>,
18621 n: usize,
18622 ) -> Result<(), Box<dyn std::error::Error>> {
18623 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18624 return Err("add3_raw geometry".into());
18625 }
18626 let f = self.func("add3_f32");
18627 let cfg = LaunchConfig {
18628 grid_dim: ((n as u32).div_ceil(256), 1, 1),
18629 block_dim: (256, 1, 1),
18630 shared_mem_bytes: 0,
18631 };
18632 let ni = n as i32;
18633 let __s_b = self.gpu.stream();
18634 let mut bld = __s_b.launch_builder(&f);
18635 bld.arg(a)
18636 .arg(b)
18637 .arg(&sh_raw)
18638 .arg(&scale_raw)
18639 .arg(dst)
18640 .arg(&ni);
18641 unsafe {
18642 bld.launch(cfg)?;
18643 }
18644 Ok(())
18645 }
18646
18647 pub fn matvec_bf16_down_addscale_into(
18650 &self,
18651 w: &CudaSlice<u8>,
18652 x: &CudaSlice<f32>,
18653 scale: &CudaSlice<f32>,
18654 dst: &mut CudaSlice<f32>,
18655 in_f: usize,
18656 out_f: usize,
18657 ) -> Result<(), Box<dyn std::error::Error>> {
18658 if w.len() != in_f * out_f * 2
18659 || x.len() < in_f
18660 || in_f % 8 != 0
18661 || dst.len() < out_f
18662 || scale.is_empty()
18663 {
18664 return Err("matvec_bf16_down_addscale geometry".into());
18665 }
18666 let f = self.func("matvec_bf16_down_addscale");
18667 let cfg = LaunchConfig {
18668 grid_dim: (out_f as u32, 1, 1),
18669 block_dim: (mmv_block(), 1, 1),
18670 shared_mem_bytes: 0,
18671 };
18672 let ini = in_f as i32;
18673 let __s_b = self.gpu.stream();
18674 let mut b = __s_b.launch_builder(&f);
18675 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18676 unsafe {
18677 b.launch(cfg)?;
18678 }
18679 Ok(())
18680 }
18681
18682 #[allow(clippy::too_many_arguments)]
18686 pub fn matvec_bf16_dual_silu_rows_into(
18687 &self,
18688 wg: &CudaSlice<u8>,
18689 wu: &CudaSlice<u8>,
18690 x: &CudaSlice<f32>,
18691 act: &mut CudaSlice<f32>,
18692 in_f: usize,
18693 out_f: usize,
18694 limit: Option<f32>,
18695 t: usize,
18696 ) -> Result<(), Box<dyn std::error::Error>> {
18697 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18698 return Err("matvec_bf16_dual_silu_rows geometry".into());
18699 }
18700 let f = self.func("matvec_bf16_dual_silu_rows");
18701 let cfg = LaunchConfig {
18702 grid_dim: (out_f as u32, t as u32, 1),
18703 block_dim: (mmv_block(), 1, 1),
18704 shared_mem_bytes: 0,
18705 };
18706 let (ini, outi) = (in_f as i32, out_f as i32);
18707 let lim = limit.unwrap_or(0.0);
18708 let __s_b = self.gpu.stream();
18709 let mut b = __s_b.launch_builder(&f);
18710 b.arg(wg)
18711 .arg(wu)
18712 .arg(x)
18713 .arg(&mut *act)
18714 .arg(&ini)
18715 .arg(&outi)
18716 .arg(&lim);
18717 unsafe {
18718 b.launch(cfg)?;
18719 }
18720 Ok(())
18721 }
18722
18723 pub fn matvec_bf16_rows_into(
18725 &self,
18726 w: &CudaSlice<u8>,
18727 x: &CudaSlice<f32>,
18728 y: &mut CudaSlice<f32>,
18729 in_f: usize,
18730 out_f: usize,
18731 t: usize,
18732 ) -> Result<(), Box<dyn std::error::Error>> {
18733 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18734 return Err("matvec_bf16_rows geometry".into());
18735 }
18736 if t == 1 && step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18742 if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
18743 return Ok(());
18744 }
18745 }
18746 let f = self.func("matvec_bf16_f32acc_x4_rows");
18747 let cfg = LaunchConfig {
18748 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18749 block_dim: (mmv_block(), 1, 1),
18750 shared_mem_bytes: 0,
18751 };
18752 let (ini, outi) = (in_f as i32, out_f as i32);
18753 let __s_b = self.gpu.stream();
18754 let mut b = __s_b.launch_builder(&f);
18755 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18756 unsafe {
18757 b.launch(cfg)?;
18758 }
18759 Ok(())
18760 }
18761
18762 pub fn matvec_bf16_dual_silu_into(
18763 &self,
18764 wg: &CudaSlice<u8>,
18765 wu: &CudaSlice<u8>,
18766 x: &CudaSlice<f32>,
18767 act: &mut CudaSlice<f32>,
18768 in_f: usize,
18769 out_f: usize,
18770 limit: Option<f32>,
18771 ) -> Result<(), Box<dyn std::error::Error>> {
18772 if wg.len() != in_f * out_f * 2
18773 || wu.len() != in_f * out_f * 2
18774 || x.len() < in_f
18775 || in_f % 8 != 0
18776 || act.len() < out_f
18777 {
18778 return Err("matvec_bf16_dual_silu geometry".into());
18779 }
18780 let f = self.func("matvec_bf16_dual_silu");
18781 let cfg = LaunchConfig {
18782 grid_dim: (out_f as u32, 1, 1),
18783 block_dim: (mmv_block(), 1, 1),
18784 shared_mem_bytes: 0,
18785 };
18786 let (ini, outi) = (in_f as i32, out_f as i32);
18787 let lim = limit.unwrap_or(0.0);
18788 let __s_b = self.gpu.stream();
18789 let mut b = __s_b.launch_builder(&f);
18790 b.arg(wg)
18791 .arg(wu)
18792 .arg(x)
18793 .arg(act)
18794 .arg(&ini)
18795 .arg(&outi)
18796 .arg(&lim);
18797 unsafe {
18798 b.launch(cfg)?;
18799 }
18800 Ok(())
18801 }
18802
18803 #[allow(clippy::too_many_arguments)]
18806 pub fn matvec_bf16_dual_view_into(
18807 &self,
18808 wg: &cudarc::driver::CudaView<'_, u8>,
18809 wu: &cudarc::driver::CudaView<'_, u8>,
18810 x: &CudaSlice<f32>,
18811 yg: &mut CudaSlice<f32>,
18812 yu: &mut CudaSlice<f32>,
18813 in_f: usize,
18814 out_f: usize,
18815 ) -> Result<(), Box<dyn std::error::Error>> {
18816 if wg.len() != in_f * out_f * 2
18817 || wu.len() != in_f * out_f * 2
18818 || x.len() < in_f
18819 || in_f % 8 != 0
18820 || yg.len() < out_f
18821 || yu.len() < out_f
18822 {
18823 return Err(format!(
18824 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18825 wg.len(),
18826 wu.len(),
18827 x.len()
18828 )
18829 .into());
18830 }
18831 let f = self.func("matvec_bf16_dual");
18832 let cfg = LaunchConfig {
18833 grid_dim: ((2 * out_f) as u32, 1, 1),
18834 block_dim: (mmv_block(), 1, 1),
18835 shared_mem_bytes: 0,
18836 };
18837 let (ini, outi) = (in_f as i32, out_f as i32);
18838 let __s_b = self.gpu.stream();
18839 let mut b = __s_b.launch_builder(&f);
18840 b.arg(wg)
18841 .arg(wu)
18842 .arg(x)
18843 .arg(yg)
18844 .arg(yu)
18845 .arg(&ini)
18846 .arg(&outi);
18847 unsafe {
18848 b.launch(cfg)?;
18849 }
18850 Ok(())
18851 }
18852
18853 #[allow(clippy::too_many_arguments)]
18855 pub fn matvec_bf16_dual_into(
18856 &self,
18857 wg: &CudaSlice<u8>,
18858 wu: &CudaSlice<u8>,
18859 x: &CudaSlice<f32>,
18860 yg: &mut CudaSlice<f32>,
18861 yu: &mut CudaSlice<f32>,
18862 in_f: usize,
18863 out_f: usize,
18864 ) -> Result<(), Box<dyn std::error::Error>> {
18865 if wg.len() != in_f * out_f * 2
18866 || wu.len() != in_f * out_f * 2
18867 || x.len() < in_f
18868 || in_f % 8 != 0
18869 || yg.len() < out_f
18870 || yu.len() < out_f
18871 {
18872 return Err(format!(
18873 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18874 wg.len(),
18875 wu.len(),
18876 x.len()
18877 )
18878 .into());
18879 }
18880 let f = self.func("matvec_bf16_dual");
18881 let cfg = LaunchConfig {
18882 grid_dim: ((2 * out_f) as u32, 1, 1),
18883 block_dim: (mmv_block(), 1, 1),
18884 shared_mem_bytes: 0,
18885 };
18886 let (ini, outi) = (in_f as i32, out_f as i32);
18887 let __s_b = self.gpu.stream();
18888 let mut b = __s_b.launch_builder(&f);
18889 b.arg(wg)
18890 .arg(wu)
18891 .arg(x)
18892 .arg(yg)
18893 .arg(yu)
18894 .arg(&ini)
18895 .arg(&outi);
18896 unsafe {
18897 b.launch(cfg)?;
18898 }
18899 Ok(())
18900 }
18901
18902 pub(crate) fn matvec_bf16_dual(
18905 &self,
18906 wg: &CudaSlice<u8>,
18907 wu: &CudaSlice<u8>,
18908 x: &CudaSlice<f32>,
18909 in_f: usize,
18910 out_f: usize,
18911 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18912 if wg.len() != in_f * out_f * 2
18913 || wu.len() != in_f * out_f * 2
18914 || x.len() < in_f
18915 || in_f % 8 != 0
18916 {
18917 return Err(format!(
18918 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18919 wg.len(),
18920 wu.len(),
18921 x.len()
18922 )
18923 .into());
18924 }
18925 let mut yg = self.alloc_uninit::<f32>(out_f)?;
18926 let mut yu = self.alloc_uninit::<f32>(out_f)?;
18927 let f = self.func("matvec_bf16_dual");
18928 let cfg = LaunchConfig {
18929 grid_dim: ((2 * out_f) as u32, 1, 1),
18930 block_dim: (mmv_block(), 1, 1),
18931 shared_mem_bytes: 0,
18932 };
18933 let (ini, outi) = (in_f as i32, out_f as i32);
18934 let __s_b = self.gpu.stream();
18935 let mut b = __s_b.launch_builder(&f);
18936 b.arg(wg)
18937 .arg(wu)
18938 .arg(x)
18939 .arg(&mut yg)
18940 .arg(&mut yu)
18941 .arg(&ini)
18942 .arg(&outi);
18943 unsafe {
18944 b.launch(cfg)?;
18945 }
18946 Ok((yg, yu))
18947 }
18948
18949 #[allow(clippy::too_many_arguments)]
18950 fn linear_bf16_chunked_inner(
18951 &self,
18952 x: &CudaSlice<f32>,
18953 data: &CudaSlice<u8>,
18954 m: usize,
18955 in_f: usize,
18956 out_f: usize,
18957 exact: bool,
18958 canonical_chunk_rows: Option<usize>,
18959 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18960 const CHUNK_BYTES: usize = 256 << 20;
18961 if m == 1
18964 && !exact
18965 && canonical_chunk_rows.is_none()
18966 && in_f % 8 == 0
18967 && Self::bf16_mmv_on()
18968 {
18969 return self.matvec_bf16(data, x, in_f, out_f);
18970 }
18971 let row_bytes = in_f
18972 .checked_mul(std::mem::size_of::<f32>())
18973 .ok_or("BF16 chunk row byte count overflow")?;
18974 if row_bytes == 0 || out_f == 0 {
18975 return Err("BF16 chunk dimensions must be nonzero".into());
18976 }
18977 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18978 let chunk_rows = match canonical_chunk_rows {
18979 Some(rows) if rows == 0 => {
18980 return Err("canonical BF16 chunk rows must be nonzero".into());
18981 }
18982 Some(rows) if rows > max_chunk_rows => {
18983 return Err(format!(
18984 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18985 )
18986 .into());
18987 }
18988 Some(rows) if out_f % rows != 0 => {
18989 return Err(format!(
18990 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18991 )
18992 .into());
18993 }
18994 Some(rows) => rows,
18995 None => max_chunk_rows,
18996 };
18997 if chunk_rows >= out_f {
18998 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18999 return if exact {
19000 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19001 } else {
19002 self.linear(x, &wf32, m, in_f, out_f)
19003 };
19004 }
19005 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19006 let mut r0 = 0usize;
19007 while r0 < out_f {
19008 let rows = chunk_rows.min(out_f - r0);
19009 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19010 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19011 let yc = if exact {
19012 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19013 } else {
19014 self.linear(x, &wf32, m, in_f, rows)?
19015 };
19016 for mi in 0..m {
19018 let src = yc.slice(mi * rows..(mi + 1) * rows);
19019 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19020 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19021 }
19022 r0 += rows;
19023 }
19024 Ok(y)
19025 }
19026
19027 pub fn linear_bf16_resident(
19031 &self,
19032 x: &CudaSlice<f32>,
19033 data: &CudaSlice<u8>,
19034 m: usize,
19035 in_f: usize,
19036 out_f: usize,
19037 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19038 if data.len() != in_f * out_f * 2 {
19039 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19040 }
19041 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19042 }
19043
19044 pub fn linear_bf16_resident_canonical_rows(
19050 &self,
19051 x: &CudaSlice<f32>,
19052 data: &CudaSlice<u8>,
19053 m: usize,
19054 in_f: usize,
19055 out_f: usize,
19056 canonical_chunk_rows: usize,
19057 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19058 if data.len() != in_f * out_f * 2 {
19059 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19060 }
19061 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19062 }
19063
19064 pub fn linear_f32_resident_canonical_rows(
19069 &self,
19070 x: &CudaSlice<f32>,
19071 data: &CudaSlice<f32>,
19072 m: usize,
19073 in_f: usize,
19074 out_f: usize,
19075 canonical_chunk_rows: usize,
19076 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19077 self.linear_f32_resident_canonical_rows_inner(
19078 x,
19079 data,
19080 m,
19081 in_f,
19082 out_f,
19083 canonical_chunk_rows,
19084 false,
19085 )
19086 }
19087
19088 pub fn linear_f32_resident_canonical_rows_strided(
19094 &self,
19095 x: &CudaSlice<f32>,
19096 data: &CudaSlice<f32>,
19097 m: usize,
19098 in_f: usize,
19099 out_f: usize,
19100 canonical_chunk_rows: usize,
19101 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19102 self.linear_f32_resident_canonical_rows_inner(
19103 x,
19104 data,
19105 m,
19106 in_f,
19107 out_f,
19108 canonical_chunk_rows,
19109 true,
19110 )
19111 }
19112
19113 fn linear_f32_resident_canonical_rows_inner(
19114 &self,
19115 x: &CudaSlice<f32>,
19116 data: &CudaSlice<f32>,
19117 m: usize,
19118 in_f: usize,
19119 out_f: usize,
19120 canonical_chunk_rows: usize,
19121 strided_output: bool,
19122 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19123 if data.len() != in_f * out_f {
19124 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19125 }
19126 if canonical_chunk_rows == 0
19127 || canonical_chunk_rows > out_f
19128 || out_f % canonical_chunk_rows != 0
19129 {
19130 return Err(format!(
19131 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19132 )
19133 .into());
19134 }
19135 if canonical_chunk_rows == out_f {
19136 return self.linear(x, data, m, in_f, out_f);
19137 }
19138
19139 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19140 let input = x.slice(0..x.len());
19141 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19142 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19143 if m == 1 {
19144 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19145 self.linear_device_into(
19146 &input,
19147 &weights,
19148 &mut destination,
19149 1,
19150 in_f,
19151 canonical_chunk_rows,
19152 )?;
19153 continue;
19154 }
19155 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19156 if strided_output {
19157 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19158 } else {
19159 for token in 0..m {
19160 let source = chunk
19161 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19162 let mut destination =
19163 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19164 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19165 }
19166 }
19167 }
19168 Ok(y)
19169 }
19170
19171 pub fn linear_f32_resident_canonical_rows_t1_into(
19176 &self,
19177 x: &CudaSlice<f32>,
19178 data: &CudaSlice<f32>,
19179 y: &mut CudaSlice<f32>,
19180 in_f: usize,
19181 out_f: usize,
19182 canonical_chunk_rows: usize,
19183 ) -> Result<(), Box<dyn std::error::Error>> {
19184 if data.len() != in_f * out_f {
19185 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19186 }
19187 if y.len() != out_f || x.len() != in_f {
19188 return Err(format!(
19189 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19190 x.len(),
19191 y.len()
19192 )
19193 .into());
19194 }
19195 if canonical_chunk_rows == 0
19196 || canonical_chunk_rows > out_f
19197 || out_f % canonical_chunk_rows != 0
19198 {
19199 return Err(format!(
19200 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19201 )
19202 .into());
19203 }
19204 let input = x.slice(0..x.len());
19205 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19206 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19207 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19208 self.linear_device_into(
19209 &input,
19210 &weights,
19211 &mut destination,
19212 1,
19213 in_f,
19214 canonical_chunk_rows,
19215 )?;
19216 }
19217 Ok(())
19218 }
19219
19220 pub fn linear_t1_into(
19223 &self,
19224 x: &cudarc::driver::CudaView<'_, f32>,
19225 w: &cudarc::driver::CudaView<'_, f32>,
19226 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19227 in_f: usize,
19228 out_f: usize,
19229 ) -> Result<(), Box<dyn std::error::Error>> {
19230 self.linear_device_into(x, w, y, 1, in_f, out_f)
19231 }
19232
19233 pub fn linear_decode_exact(
19240 &self,
19241 x: &CudaSlice<f32>,
19242 w: &CudaSlice<f32>,
19243 m_tokens: usize,
19244 in_f: usize,
19245 out_f: usize,
19246 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19247 if m_tokens == 1 {
19248 return self.linear(x, w, 1, in_f, out_f);
19249 }
19250 let xv = self.view(x, m_tokens * in_f);
19251 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19252 for t in 0..m_tokens {
19253 let row = xv.slice(t * in_f..(t + 1) * in_f);
19254 let mut xr = self.alloc_uninit::<f32>(in_f)?;
19255 self.copy_view_into(&mut xr, 0, &row, in_f)?;
19256 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19257 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19258 }
19259 Ok(y)
19260 }
19261
19262 pub fn linear(
19263 &self,
19264 x: &CudaSlice<f32>,
19265 w: &CudaSlice<f32>,
19266 m_tokens: usize,
19267 in_f: usize,
19268 out_f: usize,
19269 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19270 self.linear_device(x, w, m_tokens, in_f, out_f)
19271 }
19272
19273 fn linear_device<I>(
19274 &self,
19275 x: &I,
19276 w: &I,
19277 m_tokens: usize,
19278 in_f: usize,
19279 out_f: usize,
19280 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19281 where
19282 I: cudarc::driver::DevicePtr<f32>,
19283 {
19284 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)?;
19286 Ok(c)
19287 }
19288
19289 fn linear_device_into<I, O>(
19290 &self,
19291 x: &I,
19292 w: &I,
19293 c: &mut O,
19294 m_tokens: usize,
19295 in_f: usize,
19296 out_f: usize,
19297 ) -> Result<(), Box<dyn std::error::Error>>
19298 where
19299 I: cudarc::driver::DevicePtr<f32>,
19300 O: cudarc::driver::DevicePtrMut<f32>,
19301 {
19302 use cudarc::cublaslt::{Matmul, MatmulConfig};
19303 let cfg = MatmulConfig {
19304 transa: true,
19305 transb: false,
19306 transc: false,
19307 m: out_f as u64,
19308 n: m_tokens as u64,
19309 k: in_f as u64,
19310 alpha: 1.0,
19311 lda: in_f as i64,
19312 ldb: in_f as i64,
19313 beta: 0.0,
19314 ldc: out_f as i64,
19315 stride_a: None,
19316 stride_b: None,
19317 stride_c: None,
19318 stride_bias: None,
19319 batch_size: None,
19320 };
19321 let blas = self.gpu.blas();
19322 unsafe {
19323 blas.matmul(cfg, w, x, c, None, None)?;
19324 }
19325 Ok(())
19326 }
19327
19328 pub fn sdpa_naive(
19337 &self,
19338 q: &CudaSlice<f32>,
19339 k: &CudaSlice<f32>,
19340 v: &CudaSlice<f32>,
19341 o: &mut CudaSlice<f32>,
19342 head_dim: usize,
19343 n_head: usize,
19344 n_head_kv: usize,
19345 t: usize,
19346 t_kv: usize,
19347 scale: f32,
19348 causal: bool,
19349 ) -> Result<(), Box<dyn std::error::Error>> {
19350 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19351 return self.sdpa_naive_gmem(
19352 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19353 );
19354 }
19355 let f = self.func("sdpa_naive_f32");
19356 let cfg = LaunchConfig {
19357 grid_dim: (n_head as u32, t as u32, 1),
19358 block_dim: (128, 1, 1),
19359 shared_mem_bytes: (t_kv * 4) as u32,
19360 };
19361 let (hd, nh, nhkv, ti, tkvi, cz) = (
19362 head_dim as i32,
19363 n_head as i32,
19364 n_head_kv as i32,
19365 t as i32,
19366 t_kv as i32,
19367 causal as i32,
19368 );
19369 let __s_b = self.gpu.stream();
19370 let mut b = __s_b.launch_builder(&f);
19371 b.arg(q)
19372 .arg(k)
19373 .arg(v)
19374 .arg(o)
19375 .arg(&hd)
19376 .arg(&nh)
19377 .arg(&nhkv)
19378 .arg(&ti)
19379 .arg(&tkvi)
19380 .arg(&scale)
19381 .arg(&cz);
19382 unsafe {
19383 b.launch(cfg)?;
19384 }
19385 Ok(())
19386 }
19387
19388 #[allow(clippy::too_many_arguments)]
19397 pub fn sdpa_naive_gmem(
19398 &self,
19399 q: &CudaSlice<f32>,
19400 k: &CudaSlice<f32>,
19401 v: &CudaSlice<f32>,
19402 o: &mut CudaSlice<f32>,
19403 head_dim: usize,
19404 n_head: usize,
19405 n_head_kv: usize,
19406 t: usize,
19407 t_kv: usize,
19408 scale: f32,
19409 causal: bool,
19410 ) -> Result<(), Box<dyn std::error::Error>> {
19411 let ws_len = n_head
19412 .checked_mul(t)
19413 .and_then(|x| x.checked_mul(t_kv))
19414 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19415 let ws_bytes = ws_len
19416 .checked_mul(std::mem::size_of::<f32>())
19417 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19418 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19419 return Err(format!(
19420 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19421 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19422 needs a tiled/flash kernel, not the naive oracle"
19423 )
19424 .into());
19425 }
19426 let mut scores = self.uninit(ws_len)?;
19427 let f = self.func("sdpa_naive_gmem_f32");
19428 let cfg = LaunchConfig {
19429 grid_dim: (n_head as u32, t as u32, 1),
19430 block_dim: (128, 1, 1),
19431 shared_mem_bytes: 0,
19432 };
19433 let (hd, nh, nhkv, ti, tkvi, cz) = (
19434 head_dim as i32,
19435 n_head as i32,
19436 n_head_kv as i32,
19437 t as i32,
19438 t_kv as i32,
19439 causal as i32,
19440 );
19441 let __s_b = self.gpu.stream();
19442 let mut b = __s_b.launch_builder(&f);
19443 b.arg(q)
19444 .arg(k)
19445 .arg(v)
19446 .arg(o)
19447 .arg(&mut scores)
19448 .arg(&hd)
19449 .arg(&nh)
19450 .arg(&nhkv)
19451 .arg(&ti)
19452 .arg(&tkvi)
19453 .arg(&scale)
19454 .arg(&cz);
19455 unsafe {
19456 b.launch(cfg)?;
19457 }
19458 Ok(())
19459 }
19460
19461 #[allow(clippy::too_many_arguments)]
19466 pub fn sdpa_naive_island(
19467 &self,
19468 q: &CudaSlice<f32>,
19469 k: &CudaSlice<f32>,
19470 v: &CudaSlice<f32>,
19471 o: &mut CudaSlice<f32>,
19472 span_id: &CudaSlice<i32>,
19473 head_dim: usize,
19474 n_head: usize,
19475 n_head_kv: usize,
19476 t: usize,
19477 t_kv: usize,
19478 scale: f32,
19479 window: usize,
19480 ) -> Result<(), Box<dyn std::error::Error>> {
19481 let f = self.func("sdpa_naive_island_f32");
19482 let cfg = LaunchConfig {
19483 grid_dim: (n_head as u32, t as u32, 1),
19484 block_dim: (128, 1, 1),
19485 shared_mem_bytes: (t_kv * 4) as u32,
19486 };
19487 let (hd, nh, nhkv, ti, tkvi, wi) = (
19488 head_dim as i32,
19489 n_head as i32,
19490 n_head_kv as i32,
19491 t as i32,
19492 t_kv as i32,
19493 window as i32,
19494 );
19495 let __s_b = self.gpu.stream();
19496 let mut b = __s_b.launch_builder(&f);
19497 b.arg(q)
19498 .arg(k)
19499 .arg(v)
19500 .arg(o)
19501 .arg(span_id)
19502 .arg(&hd)
19503 .arg(&nh)
19504 .arg(&nhkv)
19505 .arg(&ti)
19506 .arg(&tkvi)
19507 .arg(&scale)
19508 .arg(&wi);
19509 unsafe {
19510 b.launch(cfg)?;
19511 }
19512 Ok(())
19513 }
19514
19515 #[allow(clippy::too_many_arguments)]
19517 pub fn sdpa_naive_w(
19518 &self,
19519 q: &CudaSlice<f32>,
19520 k: &CudaSlice<f32>,
19521 v: &CudaSlice<f32>,
19522 o: &mut CudaSlice<f32>,
19523 head_dim: usize,
19524 n_head: usize,
19525 n_head_kv: usize,
19526 t: usize,
19527 t_kv: usize,
19528 scale: f32,
19529 causal: bool,
19530 window: usize,
19531 ) -> Result<(), Box<dyn std::error::Error>> {
19532 let f = self.func("sdpa_naive_w_f32");
19533 let cfg = LaunchConfig {
19534 grid_dim: (n_head as u32, t as u32, 1),
19535 block_dim: (128, 1, 1),
19536 shared_mem_bytes: (t_kv * 4) as u32,
19537 };
19538 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19539 head_dim as i32,
19540 n_head as i32,
19541 n_head_kv as i32,
19542 t as i32,
19543 t_kv as i32,
19544 causal as i32,
19545 window as i32,
19546 );
19547 let __s_b = self.gpu.stream();
19548 let mut b = __s_b.launch_builder(&f);
19549 b.arg(q)
19550 .arg(k)
19551 .arg(v)
19552 .arg(o)
19553 .arg(&hd)
19554 .arg(&nh)
19555 .arg(&nhkv)
19556 .arg(&ti)
19557 .arg(&tkvi)
19558 .arg(&scale)
19559 .arg(&cz)
19560 .arg(&wi);
19561 unsafe {
19562 b.launch(cfg)?;
19563 }
19564 Ok(())
19565 }
19566
19567 #[allow(clippy::too_many_arguments)]
19577 pub fn sdpa_naive_w_lo(
19578 &self,
19579 q: &CudaSlice<f32>,
19580 k: &CudaSlice<f32>,
19581 v: &CudaSlice<f32>,
19582 o: &mut CudaSlice<f32>,
19583 head_dim: usize,
19584 n_head: usize,
19585 n_head_kv: usize,
19586 t: usize,
19587 t_kv: usize,
19588 scale: f32,
19589 causal: bool,
19590 window: usize,
19591 ) -> Result<(), Box<dyn std::error::Error>> {
19592 let kv_lo = if window > 0 {
19593 (t_kv - t + 1).saturating_sub(window)
19594 } else {
19595 0
19596 };
19597 let smem = (t_kv - kv_lo) * 4;
19598 if smem > 48 * 1024 {
19599 return Err(format!(
19600 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19601 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19602 a window this wide needs the multi-pass long-ctx kernel"
19603 )
19604 .into());
19605 }
19606 let f = self.func("sdpa_naive_w_lo_f32");
19607 let cfg = LaunchConfig {
19608 grid_dim: (n_head as u32, t as u32, 1),
19609 block_dim: (128, 1, 1),
19610 shared_mem_bytes: smem as u32,
19611 };
19612 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19613 head_dim as i32,
19614 n_head as i32,
19615 n_head_kv as i32,
19616 t as i32,
19617 t_kv as i32,
19618 causal as i32,
19619 window as i32,
19620 kv_lo as i32,
19621 );
19622 let __s_b = self.gpu.stream();
19623 let mut b = __s_b.launch_builder(&f);
19624 b.arg(q)
19625 .arg(k)
19626 .arg(v)
19627 .arg(o)
19628 .arg(&hd)
19629 .arg(&nh)
19630 .arg(&nhkv)
19631 .arg(&ti)
19632 .arg(&tkvi)
19633 .arg(&scale)
19634 .arg(&cz)
19635 .arg(&wi)
19636 .arg(&lo);
19637 unsafe {
19638 b.launch(cfg)?;
19639 }
19640 Ok(())
19641 }
19642
19643 pub fn sdpa_naive_view(
19645 &self,
19646 q: &CudaSlice<f32>,
19647 k: &cudarc::driver::CudaView<f32>,
19648 v: &cudarc::driver::CudaView<f32>,
19649 o: &mut CudaSlice<f32>,
19650 head_dim: usize,
19651 n_head: usize,
19652 n_head_kv: usize,
19653 t: usize,
19654 t_kv: usize,
19655 scale: f32,
19656 causal: bool,
19657 ) -> Result<(), Box<dyn std::error::Error>> {
19658 let f = self.func("sdpa_naive_f32");
19659 let cfg = LaunchConfig {
19660 grid_dim: (n_head as u32, t as u32, 1),
19661 block_dim: (128, 1, 1),
19662 shared_mem_bytes: (t_kv * 4) as u32,
19663 };
19664 let (hd, nh, nhkv, ti, tkvi, cz) = (
19665 head_dim as i32,
19666 n_head as i32,
19667 n_head_kv as i32,
19668 t as i32,
19669 t_kv as i32,
19670 causal as i32,
19671 );
19672 let __s_b = self.gpu.stream();
19673 let mut b = __s_b.launch_builder(&f);
19674 b.arg(q)
19675 .arg(k)
19676 .arg(v)
19677 .arg(o)
19678 .arg(&hd)
19679 .arg(&nh)
19680 .arg(&nhkv)
19681 .arg(&ti)
19682 .arg(&tkvi)
19683 .arg(&scale)
19684 .arg(&cz);
19685 unsafe {
19686 b.launch(cfg)?;
19687 }
19688 Ok(())
19689 }
19690
19691 #[allow(clippy::too_many_arguments)]
19699 pub fn fa_dequant_kv_view_f32(
19700 &self,
19701 k: &cudarc::driver::CudaView<u8>,
19702 v: &cudarc::driver::CudaView<u8>,
19703 kf: &mut CudaSlice<f32>,
19704 vf: &mut CudaSlice<f32>,
19705 kv_dim_k: usize,
19706 kv_dim_v: usize,
19707 t_kv: usize,
19708 k_tok_bytes: usize,
19709 v_tok_bytes: usize,
19710 g: bool,
19711 ) -> Result<(), Box<dyn std::error::Error>> {
19712 let f = if g {
19713 self.func_g("fa_dequant_kv_ws_f32")
19714 } else {
19715 self.func("fa_dequant_kv_ws_f32")
19716 };
19717 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19718 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19719 let cfg = LaunchConfig {
19720 grid_dim: (nblk.max(1), 1, 1),
19721 block_dim: (256, 1, 1),
19722 shared_mem_bytes: 0,
19723 };
19724 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19725 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19726 let __s_b = self.gpu.stream();
19727 let mut b = __s_b.launch_builder(&f);
19728 b.arg(k)
19729 .arg(v)
19730 .arg(&mut *kf)
19731 .arg(&mut *vf)
19732 .arg(&kdk)
19733 .arg(&kdv)
19734 .arg(&tkvi)
19735 .arg(&ktb)
19736 .arg(&vtb);
19737 unsafe {
19738 b.launch(cfg)?;
19739 }
19740 Ok(())
19741 }
19742
19743 #[allow(clippy::too_many_arguments)]
19744 pub fn sdpa_naive_quantized_view(
19745 &self,
19746 q: &CudaSlice<f32>,
19747 k: &cudarc::driver::CudaView<u8>,
19748 v: &cudarc::driver::CudaView<u8>,
19749 o: &mut CudaSlice<f32>,
19750 head_dim: usize,
19751 n_head: usize,
19752 n_head_kv: usize,
19753 t: usize,
19754 t_kv: usize,
19755 scale: f32,
19756 causal: bool,
19757 k_tok_bytes: usize,
19758 v_tok_bytes: usize,
19759 ) -> Result<(), Box<dyn std::error::Error>> {
19760 let kv_dim = n_head_kv * head_dim;
19761 let mut kf = self.uninit(t_kv * kv_dim)?;
19762 let mut vf = self.uninit(t_kv * kv_dim)?;
19763 let f = self.func("fa_dequant_kv_ws_f32");
19764 let total = (2 * t_kv * kv_dim) as u64;
19765 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19766 let cfg = LaunchConfig {
19767 grid_dim: (nblk.max(1), 1, 1),
19768 block_dim: (256, 1, 1),
19769 shared_mem_bytes: 0,
19770 };
19771 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19772 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19773 let __s_b = self.gpu.stream();
19774 let mut b = __s_b.launch_builder(&f);
19775 b.arg(k)
19776 .arg(v)
19777 .arg(&mut kf)
19778 .arg(&mut vf)
19779 .arg(&kv_dim_i)
19780 .arg(&kv_dim_i)
19781 .arg(&t_kv_i)
19782 .arg(&k_tok_bytes_i)
19783 .arg(&v_tok_bytes_i);
19784 unsafe { b.launch(cfg)? };
19785 self.sdpa_naive(
19786 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19787 )
19788 }
19789
19790 #[allow(clippy::too_many_arguments)]
19802 pub fn sdpa_naive_w_quantized_view(
19803 &self,
19804 q: &CudaSlice<f32>,
19805 k: &cudarc::driver::CudaView<u8>,
19806 v: &cudarc::driver::CudaView<u8>,
19807 o: &mut CudaSlice<f32>,
19808 head_dim: usize,
19809 n_head: usize,
19810 n_head_kv: usize,
19811 t: usize,
19812 t_kv: usize,
19813 scale: f32,
19814 causal: bool,
19815 window: usize,
19816 k_tok_bytes: usize,
19817 v_tok_bytes: usize,
19818 ) -> Result<(), Box<dyn std::error::Error>> {
19819 let kv_dim = n_head_kv * head_dim;
19820 let mut kf = self.uninit(t_kv * kv_dim)?;
19821 let mut vf = self.uninit(t_kv * kv_dim)?;
19822 let f = self.func("fa_dequant_kv_ws_f32");
19823 let total = (2 * t_kv * kv_dim) as u64;
19824 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19825 let cfg = LaunchConfig {
19826 grid_dim: (nblk.max(1), 1, 1),
19827 block_dim: (256, 1, 1),
19828 shared_mem_bytes: 0,
19829 };
19830 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19831 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19832 let __s_b = self.gpu.stream();
19833 let mut b = __s_b.launch_builder(&f);
19834 b.arg(k)
19835 .arg(v)
19836 .arg(&mut kf)
19837 .arg(&mut vf)
19838 .arg(&kv_dim_i)
19839 .arg(&kv_dim_i)
19840 .arg(&t_kv_i)
19841 .arg(&k_tok_bytes_i)
19842 .arg(&v_tok_bytes_i);
19843 unsafe { b.launch(cfg)? };
19844 self.sdpa_naive_w(
19845 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19846 )
19847 }
19848
19849 pub fn fa_prefill(
19853 &self,
19854 q: &CudaSlice<f32>,
19855 k: &CudaSlice<f32>,
19856 v: &CudaSlice<f32>,
19857 o: &mut CudaSlice<f32>,
19858 head_dim: usize,
19859 n_head: usize,
19860 n_head_kv: usize,
19861 t: usize,
19862 t_kv: usize,
19863 scale: f32,
19864 causal: bool,
19865 ) -> Result<(), Box<dyn std::error::Error>> {
19866 if portable_mma_gated() {
19867 return self.sdpa_naive(
19868 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19869 );
19870 }
19871 let fa3_on = head_dim == 256
19879 && causal
19880 && t == t_kv
19881 && match std::env::var("MEMRA_FA3").as_deref() {
19882 Ok("0") => false,
19883 Ok("1") => {
19887 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19888 true
19889 }
19890 _ => cfg!(memra_hopper_mma),
19891 };
19892 if fa3_on {
19893 let n = t * n_head * head_dim;
19894 let nkv = t * n_head_kv * head_dim;
19895 let mut q16 = self.alloc_u8_uninit(n * 2)?;
19896 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19897 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19898 self.f32_to_bf16_into(q, &mut q16, n)?;
19899 self.f32_to_bf16_into(k, &mut k16, nkv)?;
19900 self.f32_to_bf16_into(v, &mut v16, nkv)?;
19901 let rc = {
19902 use cudarc::driver::{DevicePtr, DevicePtrMut};
19903 let stream = self.gpu.stream();
19904 let (qp, _g1) = q16.device_ptr(&stream);
19905 let (kp, _g2) = k16.device_ptr(&stream);
19906 let (vp, _g3) = v16.device_ptr(&stream);
19907 let (op, _g4) = o.device_ptr_mut(&stream);
19908 unsafe {
19909 memra_fa3_prefill(
19910 qp as *const core::ffi::c_void,
19911 kp as *const core::ffi::c_void,
19912 vp as *const core::ffi::c_void,
19913 op as *mut f32,
19914 t as i32,
19915 n_head as i32,
19916 n_head_kv as i32,
19917 head_dim as i32,
19918 scale,
19919 stream.cu_stream() as *mut core::ffi::c_void,
19920 )
19921 }
19922 };
19923 if rc != 0 {
19924 return Err(format!("memra_fa3_prefill rc={rc}").into());
19925 }
19926 return Ok(());
19927 }
19928 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19933 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19934 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19935 const BLOCK_Q: usize = 64;
19936 const BKX: usize = 32;
19937 let f = self.func("fa_prefill_bf16_p1");
19938 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19939 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19940 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19941 f.set_attribute(
19942 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19943 shmem as i32,
19944 )?;
19945 let cfg = LaunchConfig {
19946 grid_dim: (
19947 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19948 n_head as u32,
19949 1,
19950 ),
19951 block_dim: (32, 4, 1),
19952 shared_mem_bytes: shmem,
19953 };
19954 let (hd, nh, nhkv, ti, tkvi, cz) = (
19955 head_dim as i32,
19956 n_head as i32,
19957 n_head_kv as i32,
19958 t as i32,
19959 t_kv as i32,
19960 causal as i32,
19961 );
19962 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19963 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19964 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19965 let __s_b = self.gpu.stream();
19966 let mut b = __s_b.launch_builder(&f);
19967 b.arg(&qb)
19968 .arg(&kb)
19969 .arg(&vb)
19970 .arg(o)
19971 .arg(&hd)
19972 .arg(&nh)
19973 .arg(&nhkv)
19974 .arg(&ti)
19975 .arg(&tkvi)
19976 .arg(&scale)
19977 .arg(&cz);
19978 unsafe {
19979 b.launch(cfg)?;
19980 }
19981 return Ok(());
19982 }
19983 const BK: usize = 32;
19989 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19992 let (block_q, warps, w2_sfx): (usize, u32, &str) =
19993 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19994 let hd_sfx = fa_hd_suffix(head_dim)?;
19998 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19999 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20004 let (kb16, vb16) = if bf16kv {
20005 let n = t_kv * n_head_kv * head_dim;
20006 let mut kb = self.alloc_u8_uninit(n * 2)?;
20007 let mut vb = self.alloc_u8_uninit(n * 2)?;
20008 let fcv = self.func("f32_to_bf16_bulk");
20009 let ni = n as i64;
20010 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20011 let __s_b = self.gpu.stream();
20012 let mut b = __s_b.launch_builder(&fcv);
20013 b.arg(k).arg(&mut kb).arg(&ni);
20014 unsafe {
20015 b.launch(cfgc)?;
20016 }
20017 let __s_b = self.gpu.stream();
20018 let mut b = __s_b.launch_builder(&fcv);
20019 b.arg(v).arg(&mut vb).arg(&ni);
20020 unsafe {
20021 b.launch(cfgc)?;
20022 }
20023 (Some(kb), Some(vb))
20024 } else {
20025 (None, None)
20026 };
20027 let f = self.func(&if bf16kv {
20028 format!("fa_prefill_bf16kv_pp{hd_sfx}")
20029 } else {
20030 format!(
20031 "fa_prefill_f32{}{}{hd_sfx}",
20032 if floor { "" } else { "_pp" },
20033 if floor { "" } else { w2_sfx }
20034 )
20035 });
20036 let kv_stages = if bf16kv { 2 } else { 1 };
20039 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20040 + 4 * (block_q * BK + 2 * block_q)) as u32;
20041 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20042 f.set_attribute(
20043 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20044 shmem as i32,
20045 )?;
20046 let cfg = LaunchConfig {
20047 grid_dim: (
20048 (t as u32 + block_q as u32 - 1) / block_q as u32,
20049 n_head as u32,
20050 1,
20051 ),
20052 block_dim: (32, warps, 1),
20053 shared_mem_bytes: shmem,
20054 };
20055 let (hd, nh, nhkv, ti, tkvi, cz) = (
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 );
20063 let __s_b = self.gpu.stream();
20064 let mut b = __s_b.launch_builder(&f);
20065 b.arg(q);
20066 match (&kb16, &vb16) {
20067 (Some(kb), Some(vb)) => {
20068 b.arg(kb).arg(vb);
20069 }
20070 _ => {
20071 b.arg(k).arg(v);
20072 }
20073 }
20074 b.arg(o)
20075 .arg(&hd)
20076 .arg(&nh)
20077 .arg(&nhkv)
20078 .arg(&ti)
20079 .arg(&tkvi)
20080 .arg(&scale)
20081 .arg(&cz);
20082 unsafe {
20083 b.launch(cfg)?;
20084 }
20085 Ok(())
20086 }
20087
20088 #[allow(clippy::too_many_arguments)]
20092 pub fn fa_prefill_w(
20093 &self,
20094 q: &CudaSlice<f32>,
20095 k: &CudaSlice<f32>,
20096 v: &CudaSlice<f32>,
20097 o: &mut CudaSlice<f32>,
20098 head_dim: usize,
20099 n_head: usize,
20100 n_head_kv: usize,
20101 t: usize,
20102 t_kv: usize,
20103 scale: f32,
20104 causal: bool,
20105 window: usize,
20106 ) -> Result<(), Box<dyn std::error::Error>> {
20107 if portable_mma_gated() {
20110 return self.sdpa_naive_w(
20111 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20112 );
20113 }
20114 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20118 let faw_f32 =
20119 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20120 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20121 self.fa_prefill_w_arm(
20122 q,
20123 k,
20124 v,
20125 o,
20126 head_dim,
20127 n_head,
20128 n_head_kv,
20129 t,
20130 t_kv,
20131 scale,
20132 causal,
20133 window,
20134 floor || faw_f32,
20135 floor,
20136 )
20137 }
20138
20139 #[allow(clippy::too_many_arguments)]
20142 pub fn fa_prefill_w_pre(
20143 &self,
20144 qb: &CudaSlice<u8>,
20145 kb: &CudaSlice<u8>,
20146 vb: &CudaSlice<u8>,
20147 o: &mut CudaSlice<f32>,
20148 head_dim: usize,
20149 n_head: usize,
20150 n_head_kv: usize,
20151 t: usize,
20152 t_kv: usize,
20153 scale: f32,
20154 causal: bool,
20155 window: usize,
20156 v_f16: bool,
20157 ) -> Result<(), Box<dyn std::error::Error>> {
20158 const BLOCK_Q: usize = 64;
20159 const BK: usize = 32;
20160 debug_assert_eq!(head_dim, 256);
20161 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20162 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20163 if hp {
20164 const BLOCK_QH: usize = 32;
20165 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20168 let vh: &CudaSlice<u8> = if v_f16 {
20169 vb
20170 } else {
20171 let n = t_kv * n_head_kv * head_dim;
20172 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20173 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20174 }
20175 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20176 vguard.as_ref().unwrap()
20177 };
20178 let f = self.func("fa_prefill_w_bf16_p1h2");
20179 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20180 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20181 f.set_attribute(
20182 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20183 shmem as i32,
20184 )?;
20185 let cfg = LaunchConfig {
20186 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20187 block_dim: (32, 4, 1),
20188 shared_mem_bytes: shmem,
20189 };
20190 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20191 head_dim as i32,
20192 n_head as i32,
20193 n_head_kv as i32,
20194 t as i32,
20195 t_kv as i32,
20196 causal as i32,
20197 window as i32,
20198 );
20199 let __s_b = self.gpu.stream();
20200 let mut b = __s_b.launch_builder(&f);
20201 b.arg(qb)
20202 .arg(kb)
20203 .arg(vh)
20204 .arg(o)
20205 .arg(&hd)
20206 .arg(&nh)
20207 .arg(&nhkv)
20208 .arg(&ti)
20209 .arg(&tkvi)
20210 .arg(&scale)
20211 .arg(&cz)
20212 .arg(&wi);
20213 unsafe {
20214 b.launch(cfg)?;
20215 }
20216 return Ok(());
20217 }
20218 let f = self.func("fa_prefill_w_bf16_p1");
20219 let shmem =
20220 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20221 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20222 f.set_attribute(
20223 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20224 shmem as i32,
20225 )?;
20226 let cfg = LaunchConfig {
20227 grid_dim: (
20228 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20229 n_head as u32,
20230 1,
20231 ),
20232 block_dim: (32, 4, 1),
20233 shared_mem_bytes: shmem,
20234 };
20235 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20236 head_dim as i32,
20237 n_head as i32,
20238 n_head_kv as i32,
20239 t as i32,
20240 t_kv as i32,
20241 causal as i32,
20242 window as i32,
20243 );
20244 let __s_b = self.gpu.stream();
20245 let mut b = __s_b.launch_builder(&f);
20246 b.arg(qb)
20247 .arg(kb)
20248 .arg(vb)
20249 .arg(o)
20250 .arg(&hd)
20251 .arg(&nh)
20252 .arg(&nhkv)
20253 .arg(&ti)
20254 .arg(&tkvi)
20255 .arg(&scale)
20256 .arg(&cz)
20257 .arg(&wi);
20258 unsafe {
20259 b.launch(cfg)?;
20260 }
20261 Ok(())
20262 }
20263
20264 #[allow(clippy::too_many_arguments)]
20266 pub fn fa_prefill_w_arm(
20267 &self,
20268 q: &CudaSlice<f32>,
20269 k: &CudaSlice<f32>,
20270 v: &CudaSlice<f32>,
20271 o: &mut CudaSlice<f32>,
20272 head_dim: usize,
20273 n_head: usize,
20274 n_head_kv: usize,
20275 t: usize,
20276 t_kv: usize,
20277 scale: f32,
20278 causal: bool,
20279 window: usize,
20280 f32_stage: bool,
20281 floor: bool,
20282 ) -> Result<(), Box<dyn std::error::Error>> {
20283 const BLOCK_Q: usize = 64;
20284 const BK: usize = 32;
20285 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20286 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20290 let p1 = !floor
20291 && !f32_stage
20292 && *P1_ON.get_or_init(|| {
20293 std::env::var("MEMRA_FAW_P1")
20294 .map(|v| v != "0")
20295 .unwrap_or(true)
20296 });
20297 let hp =
20298 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20299 if hp {
20300 const BLOCK_QH: usize = 32;
20301 let f = self.func("fa_prefill_w_bf16_p1h2");
20302 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20303 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20304 f.set_attribute(
20305 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20306 shmem as i32,
20307 )?;
20308 let cfg = LaunchConfig {
20309 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20310 block_dim: (32, 4, 1),
20311 shared_mem_bytes: shmem,
20312 };
20313 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20314 head_dim as i32,
20315 n_head as i32,
20316 n_head_kv as i32,
20317 t as i32,
20318 t_kv as i32,
20319 causal as i32,
20320 window as i32,
20321 );
20322 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20323 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20324 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20325 let __s_b = self.gpu.stream();
20326 let mut b = __s_b.launch_builder(&f);
20327 b.arg(&qb)
20328 .arg(&kb)
20329 .arg(&vh)
20330 .arg(o)
20331 .arg(&hd)
20332 .arg(&nh)
20333 .arg(&nhkv)
20334 .arg(&ti)
20335 .arg(&tkvi)
20336 .arg(&scale)
20337 .arg(&cz)
20338 .arg(&wi);
20339 unsafe {
20340 b.launch(cfg)?;
20341 }
20342 return Ok(());
20343 }
20344 if p1 {
20345 let f = self.func("fa_prefill_w_bf16_p1");
20346 let shmem =
20347 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20348 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20349 f.set_attribute(
20350 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20351 shmem as i32,
20352 )?;
20353 let cfg = LaunchConfig {
20354 grid_dim: (
20355 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20356 n_head as u32,
20357 1,
20358 ),
20359 block_dim: (32, 4, 1),
20360 shared_mem_bytes: shmem,
20361 };
20362 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20363 head_dim as i32,
20364 n_head as i32,
20365 n_head_kv as i32,
20366 t as i32,
20367 t_kv as i32,
20368 causal as i32,
20369 window as i32,
20370 );
20371 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20372 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20373 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20374 let __s_b = self.gpu.stream();
20375 let mut b = __s_b.launch_builder(&f);
20376 b.arg(&qb)
20377 .arg(&kb)
20378 .arg(&vb)
20379 .arg(o)
20380 .arg(&hd)
20381 .arg(&nh)
20382 .arg(&nhkv)
20383 .arg(&ti)
20384 .arg(&tkvi)
20385 .arg(&scale)
20386 .arg(&cz)
20387 .arg(&wi);
20388 unsafe {
20389 b.launch(cfg)?;
20390 }
20391 return Ok(());
20392 }
20393 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20396 let g4 = !floor
20397 && !f32_stage
20398 && n_head_kv == 1
20399 && n_head % 4 == 0
20400 && *G4_ON.get_or_init(|| {
20401 std::env::var("MEMRA_FAW_G4")
20402 .map(|v| v != "0")
20403 .unwrap_or(true)
20404 });
20405 if g4 {
20406 const SP_M: usize = 16;
20407 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20410 let o2 = *O2_ON.get_or_init(|| {
20411 std::env::var("MEMRA_FAW_O2")
20412 .map(|v| v != "0")
20413 .unwrap_or(true)
20414 });
20415 let f = self.func(if o2 {
20416 "fa_prefill_w_bf16_g4o2"
20417 } else {
20418 "fa_prefill_w_bf16_g4"
20419 });
20420 let shmem = if o2 {
20421 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20422 } else {
20423 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20424 as u32
20425 };
20426 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20427 f.set_attribute(
20428 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20429 shmem as i32,
20430 )?;
20431 let cfg = LaunchConfig {
20432 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20433 block_dim: (32, 4, 1),
20434 shared_mem_bytes: shmem,
20435 };
20436 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20437 head_dim as i32,
20438 n_head as i32,
20439 n_head_kv as i32,
20440 t as i32,
20441 t_kv as i32,
20442 causal as i32,
20443 window as i32,
20444 );
20445 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20446 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20447 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20448 let __s_b = self.gpu.stream();
20449 let mut b = __s_b.launch_builder(&f);
20450 b.arg(&qb)
20451 .arg(&kb)
20452 .arg(&vb)
20453 .arg(o)
20454 .arg(&hd)
20455 .arg(&nh)
20456 .arg(&nhkv)
20457 .arg(&ti)
20458 .arg(&tkvi)
20459 .arg(&scale)
20460 .arg(&cz)
20461 .arg(&wi);
20462 unsafe {
20463 b.launch(cfg)?;
20464 }
20465 return Ok(());
20466 }
20467 let f = self.func(if floor {
20468 "fa_prefill_w_f32"
20469 } else if f32_stage {
20470 "fa_prefill_w_f32_pp"
20471 } else {
20472 "fa_prefill_w_bf16_pp"
20473 });
20474 let shmem =
20475 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20476 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20477 f.set_attribute(
20478 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20479 shmem as i32,
20480 )?;
20481 let cfg = LaunchConfig {
20482 grid_dim: (
20483 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20484 n_head as u32,
20485 1,
20486 ),
20487 block_dim: (32, 4, 1),
20488 shared_mem_bytes: shmem,
20489 };
20490 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20491 head_dim as i32,
20492 n_head as i32,
20493 n_head_kv as i32,
20494 t as i32,
20495 t_kv as i32,
20496 causal as i32,
20497 window as i32,
20498 );
20499 if f32_stage {
20500 let __s_b = self.gpu.stream();
20501 let mut b = __s_b.launch_builder(&f);
20502 b.arg(q)
20503 .arg(k)
20504 .arg(v)
20505 .arg(o)
20506 .arg(&hd)
20507 .arg(&nh)
20508 .arg(&nhkv)
20509 .arg(&ti)
20510 .arg(&tkvi)
20511 .arg(&scale)
20512 .arg(&cz)
20513 .arg(&wi);
20514 unsafe {
20515 b.launch(cfg)?;
20516 }
20517 } else {
20518 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20519 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20520 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20521 let __s_b = self.gpu.stream();
20522 let mut b = __s_b.launch_builder(&f);
20523 b.arg(&qb)
20524 .arg(&kb)
20525 .arg(&vb)
20526 .arg(o)
20527 .arg(&hd)
20528 .arg(&nh)
20529 .arg(&nhkv)
20530 .arg(&ti)
20531 .arg(&tkvi)
20532 .arg(&scale)
20533 .arg(&cz)
20534 .arg(&wi);
20535 unsafe {
20536 b.launch(cfg)?;
20537 }
20538 }
20539 Ok(())
20540 }
20541
20542 #[allow(clippy::too_many_arguments)]
20546 pub fn fa_prefill_hd512(
20547 &self,
20548 q: &CudaSlice<f32>,
20549 k: &CudaSlice<f32>,
20550 v: &CudaSlice<f32>,
20551 o: &mut CudaSlice<f32>,
20552 head_dim: usize,
20553 n_head: usize,
20554 n_head_kv: usize,
20555 t: usize,
20556 t_kv: usize,
20557 scale: f32,
20558 causal: bool,
20559 ) -> Result<(), Box<dyn std::error::Error>> {
20560 if portable_mma_gated() {
20562 return self.sdpa_naive(
20563 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20564 );
20565 }
20566 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20572 let f32_stage =
20573 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20574 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20578 let sp = !f32_stage
20579 && *SP_ON.get_or_init(|| {
20580 std::env::var("MEMRA_FA512_SP")
20581 .map(|v| v != "0")
20582 .unwrap_or(true)
20583 });
20584 self.fa_prefill_hd512_arm(
20585 q,
20586 k,
20587 v,
20588 o,
20589 head_dim,
20590 n_head,
20591 n_head_kv,
20592 t,
20593 t_kv,
20594 scale,
20595 causal,
20596 f32_stage,
20597 sp,
20598 sp && fa_f16pv_on(),
20599 )
20600 }
20601
20602 #[allow(clippy::too_many_arguments)]
20604 pub fn fa_prefill_hd512_pre(
20605 &self,
20606 qb: &CudaSlice<u8>,
20607 kb: &CudaSlice<u8>,
20608 vb: &CudaSlice<u8>,
20609 o: &mut CudaSlice<f32>,
20610 head_dim: usize,
20611 n_head: usize,
20612 n_head_kv: usize,
20613 t: usize,
20614 t_kv: usize,
20615 scale: f32,
20616 causal: bool,
20617 v_f16: bool,
20618 ) -> Result<(), Box<dyn std::error::Error>> {
20619 debug_assert_eq!(head_dim, 512);
20620 const SP_M: usize = 16;
20621 const BKS: usize = 32;
20622 let f16pv = fa_f16pv_on();
20626 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20627 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20628 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20629 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20630 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20631 let n = t_kv * n_head_kv * head_dim;
20633 let need = n * 2;
20634 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20635 *vguard = Some(self.alloc_uninit::<u8>(need)?);
20636 }
20637 let dst = vguard.as_mut().unwrap();
20638 self.bf16_to_f16_into(vb, n, dst)?;
20639 vguard.as_ref().unwrap()
20640 } else {
20641 vb
20642 };
20643 let f = self.func(if hp {
20644 "fa_prefill_bf16_hd512_sp16h2"
20645 } else {
20646 match (f16pv, nw) {
20647 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20648 (true, _) => "fa_prefill_bf16_hd512_sp16",
20649 _ => "fa_prefill_bf16_hd512_sp",
20650 }
20651 });
20652 let (nwarp, npart) = if hp {
20653 (4usize, 4usize)
20654 } else if nw > 2 {
20655 (nw, nw)
20656 } else {
20657 (2, 1)
20658 };
20659 let shmem = if hp {
20661 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20662 as u32
20663 } else {
20664 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20665 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20666 };
20667 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20668 f.set_attribute(
20669 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20670 shmem as i32,
20671 )?;
20672 let grid_y = if hp {
20673 (n_head / 2) as u32
20674 } else {
20675 n_head as u32
20676 };
20677 let cfg = LaunchConfig {
20678 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20679 block_dim: (32, nwarp as u32, 1),
20680 shared_mem_bytes: shmem,
20681 };
20682 let (hd, nh, nhkv, ti, tkvi, cz) = (
20683 head_dim as i32,
20684 n_head as i32,
20685 n_head_kv as i32,
20686 t as i32,
20687 t_kv as i32,
20688 causal as i32,
20689 );
20690 let __s_b = self.gpu.stream();
20691 let mut b = __s_b.launch_builder(&f);
20692 b.arg(qb)
20693 .arg(kb)
20694 .arg(vref)
20695 .arg(o)
20696 .arg(&hd)
20697 .arg(&nh)
20698 .arg(&nhkv)
20699 .arg(&ti)
20700 .arg(&tkvi)
20701 .arg(&scale)
20702 .arg(&cz);
20703 unsafe {
20704 b.launch(cfg)?;
20705 }
20706 Ok(())
20707 }
20708
20709 #[allow(clippy::too_many_arguments)]
20712 pub fn fa_prefill_hd512_arm(
20713 &self,
20714 q: &CudaSlice<f32>,
20715 k: &CudaSlice<f32>,
20716 v: &CudaSlice<f32>,
20717 o: &mut CudaSlice<f32>,
20718 head_dim: usize,
20719 n_head: usize,
20720 n_head_kv: usize,
20721 t: usize,
20722 t_kv: usize,
20723 scale: f32,
20724 causal: bool,
20725 f32_stage: bool,
20726 sp: bool,
20727 f16pv: bool,
20728 ) -> Result<(), Box<dyn std::error::Error>> {
20729 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20730 if sp && !f32_stage {
20731 const SP_M: usize = 16;
20735 const BKS: usize = 32;
20736 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20737 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20738 let f = self.func(if hp {
20739 "fa_prefill_bf16_hd512_sp16h2"
20740 } else {
20741 match (f16pv, nw) {
20742 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20743 (true, _) => "fa_prefill_bf16_hd512_sp16",
20744 _ => "fa_prefill_bf16_hd512_sp",
20745 }
20746 });
20747 let (nwarp, npart) = if hp {
20748 (4usize, 4usize)
20749 } else if nw > 2 {
20750 (nw, nw)
20751 } else {
20752 (2, 1)
20753 };
20754 let shmem = if hp {
20755 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20756 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20757 } else {
20758 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20759 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20760 };
20761 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20762 f.set_attribute(
20763 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20764 shmem as i32,
20765 )?;
20766 let grid_y = if hp {
20767 (n_head / 2) as u32
20768 } else {
20769 n_head as u32
20770 };
20771 let cfg = LaunchConfig {
20772 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20773 block_dim: (32, nwarp as u32, 1),
20774 shared_mem_bytes: shmem,
20775 };
20776 let (hd, nh, nhkv, ti, tkvi, cz) = (
20777 head_dim as i32,
20778 n_head as i32,
20779 n_head_kv as i32,
20780 t as i32,
20781 t_kv as i32,
20782 causal as i32,
20783 );
20784 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20785 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20786 let vb = if f16pv {
20787 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20788 } else {
20789 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20790 };
20791 let __s_b = self.gpu.stream();
20792 let mut b = __s_b.launch_builder(&f);
20793 b.arg(&qb)
20794 .arg(&kb)
20795 .arg(&vb)
20796 .arg(o)
20797 .arg(&hd)
20798 .arg(&nh)
20799 .arg(&nhkv)
20800 .arg(&ti)
20801 .arg(&tkvi)
20802 .arg(&scale)
20803 .arg(&cz);
20804 unsafe {
20805 b.launch(cfg)?;
20806 }
20807 return Ok(());
20808 }
20809 const BLOCK_Q: usize = 32;
20810 const BK: usize = 32;
20811 const HALF: usize = 256;
20812 let f = self.func(if f32_stage {
20813 "fa_prefill_f32_hd512"
20814 } else {
20815 "fa_prefill_bf16_hd512"
20816 });
20817 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20819 + 4 * BLOCK_Q) as u32;
20820 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20821 f.set_attribute(
20822 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20823 shmem as i32,
20824 )?;
20825 let cfg = LaunchConfig {
20826 grid_dim: (
20827 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20828 n_head as u32,
20829 2,
20830 ),
20831 block_dim: (32, 2, 1),
20832 shared_mem_bytes: shmem,
20833 };
20834 let (hd, nh, nhkv, ti, tkvi, cz) = (
20835 head_dim as i32,
20836 n_head as i32,
20837 n_head_kv as i32,
20838 t as i32,
20839 t_kv as i32,
20840 causal as i32,
20841 );
20842 if f32_stage {
20843 let __s_b = self.gpu.stream();
20844 let mut b = __s_b.launch_builder(&f);
20845 b.arg(q)
20846 .arg(k)
20847 .arg(v)
20848 .arg(o)
20849 .arg(&hd)
20850 .arg(&nh)
20851 .arg(&nhkv)
20852 .arg(&ti)
20853 .arg(&tkvi)
20854 .arg(&scale)
20855 .arg(&cz);
20856 unsafe {
20857 b.launch(cfg)?;
20858 }
20859 } else {
20860 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20861 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20862 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20863 let __s_b = self.gpu.stream();
20864 let mut b = __s_b.launch_builder(&f);
20865 b.arg(&qb)
20866 .arg(&kb)
20867 .arg(&vb)
20868 .arg(o)
20869 .arg(&hd)
20870 .arg(&nh)
20871 .arg(&nhkv)
20872 .arg(&ti)
20873 .arg(&tkvi)
20874 .arg(&scale)
20875 .arg(&cz);
20876 unsafe {
20877 b.launch(cfg)?;
20878 }
20879 }
20880 Ok(())
20881 }
20882
20883 #[allow(clippy::too_many_arguments)]
20887 pub fn rope_neox2_bf16e(
20888 &self,
20889 q: &mut CudaSlice<f32>,
20890 k: &mut CudaSlice<f32>,
20891 qb: &mut CudaSlice<u8>,
20892 kb: &mut CudaSlice<u8>,
20893 pos: &CudaSlice<i32>,
20894 head_dim: usize,
20895 n_dims: usize,
20896 nh_q: usize,
20897 nh_k: usize,
20898 n_tokens: usize,
20899 base: f32,
20900 freq_scale: f32,
20901 ff: Option<&CudaSlice<f32>>,
20902 ) -> Result<(), Box<dyn std::error::Error>> {
20903 let f = self.func("rope_neox2_bf16e_f32");
20904 let rows = ((nh_q + nh_k) * n_tokens) as u32;
20905 let cfg = LaunchConfig {
20906 grid_dim: (rows, 1, 1),
20907 block_dim: ((head_dim / 2) as u32, 1, 1),
20908 shared_mem_bytes: 0,
20909 };
20910 let theta_scale = base.powf(-2.0 / n_dims as f32);
20911 let (hd, nd, nhq, nhk, nt) = (
20912 head_dim as i32,
20913 n_dims as i32,
20914 nh_q as i32,
20915 nh_k as i32,
20916 n_tokens as i32,
20917 );
20918 let __s_b = self.gpu.stream();
20919 let mut b = __s_b.launch_builder(&f);
20920 match ff {
20921 Some(t) => {
20922 b.arg(&mut *q)
20923 .arg(&mut *k)
20924 .arg(&mut *qb)
20925 .arg(&mut *kb)
20926 .arg(pos)
20927 .arg(&hd)
20928 .arg(&nd)
20929 .arg(&nhq)
20930 .arg(&nhk)
20931 .arg(&nt)
20932 .arg(&theta_scale)
20933 .arg(&freq_scale)
20934 .arg(t);
20935 unsafe {
20936 b.launch(cfg)?;
20937 }
20938 }
20939 None => {
20940 let null: u64 = 0;
20941 b.arg(&mut *q)
20942 .arg(&mut *k)
20943 .arg(&mut *qb)
20944 .arg(&mut *kb)
20945 .arg(pos)
20946 .arg(&hd)
20947 .arg(&nd)
20948 .arg(&nhq)
20949 .arg(&nhk)
20950 .arg(&nt)
20951 .arg(&theta_scale)
20952 .arg(&freq_scale)
20953 .arg(&null);
20954 unsafe {
20955 b.launch(cfg)?;
20956 }
20957 }
20958 }
20959 Ok(())
20960 }
20961
20962 pub fn f32_to_bf16(
20965 &self,
20966 x: &CudaSlice<f32>,
20967 n: usize,
20968 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20969 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
20970 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20971 let f = self.func("f32_to_bf16_flat");
20972 let n_i = n as i64;
20973 let cfg = LaunchConfig {
20974 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20975 block_dim: (256, 1, 1),
20976 shared_mem_bytes: 0,
20977 };
20978 let __s_b = self.gpu.stream();
20979 let mut b = __s_b.launch_builder(&f);
20980 b.arg(x).arg(&mut y).arg(&n_i);
20981 unsafe {
20982 b.launch(cfg)?;
20983 }
20984 Ok(y)
20985 }
20986
20987 pub fn f32_to_f16(
20988 &self,
20989 x: &CudaSlice<f32>,
20990 n: usize,
20991 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20992 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20993 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20994 let f = self.func("f32_to_f16_flat");
20995 let n_i = n as i64;
20996 let cfg = LaunchConfig {
20997 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20998 block_dim: (256, 1, 1),
20999 shared_mem_bytes: 0,
21000 };
21001 let __s_b = self.gpu.stream();
21002 let mut b = __s_b.launch_builder(&f);
21003 b.arg(x).arg(&mut y).arg(&n_i);
21004 unsafe {
21005 b.launch(cfg)?;
21006 }
21007 Ok(y)
21008 }
21009
21010 pub fn bf16_to_f16(
21012 &self,
21013 xb: &CudaSlice<u8>,
21014 n: usize,
21015 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21016 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21017 self.bf16_to_f16_into(xb, n, &mut y)?;
21018 Ok(y)
21019 }
21020
21021 pub fn bf16_to_f16_into(
21023 &self,
21024 xb: &CudaSlice<u8>,
21025 n: usize,
21026 y: &mut CudaSlice<u8>,
21027 ) -> Result<(), Box<dyn std::error::Error>> {
21028 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21029 assert!(y.len() >= n * 2);
21030 let f = self.func("bf16_to_f16_flat");
21031 let n2 = (n / 2) as i64;
21032 let cfg = LaunchConfig {
21033 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21034 block_dim: (256, 1, 1),
21035 shared_mem_bytes: 0,
21036 };
21037 let __s_b = self.gpu.stream();
21038 let mut b = __s_b.launch_builder(&f);
21039 b.arg(xb).arg(y).arg(&n2);
21040 unsafe {
21041 b.launch(cfg)?;
21042 }
21043 Ok(())
21044 }
21045
21046 #[allow(clippy::too_many_arguments)]
21051 pub fn fa_prefill_vl8(
21052 &self,
21053 seqs: &[FaSeqVl],
21054 head_dim: usize,
21055 n_head: usize,
21056 n_head_kv: usize,
21057 scale: f32,
21058 ) -> Result<(), Box<dyn std::error::Error>> {
21059 const BK: usize = 32;
21060 let b = seqs.len();
21061 assert!(b >= 1 && b <= 8);
21062 let mut packed = [FaSeqVl::default(); 8];
21063 packed[..b].copy_from_slice(seqs);
21064 let v = FaVl8(packed);
21065 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21066 let ept = (n_head_kv * head_dim) as i32;
21067 {
21068 let f = self.func("fa_mirror_vl");
21069 let max_n = (max_t as i64) * ept as i64;
21070 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21071 for which in 0..2i32 {
21072 let cfg = LaunchConfig {
21073 grid_dim: (blocks, 1, b as u32),
21074 block_dim: (256, 1, 1),
21075 shared_mem_bytes: 0,
21076 };
21077 let __s_lb = self.gpu.stream();
21078 let mut lb = __s_lb.launch_builder(&f);
21079 lb.arg(&v).arg(&ept).arg(&which);
21080 unsafe {
21081 lb.launch(cfg)?;
21082 }
21083 }
21084 }
21085 let hd_sfx = fa_hd_suffix(head_dim)?;
21086 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21087 let block_q = 64usize;
21088 let kv_stages = 2usize;
21089 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21090 + 4 * (block_q * BK + 2 * block_q)) as u32;
21091 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21092 f.set_attribute(
21093 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21094 shmem as i32,
21095 )?;
21096 let cfg = LaunchConfig {
21097 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21098 block_dim: (32, 4, 1),
21099 shared_mem_bytes: shmem,
21100 };
21101 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21102 let __s_lb = self.gpu.stream();
21103 let mut lb = __s_lb.launch_builder(&f);
21104 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21105 unsafe {
21106 lb.launch(cfg)?;
21107 }
21108 Ok(())
21109 }
21110
21111 #[allow(clippy::too_many_arguments)]
21115 pub fn attn_pre_vl8(
21116 &self,
21117 seqs: &[AttnPreVl],
21118 wq: &CudaSlice<f32>,
21119 wk: &CudaSlice<f32>,
21120 head_dim: usize,
21121 rope_dims: usize,
21122 n_head: usize,
21123 n_head_kv: usize,
21124 eps: f32,
21125 freq_base: f32,
21126 freq_scale: f32,
21127 kv_dim_k: usize,
21128 kv_dim_v: usize,
21129 k_tok_bytes: usize,
21130 v_tok_bytes: usize,
21131 ) -> Result<(), Box<dyn std::error::Error>> {
21132 let b = seqs.len();
21133 assert!(b >= 1 && b <= 8);
21134 let mut packed = [AttnPreVl::default(); 8];
21135 packed[..b].copy_from_slice(seqs);
21136 let v = AttnPreVl8(packed);
21137 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21138 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21139 {
21140 let f = self.func("q_gate_split_vl");
21141 let n = max_t * (n_head * head_dim) as u32;
21142 let cfg = LaunchConfig {
21143 grid_dim: (n.div_ceil(256), 1, b as u32),
21144 block_dim: (256, 1, 1),
21145 shared_mem_bytes: 0,
21146 };
21147 let __s_lb = self.gpu.stream();
21148 let mut lb = __s_lb.launch_builder(&f);
21149 lb.arg(&v).arg(&hd).arg(&nh);
21150 unsafe {
21151 lb.launch(cfg)?;
21152 }
21153 }
21154 {
21155 let f = self.func("attn_rms_vl");
21156 let cfg = LaunchConfig {
21157 grid_dim: (max_t * n_head as u32, 2, b as u32),
21158 block_dim: (rms_block(), 1, 1),
21159 shared_mem_bytes: 0,
21160 };
21161 let __s_lb = self.gpu.stream();
21162 let mut lb = __s_lb.launch_builder(&f);
21163 lb.arg(&v)
21164 .arg(wq)
21165 .arg(wk)
21166 .arg(&hd)
21167 .arg(&nh)
21168 .arg(&nhkv)
21169 .arg(&eps);
21170 unsafe {
21171 lb.launch(cfg)?;
21172 }
21173 }
21174 {
21175 let f = self.func("attn_rope_vl");
21176 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21177 let nd = rope_dims as i32;
21178 let cfg = LaunchConfig {
21179 grid_dim: (max_t * n_head as u32, 2, b as u32),
21180 block_dim: ((head_dim / 2) as u32, 1, 1),
21181 shared_mem_bytes: 0,
21182 };
21183 let __s_lb = self.gpu.stream();
21184 let mut lb = __s_lb.launch_builder(&f);
21185 lb.arg(&v)
21186 .arg(&hd)
21187 .arg(&nd)
21188 .arg(&nh)
21189 .arg(&nhkv)
21190 .arg(&theta_scale)
21191 .arg(&freq_scale);
21192 unsafe {
21193 lb.launch(cfg)?;
21194 }
21195 }
21196 {
21197 let f = self.func("append_kv_vl");
21198 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21199 let cfg = LaunchConfig {
21200 grid_dim: (nblk, max_t, b as u32),
21201 block_dim: (32, 1, 1),
21202 shared_mem_bytes: 0,
21203 };
21204 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21205 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21206 let __s_lb = self.gpu.stream();
21207 let mut lb = __s_lb.launch_builder(&f);
21208 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21209 unsafe {
21210 lb.launch(cfg)?;
21211 }
21212 }
21213 Ok(())
21214 }
21215
21216 pub fn fa_prefill_view(
21221 &self,
21222 q: &CudaSlice<f32>,
21223 k: &cudarc::driver::CudaView<u8>,
21224 v: &cudarc::driver::CudaView<u8>,
21225 o: &mut CudaSlice<f32>,
21226 head_dim: usize,
21227 n_head: usize,
21228 n_head_kv: usize,
21229 t: usize,
21230 t_kv: usize,
21231 scale: f32,
21232 causal: bool,
21233 k_tok_bytes: usize,
21234 v_tok_bytes: usize,
21235 g: bool,
21236 ) -> Result<(), Box<dyn std::error::Error>> {
21237 if portable_mma_gated() {
21238 return self.sdpa_naive_quantized_view(
21239 q,
21240 k,
21241 v,
21242 o,
21243 head_dim,
21244 n_head,
21245 n_head_kv,
21246 t,
21247 t_kv,
21248 scale,
21249 causal,
21250 k_tok_bytes,
21251 v_tok_bytes,
21252 );
21253 }
21254 const BLOCK_Q: usize = 64;
21255 const BK: usize = 32;
21256 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21259 let f = if g {
21260 self.func_g(&name)
21261 } else {
21262 self.func(&name)
21263 };
21264 let shmem =
21265 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21266 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21267 f.set_attribute(
21268 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21269 shmem as i32,
21270 )?;
21271 let cfg = LaunchConfig {
21272 grid_dim: (
21273 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21274 n_head as u32,
21275 1,
21276 ),
21277 block_dim: (32, 4, 1),
21278 shared_mem_bytes: shmem,
21279 };
21280 let (hd, nh, nhkv, ti, tkvi, cz) = (
21281 head_dim as i32,
21282 n_head as i32,
21283 n_head_kv as i32,
21284 t as i32,
21285 t_kv as i32,
21286 causal as i32,
21287 );
21288 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21289 let __s_b = self.gpu.stream();
21290 let mut b = __s_b.launch_builder(&f);
21291 b.arg(q)
21292 .arg(k)
21293 .arg(v)
21294 .arg(o)
21295 .arg(&hd)
21296 .arg(&nh)
21297 .arg(&nhkv)
21298 .arg(&ti)
21299 .arg(&tkvi)
21300 .arg(&scale)
21301 .arg(&cz)
21302 .arg(&ktb)
21303 .arg(&vtb);
21304 unsafe {
21305 b.launch(cfg)?;
21306 }
21307 Ok(())
21308 }
21309
21310 #[allow(clippy::too_many_arguments)]
21320 pub fn fa_prefill_view_ws(
21321 &self,
21322 q: &CudaSlice<f32>,
21323 k: &cudarc::driver::CudaView<u8>,
21324 v: &cudarc::driver::CudaView<u8>,
21325 o: &mut CudaSlice<f32>,
21326 head_dim: usize,
21327 n_head: usize,
21328 n_head_kv: usize,
21329 t: usize,
21330 t_kv: usize,
21331 scale: f32,
21332 causal: bool,
21333 k_tok_bytes: usize,
21334 v_tok_bytes: usize,
21335 g: bool,
21336 ) -> Result<(), Box<dyn std::error::Error>> {
21337 if portable_mma_gated() {
21338 return self.sdpa_naive_quantized_view(
21339 q,
21340 k,
21341 v,
21342 o,
21343 head_dim,
21344 n_head,
21345 n_head_kv,
21346 t,
21347 t_kv,
21348 scale,
21349 causal,
21350 k_tok_bytes,
21351 v_tok_bytes,
21352 );
21353 }
21354 const BLOCK_Q: usize = 64;
21355 const BK: usize = 32;
21356 let kv_dim_k = n_head_kv * head_dim;
21357 let kv_dim_v = n_head_kv * head_dim;
21358 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21360 let mut guard = self.prime_deqw_ws.lock().unwrap();
21362 let need_grow = match guard.as_ref() {
21363 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21364 None => true,
21365 };
21366 if need_grow {
21367 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21368 let (ck, cv) = guard
21369 .as_ref()
21370 .map(|(a, b)| (a.len(), b.len()))
21371 .unwrap_or((0, 0));
21372 *guard = Some((
21373 self.alloc_u8(grow(ck, k_ws_bytes))?,
21374 self.alloc_u8(grow(cv, v_ws_bytes))?,
21375 ));
21376 }
21377 let (kw, vw) = guard.as_mut().unwrap();
21378 {
21380 let f = if g {
21382 self.func_g("fa_dequant_kv_ws_bf16")
21383 } else {
21384 self.func("fa_dequant_kv_ws_bf16")
21385 };
21386 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21387 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21388 let cfg = LaunchConfig {
21389 grid_dim: (nblk.max(1), 1, 1),
21390 block_dim: (256, 1, 1),
21391 shared_mem_bytes: 0,
21392 };
21393 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21394 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21395 let __s_b = self.gpu.stream();
21396 let mut b = __s_b.launch_builder(&f);
21397 b.arg(k)
21398 .arg(v)
21399 .arg(&mut *kw)
21400 .arg(&mut *vw)
21401 .arg(&kdk)
21402 .arg(&kdv)
21403 .arg(&tkvi)
21404 .arg(&ktb)
21405 .arg(&vtb);
21406 unsafe {
21407 b.launch(cfg)?;
21408 }
21409 }
21410 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21418 .map(|v| v != "0")
21419 .unwrap_or(true);
21420 {
21421 let hd_sfx = fa_hd_suffix(head_dim)?;
21422 let f = self.func(&format!(
21423 "fa_prefill_qw{}{hd_sfx}",
21424 if db { "_db" } else { "" }
21425 ));
21426 let shmem = if db {
21427 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21429 } else {
21430 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21431 };
21432 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21433 f.set_attribute(
21434 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21435 shmem as i32,
21436 )?;
21437 let cfg = LaunchConfig {
21438 grid_dim: (
21439 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21440 n_head as u32,
21441 1,
21442 ),
21443 block_dim: (32, 4, 1),
21444 shared_mem_bytes: shmem,
21445 };
21446 let (hd, nh, nhkv, ti, tkvi, cz) = (
21447 head_dim as i32,
21448 n_head as i32,
21449 n_head_kv as i32,
21450 t as i32,
21451 t_kv as i32,
21452 causal as i32,
21453 );
21454 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21455 let __s_b = self.gpu.stream();
21456 let mut b = __s_b.launch_builder(&f);
21457 b.arg(q)
21458 .arg(&*kw)
21459 .arg(&*vw)
21460 .arg(o)
21461 .arg(&hd)
21462 .arg(&nh)
21463 .arg(&nhkv)
21464 .arg(&ti)
21465 .arg(&tkvi)
21466 .arg(&scale)
21467 .arg(&cz)
21468 .arg(&kdk)
21469 .arg(&kdv);
21470 unsafe {
21471 b.launch(cfg)?;
21472 }
21473 }
21474 Ok(())
21475 }
21476
21477 #[allow(clippy::too_many_arguments)]
21493 pub fn fa_prefill_view_ws_w_hd128(
21494 &self,
21495 q: &CudaSlice<f32>,
21496 k: &cudarc::driver::CudaView<u8>,
21497 v: &cudarc::driver::CudaView<u8>,
21498 o: &mut CudaSlice<f32>,
21499 head_dim: usize,
21500 n_head: usize,
21501 n_head_kv: usize,
21502 t: usize,
21503 t_kv: usize,
21504 scale: f32,
21505 causal: bool,
21506 window: usize,
21507 k_tok_bytes: usize,
21508 v_tok_bytes: usize,
21509 ) -> Result<(), Box<dyn std::error::Error>> {
21510 assert_eq!(
21511 head_dim, 128,
21512 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21513 );
21514 if portable_mma_gated() {
21515 return self.sdpa_naive_w_quantized_view(
21516 q,
21517 k,
21518 v,
21519 o,
21520 head_dim,
21521 n_head,
21522 n_head_kv,
21523 t,
21524 t_kv,
21525 scale,
21526 causal,
21527 window,
21528 k_tok_bytes,
21529 v_tok_bytes,
21530 );
21531 }
21532 const BLOCK_Q: usize = 64;
21533 const BK: usize = 32;
21534 let kv_dim_k = n_head_kv * head_dim;
21535 let kv_dim_v = n_head_kv * head_dim;
21536 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21538 let mut guard = self.prime_deqw_ws.lock().unwrap();
21539 let need_grow = match guard.as_ref() {
21540 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21541 None => true,
21542 };
21543 if need_grow {
21544 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21545 let (ck, cv) = guard
21546 .as_ref()
21547 .map(|(a, b)| (a.len(), b.len()))
21548 .unwrap_or((0, 0));
21549 *guard = Some((
21550 self.alloc_u8(grow(ck, k_ws_bytes))?,
21551 self.alloc_u8(grow(cv, v_ws_bytes))?,
21552 ));
21553 }
21554 let (kw, vw) = guard.as_mut().unwrap();
21555 {
21558 let f = self.func("fa_dequant_kv_ws_bf16");
21559 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21560 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21561 let cfg = LaunchConfig {
21562 grid_dim: (nblk.max(1), 1, 1),
21563 block_dim: (256, 1, 1),
21564 shared_mem_bytes: 0,
21565 };
21566 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21567 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21568 let __s_b = self.gpu.stream();
21569 let mut b = __s_b.launch_builder(&f);
21570 b.arg(k)
21571 .arg(v)
21572 .arg(&mut *kw)
21573 .arg(&mut *vw)
21574 .arg(&kdk)
21575 .arg(&kdv)
21576 .arg(&tkvi)
21577 .arg(&ktb)
21578 .arg(&vtb);
21579 unsafe {
21580 b.launch(cfg)?;
21581 }
21582 }
21583 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21585 .map(|v| v != "0")
21586 .unwrap_or(true);
21587 {
21588 let f = self.func(if db {
21589 "fa_prefill_qw_db_w_hd128"
21590 } else {
21591 "fa_prefill_qw_w_hd128"
21592 });
21593 let shmem = if db {
21594 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21595 } else {
21596 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21597 };
21598 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21599 f.set_attribute(
21600 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21601 shmem as i32,
21602 )?;
21603 let cfg = LaunchConfig {
21604 grid_dim: (
21605 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21606 n_head as u32,
21607 1,
21608 ),
21609 block_dim: (32, 4, 1),
21610 shared_mem_bytes: shmem,
21611 };
21612 let (hd, nh, nhkv, ti, tkvi, cz) = (
21613 head_dim as i32,
21614 n_head as i32,
21615 n_head_kv as i32,
21616 t as i32,
21617 t_kv as i32,
21618 causal as i32,
21619 );
21620 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21621 let __s_b = self.gpu.stream();
21622 let mut b = __s_b.launch_builder(&f);
21623 b.arg(q)
21624 .arg(&*kw)
21625 .arg(&*vw)
21626 .arg(o)
21627 .arg(&hd)
21628 .arg(&nh)
21629 .arg(&nhkv)
21630 .arg(&ti)
21631 .arg(&tkvi)
21632 .arg(&scale)
21633 .arg(&cz)
21634 .arg(&kdk)
21635 .arg(&kdv)
21636 .arg(&wnd);
21637 unsafe {
21638 b.launch(cfg)?;
21639 }
21640 }
21641 Ok(())
21642 }
21643
21644 pub fn fa_decode(
21648 &self,
21649 q: &CudaSlice<f32>,
21650 k: &cudarc::driver::CudaView<u8>,
21651 v: &cudarc::driver::CudaView<u8>,
21652 o: &mut CudaSlice<f32>,
21653 head_dim: usize,
21654 n_head: usize,
21655 n_head_kv: usize,
21656 t_kv: usize,
21657 scale: f32,
21658 k_tok_bytes: usize,
21659 v_tok_bytes: usize,
21660 ) -> Result<(), Box<dyn std::error::Error>> {
21661 self.fa_decode_kvmod(
21662 q,
21663 k,
21664 v,
21665 o,
21666 head_dim,
21667 n_head,
21668 n_head_kv,
21669 t_kv,
21670 scale,
21671 k_tok_bytes,
21672 v_tok_bytes,
21673 false,
21674 )
21675 }
21676
21677 #[allow(clippy::too_many_arguments)]
21681 #[allow(clippy::too_many_arguments)]
21685 #[allow(clippy::too_many_arguments)]
21686 fn fa_decode_scalar_unified(
21687 &self,
21688 q: &cudarc::driver::CudaView<f32>,
21689 k: &cudarc::driver::CudaView<u8>,
21690 v: &cudarc::driver::CudaView<u8>,
21691 o: &mut cudarc::driver::CudaViewMut<f32>,
21692 head_dim: usize,
21693 n_head: usize,
21694 n_head_kv: usize,
21695 t_kv_host: usize,
21696 t_kv_dev: Option<&CudaSlice<i32>>,
21697 scale: f32,
21698 n_splits: usize,
21699 split_keys: usize,
21700 k_tok_bytes: usize,
21701 v_tok_bytes: usize,
21702 g: bool,
21703 part_o: &mut CudaSlice<f32>,
21704 part_m: &mut CudaSlice<f32>,
21705 part_l: &mut CudaSlice<f32>,
21706 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21707 ) -> Result<(), Box<dyn std::error::Error>> {
21708 let f = if g {
21709 self.func_g("fa_decode_f32")
21710 } else {
21711 self.fa_func("fa_decode_f32", head_dim)
21712 };
21713 let cfg = LaunchConfig {
21714 grid_dim: (n_head as u32, n_splits as u32, 1),
21715 block_dim: (head_dim as u32, 1, 1),
21716 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21717 };
21718 let (hd, nh, nhkv, nsp) = (
21719 head_dim as i32,
21720 n_head as i32,
21721 n_head_kv as i32,
21722 n_splits as i32,
21723 );
21724 let (ktb, vtb, tkvi, ski) = (
21725 k_tok_bytes as i64,
21726 v_tok_bytes as i64,
21727 t_kv_host as i32,
21728 split_keys as i32,
21729 );
21730 let __s_b = self.gpu.stream();
21731 let mut b = __s_b.launch_builder(&f);
21732 match t_kv_dev {
21733 Some(d) => {
21734 b.arg(q)
21735 .arg(k)
21736 .arg(v)
21737 .arg(&mut *part_o)
21738 .arg(&mut *part_m)
21739 .arg(&mut *part_l)
21740 .arg(&hd)
21741 .arg(&nh)
21742 .arg(&nhkv)
21743 .arg(&tkvi)
21744 .arg(d)
21745 .arg(&scale)
21746 .arg(&nsp)
21747 .arg(&ski)
21748 .arg(&ktb)
21749 .arg(&vtb);
21750 unsafe {
21751 b.launch(cfg)?;
21752 }
21753 }
21754 None => {
21755 let null: u64 = 0;
21756 b.arg(q)
21757 .arg(k)
21758 .arg(v)
21759 .arg(&mut *part_o)
21760 .arg(&mut *part_m)
21761 .arg(&mut *part_l)
21762 .arg(&hd)
21763 .arg(&nh)
21764 .arg(&nhkv)
21765 .arg(&tkvi)
21766 .arg(&null)
21767 .arg(&scale)
21768 .arg(&nsp)
21769 .arg(&ski)
21770 .arg(&ktb)
21771 .arg(&vtb);
21772 unsafe {
21773 b.launch(cfg)?;
21774 }
21775 }
21776 }
21777 let cfg2 = LaunchConfig {
21778 grid_dim: (n_head as u32, 1, 1),
21779 block_dim: (head_dim as u32, 1, 1),
21780 shared_mem_bytes: 0,
21781 };
21782 if let Some((oq, od)) = q8_out {
21783 let fc = if g {
21785 self.func_g("fa_decode_combine_q8_1")
21786 } else {
21787 self.fa_func("fa_decode_combine_q8_1", head_dim)
21788 };
21789 let __s_b2 = self.gpu.stream();
21790 let mut b2 = __s_b2.launch_builder(&fc);
21791 b2.arg(&*part_o)
21792 .arg(&*part_m)
21793 .arg(&*part_l)
21794 .arg(oq)
21795 .arg(od)
21796 .arg(&hd)
21797 .arg(&nh)
21798 .arg(&nsp);
21799 unsafe {
21800 b2.launch(cfg2)?;
21801 }
21802 return Ok(());
21803 }
21804 let fc = if g {
21805 self.func_g("fa_decode_combine_f32")
21806 } else {
21807 self.fa_func("fa_decode_combine_f32", head_dim)
21808 };
21809 let __s_b2 = self.gpu.stream();
21810 let mut b2 = __s_b2.launch_builder(&fc);
21811 b2.arg(&*part_o)
21812 .arg(&*part_m)
21813 .arg(&*part_l)
21814 .arg(o)
21815 .arg(&hd)
21816 .arg(&nh)
21817 .arg(&nsp);
21818 unsafe {
21819 b2.launch(cfg2)?;
21820 }
21821 Ok(())
21822 }
21823
21824 pub fn fa_decode_kvmod(
21825 &self,
21826 q: &CudaSlice<f32>,
21827 k: &cudarc::driver::CudaView<u8>,
21828 v: &cudarc::driver::CudaView<u8>,
21829 o: &mut CudaSlice<f32>,
21830 head_dim: usize,
21831 n_head: usize,
21832 n_head_kv: usize,
21833 t_kv: usize,
21834 scale: f32,
21835 k_tok_bytes: usize,
21836 v_tok_bytes: usize,
21837 g: bool,
21838 ) -> Result<(), Box<dyn std::error::Error>> {
21839 let q_view = q.as_view();
21840 let mut o_view = o.as_view_mut();
21841 self.fa_decode_kvmod_view(
21842 &q_view,
21843 k,
21844 v,
21845 &mut o_view,
21846 head_dim,
21847 n_head,
21848 n_head_kv,
21849 t_kv,
21850 scale,
21851 k_tok_bytes,
21852 v_tok_bytes,
21853 g,
21854 )
21855 }
21856
21857 #[allow(clippy::too_many_arguments)]
21862 pub fn fa_decode_kvmod_view(
21863 &self,
21864 q: &cudarc::driver::CudaView<f32>,
21865 k: &cudarc::driver::CudaView<u8>,
21866 v: &cudarc::driver::CudaView<u8>,
21867 o: &mut cudarc::driver::CudaViewMut<f32>,
21868 head_dim: usize,
21869 n_head: usize,
21870 n_head_kv: usize,
21871 t_kv: usize,
21872 scale: f32,
21873 k_tok_bytes: usize,
21874 v_tok_bytes: usize,
21875 g: bool,
21876 ) -> Result<(), Box<dyn std::error::Error>> {
21877 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21898 if g && head_dim == 256 && !fa_v4_at(t_kv) {
21902 fa_vec = false;
21903 }
21904 let sp = fa_split_keys(t_kv, n_head_kv);
21905 let n_splits = if fa_vec {
21906 ((t_kv + sp - 1) / sp).max(1)
21907 } else {
21908 ((t_kv + 255) / 256).max(1)
21909 };
21910 let o_len = n_head * n_splits * head_dim;
21911 let ml_len = n_head * n_splits;
21912 let mut part_guard = self.fa_part_pool.lock().unwrap();
21913 if part_guard
21914 .as_ref()
21915 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21916 .unwrap_or(true)
21917 {
21918 let old = part_guard.take();
21929 let (co, cm) = old
21930 .as_ref()
21931 .map(|pp| (pp.0.len(), pp.1.len()))
21932 .unwrap_or((0, 0));
21933 if let Some(old) = old {
21934 self.fa_part_retired.lock().unwrap().push(old);
21935 }
21936 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21937 eprintln!(
21938 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21939 co, o_len, cm, ml_len
21940 );
21941 }
21942 *part_guard = Some((
21943 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21944 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21945 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21946 ));
21947 }
21948 let pg = part_guard.as_mut().unwrap();
21949 self.gpu
21950 .stream()
21951 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21952 self.gpu
21953 .stream()
21954 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21955 self.gpu
21956 .stream()
21957 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21958 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21959 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21960 let (hd, nh, nhkv, tkvi, nsp) = (
21961 head_dim as i32,
21962 n_head as i32,
21963 n_head_kv as i32,
21964 t_kv as i32,
21965 n_splits as i32,
21966 );
21967 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21968 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
21972 let fa512_min = fa512_min_tkv();
21977 let deep = fa_vec
21980 && head_dim == 256
21981 && fa_v4_at(t_kv)
21982 && !g
21983 && fa_deep_at(t_kv)
21984 && !matches!(fa_v4_mode(), "noB3" | "stage");
21985 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21986 let gqa = (n_head / n_head_kv).max(1) as u32;
21989 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21990 (
21991 fv,
21992 LaunchConfig {
21993 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21994 block_dim: (32, gqa, 1),
21995 shared_mem_bytes: 0,
21996 },
21997 )
21998 } else if fa_vec && head_dim <= 256 {
21999 let gqa = (n_head / n_head_kv).max(1) as u32;
22000 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22011 let smem_tkv = *SMEM_TKV.get_or_init(|| {
22012 std::env::var("MEMRA_FA_SMEM_TKV")
22013 .ok()
22014 .and_then(|v| v.parse().ok())
22015 .unwrap_or_else(|| {
22016 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22017 })
22018 });
22019 if fa_v4_at(t_kv) && head_dim == 256 {
22020 let v4name = match fa_v4_mode() {
22024 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
22027 _ => "fa_decode_vec_q_v4",
22028 };
22029 let fv = if g {
22030 self.func_g(v4name)
22031 } else {
22032 self.func(v4name)
22033 };
22034 let shmem = (if deep { 12160 } else { 11520 }
22037 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22038 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22039 fv.set_attribute(
22040 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22041 shmem as i32,
22042 )?;
22043 (
22044 fv,
22045 LaunchConfig {
22046 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22047 block_dim: (32, gqa, 1),
22048 shared_mem_bytes: shmem,
22049 },
22050 )
22051 } else if fa_v3_active(head_dim) {
22052 let fv = if g {
22055 self.func_g("fa_decode_vec_q_v3")
22056 } else {
22057 self.func("fa_decode_vec_q_v3")
22058 };
22059 let shmem = (32 * head_dim * 2) as u32; (
22061 fv,
22062 LaunchConfig {
22063 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22064 block_dim: (32, gqa, 1),
22065 shared_mem_bytes: shmem,
22066 },
22067 )
22068 } else if fa_v2_on() {
22069 let fv = if g {
22073 self.func_g("fa_decode_vec_q_v2")
22074 } else {
22075 self.func("fa_decode_vec_q_v2")
22076 };
22077 let shmem = (2 * 32 * head_dim * 2) as u32; (
22079 fv,
22080 LaunchConfig {
22081 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22082 block_dim: (32, gqa, 1),
22083 shared_mem_bytes: shmem,
22084 },
22085 )
22086 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22087 {
22088 let fv = if g {
22092 self.func_g("fa_decode_vec_q_smem")
22093 } else {
22094 self.func("fa_decode_vec_q_smem")
22095 };
22096 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
22098 fv.set_attribute(
22099 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22100 shmem as i32,
22101 )?;
22102 (
22103 fv,
22104 LaunchConfig {
22105 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22106 block_dim: (32, gqa, 1),
22107 shared_mem_bytes: shmem,
22108 },
22109 )
22110 } else {
22111 let fv = if g {
22114 self.func_g("fa_decode_vec_q")
22115 } else {
22116 self.func("fa_decode_vec_q")
22117 };
22118 (
22119 fv,
22120 LaunchConfig {
22121 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22122 block_dim: (32, gqa, 1),
22123 shared_mem_bytes: 0,
22124 },
22125 )
22126 }
22127 } else {
22128 return self.fa_decode_scalar_unified(
22131 q,
22132 k,
22133 v,
22134 o,
22135 head_dim,
22136 n_head,
22137 n_head_kv,
22138 t_kv,
22139 None,
22140 scale,
22141 n_splits,
22142 if fa_vec { sp } else { 256 },
22143 k_tok_bytes,
22144 v_tok_bytes,
22145 g,
22146 part_o,
22147 part_m,
22148 part_l,
22149 None,
22150 );
22151 };
22152 let __s_b = self.gpu.stream();
22153 let mut b = __s_b.launch_builder(&f);
22154 b.arg(q)
22155 .arg(k)
22156 .arg(v)
22157 .arg(&mut *part_o)
22158 .arg(&mut *part_m)
22159 .arg(&mut *part_l)
22160 .arg(&hd)
22161 .arg(&nh)
22162 .arg(&nhkv)
22163 .arg(&tkvi)
22164 .arg(&scale)
22165 .arg(&nsp)
22166 .arg(&ktb)
22167 .arg(&vtb);
22168 unsafe {
22169 b.launch(cfg)?;
22170 }
22171 let (fc, cfg2) = (
22174 if g {
22175 self.func_g("fa_decode_combine_f32")
22176 } else {
22177 self.fa_func("fa_decode_combine_f32", head_dim)
22178 },
22179 LaunchConfig {
22180 grid_dim: (n_head as u32, 1, 1),
22181 block_dim: (head_dim as u32, 1, 1),
22182 shared_mem_bytes: 0,
22183 },
22184 );
22185 let __s_b2 = self.gpu.stream();
22186 let mut b2 = __s_b2.launch_builder(&fc);
22187 b2.arg(&*part_o)
22188 .arg(&*part_m)
22189 .arg(&*part_l)
22190 .arg(o)
22191 .arg(&hd)
22192 .arg(&nh)
22193 .arg(&nsp);
22194 unsafe {
22195 b2.launch(cfg2)?;
22196 }
22197 Ok(())
22198 }
22199
22200 #[allow(clippy::too_many_arguments)]
22211 pub fn fa_decode_batch_seqs_v4(
22212 &self,
22213 q: &CudaSlice<f32>,
22214 kv_ptrs: &cudarc::driver::CudaView<u64>,
22215 pos_seq: &CudaSlice<i32>,
22216 o: &mut CudaSlice<f32>,
22217 head_dim: usize,
22218 n_head: usize,
22219 n_head_kv: usize,
22220 b_n: usize,
22221 t_kv_max: usize,
22222 scale: f32,
22223 split_keys: usize,
22224 k_tok_bytes: usize,
22225 v_tok_bytes: usize,
22226 ) -> Result<(), Box<dyn std::error::Error>> {
22227 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22228 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22229 let o_len = b_n * n_head * n_splits_max * head_dim;
22230 let ml_len = b_n * n_head * n_splits_max;
22231 let mut part_guard = self.fa_part_pool.lock().unwrap();
22232 if part_guard
22233 .as_ref()
22234 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22235 .unwrap_or(true)
22236 {
22237 let old = part_guard.take();
22248 let (co, cm) = old
22249 .as_ref()
22250 .map(|pp| (pp.0.len(), pp.1.len()))
22251 .unwrap_or((0, 0));
22252 if let Some(old) = old {
22253 self.fa_part_retired.lock().unwrap().push(old);
22254 }
22255 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22256 eprintln!(
22257 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22258 co, o_len, cm, ml_len
22259 );
22260 }
22261 *part_guard = Some((
22262 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22263 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22264 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22265 ));
22266 }
22267 let pg = part_guard.as_mut().unwrap();
22268 self.gpu
22269 .stream()
22270 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22271 self.gpu
22272 .stream()
22273 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22274 self.gpu
22275 .stream()
22276 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22277 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22278 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22279 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22280 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22281 let gqa = (n_head / n_head_kv).max(1) as u32;
22282 let f = self.func("fa_decode_vec_q_seqs_v4");
22283 let shmem = (11520 + 32 * head_dim * 2) as u32;
22285 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22286 f.set_attribute(
22287 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22288 shmem as i32,
22289 )?;
22290 let cfg = LaunchConfig {
22291 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22292 block_dim: (32, gqa, 1),
22293 shared_mem_bytes: shmem,
22294 };
22295 {
22296 let __s_b = self.gpu.stream();
22297 let mut b = __s_b.launch_builder(&f);
22298 b.arg(q)
22299 .arg(kv_ptrs)
22300 .arg(pos_seq)
22301 .arg(&mut *part_o)
22302 .arg(&mut *part_m)
22303 .arg(&mut *part_l)
22304 .arg(&hd)
22305 .arg(&nh)
22306 .arg(&nhkv)
22307 .arg(&scale)
22308 .arg(&nspm)
22309 .arg(&spk)
22310 .arg(&ktb)
22311 .arg(&vtb);
22312 unsafe {
22313 b.launch(cfg)?;
22314 }
22315 }
22316 let fc = self.func("fa_decode_combine_seqs");
22317 let cfg2 = LaunchConfig {
22318 grid_dim: (n_head as u32, b_n as u32, 1),
22319 block_dim: (head_dim as u32, 1, 1),
22320 shared_mem_bytes: 0,
22321 };
22322 let __s_b2 = self.gpu.stream();
22323 let mut b2 = __s_b2.launch_builder(&fc);
22324 b2.arg(&*part_o)
22325 .arg(&*part_m)
22326 .arg(&*part_l)
22327 .arg(o)
22328 .arg(&hd)
22329 .arg(&nh)
22330 .arg(pos_seq)
22331 .arg(&nspm)
22332 .arg(&spk);
22333 unsafe {
22334 b2.launch(cfg2)?;
22335 }
22336 Ok(())
22337 }
22338
22339 #[allow(clippy::too_many_arguments)]
22346 pub fn append_kv_quantized_seqs(
22347 &self,
22348 k_rows: &CudaSlice<f32>,
22349 v_rows: &CudaSlice<f32>,
22350 kv_ptrs: &cudarc::driver::CudaView<u64>,
22351 pos_seq: &CudaSlice<i32>,
22352 b_n: usize,
22353 kv_dim_k: usize,
22354 kv_dim_v: usize,
22355 k_tok_bytes: usize,
22356 v_tok_bytes: usize,
22357 ) -> Result<(), Box<dyn std::error::Error>> {
22358 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22359 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22360 let cfg = LaunchConfig {
22361 grid_dim: (nblk, b_n as u32, 1),
22362 block_dim: (32, 1, 1),
22363 shared_mem_bytes: 0,
22364 };
22365 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22366 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22367 let __s_b = self.gpu.stream();
22368 let mut b = __s_b.launch_builder(&f);
22369 b.arg(k_rows)
22370 .arg(v_rows)
22371 .arg(kv_ptrs)
22372 .arg(pos_seq)
22373 .arg(&kdk)
22374 .arg(&kdv)
22375 .arg(&ktb)
22376 .arg(&vtb);
22377 unsafe {
22378 b.launch(cfg)?;
22379 }
22380 Ok(())
22381 }
22382
22383 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22389 std::env::var("MEMRA_NO_FA_VEC").is_err()
22390 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22391 && base_len + 1 >= fa_vec_min_tkv()
22392 && head_dim <= 256
22393 && head_dim % 32 == 0
22394 }
22395
22396 #[allow(clippy::too_many_arguments)]
22405 pub fn fa_decode_rows(
22406 &self,
22407 q: &CudaSlice<f32>,
22408 k: &cudarc::driver::CudaView<u8>,
22409 v: &cudarc::driver::CudaView<u8>,
22410 o: &mut CudaSlice<f32>,
22411 head_dim: usize,
22412 n_head: usize,
22413 n_head_kv: usize,
22414 base_len: usize,
22415 t: usize,
22416 scale: f32,
22417 k_tok_bytes: usize,
22418 v_tok_bytes: usize,
22419 base_dev: Option<(&CudaSlice<i32>, i32)>,
22423 kv_shared: bool,
22426 g: bool,
22430 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22433 ) -> Result<(), Box<dyn std::error::Error>> {
22434 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22435 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
22442 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22443 let v = *SP512.get_or_init(|| {
22446 std::env::var("MEMRA_FA_SP512")
22447 .ok()
22448 .and_then(|x| x.parse().ok())
22449 .unwrap_or(0)
22450 });
22451 sp = if v >= 8 {
22452 v
22453 } else {
22454 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22455 };
22456 }
22457 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22458 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22459 let gqa = (n_head / n_head_kv).max(1) as u32;
22460 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22471 groups.push((0, t, sp));
22472 } else {
22473 let mut r0 = 0usize;
22474 while r0 < t {
22475 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22476 let mut r1 = r0 + 1;
22477 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22478 r1 += 1;
22479 }
22480 groups.push((r0, r1 - r0, sp_g));
22481 r0 = r1;
22482 }
22483 }
22484 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22488 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22489 std::env::var("MEMRA_FA_SMEM_TKV")
22490 .ok()
22491 .and_then(|v| v.parse().ok())
22492 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22493 });
22494 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22495 let v3 = fa_v3_active(head_dim);
22496 let smem_rows =
22497 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22498 let _ = kv_shared;
22503 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22506 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22520 let tb512 = head_dim == 512
22522 && sp <= 32
22523 && n_head / n_head_kv.max(1) <= 16
22524 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22525 let fname = if tb512 {
22526 "fa_decode_vec_q_rows_v4_512_tb"
22527 } else if i2 {
22528 "fa_decode_vec_q_rows_dpl16_i2"
22529 } else if head_dim == 512 {
22530 "fa_decode_vec_q_rows_dpl16"
22531 }
22532 else if v4 {
22534 "fa_decode_vec_q_rows_v4"
22535 } else if v3 {
22536 "fa_decode_vec_q_rows_v3"
22537 } else if fa_v2_on() {
22538 "fa_decode_vec_q_rows_v2"
22539 } else if smem_rows {
22540 "fa_decode_vec_q_rows_smem"
22541 } else {
22542 "fa_decode_vec_q_rows"
22543 };
22544 let f = if head_dim == 512 {
22545 self.fa_func(fname, head_dim)
22546 } else if g {
22547 self.func_g(if smem_rows {
22555 "fa_decode_vec_q_rows"
22556 } else {
22557 fname
22558 })
22559 } else {
22560 self.func(fname)
22561 };
22562 let shmem = if tb512 {
22563 let gk = Self::gkv_on();
22565 let sh =
22566 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22567 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22568 f.set_attribute(
22569 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22570 sh as i32,
22571 )?;
22572 sh
22573 } else if v4 || v3 || smem_rows || fa_v2_on() {
22574 let sh = (if v4 {
22576 11520 + 32 * head_dim * if g { 1 } else { 2 }
22577 } else if v3 {
22578 32 * head_dim * 2
22579 } else {
22580 2 * 32 * head_dim * 2
22581 }) as u32;
22582 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22583 f.set_attribute(
22584 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22585 sh as i32,
22586 )?;
22587 sh
22588 } else {
22589 0
22590 };
22591 for &(r0, t_g, sp_g) in &groups {
22595 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22596 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22597 let base_i = (base_len + r0) as i32;
22598 let o_len = t_g * n_head * n_splits_g * head_dim;
22599 let ml_len = t_g * n_head * n_splits_g;
22600 let mut part_guard = self.fa_part_pool.lock().unwrap();
22601 if part_guard
22602 .as_ref()
22603 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22604 .unwrap_or(true)
22605 {
22606 let old = part_guard.take();
22617 let (co, cm) = old
22618 .as_ref()
22619 .map(|pp| (pp.0.len(), pp.1.len()))
22620 .unwrap_or((0, 0));
22621 if let Some(old) = old {
22622 self.fa_part_retired.lock().unwrap().push(old);
22623 }
22624 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22625 eprintln!(
22626 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22627 co, o_len, cm, ml_len
22628 );
22629 }
22630 *part_guard = Some((
22631 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22632 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22633 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22634 ));
22635 }
22636 let pg = part_guard.as_mut().unwrap();
22637 self.gpu
22638 .stream()
22639 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22640 self.gpu
22641 .stream()
22642 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22643 self.gpu
22644 .stream()
22645 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22646 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22647 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22648 let qv = self.view(q, t * n_head * head_dim);
22649 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22650 let cfg = LaunchConfig {
22651 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22652 block_dim: (32, gqa, 1),
22653 shared_mem_bytes: shmem,
22654 };
22655 {
22656 let __s_b = self.gpu.stream();
22657 let mut b = __s_b.launch_builder(&f);
22658 if tb512 {
22659 let (bd, plus) =
22661 base_dev.expect("hd512 rows twin requires a device base counter");
22662 let plus_g = plus + r0 as i32;
22663 let nr = t_g as i32;
22664 if Self::pdl_on() && Self::pdl_wb_on() {
22665 use cudarc::driver::{DevicePtr, DevicePtrMut};
22667 let s = &self.gpu.stream();
22668 let (pq, _b0) = q_g.device_ptr(s);
22669 let (pk, _b1) = k.device_ptr(s);
22670 let (pv, _b2) = v.device_ptr(s);
22671 let (po, _b3) = part_o.device_ptr_mut(s);
22672 let (pm, _b4) = part_m.device_ptr_mut(s);
22673 let (pl, _b5) = part_l.device_ptr_mut(s);
22674 let (pb, _b6) = bd.device_ptr(s);
22675 let mut ps = [
22676 &pq as *const _ as *mut std::ffi::c_void,
22677 &pk as *const _ as *mut _,
22678 &pv as *const _ as *mut _,
22679 &po as *const _ as *mut _,
22680 &pm as *const _ as *mut _,
22681 &pl as *const _ as *mut _,
22682 &hd as *const _ as *mut _,
22683 &nh as *const _ as *mut _,
22684 &nhkv as *const _ as *mut _,
22685 &pb as *const _ as *mut _,
22686 &plus_g as *const _ as *mut _,
22687 &scale as *const _ as *mut _,
22688 &nspm as *const _ as *mut _,
22689 &spk as *const _ as *mut _,
22690 &ktb as *const _ as *mut _,
22691 &vtb as *const _ as *mut _,
22692 &nr as *const _ as *mut _,
22693 ];
22694 unsafe {
22695 self.launch_pdl_flash(
22696 Self::gkv_on(),
22697 "fa_decode_vec_q_rows_v4_512_tb",
22698 (n_head_kv as u32, n_splits_g as u32, 1),
22699 (32, gqa, 1),
22700 shmem,
22701 &mut ps,
22702 )?;
22703 }
22704 } else {
22705 let cfg_tb = LaunchConfig {
22706 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22707 block_dim: (32, gqa, 1),
22708 shared_mem_bytes: shmem,
22709 };
22710 b.arg(&q_g)
22711 .arg(k)
22712 .arg(v)
22713 .arg(&mut *part_o)
22714 .arg(&mut *part_m)
22715 .arg(&mut *part_l)
22716 .arg(&hd)
22717 .arg(&nh)
22718 .arg(&nhkv)
22719 .arg(bd)
22720 .arg(&plus_g)
22721 .arg(&scale)
22722 .arg(&nspm)
22723 .arg(&spk)
22724 .arg(&ktb)
22725 .arg(&vtb)
22726 .arg(&nr);
22727 unsafe {
22728 b.launch(cfg_tb)?;
22729 }
22730 }
22731 } else if head_dim == 512 {
22732 let (bd, plus) =
22733 base_dev.expect("hd512 rows twin requires a device base counter");
22734 let plus_g = plus + r0 as i32;
22735 b.arg(&q_g)
22736 .arg(k)
22737 .arg(v)
22738 .arg(&mut *part_o)
22739 .arg(&mut *part_m)
22740 .arg(&mut *part_l)
22741 .arg(&hd)
22742 .arg(&nh)
22743 .arg(&nhkv)
22744 .arg(bd)
22745 .arg(&plus_g)
22746 .arg(&scale)
22747 .arg(&nspm)
22748 .arg(&spk)
22749 .arg(&ktb)
22750 .arg(&vtb);
22751 unsafe {
22752 b.launch(cfg)?;
22753 }
22754 } else {
22755 b.arg(&q_g)
22756 .arg(k)
22757 .arg(v)
22758 .arg(&mut *part_o)
22759 .arg(&mut *part_m)
22760 .arg(&mut *part_l)
22761 .arg(&hd)
22762 .arg(&nh)
22763 .arg(&nhkv)
22764 .arg(&base_i)
22765 .arg(&scale)
22766 .arg(&nspm)
22767 .arg(&spk)
22768 .arg(&ktb)
22769 .arg(&vtb);
22770 unsafe {
22771 b.launch(cfg)?;
22772 }
22773 }
22774 }
22775 let cfg2 = LaunchConfig {
22776 grid_dim: (n_head as u32, t_g as u32, 1),
22777 block_dim: (head_dim as u32, 1, 1),
22778 shared_mem_bytes: 0,
22779 };
22780 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22781 if head_dim == 512 {
22782 let (bd, plus) = base_dev.unwrap();
22785 let plus_g = plus + r0 as i32;
22786 if let Some((oq, od)) = q8_out.as_mut() {
22787 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22789 if Self::pdl_on() && Self::pdl_wb_on() {
22790 use cudarc::driver::{DevicePtr, DevicePtrMut};
22792 let s = &self.gpu.stream();
22793 let (po, _g0) = part_o.device_ptr(s);
22794 let (pm, _g1) = part_m.device_ptr(s);
22795 let (pl, _g2) = part_l.device_ptr(s);
22796 let (pq, _g3) = oq.device_ptr_mut(s);
22797 let (pd, _g4) = od.device_ptr_mut(s);
22798 let (pb, _g5) = bd.device_ptr(s);
22799 let mut ps = [
22800 &po as *const _ as *mut std::ffi::c_void,
22801 &pm as *const _ as *mut _,
22802 &pl as *const _ as *mut _,
22803 &pq as *const _ as *mut _,
22804 &pd as *const _ as *mut _,
22805 &hd as *const _ as *mut _,
22806 &nh as *const _ as *mut _,
22807 &pb as *const _ as *mut _,
22808 &plus_g as *const _ as *mut _,
22809 &nspm as *const _ as *mut _,
22810 &spk as *const _ as *mut _,
22811 ];
22812 unsafe {
22813 self.launch_pdl_flash(
22814 Self::gkv_on(),
22815 "fa_decode_combine_rows_dc_q8_1",
22816 cfg2.grid_dim,
22817 cfg2.block_dim,
22818 0,
22819 &mut ps,
22820 )?;
22821 }
22822 continue;
22823 }
22824 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22825 let __s_b2 = self.gpu.stream();
22826 let mut b2 = __s_b2.launch_builder(&fc);
22827 b2.arg(&*part_o)
22828 .arg(&*part_m)
22829 .arg(&*part_l)
22830 .arg(&mut **oq)
22831 .arg(&mut **od)
22832 .arg(&hd)
22833 .arg(&nh)
22834 .arg(bd)
22835 .arg(&plus_g)
22836 .arg(&nspm)
22837 .arg(&spk);
22838 unsafe {
22839 b2.launch(cfg2)?;
22840 }
22841 continue;
22842 }
22843 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22844 let __s_b2 = self.gpu.stream();
22845 let mut b2 = __s_b2.launch_builder(&fc);
22846 b2.arg(&*part_o)
22847 .arg(&*part_m)
22848 .arg(&*part_l)
22849 .arg(&mut o_g)
22850 .arg(&hd)
22851 .arg(&nh)
22852 .arg(bd)
22853 .arg(&plus_g)
22854 .arg(&nspm)
22855 .arg(&spk);
22856 unsafe {
22857 b2.launch(cfg2)?;
22858 }
22859 } else {
22860 assert!(
22863 q8_out.is_none(),
22864 "rows q8 emit requires the hd512 dc combine"
22865 );
22866 let fc = self.func("fa_decode_combine_rows");
22867 let __s_b2 = self.gpu.stream();
22868 let mut b2 = __s_b2.launch_builder(&fc);
22869 b2.arg(&*part_o)
22870 .arg(&*part_m)
22871 .arg(&*part_l)
22872 .arg(&mut o_g)
22873 .arg(&hd)
22874 .arg(&nh)
22875 .arg(&base_i)
22876 .arg(&nspm)
22877 .arg(&spk);
22878 unsafe {
22879 b2.launch(cfg2)?;
22880 }
22881 }
22882 }
22883 Ok(())
22884 }
22885
22886 #[allow(clippy::too_many_arguments)]
22890 pub fn fa_decode_rows_w(
22891 &self,
22892 q: &CudaSlice<f32>,
22893 k: &cudarc::driver::CudaView<u8>,
22894 v: &cudarc::driver::CudaView<u8>,
22895 o: &mut CudaSlice<f32>,
22896 head_dim: usize,
22897 n_head: usize,
22898 n_head_kv: usize,
22899 base_dev: &CudaSlice<i32>,
22900 base_plus: i32,
22901 t: usize,
22902 scale: f32,
22903 window: usize,
22904 k_tok_bytes: usize,
22905 v_tok_bytes: usize,
22906 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22907 ) -> Result<(), Box<dyn std::error::Error>> {
22908 debug_assert!(head_dim == 256);
22913 let sp = {
22921 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22922 let v = *SPW.get_or_init(|| {
22923 std::env::var("MEMRA_FA_SPW")
22924 .ok()
22925 .and_then(|x| x.parse().ok())
22926 .unwrap_or(0)
22927 });
22928 if v >= 8 {
22929 v
22930 } else {
22931 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22932 }
22933 };
22934 let n_splits_max = (window + sp - 1) / sp;
22935 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22936 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22937 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22938 let gqa = (n_head / n_head_kv).max(1) as u32;
22939 let o_len = t * n_head * n_splits_max * head_dim;
22940 let ml_len = t * n_head * n_splits_max;
22941 let mut part_guard = self.fa_part_pool.lock().unwrap();
22942 if part_guard
22943 .as_ref()
22944 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22945 .unwrap_or(true)
22946 {
22947 let old = part_guard.take();
22958 let (co, cm) = old
22959 .as_ref()
22960 .map(|pp| (pp.0.len(), pp.1.len()))
22961 .unwrap_or((0, 0));
22962 if let Some(old) = old {
22963 self.fa_part_retired.lock().unwrap().push(old);
22964 }
22965 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22966 eprintln!(
22967 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22968 co, o_len, cm, ml_len
22969 );
22970 }
22971 *part_guard = Some((
22972 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22973 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22974 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22975 ));
22976 }
22977 let pg = part_guard.as_mut().unwrap();
22978 self.gpu
22979 .stream()
22980 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22981 self.gpu
22982 .stream()
22983 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22984 self.gpu
22985 .stream()
22986 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22987 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22988 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22994 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22995 std::env::var("MEMRA_FA_SMEM_TKV")
22996 .ok()
22997 .and_then(|v| v.parse().ok())
22998 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22999 });
23000 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23006 let wg = Self::wkv_on();
23011 let sp2 =
23014 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23015 if sp2 {
23016 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23017 if Self::pdl_on() && Self::pdl_wb_on() {
23018 use cudarc::driver::{DevicePtr, DevicePtrMut};
23020 let s = &self.gpu.stream();
23021 let (pq, _b0) = q.device_ptr(s);
23022 let (pk, _b1) = k.device_ptr(s);
23023 let (pv, _b2) = v.device_ptr(s);
23024 let (po, _b3) = part_o.device_ptr_mut(s);
23025 let (pm, _b4) = part_m.device_ptr_mut(s);
23026 let (pl, _b5) = part_l.device_ptr_mut(s);
23027 let (pb, _b6) = base_dev.device_ptr(s);
23028 let mut ps = [
23029 &pq as *const _ as *mut std::ffi::c_void,
23030 &pk as *const _ as *mut _,
23031 &pv as *const _ as *mut _,
23032 &po as *const _ as *mut _,
23033 &pm as *const _ as *mut _,
23034 &pl as *const _ as *mut _,
23035 &hd as *const _ as *mut _,
23036 &nh as *const _ as *mut _,
23037 &nhkv as *const _ as *mut _,
23038 &pb as *const _ as *mut _,
23039 &base_plus as *const _ as *mut _,
23040 &scale as *const _ as *mut _,
23041 &nspm as *const _ as *mut _,
23042 &spk as *const _ as *mut _,
23043 &ktb as *const _ as *mut _,
23044 &vtb as *const _ as *mut _,
23045 &wini as *const _ as *mut _,
23046 ];
23047 unsafe {
23048 self.launch_pdl_flash(
23049 wg,
23050 "fa_decode_vec_q_rows_v4_w_sp",
23051 (n_head_kv as u32, n_splits_max as u32, t as u32),
23052 (32, gqa + 1, 1),
23053 sh,
23054 &mut ps,
23055 )?;
23056 }
23057 } else {
23058 let f = if wg {
23059 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23060 } else {
23061 self.func("fa_decode_vec_q_rows_v4_w_sp")
23062 };
23063 f.set_attribute(
23064 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23065 sh as i32,
23066 )?;
23067 let cfg = LaunchConfig {
23068 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23069 block_dim: (32, gqa + 1, 1),
23070 shared_mem_bytes: sh,
23071 };
23072 let __s_b = self.gpu.stream();
23073 let mut b = __s_b.launch_builder(&f);
23074 b.arg(q)
23075 .arg(k)
23076 .arg(v)
23077 .arg(&mut *part_o)
23078 .arg(&mut *part_m)
23079 .arg(&mut *part_l)
23080 .arg(&hd)
23081 .arg(&nh)
23082 .arg(&nhkv)
23083 .arg(base_dev)
23084 .arg(&base_plus)
23085 .arg(&scale)
23086 .arg(&nspm)
23087 .arg(&spk)
23088 .arg(&ktb)
23089 .arg(&vtb)
23090 .arg(&wini);
23091 unsafe {
23092 b.launch(cfg)?;
23093 }
23094 }
23095 } else {
23096 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23097 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23099 use cudarc::driver::{DevicePtr, DevicePtrMut};
23100 let s = &self.gpu.stream();
23101 let (pq, _b0) = q.device_ptr(s);
23102 let (pk, _b1) = k.device_ptr(s);
23103 let (pv, _b2) = v.device_ptr(s);
23104 let (po, _b3) = part_o.device_ptr_mut(s);
23105 let (pm, _b4) = part_m.device_ptr_mut(s);
23106 let (pl, _b5) = part_l.device_ptr_mut(s);
23107 let (pb, _b6) = base_dev.device_ptr(s);
23108 let mut ps = [
23109 &pq as *const _ as *mut std::ffi::c_void,
23110 &pk as *const _ as *mut _,
23111 &pv as *const _ as *mut _,
23112 &po as *const _ as *mut _,
23113 &pm as *const _ as *mut _,
23114 &pl as *const _ as *mut _,
23115 &hd as *const _ as *mut _,
23116 &nh as *const _ as *mut _,
23117 &nhkv as *const _ as *mut _,
23118 &pb as *const _ as *mut _,
23119 &base_plus as *const _ as *mut _,
23120 &scale as *const _ as *mut _,
23121 &nspm as *const _ as *mut _,
23122 &spk as *const _ as *mut _,
23123 &ktb as *const _ as *mut _,
23124 &vtb as *const _ as *mut _,
23125 &wini as *const _ as *mut _,
23126 ];
23127 unsafe {
23128 self.launch_pdl_flash(
23129 wg,
23130 "fa_decode_vec_q_rows_v4_w",
23131 (n_head_kv as u32, n_splits_max as u32, t as u32),
23132 (32, gqa, 1),
23133 sh,
23134 &mut ps,
23135 )?;
23136 }
23137 } else {
23138 let pick = |name: &str| {
23139 if wg {
23140 self.func_g(name)
23141 } else {
23142 self.func(name)
23143 }
23144 };
23145 let (f, sh) = if fa_v4_at(window) {
23146 let f = pick("fa_decode_vec_q_rows_v4_w");
23147 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23148 } else if smem_tkv > 0 && window >= smem_tkv {
23149 (
23152 pick("fa_decode_vec_q_rows_smem_w"),
23153 (2 * 32 * head_dim * 2) as u32,
23154 )
23155 } else {
23156 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23157 };
23158 f.set_attribute(
23159 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23160 sh as i32,
23161 )?;
23162 let cfg = LaunchConfig {
23163 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23164 block_dim: (32, gqa, 1),
23165 shared_mem_bytes: sh,
23166 };
23167 let __s_b = self.gpu.stream();
23168 let mut b = __s_b.launch_builder(&f);
23169 b.arg(q)
23170 .arg(k)
23171 .arg(v)
23172 .arg(&mut *part_o)
23173 .arg(&mut *part_m)
23174 .arg(&mut *part_l)
23175 .arg(&hd)
23176 .arg(&nh)
23177 .arg(&nhkv)
23178 .arg(base_dev)
23179 .arg(&base_plus)
23180 .arg(&scale)
23181 .arg(&nspm)
23182 .arg(&spk)
23183 .arg(&ktb)
23184 .arg(&vtb)
23185 .arg(&wini);
23186 unsafe {
23187 b.launch(cfg)?;
23188 }
23189 }
23190 }
23191 let cfg2 = LaunchConfig {
23192 grid_dim: (n_head as u32, t as u32, 1),
23193 block_dim: (head_dim as u32, 1, 1),
23194 shared_mem_bytes: 0,
23195 };
23196 if let Some((oq, od)) = q8_out {
23197 if Self::pdl_on() && Self::pdl_wb_on() {
23200 use cudarc::driver::{DevicePtr, DevicePtrMut};
23202 let s = &self.gpu.stream();
23203 let (po, _g0) = part_o.device_ptr(s);
23204 let (pm, _g1) = part_m.device_ptr(s);
23205 let (pl, _g2) = part_l.device_ptr(s);
23206 let (pq, _g3) = oq.device_ptr_mut(s);
23207 let (pd, _g4) = od.device_ptr_mut(s);
23208 let mut ps = [
23209 &po as *const _ as *mut std::ffi::c_void,
23210 &pm as *const _ as *mut _,
23211 &pl as *const _ as *mut _,
23212 &pq as *const _ as *mut _,
23213 &pd as *const _ as *mut _,
23214 &hd as *const _ as *mut _,
23215 &nh as *const _ as *mut _,
23216 &nspm as *const _ as *mut _,
23217 &spk as *const _ as *mut _,
23218 &wini as *const _ as *mut _,
23219 ];
23220 unsafe {
23221 self.launch_pdl_flash(
23222 wg,
23223 "fa_decode_combine_rows_w_q8_1",
23224 cfg2.grid_dim,
23225 cfg2.block_dim,
23226 0,
23227 &mut ps,
23228 )?;
23229 }
23230 return Ok(());
23231 }
23232 let fc = if wg {
23233 self.func_g("fa_decode_combine_rows_w_q8_1")
23234 } else {
23235 self.func("fa_decode_combine_rows_w_q8_1")
23236 };
23237 let __s_b2 = self.gpu.stream();
23238 let mut b2 = __s_b2.launch_builder(&fc);
23239 b2.arg(&*part_o)
23240 .arg(&*part_m)
23241 .arg(&*part_l)
23242 .arg(oq)
23243 .arg(od)
23244 .arg(&hd)
23245 .arg(&nh)
23246 .arg(&nspm)
23247 .arg(&spk)
23248 .arg(&wini);
23249 unsafe {
23250 b2.launch(cfg2)?;
23251 }
23252 return Ok(());
23253 }
23254 let fc = if wg {
23255 self.func_g("fa_decode_combine_rows_w")
23256 } else {
23257 self.func("fa_decode_combine_rows_w")
23258 };
23259 let __s_b2 = self.gpu.stream();
23260 let mut b2 = __s_b2.launch_builder(&fc);
23261 b2.arg(&*part_o)
23262 .arg(&*part_m)
23263 .arg(&*part_l)
23264 .arg(o)
23265 .arg(&hd)
23266 .arg(&nh)
23267 .arg(&nspm)
23268 .arg(&spk)
23269 .arg(&wini);
23270 unsafe {
23271 b2.launch(cfg2)?;
23272 }
23273 Ok(())
23274 }
23275
23276 #[allow(clippy::too_many_arguments)]
23282 pub fn fa_decode_rows_dc(
23283 &self,
23284 q: &CudaSlice<f32>,
23285 k: &cudarc::driver::CudaView<u8>,
23286 v: &cudarc::driver::CudaView<u8>,
23287 o: &mut CudaSlice<f32>,
23288 head_dim: usize,
23289 n_head: usize,
23290 n_head_kv: usize,
23291 base_dev: &CudaSlice<i32>,
23292 t_kv_upper: usize,
23293 t: usize,
23294 scale: f32,
23295 k_tok_bytes: usize,
23296 v_tok_bytes: usize,
23297 base_plus: i32,
23298 g: bool,
23299 ) -> Result<(), Box<dyn std::error::Error>> {
23300 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23301 assert!(
23302 v4 || fa_v3_active(head_dim),
23303 "stream fa rows requires the v3 or v4 lane"
23304 );
23305 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23306 if v4 {
23307 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23308 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23309 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23310 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23311 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23312 let gqa = (n_head / n_head_kv).max(1) as u32;
23313 let o_len = t * n_head * n_splits_max * head_dim;
23314 let ml_len = t * n_head * n_splits_max;
23315 let mut part_guard = self.fa_part_pool.lock().unwrap();
23316 if part_guard
23317 .as_ref()
23318 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23319 .unwrap_or(true)
23320 {
23321 let old = part_guard.take();
23332 let (co, cm) = old
23333 .as_ref()
23334 .map(|pp| (pp.0.len(), pp.1.len()))
23335 .unwrap_or((0, 0));
23336 if let Some(old) = old {
23337 self.fa_part_retired.lock().unwrap().push(old);
23338 }
23339 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23340 eprintln!(
23341 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23342 co, o_len, cm, ml_len
23343 );
23344 }
23345 *part_guard = Some((
23346 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23347 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23348 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23349 ));
23350 }
23351 let pg = part_guard.as_mut().unwrap();
23352 self.gpu
23353 .stream()
23354 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23355 self.gpu
23356 .stream()
23357 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23358 self.gpu
23359 .stream()
23360 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23361 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23362 let f = if g {
23363 self.func_g("fa_decode_vec_q_rows_v4_dc")
23364 } else {
23365 self.func("fa_decode_vec_q_rows_v4_dc")
23366 };
23367 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23368 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23369 f.set_attribute(
23370 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23371 sh as i32,
23372 )?;
23373 let cfg = LaunchConfig {
23374 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23375 block_dim: (32, gqa, 1),
23376 shared_mem_bytes: sh,
23377 };
23378 let __s_b = self.gpu.stream();
23379 let mut b = __s_b.launch_builder(&f);
23380 b.arg(q)
23381 .arg(k)
23382 .arg(v)
23383 .arg(&mut *part_o)
23384 .arg(&mut *part_m)
23385 .arg(&mut *part_l)
23386 .arg(&hd)
23387 .arg(&nh)
23388 .arg(&nhkv)
23389 .arg(base_dev)
23390 .arg(&base_plus)
23391 .arg(&scale)
23392 .arg(&nspm)
23393 .arg(&spk)
23394 .arg(&ktb)
23395 .arg(&vtb);
23396 unsafe {
23397 b.launch(cfg)?;
23398 }
23399 let fc = self.func("fa_decode_combine_rows_dc");
23400 let cfg2 = LaunchConfig {
23401 grid_dim: (n_head as u32, t as u32, 1),
23402 block_dim: (head_dim as u32, 1, 1),
23403 shared_mem_bytes: 0,
23404 };
23405 let __s_b2 = self.gpu.stream();
23406 let mut b2 = __s_b2.launch_builder(&fc);
23407 b2.arg(&*part_o)
23408 .arg(&*part_m)
23409 .arg(&*part_l)
23410 .arg(o)
23411 .arg(&hd)
23412 .arg(&nh)
23413 .arg(base_dev)
23414 .arg(&base_plus)
23415 .arg(&nspm)
23416 .arg(&spk);
23417 unsafe {
23418 b2.launch(cfg2)?;
23419 }
23420 return Ok(());
23421 }
23422 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23423 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23424 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23425 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23426 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23427 let gqa = (n_head / n_head_kv).max(1) as u32;
23428 let o_len = t * n_head * n_splits_max * head_dim;
23429 let ml_len = t * n_head * n_splits_max;
23430 let mut part_guard = self.fa_part_pool.lock().unwrap();
23431 if part_guard
23432 .as_ref()
23433 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23434 .unwrap_or(true)
23435 {
23436 let old = part_guard.take();
23447 let (co, cm) = old
23448 .as_ref()
23449 .map(|pp| (pp.0.len(), pp.1.len()))
23450 .unwrap_or((0, 0));
23451 if let Some(old) = old {
23452 self.fa_part_retired.lock().unwrap().push(old);
23453 }
23454 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23455 eprintln!(
23456 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23457 co, o_len, cm, ml_len
23458 );
23459 }
23460 *part_guard = Some((
23461 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23462 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23463 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23464 ));
23465 }
23466 let pg = part_guard.as_mut().unwrap();
23467 self.gpu
23468 .stream()
23469 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23470 self.gpu
23471 .stream()
23472 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23473 self.gpu
23474 .stream()
23475 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23476 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23477 let f = self.func("fa_decode_vec_q_rows_v3_dc");
23478 let sh = (32 * head_dim * 2) as u32;
23479 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23480 f.set_attribute(
23481 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23482 sh as i32,
23483 )?;
23484 let cfg = LaunchConfig {
23485 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23486 block_dim: (32, gqa, 1),
23487 shared_mem_bytes: sh,
23488 };
23489 let __s_b = self.gpu.stream();
23490 let mut b = __s_b.launch_builder(&f);
23491 b.arg(q)
23492 .arg(k)
23493 .arg(v)
23494 .arg(&mut *part_o)
23495 .arg(&mut *part_m)
23496 .arg(&mut *part_l)
23497 .arg(&hd)
23498 .arg(&nh)
23499 .arg(&nhkv)
23500 .arg(base_dev)
23501 .arg(&scale)
23502 .arg(&nspm)
23503 .arg(&spk)
23504 .arg(&ktb)
23505 .arg(&vtb);
23506 unsafe {
23507 b.launch(cfg)?;
23508 }
23509 let fc = self.func("fa_decode_combine_rows_dc");
23510 let cfg2 = LaunchConfig {
23511 grid_dim: (n_head as u32, t as u32, 1),
23512 block_dim: (head_dim as u32, 1, 1),
23513 shared_mem_bytes: 0,
23514 };
23515 let plus0 = 0i32;
23516 let __s_b2 = self.gpu.stream();
23517 let mut b2 = __s_b2.launch_builder(&fc);
23518 b2.arg(&*part_o)
23519 .arg(&*part_m)
23520 .arg(&*part_l)
23521 .arg(o)
23522 .arg(&hd)
23523 .arg(&nh)
23524 .arg(base_dev)
23525 .arg(&plus0)
23526 .arg(&nspm)
23527 .arg(&spk);
23528 unsafe {
23529 b2.launch(cfg2)?;
23530 }
23531 Ok(())
23532 }
23533
23534 pub fn fa_decode_dc(
23545 &self,
23546 q: &CudaSlice<f32>,
23547 k: &cudarc::driver::CudaView<u8>,
23548 v: &cudarc::driver::CudaView<u8>,
23549 o: &mut CudaSlice<f32>,
23550 head_dim: usize,
23551 n_head: usize,
23552 n_head_kv: usize,
23553 t_kv_dev: &CudaSlice<i32>,
23554 bucket_max: usize,
23555 scale: f32,
23556 k_tok_bytes: usize,
23557 v_tok_bytes: usize,
23558 g: bool,
23559 ) -> Result<(), Box<dyn std::error::Error>> {
23560 self.fa_decode_dc_q8(
23561 q,
23562 k,
23563 v,
23564 o,
23565 head_dim,
23566 n_head,
23567 n_head_kv,
23568 t_kv_dev,
23569 bucket_max,
23570 scale,
23571 k_tok_bytes,
23572 v_tok_bytes,
23573 g,
23574 None,
23575 )
23576 }
23577
23578 #[allow(clippy::too_many_arguments)]
23581 pub fn fa_decode_dc_q8(
23582 &self,
23583 q: &CudaSlice<f32>,
23584 k: &cudarc::driver::CudaView<u8>,
23585 v: &cudarc::driver::CudaView<u8>,
23586 o: &mut CudaSlice<f32>,
23587 head_dim: usize,
23588 n_head: usize,
23589 n_head_kv: usize,
23590 t_kv_dev: &CudaSlice<i32>,
23591 bucket_max: usize,
23592 scale: f32,
23593 k_tok_bytes: usize,
23594 v_tok_bytes: usize,
23595 g: bool,
23596 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23597 ) -> Result<(), Box<dyn std::error::Error>> {
23598 let mut fa_vec =
23606 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23607 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23608 fa_vec = false;
23609 } let sp = fa_split_keys(bucket_max, n_head_kv);
23611 let n_splits = if fa_vec {
23612 ((bucket_max + sp - 1) / sp).max(1)
23613 } else {
23614 ((bucket_max + 255) / 256).max(1)
23615 };
23616 let o_len = n_head * n_splits * head_dim;
23617 let ml_len = n_head * n_splits;
23618 let mut part_guard = self.fa_part_pool.lock().unwrap();
23619 if part_guard
23620 .as_ref()
23621 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23622 .unwrap_or(true)
23623 {
23624 let old = part_guard.take();
23635 let (co, cm) = old
23636 .as_ref()
23637 .map(|pp| (pp.0.len(), pp.1.len()))
23638 .unwrap_or((0, 0));
23639 if let Some(old) = old {
23640 self.fa_part_retired.lock().unwrap().push(old);
23641 }
23642 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23643 eprintln!(
23644 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23645 co, o_len, cm, ml_len
23646 );
23647 }
23648 *part_guard = Some((
23649 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23650 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23651 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23652 ));
23653 }
23654 let pg = part_guard.as_mut().unwrap();
23655 self.gpu
23656 .stream()
23657 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23658 self.gpu
23659 .stream()
23660 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23661 self.gpu
23662 .stream()
23663 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23664 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23665 let (hd, nh, nhkv, nsp) = (
23666 head_dim as i32,
23667 n_head as i32,
23668 n_head_kv as i32,
23669 n_splits as i32,
23670 );
23671 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23672 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23673 let deep = fa_vec
23676 && head_dim == 256
23677 && fa_v4_at(bucket_max)
23678 && !g
23679 && fa_deep_at(bucket_max)
23680 && !matches!(fa_v4_mode(), "noB3" | "stage");
23681 let (f, cfg) = if fa_vec
23682 && head_dim == 512
23683 && bucket_max >= {
23684 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23685 *FA512_MIN_DC.get_or_init(|| {
23686 std::env::var("MEMRA_FA512_MIN")
23687 .ok()
23688 .and_then(|v| v.parse().ok())
23689 .unwrap_or(512)
23690 })
23691 } {
23692 let gqa = (n_head / n_head_kv).max(1) as u32;
23694 (
23695 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23696 LaunchConfig {
23697 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23698 block_dim: (32, gqa, 1),
23699 shared_mem_bytes: 0,
23700 },
23701 )
23702 } else if fa_vec && head_dim == 512 {
23703 let q_view = q.as_view();
23706 let mut o_view = o.as_view_mut();
23707 return self.fa_decode_scalar_unified(
23708 &q_view,
23709 k,
23710 v,
23711 &mut o_view,
23712 head_dim,
23713 n_head,
23714 n_head_kv,
23715 0,
23716 Some(t_kv_dev),
23717 scale,
23718 n_splits,
23719 sp,
23720 k_tok_bytes,
23721 v_tok_bytes,
23722 g,
23723 &mut *part_o,
23724 &mut *part_m,
23725 &mut *part_l,
23726 q8_out,
23727 );
23728 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23729 let gqa = (n_head / n_head_kv).max(1) as u32;
23732 let fv = if g {
23733 self.func_g("fa_decode_vec_q_v4_dc")
23734 } else if deep {
23735 self.func("fa_decode_vec_q_v4_deep_dc")
23736 } else {
23737 self.func("fa_decode_vec_q_v4_dc")
23738 };
23739 let shmem =
23740 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23741 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23742 fv.set_attribute(
23743 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23744 shmem as i32,
23745 )?;
23746 (
23747 fv,
23748 LaunchConfig {
23749 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23750 block_dim: (32, gqa, 1),
23751 shared_mem_bytes: shmem,
23752 },
23753 )
23754 } else if fa_vec && fa_v3_active(head_dim) {
23755 let gqa = (n_head / n_head_kv).max(1) as u32;
23758 let fv = if g {
23759 self.func_g("fa_decode_vec_q_v3_dc")
23760 } else {
23761 self.func("fa_decode_vec_q_v3_dc")
23762 };
23763 let shmem = (32 * head_dim * 2) as u32; (
23765 fv,
23766 LaunchConfig {
23767 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23768 block_dim: (32, gqa, 1),
23769 shared_mem_bytes: shmem,
23770 },
23771 )
23772 } else if fa_vec && fa_v2_on() {
23773 let gqa = (n_head / n_head_kv).max(1) as u32;
23777 let fv = if g {
23778 self.func_g("fa_decode_vec_q_v2_dc")
23779 } else {
23780 self.func("fa_decode_vec_q_v2_dc")
23781 };
23782 let shmem = (2 * 32 * head_dim * 2) as u32; (
23784 fv,
23785 LaunchConfig {
23786 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23787 block_dim: (32, gqa, 1),
23788 shared_mem_bytes: shmem,
23789 },
23790 )
23791 } else if fa_vec {
23792 let gqa = (n_head / n_head_kv).max(1) as u32;
23793 let fv = if g {
23795 self.func_g("fa_decode_vec_q_dc")
23796 } else {
23797 self.func("fa_decode_vec_q_dc")
23798 };
23799 (
23800 fv,
23801 LaunchConfig {
23802 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23803 block_dim: (32, gqa, 1),
23804 shared_mem_bytes: 0,
23805 },
23806 )
23807 } else {
23808 let q_view = q.as_view();
23809 let mut o_view = o.as_view_mut();
23810 return self.fa_decode_scalar_unified(
23811 &q_view,
23812 k,
23813 v,
23814 &mut o_view,
23815 head_dim,
23816 n_head,
23817 n_head_kv,
23818 0,
23819 Some(t_kv_dev),
23820 scale,
23821 n_splits,
23822 if fa_vec { sp } else { 256 },
23823 k_tok_bytes,
23824 v_tok_bytes,
23825 g,
23826 &mut *part_o,
23827 &mut *part_m,
23828 &mut *part_l,
23829 q8_out,
23830 );
23831 };
23832 let ski = sp as i32; let __s_b = self.gpu.stream();
23834 let mut b = __s_b.launch_builder(&f);
23835 b.arg(q)
23836 .arg(k)
23837 .arg(v)
23838 .arg(&mut *part_o)
23839 .arg(&mut *part_m)
23840 .arg(&mut *part_l)
23841 .arg(&hd)
23842 .arg(&nh)
23843 .arg(&nhkv)
23844 .arg(t_kv_dev)
23845 .arg(&scale)
23846 .arg(&nsp)
23847 .arg(&ski)
23848 .arg(&ktb)
23849 .arg(&vtb);
23850 unsafe {
23851 b.launch(cfg)?;
23852 }
23853 let cfg2 = LaunchConfig {
23854 grid_dim: (n_head as u32, 1, 1),
23855 block_dim: (head_dim as u32, 1, 1),
23856 shared_mem_bytes: 0,
23857 };
23858 if let Some((oq, od)) = q8_out {
23859 let fc = if g {
23860 self.func_g("fa_decode_combine_q8_1")
23861 } else {
23862 self.fa_func("fa_decode_combine_q8_1", head_dim)
23863 };
23864 let __s_b2 = self.gpu.stream();
23865 let mut b2 = __s_b2.launch_builder(&fc);
23866 b2.arg(&*part_o)
23867 .arg(&*part_m)
23868 .arg(&*part_l)
23869 .arg(oq)
23870 .arg(od)
23871 .arg(&hd)
23872 .arg(&nh)
23873 .arg(&nsp);
23874 unsafe {
23875 b2.launch(cfg2)?;
23876 }
23877 return Ok(());
23878 }
23879 let fc = if g {
23880 self.func_g("fa_decode_combine_f32")
23881 } else {
23882 self.fa_func("fa_decode_combine_f32", head_dim)
23883 };
23884 let __s_b2 = self.gpu.stream();
23885 let mut b2 = __s_b2.launch_builder(&fc);
23886 b2.arg(&*part_o)
23887 .arg(&*part_m)
23888 .arg(&*part_l)
23889 .arg(o)
23890 .arg(&hd)
23891 .arg(&nh)
23892 .arg(&nsp);
23893 unsafe {
23894 b2.launch(cfg2)?;
23895 }
23896 Ok(())
23897 }
23898
23899 #[allow(clippy::too_many_arguments)]
23903 pub fn append_kv_quantized_dcw(
23904 &self,
23905 k_row: &CudaSlice<f32>,
23906 v_row: &CudaSlice<f32>,
23907 kc: &mut CudaSlice<u8>,
23908 vc: &mut CudaSlice<u8>,
23909 len_dev: &CudaSlice<i32>,
23910 base_dev: Option<&CudaSlice<i32>>,
23911 kv_dim_k: usize,
23912 kv_dim_v: usize,
23913 k_tok_bytes: usize,
23914 v_tok_bytes: usize,
23915 ) -> Result<(), Box<dyn std::error::Error>> {
23916 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23917 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23918 let cfg = LaunchConfig {
23919 grid_dim: (nblk, 1, 1),
23920 block_dim: (32, 1, 1),
23921 shared_mem_bytes: 0,
23922 };
23923 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23924 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23925 let null: u64 = 0;
23926 let __s_b = self.gpu.stream();
23927 let mut b = __s_b.launch_builder(&f);
23928 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23929 match base_dev {
23930 Some(base) => {
23931 b.arg(base);
23932 }
23933 None => {
23934 b.arg(&null);
23935 }
23936 }
23937 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23938 unsafe {
23939 b.launch(cfg)?;
23940 }
23941 Ok(())
23942 }
23943
23944 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23946 let f = self.func("inc_i32");
23947 let cfg = LaunchConfig {
23948 grid_dim: (1, 1, 1),
23949 block_dim: (1, 1, 1),
23950 shared_mem_bytes: 0,
23951 };
23952 let __s_b = self.gpu.stream();
23953 let mut b = __s_b.launch_builder(&f);
23954 b.arg(counter);
23955 unsafe {
23956 b.launch(cfg)?;
23957 }
23958 Ok(())
23959 }
23960
23961 #[allow(clippy::too_many_arguments)]
23970 fn fa_part_pool_grow(
23974 &self,
23975 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
23976 o_len: usize,
23977 ml_len: usize,
23978 ) -> Result<(), Box<dyn std::error::Error>> {
23979 if part_guard
23980 .as_ref()
23981 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23982 .unwrap_or(true)
23983 {
23984 let old = part_guard.take();
23985 let (co, cm) = old
23986 .as_ref()
23987 .map(|pp| (pp.0.len(), pp.1.len()))
23988 .unwrap_or((0, 0));
23989 if let Some(old) = old {
23990 self.fa_part_retired.lock().unwrap().push(old);
23991 }
23992 *part_guard = Some((
23993 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23994 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23995 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23996 ));
23997 }
23998 Ok(())
23999 }
24000
24001 pub fn fa_dcw_pool_ensure(
24004 &self,
24005 head_dim: usize,
24006 n_head: usize,
24007 n_head_kv: usize,
24008 bucket_max: usize,
24009 ) -> Result<(), Box<dyn std::error::Error>> {
24010 let sp = fa_split_keys(bucket_max, n_head_kv);
24011 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24012 let o_len = n_head * n_splits * head_dim;
24013 let ml_len = n_head * n_splits;
24014 let mut part_guard = self.fa_part_pool.lock().unwrap();
24015 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24016 }
24017
24018 #[allow(clippy::too_many_arguments)]
24026 pub fn fa_decode_dcw2(
24027 &self,
24028 q2: &CudaSlice<f32>,
24029 k_ring: &cudarc::driver::CudaView<u8>,
24030 v_ring: &cudarc::driver::CudaView<u8>,
24031 o2: &mut CudaSlice<f32>,
24032 head_dim: usize,
24033 n_head: usize,
24034 n_head_kv: usize,
24035 len_dev: &CudaSlice<i32>,
24036 base_dev: Option<&CudaSlice<i32>>,
24037 window: usize,
24038 bucket_max: usize,
24039 scale: f32,
24040 k_tok_bytes: usize,
24041 v_tok_bytes: usize,
24042 gate2: &CudaSlice<f32>,
24043 ) -> Result<(), Box<dyn std::error::Error>> {
24044 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24045 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24046 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24047 }
24048 let sp = fa_split_keys(bucket_max, n_head_kv);
24049 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24050 let o_len = 2 * n_head * n_splits * head_dim;
24052 let ml_len = 2 * n_head * n_splits;
24053 let mut part_guard = self.fa_part_pool.lock().unwrap();
24054 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24055 let pg = part_guard.as_mut().unwrap();
24056 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24057 let (hd, nh, nhkv, nsp) = (
24058 head_dim as i32,
24059 n_head as i32,
24060 n_head_kv as i32,
24061 n_splits as i32,
24062 );
24063 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24064 let (ski, win) = (sp as i32, window as i32);
24065 let gqa = (n_head / n_head_kv).max(1) as u32;
24066 let smem = (32 * head_dim * 2) as u32;
24067 let f = self.func("fa_decode_vec_q_v3_dcw2");
24068 let cfg = LaunchConfig {
24069 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24070 block_dim: (32, gqa, 1),
24071 shared_mem_bytes: smem,
24072 };
24073 let null: u64 = 0;
24074 {
24075 let __s_b = self.gpu.stream();
24076 let mut b = __s_b.launch_builder(&f);
24077 b.arg(q2)
24078 .arg(k_ring)
24079 .arg(v_ring)
24080 .arg(&mut *part_o)
24081 .arg(&mut *part_m)
24082 .arg(&mut *part_l)
24083 .arg(&hd)
24084 .arg(&nh)
24085 .arg(&nhkv)
24086 .arg(len_dev);
24087 match base_dev {
24088 Some(base) => {
24089 b.arg(base);
24090 }
24091 None => {
24092 b.arg(&null);
24093 }
24094 }
24095 b.arg(&win)
24096 .arg(&scale)
24097 .arg(&nsp)
24098 .arg(&ski)
24099 .arg(&ktb)
24100 .arg(&vtb);
24101 unsafe {
24102 b.launch(cfg)?;
24103 }
24104 }
24105 let fc = {
24109 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24110 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24111 self.func("fa_decode_combine_gate_f32_s")
24112 } else {
24113 self.func("fa_decode_combine_gate_f32")
24114 }
24115 };
24116 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24117 let nh2 = (2 * n_head) as i32;
24118 let cfg2 = LaunchConfig {
24119 grid_dim: ((2 * n_head) as u32, 1, 1),
24120 block_dim: (head_dim as u32, 1, 1),
24121 shared_mem_bytes: if combine_shared {
24122 (2 * n_splits * 4) as u32
24123 } else {
24124 0
24125 },
24126 };
24127 let __s_b2 = self.gpu.stream();
24128 let mut b2 = __s_b2.launch_builder(&fc);
24129 b2.arg(&*part_o)
24130 .arg(&*part_m)
24131 .arg(&*part_l)
24132 .arg(gate2)
24133 .arg(o2)
24134 .arg(&hd)
24135 .arg(&nh2)
24136 .arg(&nsp);
24137 unsafe {
24138 b2.launch(cfg2)?;
24139 }
24140 Ok(())
24141 }
24142
24143 #[allow(clippy::too_many_arguments)]
24152 pub fn fa_decode_dcw_rows(
24153 &self,
24154 q_rows: &CudaSlice<f32>,
24155 tab: &CudaSlice<u64>,
24156 o_rows: &mut CudaSlice<f32>,
24157 t: usize,
24158 head_dim: usize,
24159 n_head: usize,
24160 n_head_kv: usize,
24161 window: usize,
24162 max_ns: usize,
24163 scale: f32,
24164 k_tok_bytes: usize,
24165 v_tok_bytes: usize,
24166 gate_rows: &CudaSlice<f32>,
24167 ) -> Result<(), Box<dyn std::error::Error>> {
24168 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24169 || head_dim > 256
24170 || head_dim % 32 != 0
24171 || !fa_v3_on()
24172 {
24173 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24174 }
24175 if fa_sm_count() < 128
24176 || std::env::var("MEMRA_FA_SPLIT").is_ok()
24177 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24178 || std::env::var("MEMRA_FA_SP16").is_ok()
24179 {
24180 return Err(
24181 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24182 (or a <128-SM rig) keep the per-row path"
24183 .into(),
24184 );
24185 }
24186 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24187 return Err("fa_decode_dcw_rows geometry".into());
24188 }
24189 let o_len = t * n_head * max_ns * head_dim;
24190 let ml_len = t * n_head * max_ns;
24191 let mut part_guard = self.fa_part_pool.lock().unwrap();
24192 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24193 let pg = part_guard.as_mut().unwrap();
24194 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24195 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24196 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24197 let (win, mns) = (window as i32, max_ns as i32);
24198 let gqa = (n_head / n_head_kv).max(1) as u32;
24199 let smem = (32 * head_dim * 2) as u32;
24200 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24201 let cfg = LaunchConfig {
24202 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24203 block_dim: (32, gqa, 1),
24204 shared_mem_bytes: smem,
24205 };
24206 {
24207 let __s_b = self.gpu.stream();
24208 let mut b = __s_b.launch_builder(&f);
24209 b.arg(q_rows)
24210 .arg(tab)
24211 .arg(&mut *part_o)
24212 .arg(&mut *part_m)
24213 .arg(&mut *part_l)
24214 .arg(&hd)
24215 .arg(&nh)
24216 .arg(&nhkv)
24217 .arg(&win)
24218 .arg(&scale)
24219 .arg(&mns)
24220 .arg(&ktb)
24221 .arg(&vtb);
24222 unsafe {
24223 b.launch(cfg)?;
24224 }
24225 }
24226 let fc = {
24230 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24231 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24232 self.func("fa_decode_combine_gate_f32_s")
24233 } else {
24234 self.func("fa_decode_combine_gate_f32")
24235 }
24236 };
24237 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24238 let nht = (t * n_head) as i32;
24239 let cfg2 = LaunchConfig {
24240 grid_dim: ((t * n_head) as u32, 1, 1),
24241 block_dim: (head_dim as u32, 1, 1),
24242 shared_mem_bytes: if combine_shared {
24243 (2 * max_ns * 4) as u32
24244 } else {
24245 0
24246 },
24247 };
24248 let __s_b2 = self.gpu.stream();
24249 let mut b2 = __s_b2.launch_builder(&fc);
24250 b2.arg(&*part_o)
24251 .arg(&*part_m)
24252 .arg(&*part_l)
24253 .arg(gate_rows)
24254 .arg(o_rows)
24255 .arg(&hd)
24256 .arg(&nht)
24257 .arg(&mns);
24258 unsafe {
24259 b2.launch(cfg2)?;
24260 }
24261 Ok(())
24262 }
24263
24264 pub fn fa_decode_dcw(
24265 &self,
24266 q: &CudaSlice<f32>,
24267 k_ring: &cudarc::driver::CudaView<u8>,
24268 v_ring: &cudarc::driver::CudaView<u8>,
24269 o: &mut CudaSlice<f32>,
24270 head_dim: usize,
24271 n_head: usize,
24272 n_head_kv: usize,
24273 len_dev: &CudaSlice<i32>,
24274 base_dev: Option<&CudaSlice<i32>>,
24275 window: usize,
24276 bucket_max: usize,
24277 scale: f32,
24278 k_tok_bytes: usize,
24279 v_tok_bytes: usize,
24280 fused_gate: Option<&CudaSlice<f32>>,
24284 ) -> Result<(), Box<dyn std::error::Error>> {
24285 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24286 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24287 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"
24288 .into());
24289 }
24290 let sp = fa_split_keys(bucket_max, n_head_kv);
24291 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24292 let o_len = n_head * n_splits * head_dim;
24293 let ml_len = n_head * n_splits;
24294 let mut part_guard = self.fa_part_pool.lock().unwrap();
24295 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24296 let pg = part_guard.as_mut().unwrap();
24297 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24302 let memset_on = *MEMSET_ON
24307 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24308 || crate::tp::token_graph_building();
24309 if memset_on {
24310 self.gpu
24311 .stream()
24312 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24313 self.gpu
24314 .stream()
24315 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24316 self.gpu
24317 .stream()
24318 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24319 }
24320 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24321 let (hd, nh, nhkv, nsp) = (
24322 head_dim as i32,
24323 n_head as i32,
24324 n_head_kv as i32,
24325 n_splits as i32,
24326 );
24327 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24328 let (ski, win) = (sp as i32, window as i32);
24329 let gqa = (n_head / n_head_kv).max(1) as u32;
24330 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24334 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24335 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24336 Ok("2") => 2,
24337 Ok("1") => 1,
24338 _ => 0,
24339 });
24340 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24345 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24346 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24347 std::sync::Mutex::new(None);
24348 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24352 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24353 && (n_head / n_head_kv) % 2 == 0
24354 && (n_head / n_head_kv) >= 2;
24355 let f = if fprof {
24356 self.func("fa_decode_vec_q_v3_dcw_prof")
24357 } else if hs2 {
24358 self.func("fa_decode_vec_q_v3_dcw_hs2")
24359 } else if hoist == 2 {
24360 self.func("fa_decode_vec_q_v3_dcw_hc")
24362 } else if hoist == 1 {
24363 self.func("fa_decode_vec_q_v3_dcw_h")
24365 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24366 self.func("fa_decode_vec_q_v3_dcw_u8")
24367 } else {
24368 self.func("fa_decode_vec_q_v3_dcw")
24369 };
24370 let cfg = LaunchConfig {
24371 grid_dim: if hs2 {
24372 ((2 * n_head_kv) as u32, n_splits as u32, 1)
24373 } else {
24374 (n_head_kv as u32, n_splits as u32, 1)
24375 },
24376 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24377 shared_mem_bytes: smem,
24378 };
24379 let null: u64 = 0;
24380 let __s_b = self.gpu.stream();
24381 let mut b = __s_b.launch_builder(&f);
24382 b.arg(q)
24383 .arg(k_ring)
24384 .arg(v_ring)
24385 .arg(&mut *part_o)
24386 .arg(&mut *part_m)
24387 .arg(&mut *part_l)
24388 .arg(&hd)
24389 .arg(&nh)
24390 .arg(&nhkv)
24391 .arg(len_dev);
24392 match base_dev {
24393 Some(base) => {
24394 b.arg(base);
24395 }
24396 None => {
24397 b.arg(&null);
24398 }
24399 }
24400 b.arg(&win)
24401 .arg(&scale)
24402 .arg(&nsp)
24403 .arg(&ski)
24404 .arg(&ktb)
24405 .arg(&vtb);
24406 if fprof {
24407 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24408 if guard
24409 .as_ref()
24410 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24411 {
24412 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24413 }
24414 let (_, buf) = guard.as_mut().expect("armed above");
24415 b.arg(&*buf);
24416 unsafe {
24417 b.launch(cfg)?;
24418 }
24419 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24420 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24421 if n % 430 == 0 {
24422 self.stream().synchronize()?;
24423 let h = self.dtoh_u64(buf)?;
24424 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24425 let tot: u64 = h[..6].iter().sum();
24426 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24427 for (i, name) in phases.iter().enumerate() {
24428 let pct = if tot > 0 {
24429 h[i] as f64 / tot as f64 * 100.0
24430 } else {
24431 0.0
24432 };
24433 line.push_str(&format!(" {name}={pct:.1}%"));
24434 }
24435 if h[6] > 0 {
24436 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24437 }
24438 eprintln!("{line}");
24439 }
24440 } else {
24441 unsafe {
24442 b.launch(cfg)?;
24443 }
24444 }
24445 let mut combine_shared = false;
24446 let fc = if fused_gate.is_some() {
24447 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24450 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24451 combine_shared = true;
24452 self.func("fa_decode_combine_gate_f32_s")
24453 } else {
24454 self.func("fa_decode_combine_gate_f32")
24455 }
24456 } else {
24457 self.fa_func("fa_decode_combine_f32", head_dim)
24458 };
24459 let cfg2 = LaunchConfig {
24460 grid_dim: (n_head as u32, 1, 1),
24461 block_dim: (head_dim as u32, 1, 1),
24462 shared_mem_bytes: if combine_shared {
24463 (2 * n_splits * 4) as u32
24464 } else {
24465 0
24466 },
24467 };
24468 let __s_b2 = self.gpu.stream();
24469 let mut b2 = __s_b2.launch_builder(&fc);
24470 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24471 if let Some(gate_row) = fused_gate {
24472 b2.arg(gate_row);
24473 }
24474 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24475 unsafe {
24476 b2.launch(cfg2)?;
24477 }
24478 Ok(())
24479 }
24480
24481 pub fn fa_geom_eager(
24487 &self,
24488 t_kv: usize,
24489 head_dim: usize,
24490 n_head_kv: usize,
24491 g: bool,
24492 ) -> (bool, usize) {
24493 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24497 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24503 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24504 if g && head_dim == 256 && !fa_v4_at(t_kv) {
24510 fa_vec = false;
24511 }
24512 let sp = fa_split_keys(t_kv, n_head_kv);
24513 let n_splits = if fa_vec {
24514 ((t_kv + sp - 1) / sp).max(1)
24515 } else {
24516 ((t_kv + 255) / 256).max(1)
24517 };
24518 (fa_vec, n_splits)
24519 }
24520
24521 pub fn fa_bucket_key(
24527 &self,
24528 t_kv: usize,
24529 head_dim: usize,
24530 n_head_kv: usize,
24531 g: bool,
24532 ) -> (bool, usize) {
24533 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24534 }
24535
24536 pub fn capture_graph_retained<F>(
24548 &self,
24549 step: F,
24550 ) -> Result<
24551 (
24552 cudarc::driver::CudaGraph,
24553 Vec<Box<dyn std::any::Any + Send>>,
24554 ),
24555 Box<dyn std::error::Error>,
24556 >
24557 where
24558 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24559 {
24560 use cudarc::driver::sys::CUgraphInstantiate_flags;
24561 self.capture_graph_retained_flags(
24562 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24563 step,
24564 )
24565 }
24566
24567 pub fn capture_graph_retained_flags<F>(
24572 &self,
24573 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24574 mut step: F,
24575 ) -> Result<
24576 (
24577 cudarc::driver::CudaGraph,
24578 Vec<Box<dyn std::any::Any + Send>>,
24579 ),
24580 Box<dyn std::error::Error>,
24581 >
24582 where
24583 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24584 {
24585 use cudarc::driver::sys::CUstreamCaptureMode;
24586 self.capture_keep.lock().unwrap().clear();
24594 let was_tracking = self.gpu.ctx.is_event_tracking();
24595 if was_tracking {
24596 unsafe {
24597 self.gpu.ctx.disable_event_tracking();
24598 }
24599 }
24600 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24601 self.capture_keep_on
24602 .store(true, std::sync::atomic::Ordering::Relaxed);
24603 let w = (|| {
24604 step(self)?;
24605 step(self)
24606 })();
24607 self.capture_keep_on
24608 .store(false, std::sync::atomic::Ordering::Relaxed);
24609 w?;
24610 self.gpu.stream().synchronize()?;
24611 self.gpu
24612 .stream()
24613 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24614 let r = step(self);
24615 let g = self.gpu.stream().end_capture(flags);
24616 r?;
24617 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24618 graph.upload()?;
24619 Ok(graph)
24620 };
24621 let result = run();
24622 self.capture_keep_on
24623 .store(false, std::sync::atomic::Ordering::Relaxed);
24624 if was_tracking {
24625 unsafe {
24626 self.gpu.ctx.enable_event_tracking();
24627 }
24628 }
24629 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24630 Ok((result?, keeper))
24631 }
24632
24633 pub fn capture_graph_retained_nowarm<F>(
24639 &self,
24640 mut step: F,
24641 ) -> Result<
24642 (
24643 cudarc::driver::CudaGraph,
24644 Vec<Box<dyn std::any::Any + Send>>,
24645 ),
24646 Box<dyn std::error::Error>,
24647 >
24648 where
24649 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24650 {
24651 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24652 let was_tracking = self.gpu.ctx.is_event_tracking();
24653 if was_tracking {
24654 unsafe {
24655 self.gpu.ctx.disable_event_tracking();
24656 }
24657 }
24658 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24659 self.gpu.stream().synchronize()?;
24660 self.gpu
24661 .stream()
24662 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24663 let r = step(self);
24664 let g = self.gpu.stream().end_capture(
24665 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24666 );
24667 r?;
24668 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24669 graph.upload()?;
24670 Ok(graph)
24671 };
24672 let result = run();
24673 if was_tracking {
24674 unsafe {
24675 self.gpu.ctx.enable_event_tracking();
24676 }
24677 }
24678 Ok((result?, Vec::new()))
24679 }
24680
24681 pub fn capture_graph<F>(
24682 &self,
24683 mut step: F,
24684 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24685 where
24686 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24687 {
24688 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24689 let was_tracking = self.gpu.ctx.is_event_tracking();
24697 if was_tracking {
24698 unsafe {
24699 self.gpu.ctx.disable_event_tracking();
24700 }
24701 }
24702 let iflag = {
24709 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24710 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24711 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24714 Ok("priority") => {
24715 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24716 }
24717 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24718 })
24719 };
24720 let ct = {
24727 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24728 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24729 };
24730 let warmups = {
24753 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24754 *W.get_or_init(|| {
24755 std::env::var("MEMRA_GRAPH_WARMUPS")
24756 .ok()
24757 .and_then(|v| v.parse().ok())
24758 .filter(|n| *n >= 1)
24759 .unwrap_or(1)
24760 })
24761 };
24762 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24763 let t_w = std::time::Instant::now();
24764 for _ in 0..warmups {
24766 step(self)?;
24767 }
24768 self.gpu.stream().synchronize()?;
24769 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24770 let t_c = std::time::Instant::now();
24772 self.gpu
24773 .stream()
24774 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24775 let r = step(self);
24778 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24779 let t_i = std::time::Instant::now();
24780 let g = self.gpu.stream().end_capture(iflag);
24781 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24782 r?;
24783 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24784 let t_u = std::time::Instant::now();
24785 graph.upload()?;
24786 if ct {
24787 println!(
24788 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
24789 instantiate {ms_inst:.2} ms upload {:.2} ms",
24790 t_u.elapsed().as_secs_f64() * 1e3
24791 );
24792 }
24793 Ok(graph)
24794 };
24795 let result = run();
24796 if was_tracking {
24797 unsafe {
24798 self.gpu.ctx.enable_event_tracking();
24799 }
24800 }
24801 result
24802 }
24803
24804 pub fn gdn_scan_s128_view(
24806 &self,
24807 q: &CudaSlice<f32>,
24808 k: &CudaSlice<f32>,
24809 v: &CudaSlice<f32>,
24810 g: &CudaSlice<f32>,
24811 beta: &CudaSlice<f32>,
24812 state_in: &cudarc::driver::CudaView<f32>,
24813 state_out: &mut cudarc::driver::CudaViewMut<f32>,
24814 o: &mut CudaSlice<f32>,
24815 n_head: usize,
24816 t: usize,
24817 scale: f32,
24818 ) -> Result<(), Box<dyn std::error::Error>> {
24819 let f = self.func("gdn_scan_s128");
24820 const S_V: u32 = 128;
24821 const WARP: u32 = 32;
24822 const COLS: u32 = 4;
24823 let cfg = LaunchConfig {
24824 grid_dim: (n_head as u32, 1, S_V / COLS),
24825 block_dim: (WARP, COLS, 1),
24826 shared_mem_bytes: 0,
24827 };
24828 let (h, ti) = (n_head as i32, t as i32);
24829 let __s_b = self.gpu.stream();
24830 let mut b = __s_b.launch_builder(&f);
24831 b.arg(q)
24832 .arg(k)
24833 .arg(v)
24834 .arg(g)
24835 .arg(beta)
24836 .arg(state_in)
24837 .arg(state_out)
24838 .arg(o)
24839 .arg(&h)
24840 .arg(&ti)
24841 .arg(&scale);
24842 unsafe {
24843 b.launch(cfg)?;
24844 }
24845 Ok(())
24846 }
24847
24848 pub fn ssm_conv1d_view(
24850 &self,
24851 x: &cudarc::driver::CudaView<f32>,
24852 w: &CudaSlice<f32>,
24853 y: &mut CudaSlice<f32>,
24854 conv_dim: usize,
24855 t: usize,
24856 d_conv: usize,
24857 silu: bool,
24858 ) -> Result<(), Box<dyn std::error::Error>> {
24859 let f = self.func("ssm_conv1d_silu_f32");
24860 let cfg = LaunchConfig {
24862 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24863 block_dim: (256, 1, 1),
24864 shared_mem_bytes: 0,
24865 };
24866 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24867 let __s_b = self.gpu.stream();
24868 let mut b = __s_b.launch_builder(&f);
24869 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24870 unsafe {
24871 b.launch(cfg)?;
24872 }
24873 Ok(())
24874 }
24875
24876 pub fn ssm_conv1d_tm(
24883 &self,
24884 qkv_tm: &CudaSlice<f32>,
24885 w: &CudaSlice<f32>,
24886 y: &mut CudaSlice<f32>,
24887 conv_dim: usize,
24888 t: usize,
24889 d_conv: usize,
24890 ) -> Result<(), Box<dyn std::error::Error>> {
24891 let f = self.func("ssm_conv1d_tm_f32");
24892 let cfg = LaunchConfig {
24893 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24894 block_dim: (256, 1, 1),
24895 shared_mem_bytes: 0,
24896 };
24897 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24898 let __s_b = self.gpu.stream();
24899 let mut b = __s_b.launch_builder(&f);
24900 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24901 unsafe {
24902 b.launch(cfg)?;
24903 }
24904 Ok(())
24905 }
24906
24907 pub fn ssm_conv1d_tm_state(
24915 &self,
24916 qkv_tm: &CudaSlice<f32>,
24917 conv_state: &mut CudaSlice<f32>,
24918 w: &CudaSlice<f32>,
24919 y: &mut CudaSlice<f32>,
24920 conv_dim: usize,
24921 t: usize,
24922 d_conv: usize,
24923 ) -> Result<(), Box<dyn std::error::Error>> {
24924 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24925 }
24926
24927 #[allow(clippy::too_many_arguments)]
24930 pub fn ssm_conv1d_tm_state_pad(
24931 &self,
24932 qkv_tm: &CudaSlice<f32>,
24933 conv_state: &mut CudaSlice<f32>,
24934 w: &CudaSlice<f32>,
24935 y: &mut CudaSlice<f32>,
24936 conv_dim: usize,
24937 t: usize,
24938 d_conv: usize,
24939 pad_len: Option<&CudaSlice<i32>>,
24940 ) -> Result<(), Box<dyn std::error::Error>> {
24941 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24942 let ring_old = if t < d_conv - 1 {
24946 Some(self.clone_dtod(conv_state)?)
24947 } else {
24948 None
24949 };
24950 {
24951 let f = self.func("ssm_conv1d_tm_state_f32");
24952 let cfg = LaunchConfig {
24953 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24954 block_dim: (256, 1, 1),
24955 shared_mem_bytes: 0,
24956 };
24957 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24958 let __s_b = self.gpu.stream();
24959 let mut b = __s_b.launch_builder(&f);
24960 b.arg(qkv_tm)
24961 .arg(&*conv_state)
24962 .arg(w)
24963 .arg(y)
24964 .arg(&cd)
24965 .arg(&ti)
24966 .arg(&dc);
24967 unsafe {
24968 b.launch(cfg)?;
24969 }
24970 }
24971 match (ring_old, pad_len) {
24972 (None, Some(len_d)) => {
24973 let f = self.func("ssm_conv_ring_update_dev_f32");
24974 let n = conv_dim * (d_conv - 1);
24975 let cfg = LaunchConfig::for_num_elems(n as u32);
24976 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24977 let __s_b = self.gpu.stream();
24978 let mut b = __s_b.launch_builder(&f);
24979 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24980 unsafe {
24981 b.launch(cfg)?;
24982 }
24983 }
24984 (None, None) => {
24985 let f = self.func("ssm_conv_ring_update_f32");
24986 let n = conv_dim * (d_conv - 1);
24987 let cfg = LaunchConfig::for_num_elems(n as u32);
24988 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24989 let __s_b = self.gpu.stream();
24990 let mut b = __s_b.launch_builder(&f);
24991 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24992 unsafe {
24993 b.launch(cfg)?;
24994 }
24995 }
24996 (Some(old), _) => {
24997 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
24998 }
24999 }
25000 Ok(())
25001 }
25002
25003 pub fn ssm_conv1d_tm_state_pad_v(
25005 &self,
25006 qkv_tm: &cudarc::driver::CudaView<f32>,
25007 conv_state: &mut CudaSlice<f32>,
25008 w: &CudaSlice<f32>,
25009 y: &mut CudaSlice<f32>,
25010 conv_dim: usize,
25011 t: usize,
25012 d_conv: usize,
25013 pad_len: Option<&CudaSlice<i32>>,
25014 ) -> Result<(), Box<dyn std::error::Error>> {
25015 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25016 let ring_old = if t < d_conv - 1 {
25020 Some(self.clone_dtod(conv_state)?)
25021 } else {
25022 None
25023 };
25024 {
25025 let f = self.func("ssm_conv1d_tm_state_f32");
25026 let cfg = LaunchConfig {
25027 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25028 block_dim: (256, 1, 1),
25029 shared_mem_bytes: 0,
25030 };
25031 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25032 let __s_b = self.gpu.stream();
25033 let mut b = __s_b.launch_builder(&f);
25034 b.arg(qkv_tm)
25035 .arg(&*conv_state)
25036 .arg(w)
25037 .arg(y)
25038 .arg(&cd)
25039 .arg(&ti)
25040 .arg(&dc);
25041 unsafe {
25042 b.launch(cfg)?;
25043 }
25044 }
25045 match (ring_old, pad_len) {
25046 (None, Some(len_d)) => {
25047 let f = self.func("ssm_conv_ring_update_dev_f32");
25048 let n = conv_dim * (d_conv - 1);
25049 let cfg = LaunchConfig::for_num_elems(n as u32);
25050 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25051 let __s_b = self.gpu.stream();
25052 let mut b = __s_b.launch_builder(&f);
25053 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25054 unsafe {
25055 b.launch(cfg)?;
25056 }
25057 }
25058 (None, None) => {
25059 let f = self.func("ssm_conv_ring_update_f32");
25060 let n = conv_dim * (d_conv - 1);
25061 let cfg = LaunchConfig::for_num_elems(n as u32);
25062 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25063 let __s_b = self.gpu.stream();
25064 let mut b = __s_b.launch_builder(&f);
25065 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25066 unsafe {
25067 b.launch(cfg)?;
25068 }
25069 }
25070 (Some(_), _) => unreachable!(
25071 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25072 ),
25073 }
25074 Ok(())
25075 }
25076
25077 pub fn ssm_conv_ring_rebuild(
25082 &self,
25083 qkv_tm: &CudaSlice<f32>,
25084 ring_old: &CudaSlice<f32>,
25085 conv_state: &mut CudaSlice<f32>,
25086 conv_dim: usize,
25087 tc: usize,
25088 d_conv: usize,
25089 ) -> Result<(), Box<dyn std::error::Error>> {
25090 let f = self.func("ssm_conv_ring_rebuild_f32");
25091 let n = conv_dim * (d_conv - 1);
25092 let cfg = LaunchConfig::for_num_elems(n as u32);
25093 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25094 let __s_b = self.gpu.stream();
25095 let mut b = __s_b.launch_builder(&f);
25096 b.arg(qkv_tm)
25097 .arg(ring_old)
25098 .arg(conv_state)
25099 .arg(&cd)
25100 .arg(&ti)
25101 .arg(&dc);
25102 unsafe {
25103 b.launch(cfg)?;
25104 }
25105 Ok(())
25106 }
25107
25108 #[allow(clippy::too_many_arguments)]
25113 pub fn gdn_prep_decode(
25114 &self,
25115 conv_out: &CudaSlice<f32>,
25116 beta_raw: &CudaSlice<f32>,
25117 alpha: &CudaSlice<f32>,
25118 dt_bias: &CudaSlice<f32>,
25119 a: &CudaSlice<f32>,
25120 q_l2: &mut CudaSlice<f32>,
25121 k_l2: &mut CudaSlice<f32>,
25122 v_g: &mut CudaSlice<f32>,
25123 beta: &mut CudaSlice<f32>,
25124 g_log: &mut CudaSlice<f32>,
25125 d_state: usize,
25126 num_v: usize,
25127 num_k: usize,
25128 key_dim: usize,
25129 eps: f32,
25130 ) -> Result<(), Box<dyn std::error::Error>> {
25131 let f = self.func("gdn_prep_decode_f32");
25132 let cfg = LaunchConfig {
25133 grid_dim: (num_v as u32, 1, 1),
25134 block_dim: (32, 4, 1),
25135 shared_mem_bytes: 0,
25136 };
25137 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25138 let __s_b = self.gpu.stream();
25139 let mut b = __s_b.launch_builder(&f);
25140 b.arg(conv_out)
25141 .arg(beta_raw)
25142 .arg(alpha)
25143 .arg(dt_bias)
25144 .arg(a)
25145 .arg(q_l2)
25146 .arg(k_l2)
25147 .arg(v_g)
25148 .arg(beta)
25149 .arg(g_log)
25150 .arg(&ds)
25151 .arg(&nv)
25152 .arg(&nk)
25153 .arg(&kd)
25154 .arg(&eps);
25155 unsafe {
25156 b.launch(cfg)?;
25157 }
25158 Ok(())
25159 }
25160
25161 #[allow(clippy::too_many_arguments)]
25165 pub fn ssm_conv1d_gdn(
25166 &self,
25167 qkv_tm: &CudaSlice<f32>,
25168 w: &CudaSlice<f32>,
25169 q_g: &mut CudaSlice<f32>,
25170 k_g: &mut CudaSlice<f32>,
25171 v_g: &mut CudaSlice<f32>,
25172 conv_dim: usize,
25173 t: usize,
25174 d_conv: usize,
25175 d_state: usize,
25176 num_v: usize,
25177 num_k: usize,
25178 key_dim: usize,
25179 ) -> Result<(), Box<dyn std::error::Error>> {
25180 let f = self.func("ssm_conv1d_gdn_f32");
25181 let cfg = LaunchConfig {
25182 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25183 block_dim: (256, 1, 1),
25184 shared_mem_bytes: 0,
25185 };
25186 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25187 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25188 let __s_b = self.gpu.stream();
25189 let mut b = __s_b.launch_builder(&f);
25190 b.arg(qkv_tm)
25191 .arg(w)
25192 .arg(q_g)
25193 .arg(k_g)
25194 .arg(v_g)
25195 .arg(&cd)
25196 .arg(&ti)
25197 .arg(&dc)
25198 .arg(&ds)
25199 .arg(&nv)
25200 .arg(&nk)
25201 .arg(&kd);
25202 unsafe {
25203 b.launch(cfg)?;
25204 }
25205 Ok(())
25206 }
25207
25208 pub fn ssm_conv1d(
25209 &self,
25210 x: &CudaSlice<f32>,
25211 w: &CudaSlice<f32>,
25212 y: &mut CudaSlice<f32>,
25213 conv_dim: usize,
25214 t: usize,
25215 d_conv: usize,
25216 silu: bool,
25217 ) -> Result<(), Box<dyn std::error::Error>> {
25218 let f = self.func("ssm_conv1d_silu_f32");
25219 let cfg = LaunchConfig {
25220 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25221 block_dim: (256, 1, 1),
25222 shared_mem_bytes: 0,
25223 };
25224 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25225 let __s_b = self.gpu.stream();
25226 let mut b = __s_b.launch_builder(&f);
25227 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25228 unsafe {
25229 b.launch(cfg)?;
25230 }
25231 Ok(())
25232 }
25233
25234 pub fn gdn_scan_s128(
25237 &self,
25238 q: &CudaSlice<f32>,
25239 k: &CudaSlice<f32>,
25240 v: &CudaSlice<f32>,
25241 g: &CudaSlice<f32>,
25242 beta: &CudaSlice<f32>,
25243 state_in: &CudaSlice<f32>,
25244 state_out: &mut CudaSlice<f32>,
25245 o: &mut CudaSlice<f32>,
25246 n_head: usize,
25247 t: usize,
25248 scale: f32,
25249 ) -> Result<(), Box<dyn std::error::Error>> {
25250 let f = self.func("gdn_scan_s128");
25251 const S_V: u32 = 128;
25252 const WARP: u32 = 32;
25253 const COLS_PER_BLOCK: u32 = 4;
25254 let cfg = LaunchConfig {
25255 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25256 block_dim: (WARP, COLS_PER_BLOCK, 1),
25257 shared_mem_bytes: 0,
25258 };
25259 let (h, ti) = (n_head as i32, t as i32);
25260 let __s_b = self.gpu.stream();
25261 let mut b = __s_b.launch_builder(&f);
25262 b.arg(q)
25263 .arg(k)
25264 .arg(v)
25265 .arg(g)
25266 .arg(beta)
25267 .arg(state_in)
25268 .arg(state_out)
25269 .arg(o)
25270 .arg(&h)
25271 .arg(&ti)
25272 .arg(&scale);
25273 unsafe {
25274 b.launch(cfg)?;
25275 }
25276 Ok(())
25277 }
25278
25279 #[allow(clippy::too_many_arguments)]
25284 pub fn ssm_conv1d_fused_decode_b(
25285 &self,
25286 qkv_cols: &CudaSlice<f32>,
25287 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25288 w: &CudaSlice<f32>,
25289 conv_outs: &mut CudaSlice<f32>,
25290 conv_dim: usize,
25291 d_conv: usize,
25292 b_n: usize,
25293 ) -> Result<(), Box<dyn std::error::Error>> {
25294 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25295 let cfg = LaunchConfig {
25296 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25297 block_dim: (256, 1, 1),
25298 shared_mem_bytes: 0,
25299 };
25300 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25301 let __s_b = self.gpu.stream();
25302 let mut b = __s_b.launch_builder(&f);
25303 b.arg(qkv_cols)
25304 .arg(conv_state_ptrs)
25305 .arg(w)
25306 .arg(conv_outs)
25307 .arg(&cd)
25308 .arg(&dc);
25309 unsafe {
25310 b.launch(cfg)?;
25311 }
25312 Ok(())
25313 }
25314
25315 #[allow(clippy::too_many_arguments)]
25316 pub fn gdn_prep_decode_b(
25317 &self,
25318 conv_outs: &CudaSlice<f32>,
25319 beta_raws: &CudaSlice<f32>,
25320 alphas: &CudaSlice<f32>,
25321 dt_bias: &CudaSlice<f32>,
25322 a: &CudaSlice<f32>,
25323 q_l2: &mut CudaSlice<f32>,
25324 k_l2: &mut CudaSlice<f32>,
25325 v_g: &mut CudaSlice<f32>,
25326 beta: &mut CudaSlice<f32>,
25327 g_log: &mut CudaSlice<f32>,
25328 d_state: usize,
25329 num_v: usize,
25330 num_k: usize,
25331 key_dim: usize,
25332 eps: f32,
25333 conv_dim: usize,
25334 b_n: usize,
25335 ) -> Result<(), Box<dyn std::error::Error>> {
25336 let f = self.func("gdn_prep_decode_b_f32");
25337 let cfg = LaunchConfig {
25338 grid_dim: (num_v as u32, 1, b_n as u32),
25339 block_dim: (32, 4, 1),
25340 shared_mem_bytes: 0,
25341 };
25342 let (ds, nv, nk, kd, cd) = (
25343 d_state as i32,
25344 num_v as i32,
25345 num_k as i32,
25346 key_dim as i32,
25347 conv_dim as i32,
25348 );
25349 let __s_b = self.gpu.stream();
25350 let mut b = __s_b.launch_builder(&f);
25351 b.arg(conv_outs)
25352 .arg(beta_raws)
25353 .arg(alphas)
25354 .arg(dt_bias)
25355 .arg(a)
25356 .arg(q_l2)
25357 .arg(k_l2)
25358 .arg(v_g)
25359 .arg(beta)
25360 .arg(g_log)
25361 .arg(&ds)
25362 .arg(&nv)
25363 .arg(&nk)
25364 .arg(&kd)
25365 .arg(&eps)
25366 .arg(&cd);
25367 unsafe {
25368 b.launch(cfg)?;
25369 }
25370 Ok(())
25371 }
25372
25373 #[allow(clippy::too_many_arguments)]
25374 pub fn gdn_scan_s128_batched(
25375 &self,
25376 q: &CudaSlice<f32>,
25377 k: &CudaSlice<f32>,
25378 v: &CudaSlice<f32>,
25379 g: &CudaSlice<f32>,
25380 beta: &CudaSlice<f32>,
25381 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25382 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25383 o: &mut CudaSlice<f32>,
25384 n_head: usize,
25385 b_n: usize,
25386 scale: f32,
25387 ) -> Result<(), Box<dyn std::error::Error>> {
25388 let f = self.func("gdn_scan_s128_b");
25389 const S_V: u32 = 128;
25390 const WARP: u32 = 32;
25391 const COLS_PER_BLOCK: u32 = 4;
25392 let cfg = LaunchConfig {
25393 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25394 block_dim: (WARP, COLS_PER_BLOCK, 1),
25395 shared_mem_bytes: 0,
25396 };
25397 let h = n_head as i32;
25398 let __s_b = self.gpu.stream();
25399 let mut b = __s_b.launch_builder(&f);
25400 b.arg(q)
25401 .arg(k)
25402 .arg(v)
25403 .arg(g)
25404 .arg(beta)
25405 .arg(state_in_ptrs)
25406 .arg(state_out_ptrs)
25407 .arg(o)
25408 .arg(&h)
25409 .arg(&scale);
25410 unsafe {
25411 b.launch(cfg)?;
25412 }
25413 Ok(())
25414 }
25415
25416 #[allow(clippy::too_many_arguments)]
25422 pub fn ssm_conv1d_fused_decode_b_view(
25423 &self,
25424 qkv_cols: &cudarc::driver::CudaView<f32>,
25425 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25426 w: &CudaSlice<f32>,
25427 conv_outs: &mut CudaSlice<f32>,
25428 conv_dim: usize,
25429 d_conv: usize,
25430 b_n: usize,
25431 ) -> Result<(), Box<dyn std::error::Error>> {
25432 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25433 let cfg = LaunchConfig {
25434 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25435 block_dim: (256, 1, 1),
25436 shared_mem_bytes: 0,
25437 };
25438 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25439 let __s_b = self.gpu.stream();
25440 let mut b = __s_b.launch_builder(&f);
25441 b.arg(qkv_cols)
25442 .arg(conv_state_ptrs)
25443 .arg(w)
25444 .arg(conv_outs)
25445 .arg(&cd)
25446 .arg(&dc);
25447 unsafe {
25448 b.launch(cfg)?;
25449 }
25450 Ok(())
25451 }
25452
25453 #[allow(clippy::too_many_arguments)]
25454 pub fn gdn_prep_decode_b_view(
25455 &self,
25456 conv_outs: &CudaSlice<f32>,
25457 beta_raws: &cudarc::driver::CudaView<f32>,
25458 alphas: &cudarc::driver::CudaView<f32>,
25459 dt_bias: &CudaSlice<f32>,
25460 a: &CudaSlice<f32>,
25461 q_l2: &mut CudaSlice<f32>,
25462 k_l2: &mut CudaSlice<f32>,
25463 v_g: &mut CudaSlice<f32>,
25464 beta: &mut CudaSlice<f32>,
25465 g_log: &mut CudaSlice<f32>,
25466 d_state: usize,
25467 num_v: usize,
25468 num_k: usize,
25469 key_dim: usize,
25470 eps: f32,
25471 conv_dim: usize,
25472 b_n: usize,
25473 ) -> Result<(), Box<dyn std::error::Error>> {
25474 let f = self.func("gdn_prep_decode_b_f32");
25475 let cfg = LaunchConfig {
25476 grid_dim: (num_v as u32, 1, b_n as u32),
25477 block_dim: (32, 4, 1),
25478 shared_mem_bytes: 0,
25479 };
25480 let (ds, nv, nk, kd, cd) = (
25481 d_state as i32,
25482 num_v as i32,
25483 num_k as i32,
25484 key_dim as i32,
25485 conv_dim as i32,
25486 );
25487 let __s_b = self.gpu.stream();
25488 let mut b = __s_b.launch_builder(&f);
25489 b.arg(conv_outs)
25490 .arg(beta_raws)
25491 .arg(alphas)
25492 .arg(dt_bias)
25493 .arg(a)
25494 .arg(q_l2)
25495 .arg(k_l2)
25496 .arg(v_g)
25497 .arg(beta)
25498 .arg(g_log)
25499 .arg(&ds)
25500 .arg(&nv)
25501 .arg(&nk)
25502 .arg(&kd)
25503 .arg(&eps)
25504 .arg(&cd);
25505 unsafe {
25506 b.launch(cfg)?;
25507 }
25508 Ok(())
25509 }
25510
25511 #[allow(clippy::too_many_arguments)]
25512 pub fn gdn_scan_s128_batched_view(
25513 &self,
25514 q: &CudaSlice<f32>,
25515 k: &CudaSlice<f32>,
25516 v: &CudaSlice<f32>,
25517 g: &CudaSlice<f32>,
25518 beta: &CudaSlice<f32>,
25519 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25520 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25521 o: &mut cudarc::driver::CudaViewMut<f32>,
25522 n_head: usize,
25523 b_n: usize,
25524 scale: f32,
25525 ) -> Result<(), Box<dyn std::error::Error>> {
25526 let f = self.func("gdn_scan_s128_b");
25527 const S_V: u32 = 128;
25528 const WARP: u32 = 32;
25529 const COLS_PER_BLOCK: u32 = 4;
25530 let cfg = LaunchConfig {
25531 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25532 block_dim: (WARP, COLS_PER_BLOCK, 1),
25533 shared_mem_bytes: 0,
25534 };
25535 let h = n_head as i32;
25536 let __s_b = self.gpu.stream();
25537 let mut b = __s_b.launch_builder(&f);
25538 b.arg(q)
25539 .arg(k)
25540 .arg(v)
25541 .arg(g)
25542 .arg(beta)
25543 .arg(state_in_ptrs)
25544 .arg(state_out_ptrs)
25545 .arg(o)
25546 .arg(&h)
25547 .arg(&scale);
25548 unsafe {
25549 b.launch(cfg)?;
25550 }
25551 Ok(())
25552 }
25553
25554 pub fn gdn_chunked_enabled() -> bool {
25563 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25564 *E.get_or_init(|| {
25565 std::env::var("MEMRA_GDN_CHUNKED")
25566 .map(|v| v != "0")
25567 .unwrap_or(true)
25568 })
25569 }
25570
25571 pub fn gdn_chunk_size() -> usize {
25576 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25577 *C.get_or_init(|| {
25578 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25579 .ok()
25580 .and_then(|v| v.parse().ok())
25581 .unwrap_or(32);
25582 c.clamp(32, 128) / 32 * 32
25583 })
25584 }
25585
25586 #[allow(clippy::too_many_arguments)]
25591 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25594 #[allow(clippy::too_many_arguments)]
25595 pub fn gdn_chunk_k123(
25596 &self,
25597 q: &CudaSlice<f32>,
25598 k: &CudaSlice<f32>,
25599 v: &CudaSlice<f32>,
25600 g: &CudaSlice<f32>,
25601 beta: &CudaSlice<f32>,
25602 wb16: Option<&mut CudaSlice<u8>>,
25603 n_head: usize,
25604 t: usize,
25605 c: usize,
25606 hk: usize,
25607 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25608 ) -> Result<
25609 (
25610 CudaSlice<f32>,
25611 CudaSlice<f32>,
25612 CudaSlice<f32>,
25613 CudaSlice<f32>,
25614 ),
25615 Box<dyn std::error::Error>,
25616 > {
25617 const D: usize = 128;
25618 let h = n_head;
25619 let nc = (t + c - 1) / c;
25620 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25621 let mut gcum = self.uninit(t * h)?;
25622 let mut a = self.uninit(nc * h * c * c)?;
25623 let mut p = self.uninit(nc * h * c * c)?;
25624 let mut u = self.uninit(nc * h * c * D)?;
25625 let mut w = self.uninit(nc * h * c * D)?;
25626 {
25627 let f = self.func("gdn_chunk_cumgate_f32");
25629 let cfg = LaunchConfig {
25630 grid_dim: (nc as u32, h as u32, 1),
25631 block_dim: (32, 1, 1),
25632 shared_mem_bytes: 0,
25633 };
25634 let __s_b = self.gpu.stream();
25635 let mut b = __s_b.launch_builder(&f);
25636 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25637 unsafe {
25638 b.launch(cfg)?;
25639 }
25640 }
25641 if let Some((qb, kb, pb)) = k2w {
25642 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25645 let f = self.func("gdn_k2_wgmma");
25646 let cfg = LaunchConfig {
25647 grid_dim: (nc as u32, h as u32, 1),
25648 block_dim: (128, 1, 1),
25649 shared_mem_bytes: 0,
25650 };
25651 let hki = hk as i32;
25652 let __s_b = self.gpu.stream();
25653 let mut b = __s_b.launch_builder(&f);
25654 b.arg(qb)
25655 .arg(kb)
25656 .arg(&gcum)
25657 .arg(beta)
25658 .arg(&mut a)
25659 .arg(&mut *pb)
25660 .arg(&hi)
25661 .arg(&ti)
25662 .arg(&ci)
25663 .arg(&hki);
25664 unsafe {
25665 b.launch(cfg)?;
25666 }
25667 } else if c <= 64 && !portable_mma_gated() {
25668 let f = self.func("gdn_chunk_attn_f32");
25670 f.set_attribute(
25671 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25672 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25673 )?;
25674 let jt = ((c + 31) / 32) as u32;
25675 let cfg = LaunchConfig {
25676 grid_dim: (nc as u32, h as u32, jt),
25677 block_dim: (256, 1, 1),
25678 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25679 };
25680 let hki = hk as i32;
25681 let __s_b = self.gpu.stream();
25682 let mut b = __s_b.launch_builder(&f);
25683 b.arg(q)
25684 .arg(k)
25685 .arg(&gcum)
25686 .arg(beta)
25687 .arg(&mut a)
25688 .arg(&mut p)
25689 .arg(&hi)
25690 .arg(&ti)
25691 .arg(&ci)
25692 .arg(&hki);
25693 unsafe {
25694 b.launch(cfg)?;
25695 }
25696 } else {
25697 assert!(
25699 hk == h,
25700 "generic K2 is broadcast-only (de-broadcast rides C==32)"
25701 );
25702 let f = self.func("gdn_chunk_attn_g_f32");
25703 let cfg = LaunchConfig {
25704 grid_dim: (nc as u32, h as u32, 1),
25705 block_dim: (32, 8, 1),
25706 shared_mem_bytes: 0,
25707 };
25708 let __s_b = self.gpu.stream();
25709 let mut b = __s_b.launch_builder(&f);
25710 b.arg(q)
25711 .arg(k)
25712 .arg(&gcum)
25713 .arg(beta)
25714 .arg(&mut a)
25715 .arg(&mut p)
25716 .arg(&hi)
25717 .arg(&ti)
25718 .arg(&ci);
25719 unsafe {
25720 b.launch(cfg)?;
25721 }
25722 }
25723 {
25724 let cfg = LaunchConfig {
25726 grid_dim: (nc as u32, h as u32, 1),
25727 block_dim: (256, 1, 1),
25728 shared_mem_bytes: 0,
25729 };
25730 match c {
25731 32 | 64 => {
25732 let f = self.func(if c == 32 {
25733 "gdn_chunk_solve32_f32"
25734 } else {
25735 "gdn_chunk_solve64_f32"
25736 });
25737 let wb: u64 = match wb16 {
25739 Some(d) => self.addr_u8(d),
25740 None => 0,
25741 };
25742 let hki = hk as i32;
25743 let __s_b = self.gpu.stream();
25744 let mut b = __s_b.launch_builder(&f);
25745 b.arg(v)
25746 .arg(k)
25747 .arg(&a)
25748 .arg(&gcum)
25749 .arg(&mut u)
25750 .arg(&mut w)
25751 .arg(&wb)
25752 .arg(&hi)
25753 .arg(&ti)
25754 .arg(&hki);
25755 unsafe {
25756 b.launch(cfg)?;
25757 }
25758 }
25759 _ => {
25760 assert!(hk == h, "generic K3 is broadcast-only");
25761 let f = self.func("gdn_chunk_solve_f32");
25762 let __s_b = self.gpu.stream();
25763 let mut b = __s_b.launch_builder(&f);
25764 b.arg(v)
25765 .arg(k)
25766 .arg(&a)
25767 .arg(&gcum)
25768 .arg(&mut u)
25769 .arg(&mut w)
25770 .arg(&hi)
25771 .arg(&ti)
25772 .arg(&ci);
25773 unsafe {
25774 b.launch(cfg)?;
25775 }
25776 }
25777 }
25778 }
25779 Ok((gcum, p, u, w))
25780 }
25781
25782 pub fn gdn_db_on() -> bool {
25786 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25787 }
25788
25789 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25798 !portable_mma_gated()
25799 && c == 32
25800 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25801 Ok("1") => true,
25802 Ok("0") => false,
25803 _ => gdn_mma_default_on(),
25804 }
25805 }
25806
25807 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25814 cfg!(memra_hopper_mma)
25815 && self.gdn_mma_enabled(c)
25816 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25817 }
25818
25819 #[allow(clippy::too_many_arguments)]
25824 pub fn ssm_conv1d_gdn_state_pad(
25825 &self,
25826 qkv_tm: &cudarc::driver::CudaView<f32>,
25827 conv_state: &mut CudaSlice<f32>,
25828 w: &CudaSlice<f32>,
25829 q_g: &mut CudaSlice<f32>,
25830 k_g: &mut CudaSlice<f32>,
25831 v_g: &mut CudaSlice<f32>,
25832 conv_dim: usize,
25833 t: usize,
25834 d_conv: usize,
25835 d_state: usize,
25836 num_v: usize,
25837 num_k: usize,
25838 key_dim: usize,
25839 hk: usize,
25840 pad_len: Option<&CudaSlice<i32>>,
25841 ) -> Result<(), Box<dyn std::error::Error>> {
25842 assert!(
25843 t >= d_conv - 1,
25844 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25845 );
25846 {
25847 let f = self.func("ssm_conv1d_gdn_state_f32");
25848 let cfg = LaunchConfig {
25849 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25850 block_dim: (256, 1, 1),
25851 shared_mem_bytes: 0,
25852 };
25853 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25854 let (ds, nv, nk, kd, hki) = (
25855 d_state as i32,
25856 num_v as i32,
25857 num_k as i32,
25858 key_dim as i32,
25859 hk as i32,
25860 );
25861 let __s_b = self.gpu.stream();
25862 let mut b = __s_b.launch_builder(&f);
25863 b.arg(qkv_tm)
25864 .arg(&*conv_state)
25865 .arg(w)
25866 .arg(q_g)
25867 .arg(k_g)
25868 .arg(v_g)
25869 .arg(&cd)
25870 .arg(&ti)
25871 .arg(&dc)
25872 .arg(&ds)
25873 .arg(&nv)
25874 .arg(&nk)
25875 .arg(&kd)
25876 .arg(&hki);
25877 unsafe {
25878 b.launch(cfg)?;
25879 }
25880 }
25881 match pad_len {
25882 Some(len_d) => {
25883 let f = self.func("ssm_conv_ring_update_dev_f32");
25884 let n = conv_dim * (d_conv - 1);
25885 let cfg = LaunchConfig::for_num_elems(n as u32);
25886 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25887 let __s_b = self.gpu.stream();
25888 let mut b = __s_b.launch_builder(&f);
25889 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25890 unsafe {
25891 b.launch(cfg)?;
25892 }
25893 }
25894 None => {
25895 let f = self.func("ssm_conv_ring_update_f32");
25896 let n = conv_dim * (d_conv - 1);
25897 let cfg = LaunchConfig::for_num_elems(n as u32);
25898 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25899 let __s_b = self.gpu.stream();
25900 let mut b = __s_b.launch_builder(&f);
25901 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25902 unsafe {
25903 b.launch(cfg)?;
25904 }
25905 }
25906 }
25907 Ok(())
25908 }
25909
25910 pub fn gdn_chunk_alloc(
25914 &self,
25915 n_head: usize,
25916 t: usize,
25917 c: usize,
25918 hk: usize,
25919 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25920 const D: usize = 128;
25921 assert!(
25922 c == 32,
25923 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25924 );
25925 let h = n_head;
25926 let nc = (t + c - 1) / c;
25927 Ok(GdnChunkBufs {
25928 gcum: self.uninit(t * h)?,
25929 a: self.uninit(nc * h * c * c)?,
25930 p: self.uninit(nc * h * c * c)?,
25931 u: self.uninit(nc * h * c * D)?,
25932 w: self.uninit(nc * h * c * D)?,
25933 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25934 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25935 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25936 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25937 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25938 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25939 o: self.uninit(D * h * t)?,
25940 t,
25941 nc,
25942 })
25943 }
25944
25945 pub fn f32_to_bf16_v(
25947 &self,
25948 x: &cudarc::driver::CudaView<f32>,
25949 dst: &mut CudaSlice<u8>,
25950 n: usize,
25951 ) -> Result<(), Box<dyn std::error::Error>> {
25952 let f = self.func("f32_to_bf16_bulk");
25953 let ni = n as i64;
25954 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25955 let __s_b = self.gpu.stream();
25956 let mut b = __s_b.launch_builder(&f);
25957 b.arg(x).arg(dst).arg(&ni);
25958 unsafe {
25959 b.launch(cfg)?;
25960 }
25961 Ok(())
25962 }
25963
25964 pub fn f32_to_bf16_into(
25966 &self,
25967 x: &CudaSlice<f32>,
25968 dst: &mut CudaSlice<u8>,
25969 n: usize,
25970 ) -> Result<(), Box<dyn std::error::Error>> {
25971 let f = self.func("f32_to_bf16_bulk");
25972 let ni = n as i64;
25973 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25974 let __s_b = self.gpu.stream();
25975 let mut b = __s_b.launch_builder(&f);
25976 b.arg(x).arg(dst).arg(&ni);
25977 unsafe {
25978 b.launch(cfg)?;
25979 }
25980 Ok(())
25981 }
25982
25983 pub fn gdn_chunk_k123_vl8(
25986 &self,
25987 seqs: &[GdnSeqVl],
25988 n_head: usize,
25989 hk: usize,
25990 wq: Option<&GdnWVl8>,
25991 ) -> Result<(), Box<dyn std::error::Error>> {
25992 let b = seqs.len();
25993 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
25994 let mut packed = [GdnSeqVl::default(); 8];
25995 packed[..b].copy_from_slice(seqs);
25996 let v = GdnVl8(packed);
25997 let (hi, ci) = (n_head as i32, 32i32);
25998 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25999 {
26000 let f = self.func("gdn_chunk_cumgate_vl");
26001 let cfg = LaunchConfig {
26002 grid_dim: (max_nc, n_head as u32, b as u32),
26003 block_dim: (32, 1, 1),
26004 shared_mem_bytes: 0,
26005 };
26006 let __s_lb = self.gpu.stream();
26007 let mut lb = __s_lb.launch_builder(&f);
26008 lb.arg(&v).arg(&hi).arg(&ci);
26009 unsafe {
26010 lb.launch(cfg)?;
26011 }
26012 }
26013 let hki = hk as i32;
26014 if let Some(w) = wq {
26015 let f = self.func("gdn_k2_wgmma_vl");
26017 let cfg = LaunchConfig {
26018 grid_dim: (max_nc, n_head as u32, b as u32),
26019 block_dim: (128, 1, 1),
26020 shared_mem_bytes: 0,
26021 };
26022 let __s_lb = self.gpu.stream();
26023 let mut lb = __s_lb.launch_builder(&f);
26024 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26025 unsafe {
26026 lb.launch(cfg)?;
26027 }
26028 } else {
26029 let f = self.func("gdn_chunk_attn_vl");
26030 f.set_attribute(
26031 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26032 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26033 )?;
26034 let cfg = LaunchConfig {
26035 grid_dim: (max_nc, n_head as u32, b as u32),
26036 block_dim: (256, 1, 1),
26037 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26038 };
26039 let __s_lb = self.gpu.stream();
26040 let mut lb = __s_lb.launch_builder(&f);
26041 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26042 unsafe {
26043 lb.launch(cfg)?;
26044 }
26045 }
26046 {
26047 let f = self.func("gdn_chunk_solve32_vl");
26048 let cfg = LaunchConfig {
26049 grid_dim: (max_nc, n_head as u32, b as u32),
26050 block_dim: (256, 1, 1),
26051 shared_mem_bytes: 0,
26052 };
26053 let __s_lb = self.gpu.stream();
26054 let mut lb = __s_lb.launch_builder(&f);
26055 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26056 unsafe {
26057 lb.launch(cfg)?;
26058 }
26059 }
26060 Ok(())
26061 }
26062
26063 #[allow(clippy::too_many_arguments)]
26067 pub fn gdn_prep_vl8(
26068 &self,
26069 seqs: &[GdnPrepVl],
26070 conv_w: &CudaSlice<f32>,
26071 dt_bias: &CudaSlice<f32>,
26072 a: &CudaSlice<f32>,
26073 conv_dim: usize,
26074 d_conv: usize,
26075 d_state: usize,
26076 num_v: usize,
26077 num_k: usize,
26078 key_dim: usize,
26079 hk: usize,
26080 eps: f32,
26081 ) -> Result<(), Box<dyn std::error::Error>> {
26082 let b = seqs.len();
26083 assert!(b >= 1 && b <= 8);
26084 let mut packed = [GdnPrepVl::default(); 8];
26085 packed[..b].copy_from_slice(seqs);
26086 let v = GdnPrepVl8(packed);
26087 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26088 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26089 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26090 assert!(
26091 conv_fuse || hk == num_v,
26092 "de-broadcast requires the fused conv"
26093 );
26094 if conv_fuse {
26095 let f = self.func("ssm_conv1d_gdn_state_vl");
26096 let cfg = LaunchConfig {
26097 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26098 block_dim: (256, 1, 1),
26099 shared_mem_bytes: 0,
26100 };
26101 let (dsi, nvi, nki, kdi, hki) = (
26102 d_state as i32,
26103 num_v as i32,
26104 num_k as i32,
26105 key_dim as i32,
26106 hk as i32,
26107 );
26108 let __s_lb = self.gpu.stream();
26109 let mut lb = __s_lb.launch_builder(&f);
26110 lb.arg(&v)
26111 .arg(conv_w)
26112 .arg(&cdi)
26113 .arg(&dci)
26114 .arg(&dsi)
26115 .arg(&nvi)
26116 .arg(&nki)
26117 .arg(&kdi)
26118 .arg(&hki);
26119 unsafe {
26120 lb.launch(cfg)?;
26121 }
26122 } else {
26123 let f = self.func("ssm_conv1d_tm_state_vl");
26124 let cfg = LaunchConfig {
26125 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26126 block_dim: (256, 1, 1),
26127 shared_mem_bytes: 0,
26128 };
26129 let __s_lb = self.gpu.stream();
26130 let mut lb = __s_lb.launch_builder(&f);
26131 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26132 unsafe {
26133 lb.launch(cfg)?;
26134 }
26135 }
26136 {
26137 let f = self.func("ssm_conv_ring_update_vl");
26138 let n = (conv_dim * (d_conv - 1)) as u32;
26139 let cfg = LaunchConfig {
26140 grid_dim: (n.div_ceil(256), 1, b as u32),
26141 block_dim: (256, 1, 1),
26142 shared_mem_bytes: 0,
26143 };
26144 let __s_lb = self.gpu.stream();
26145 let mut lb = __s_lb.launch_builder(&f);
26146 lb.arg(&v).arg(&cdi).arg(&dci);
26147 unsafe {
26148 lb.launch(cfg)?;
26149 }
26150 }
26151 if !conv_fuse {
26152 let f = self.func("qkv_to_gdn_repack_vl");
26153 let n = max_t * (num_v * d_state) as u32;
26154 let cfg = LaunchConfig {
26155 grid_dim: (n.div_ceil(256), 1, b as u32),
26156 block_dim: (256, 1, 1),
26157 shared_mem_bytes: 0,
26158 };
26159 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26160 let __s_lb = self.gpu.stream();
26161 let mut lb = __s_lb.launch_builder(&f);
26162 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26163 unsafe {
26164 lb.launch(cfg)?;
26165 }
26166 }
26167 if Self::l2_v2_on(d_state) {
26168 let f = self.func("gdn_l2_v2_vl");
26169 let cfg = LaunchConfig {
26170 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26171 block_dim: (256, 1, 1),
26172 shared_mem_bytes: 0,
26173 };
26174 let (dsi, nvi) = (d_state as i32, hk as i32);
26175 let __s_lb = self.gpu.stream();
26176 let mut lb = __s_lb.launch_builder(&f);
26177 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26178 unsafe {
26179 lb.launch(cfg)?;
26180 }
26181 } else {
26182 let f = self.func("gdn_l2_vl");
26183 let cfg = LaunchConfig {
26184 grid_dim: (max_t * hk as u32, 2, b as u32),
26185 block_dim: (256, 1, 1),
26186 shared_mem_bytes: 0,
26187 };
26188 let (dsi, nvi) = (d_state as i32, hk as i32);
26189 let __s_lb = self.gpu.stream();
26190 let mut lb = __s_lb.launch_builder(&f);
26191 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26192 unsafe {
26193 lb.launch(cfg)?;
26194 }
26195 }
26196 {
26197 let f = self.func("gdn_gate_prep_vl");
26198 let n = max_t * num_v as u32;
26199 let cfg = LaunchConfig {
26200 grid_dim: (n.div_ceil(256), 1, b as u32),
26201 block_dim: (256, 1, 1),
26202 shared_mem_bytes: 0,
26203 };
26204 let nvi = num_v as i32;
26205 let __s_lb = self.gpu.stream();
26206 let mut lb = __s_lb.launch_builder(&f);
26207 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26208 unsafe {
26209 lb.launch(cfg)?;
26210 }
26211 }
26212 Ok(())
26213 }
26214
26215 pub fn gdn_mirror_vl8(
26217 &self,
26218 seqs: &[GdnSeqVl],
26219 n_head: usize,
26220 which: i32,
26221 hk: usize,
26222 ) -> Result<(), Box<dyn std::error::Error>> {
26223 let b = seqs.len();
26224 assert!(b >= 1 && b <= 8);
26225 let mut packed = [GdnSeqVl::default(); 8];
26226 packed[..b].copy_from_slice(seqs);
26227 let v = GdnVl8(packed);
26228 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26229 let max_n = seqs
26230 .iter()
26231 .map(|s| {
26232 if which == 0 {
26233 s.t as i64 * ept as i64
26234 } else {
26235 s.nc as i64 * ept as i64 * 32
26236 }
26237 })
26238 .max()
26239 .unwrap();
26240 let f = self.func("gdn_mirror_vl");
26241 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26242 let cfg = LaunchConfig {
26243 grid_dim: (blocks, 1, b as u32),
26244 block_dim: (256, 1, 1),
26245 shared_mem_bytes: 0,
26246 };
26247 let __s_lb = self.gpu.stream();
26248 let mut lb = __s_lb.launch_builder(&f);
26249 lb.arg(&v).arg(&ept).arg(&which);
26250 unsafe {
26251 lb.launch(cfg)?;
26252 }
26253 Ok(())
26254 }
26255
26256 pub fn gdn_tail_vl8(
26258 &self,
26259 seqs: &[GdnPrepVl],
26260 norm_w: &CudaSlice<f32>,
26261 d_state: usize,
26262 num_v: usize,
26263 eps: f32,
26264 ) -> Result<(), Box<dyn std::error::Error>> {
26265 let b = seqs.len();
26266 assert!(b >= 1 && b <= 8);
26267 let mut packed = [GdnPrepVl::default(); 8];
26268 packed[..b].copy_from_slice(seqs);
26269 let v = GdnPrepVl8(packed);
26270 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26271 let f = self.func("gated_rmsnorm_f16out_vl");
26272 let cfg = LaunchConfig {
26274 grid_dim: (max_t * num_v as u32, 1, b as u32),
26275 block_dim: (128, 1, 1),
26276 shared_mem_bytes: 0,
26277 };
26278 let (dsi, nvi) = (d_state as i32, num_v as i32);
26279 let __s_lb = self.gpu.stream();
26280 let mut lb = __s_lb.launch_builder(&f);
26281 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26282 unsafe {
26283 lb.launch(cfg)?;
26284 }
26285 Ok(())
26286 }
26287
26288 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26291 use cudarc::driver::DevicePtr;
26292 let s = self.gpu.stream();
26293 let (p, _g) = x.device_ptr(&s);
26294 p as u64
26295 }
26296 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26297 use cudarc::driver::DevicePtrMut;
26298 let s = self.gpu.stream();
26299 let (p, _g) = x.device_ptr_mut(&s);
26300 p as u64
26301 }
26302 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26303 use cudarc::driver::DevicePtr;
26304 let s = self.gpu.stream();
26305 let (p, _g) = x.device_ptr(&s);
26306 p as u64
26307 }
26308 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26309 use cudarc::driver::DevicePtr;
26310 let s = self.gpu.stream();
26311 let (p, _g) = x.device_ptr(&s);
26312 p as u64
26313 }
26314
26315 pub fn gdn_chunk_vl8(
26319 &self,
26320 seqs: &[GdnSeqVl],
26321 n_head: usize,
26322 scale: f32,
26323 hk: usize,
26324 wq: Option<&GdnWVl8>,
26325 ) -> Result<(), Box<dyn std::error::Error>> {
26326 const NSPLIT: u32 = 4;
26327 let b = seqs.len();
26328 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26329 let mut packed = [GdnSeqVl::default(); 8];
26330 packed[..b].copy_from_slice(seqs);
26331 let v = GdnVl8(packed);
26332 let (hi, ci) = (n_head as i32, 32i32);
26333 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26334 let hki = hk as i32;
26335 if let Some(w) = wq {
26336 let f = self.func("gdn_k45_wgmma_vl");
26338 let cfg = LaunchConfig {
26339 grid_dim: (n_head as u32, NSPLIT, b as u32),
26340 block_dim: (256, 1, 1),
26341 shared_mem_bytes: 0,
26342 };
26343 let __s_lb = self.gpu.stream();
26344 let mut lb = __s_lb.launch_builder(&f);
26345 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26346 unsafe {
26347 lb.launch(cfg)?;
26348 }
26349 let _ = max_nc;
26350 return Ok(());
26351 }
26352 {
26353 let f = self.func("gdn_chunk_state_mma_vl");
26354 let cfg = LaunchConfig {
26355 grid_dim: (n_head as u32, NSPLIT, b as u32),
26356 block_dim: (256, 1, 1),
26357 shared_mem_bytes: 0,
26358 };
26359 let __s_lb = self.gpu.stream();
26360 let mut lb = __s_lb.launch_builder(&f);
26361 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26362 unsafe {
26363 lb.launch(cfg)?;
26364 }
26365 }
26366 {
26367 let f = self.func("gdn_chunk_output_mma_vl");
26368 let cfg = LaunchConfig {
26369 grid_dim: (max_nc, n_head as u32, b as u32),
26370 block_dim: (256, 1, 1),
26371 shared_mem_bytes: 0,
26372 };
26373 let __s_lb = self.gpu.stream();
26374 let mut lb = __s_lb.launch_builder(&f);
26375 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26376 unsafe {
26377 lb.launch(cfg)?;
26378 }
26379 }
26380 Ok(())
26381 }
26382 pub fn gdn_scan_chunked(
26383 &self,
26384 q: &CudaSlice<f32>,
26385 k: &CudaSlice<f32>,
26386 v: &CudaSlice<f32>,
26387 g: &CudaSlice<f32>,
26388 beta: &CudaSlice<f32>,
26389 kb16_pre: Option<&CudaSlice<u8>>,
26390 qb16_pre: Option<&CudaSlice<u8>>,
26391 state_in: &CudaSlice<f32>,
26392 state_out: &mut CudaSlice<f32>,
26393 o: &mut CudaSlice<f32>,
26394 n_head: usize,
26395 t: usize,
26396 scale: f32,
26397 c: usize,
26398 hk: usize,
26399 ) -> Result<(), Box<dyn std::error::Error>> {
26400 const D: usize = 128;
26401 const NSPLIT: u32 = 4;
26402 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26403 let h = n_head;
26404 let nc = (t + c - 1) / c;
26405 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26406 let gdn_mma_pre = !portable_mma_gated()
26411 && c == 32
26412 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26413 Ok("1") => true,
26414 Ok("0") => false,
26415 _ => gdn_mma_default_on(),
26416 };
26417 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26418 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26419 } else {
26420 None
26421 };
26422 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26427 && gdn_mma_pre
26428 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26429 let nk = t * hk * D;
26430 let mut kb16_local: Option<CudaSlice<u8>> = None;
26431 if gdn_mma_pre && kb16_pre.is_none() {
26432 let mut kb = self.alloc_u8_uninit(nk * 2)?;
26433 let f = self.func("f32_to_bf16_bulk");
26434 let n2 = nk as i64;
26435 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26436 let __s_b = self.gpu.stream();
26437 let mut b = __s_b.launch_builder(&f);
26438 b.arg(k).arg(&mut kb).arg(&n2);
26439 unsafe {
26440 b.launch(cfg2)?;
26441 }
26442 kb16_local = Some(kb);
26443 }
26444 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26445 if let Some(kb) = kb16_pre {
26446 assert!(kb.len() >= nk * 2, "kb16_pre too small");
26447 }
26448 let mut qb16: Option<CudaSlice<u8>> = None;
26449 let mut pb16: Option<CudaSlice<u8>> = None;
26450 if gdn_wgmma_pre {
26451 if qb16_pre.is_none() {
26454 let mut qb = self.alloc_u8_uninit(nk * 2)?;
26455 let f = self.func("f32_to_bf16_bulk");
26456 let n2 = nk as i64;
26457 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26458 let __s_b = self.gpu.stream();
26459 let mut b = __s_b.launch_builder(&f);
26460 b.arg(q).arg(&mut qb).arg(&n2);
26461 unsafe {
26462 b.launch(cfg2)?;
26463 }
26464 qb16 = Some(qb);
26465 } else if let Some(qb) = qb16_pre {
26466 assert!(qb.len() >= nk * 2, "qb16_pre too small");
26467 }
26468 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26469 }
26470 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26471 let k2w = if gdn_wgmma_pre {
26472 Some((
26473 *qb16_ref0.as_ref().unwrap(),
26474 *kb16_ref0.as_ref().unwrap(),
26475 pb16.as_mut().unwrap(),
26476 ))
26477 } else {
26478 None
26479 };
26480 let (gcum, p, u, w) =
26481 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26482 let _ = &w;
26483 let mut y = self.uninit(nc * h * c * D)?;
26484 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
26498 && c == 32
26499 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26500 Ok("1") => true,
26501 Ok("0") => false,
26502 _ => gdn_mma_default_on(),
26503 };
26504 if gdn_mma {
26505 let wb16 = wb16_pre
26506 .take()
26507 .expect("mma path pre-allocates wb16 (K3 store fold)");
26508 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26509 if gdn_wgmma_pre {
26521 let qb16 = qb16_ref0.unwrap();
26523 let pb16 = pb16.as_ref().unwrap();
26524 {
26525 let f = self.func("gdn_k45_wgmma");
26526 let cfg = LaunchConfig {
26527 grid_dim: (h as u32, 4, 1),
26528 block_dim: (256, 1, 1),
26529 shared_mem_bytes: 0,
26530 };
26531 let hki = hk as i32;
26532 let __s_b = self.gpu.stream();
26533 let mut b = __s_b.launch_builder(&f);
26534 b.arg(kb16_ref)
26535 .arg(&gcum)
26536 .arg(beta)
26537 .arg(&u)
26538 .arg(&wb16)
26539 .arg(qb16)
26540 .arg(pb16)
26541 .arg(o)
26542 .arg(&scale)
26543 .arg(state_in)
26544 .arg(&mut *state_out)
26545 .arg(&hi)
26546 .arg(&ti)
26547 .arg(&ci)
26548 .arg(&hki);
26549 unsafe {
26550 b.launch(cfg)?;
26551 }
26552 }
26553 return Ok(());
26554 }
26555 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26559 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26560 {
26561 let f = self.func("gdn_chunk_state_mma");
26562 let cfg = LaunchConfig {
26563 grid_dim: (h as u32, NSPLIT, 1),
26564 block_dim: (256, 1, 1),
26565 shared_mem_bytes: 0,
26566 };
26567 let hki = hk as i32;
26568 let __s_b = self.gpu.stream();
26569 let mut b = __s_b.launch_builder(&f);
26570 b.arg(kb16_ref)
26571 .arg(&gcum)
26572 .arg(beta)
26573 .arg(&u)
26574 .arg(&wb16)
26575 .arg(&mut y16)
26576 .arg(&mut ssnap16)
26577 .arg(state_in)
26578 .arg(&mut *state_out)
26579 .arg(&hi)
26580 .arg(&ti)
26581 .arg(&ci)
26582 .arg(&hki);
26583 unsafe {
26584 b.launch(cfg)?;
26585 }
26586 }
26587 {
26588 let f = self.func("gdn_chunk_output_mma");
26590 let jt = ((c + 31) / 32) as u32;
26591 let cfg = LaunchConfig {
26592 grid_dim: (nc as u32, h as u32, jt),
26593 block_dim: (256, 1, 1),
26594 shared_mem_bytes: 0,
26595 };
26596 let hki = hk as i32;
26597 let __s_b = self.gpu.stream();
26598 let mut b = __s_b.launch_builder(&f);
26599 b.arg(q)
26600 .arg(&gcum)
26601 .arg(&p)
26602 .arg(&y16)
26603 .arg(&ssnap16)
26604 .arg(o)
26605 .arg(&hi)
26606 .arg(&ti)
26607 .arg(&ci)
26608 .arg(&scale)
26609 .arg(&hki);
26610 unsafe {
26611 b.launch(cfg)?;
26612 }
26613 }
26614 return Ok(());
26615 }
26616 {
26617 let f = self.func("gdn_chunk_state_f32");
26619 let cfg = LaunchConfig {
26620 grid_dim: (h as u32, NSPLIT, 1),
26621 block_dim: (256, 1, 1),
26622 shared_mem_bytes: 0,
26623 };
26624 let __s_b = self.gpu.stream();
26625 let mut b = __s_b.launch_builder(&f);
26626 b.arg(k)
26627 .arg(&gcum)
26628 .arg(beta)
26629 .arg(&u)
26630 .arg(&w)
26631 .arg(&mut y)
26632 .arg(&mut ssnap)
26633 .arg(state_in)
26634 .arg(&mut *state_out)
26635 .arg(&hi)
26636 .arg(&ti)
26637 .arg(&ci);
26638 unsafe {
26639 b.launch(cfg)?;
26640 }
26641 }
26642 {
26643 let f = self.func("gdn_chunk_output_f32");
26645 let jt = ((c + 31) / 32) as u32;
26646 let cfg = LaunchConfig {
26647 grid_dim: (nc as u32, h as u32, jt),
26648 block_dim: (256, 1, 1),
26649 shared_mem_bytes: 0,
26650 };
26651 let __s_b = self.gpu.stream();
26652 let mut b = __s_b.launch_builder(&f);
26653 b.arg(q)
26654 .arg(&gcum)
26655 .arg(&p)
26656 .arg(&y)
26657 .arg(&ssnap)
26658 .arg(o)
26659 .arg(&hi)
26660 .arg(&ti)
26661 .arg(&ci)
26662 .arg(&scale);
26663 unsafe {
26664 b.launch(cfg)?;
26665 }
26666 }
26667 Ok(())
26668 }
26669
26670 #[allow(clippy::too_many_arguments)]
26679 #[allow(clippy::too_many_arguments)]
26680 pub fn gdn_scan_prefill(
26681 &self,
26682 q: &CudaSlice<f32>,
26683 k: &CudaSlice<f32>,
26684 v: &CudaSlice<f32>,
26685 g: &CudaSlice<f32>,
26686 beta: &CudaSlice<f32>,
26687 kb16_pre: Option<&CudaSlice<u8>>,
26688 qb16_pre: Option<&CudaSlice<u8>>,
26689 state_in: &CudaSlice<f32>,
26690 state_out: &mut CudaSlice<f32>,
26691 o: &mut CudaSlice<f32>,
26692 n_head: usize,
26693 t: usize,
26694 scale: f32,
26695 hk: usize,
26696 ) -> Result<(), Box<dyn std::error::Error>> {
26697 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26698 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26699 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26700 }
26701 if Self::gdn_chunked_enabled() && t >= 16 {
26702 self.gdn_scan_chunked(
26703 q,
26704 k,
26705 v,
26706 g,
26707 beta,
26708 kb16_pre,
26709 qb16_pre,
26710 state_in,
26711 state_out,
26712 o,
26713 n_head,
26714 t,
26715 scale,
26716 Self::gdn_chunk_size(),
26717 hk,
26718 )
26719 } else {
26720 assert!(
26721 hk == n_head,
26722 "s128 scan is broadcast-only (prep guarantees by predicate)"
26723 );
26724 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26725 }
26726 }
26727
26728 #[allow(clippy::too_many_arguments)]
26730 fn gdn_scan_diff(
26731 &self,
26732 q: &CudaSlice<f32>,
26733 k: &CudaSlice<f32>,
26734 v: &CudaSlice<f32>,
26735 g: &CudaSlice<f32>,
26736 beta: &CudaSlice<f32>,
26737 state_in: &CudaSlice<f32>,
26738 state_out: &mut CudaSlice<f32>,
26739 o: &mut CudaSlice<f32>,
26740 n_head: usize,
26741 t: usize,
26742 scale: f32,
26743 ) -> Result<(), Box<dyn std::error::Error>> {
26744 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26745 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26746 let mut o_c = self.uninit(o.len())?;
26747 let mut st_c = self.uninit(state_out.len())?;
26748 self.gdn_scan_chunked(
26749 q,
26750 k,
26751 v,
26752 g,
26753 beta,
26754 None,
26755 None,
26756 state_in,
26757 &mut st_c,
26758 &mut o_c,
26759 n_head,
26760 t,
26761 scale,
26762 Self::gdn_chunk_size(),
26763 n_head,
26764 )?;
26765 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26766 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26767 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26768 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26769 let mut max_abs = 0f32;
26770 let mut max_rel = 0f32;
26771 let mut sum_rel = 0f64;
26772 for (x, y) in a.iter().zip(b) {
26773 let ad = (x - y).abs();
26774 let rel = ad / x.abs().max(y.abs()).max(1e-3);
26775 if ad > max_abs {
26776 max_abs = ad;
26777 }
26778 if rel > max_rel {
26779 max_rel = rel;
26780 }
26781 sum_rel += rel as f64;
26782 }
26783 (max_abs, max_rel, sum_rel / a.len() as f64)
26784 };
26785 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26786 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26787 println!(
26788 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26789 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26790 Self::gdn_chunk_size()
26791 );
26792 Ok(())
26793 }
26794
26795 pub fn gdn_glog(
26797 &self,
26798 alpha: &CudaSlice<f32>,
26799 dt_bias: &CudaSlice<f32>,
26800 a: &CudaSlice<f32>,
26801 g_log: &mut CudaSlice<f32>,
26802 n_head: usize,
26803 t: usize,
26804 ) -> Result<(), Box<dyn std::error::Error>> {
26805 let f = self.func("gdn_glog_f32");
26806 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26807 let (h, ti) = (n_head as i32, t as i32);
26808 let __s_b = self.gpu.stream();
26809 let mut b = __s_b.launch_builder(&f);
26810 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26811 unsafe {
26812 b.launch(cfg)?;
26813 }
26814 Ok(())
26815 }
26816
26817 pub fn sigmoid_v(
26820 &self,
26821 x: &cudarc::driver::CudaView<f32>,
26822 y: &mut CudaSlice<f32>,
26823 n: usize,
26824 ) -> Result<(), Box<dyn std::error::Error>> {
26825 let f = self.func("sigmoid_f32");
26826 let cfg = LaunchConfig::for_num_elems(n as u32);
26827 let ni = n as i32;
26828 let __s_b = self.gpu.stream();
26829 let mut b = __s_b.launch_builder(&f);
26830 b.arg(x).arg(y).arg(&ni);
26831 unsafe {
26832 b.launch(cfg)?;
26833 }
26834 Ok(())
26835 }
26836
26837 pub fn gdn_glog_v(
26838 &self,
26839 alpha: &cudarc::driver::CudaView<f32>,
26840 dt_bias: &CudaSlice<f32>,
26841 a: &CudaSlice<f32>,
26842 g_log: &mut CudaSlice<f32>,
26843 n_head: usize,
26844 t: usize,
26845 ) -> Result<(), Box<dyn std::error::Error>> {
26846 let f = self.func("gdn_glog_f32");
26847 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26848 let (h, ti) = (n_head as i32, t as i32);
26849 let __s_b = self.gpu.stream();
26850 let mut b = __s_b.launch_builder(&f);
26851 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26852 unsafe {
26853 b.launch(cfg)?;
26854 }
26855 Ok(())
26856 }
26857
26858 pub fn sigmoid(
26859 &self,
26860 x: &CudaSlice<f32>,
26861 y: &mut CudaSlice<f32>,
26862 n: usize,
26863 ) -> Result<(), Box<dyn std::error::Error>> {
26864 let f = self.func("sigmoid_f32");
26865 let cfg = LaunchConfig::for_num_elems(n as u32);
26866 let ni = n as i32;
26867 let __s_b = self.gpu.stream();
26868 let mut b = __s_b.launch_builder(&f);
26869 b.arg(x).arg(y).arg(&ni);
26870 unsafe {
26871 b.launch(cfg)?;
26872 }
26873 Ok(())
26874 }
26875
26876 pub fn sig_mul_f16out(
26879 &self,
26880 a: &CudaSlice<f32>,
26881 g: &CudaSlice<f32>,
26882 dst: &mut CudaSlice<f32>,
26883 dst16: &mut CudaSlice<u8>,
26884 n: usize,
26885 ) -> Result<(), Box<dyn std::error::Error>> {
26886 let f = self.func("sig_mul_f16out_f32");
26887 let cfg = LaunchConfig::for_num_elems(n as u32);
26888 let ni = n as i32;
26889 let __s_b = self.gpu.stream();
26890 let mut b = __s_b.launch_builder(&f);
26891 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26892 unsafe {
26893 b.launch(cfg)?;
26894 }
26895 Ok(())
26896 }
26897
26898 #[allow(clippy::too_many_arguments)]
26907 pub fn attn_head_gate(
26908 &self,
26909 a: &CudaSlice<f32>,
26910 g: &CudaSlice<f32>,
26911 dst: &mut CudaSlice<f32>,
26912 dst16: Option<&mut CudaSlice<u8>>,
26913 head_dim: usize,
26914 n_head: usize,
26915 t: usize,
26916 ) -> Result<(), Box<dyn std::error::Error>> {
26917 let f = self.func("attn_head_gate_f32");
26918 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26919 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26920 let d16: u64 = match dst16 {
26922 Some(d) => self.addr_u8(d),
26923 None => 0,
26924 };
26925 let __s_b = self.gpu.stream();
26926 let mut b = __s_b.launch_builder(&f);
26927 b.arg(a)
26928 .arg(g)
26929 .arg(dst)
26930 .arg(&d16)
26931 .arg(&hd)
26932 .arg(&nh)
26933 .arg(&ti);
26934 unsafe {
26935 b.launch(cfg)?;
26936 }
26937 Ok(())
26938 }
26939
26940 #[allow(clippy::too_many_arguments)]
26949 pub fn swiglu_clamped_mul_scaled(
26950 &self,
26951 gate: &CudaSlice<f32>,
26952 up: &CudaSlice<f32>,
26953 gs: f32,
26954 us: f32,
26955 limit: f32,
26956 dst: &mut CudaSlice<f32>,
26957 n: usize,
26958 ) -> Result<(), Box<dyn std::error::Error>> {
26959 debug_assert!(
26960 limit > 1e-6,
26961 "swiglu_clamped needs a live limit; use silu_mul_scaled"
26962 );
26963 let f = self.func("swiglu_clamped_mul_scaled_f32");
26964 let cfg = LaunchConfig::for_num_elems(n as u32);
26965 let ni = n as i32;
26966 let __s_b = self.gpu.stream();
26967 let mut b = __s_b.launch_builder(&f);
26968 b.arg(gate)
26969 .arg(up)
26970 .arg(&gs)
26971 .arg(&us)
26972 .arg(&limit)
26973 .arg(dst)
26974 .arg(&ni);
26975 unsafe {
26976 b.launch(cfg)?;
26977 }
26978 Ok(())
26979 }
26980
26981 pub fn gated_rmsnorm(
26983 &self,
26984 o: &CudaSlice<f32>,
26985 w: &CudaSlice<f32>,
26986 z: &CudaSlice<f32>,
26987 dst: &mut CudaSlice<f32>,
26988 ncols: usize,
26989 nrows: usize,
26990 eps: f32,
26991 ) -> Result<(), Box<dyn std::error::Error>> {
26992 let f = self.func("gated_rmsnorm_f32");
26993 let cfg = LaunchConfig {
26994 grid_dim: (nrows as u32, 1, 1),
26995 block_dim: (128, 1, 1),
26996 shared_mem_bytes: 0,
26997 };
26998 let (nc, e) = (ncols as i32, eps);
26999 let __s_b = self.gpu.stream();
27000 let mut b = __s_b.launch_builder(&f);
27001 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27002 unsafe {
27003 b.launch(cfg)?;
27004 }
27005 Ok(())
27006 }
27007
27008 pub fn gated_rmsnorm_f16out(
27011 &self,
27012 o: &CudaSlice<f32>,
27013 w: &CudaSlice<f32>,
27014 z: &CudaSlice<f32>,
27015 dst: &mut CudaSlice<f32>,
27016 dst16: &mut CudaSlice<u8>,
27017 ncols: usize,
27018 nrows: usize,
27019 eps: f32,
27020 ) -> Result<(), Box<dyn std::error::Error>> {
27021 let f = self.func("gated_rmsnorm_f16out_f32");
27022 let cfg = LaunchConfig {
27024 grid_dim: (nrows as u32, 1, 1),
27025 block_dim: (128, 1, 1),
27026 shared_mem_bytes: 0,
27027 };
27028 let (nc, e) = (ncols as i32, eps);
27029 let __s_b = self.gpu.stream();
27030 let mut b = __s_b.launch_builder(&f);
27031 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27032 unsafe {
27033 b.launch(cfg)?;
27034 }
27035 Ok(())
27036 }
27037
27038 #[allow(clippy::too_many_arguments)]
27042 pub fn add_rms_norm_zq8(
27043 &self,
27044 a: &CudaSlice<f32>,
27045 b_in: &CudaSlice<f32>,
27046 w: &CudaSlice<f32>,
27047 res: &mut CudaSlice<f32>,
27048 z: &mut CudaSlice<f32>,
27049 ncols: usize,
27050 nrows: usize,
27051 eps: f32,
27052 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27053 assert!(ncols % 32 == 0);
27054 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27055 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27056 let f = self.func("add_rms_norm_zq8");
27057 let cfg = LaunchConfig {
27058 grid_dim: (nrows as u32, 1, 1),
27059 block_dim: (1024, 1, 1),
27060 shared_mem_bytes: 0,
27061 };
27062 let (nc, ep) = (ncols as i32, eps);
27063 let __s_b = self.gpu.stream();
27064 let mut b = __s_b.launch_builder(&f);
27065 b.arg(a)
27066 .arg(b_in)
27067 .arg(w)
27068 .arg(res)
27069 .arg(z)
27070 .arg(&mut q)
27071 .arg(&mut d)
27072 .arg(&nc)
27073 .arg(&ep);
27074 unsafe {
27075 b.launch(cfg)?;
27076 }
27077 Ok((q, d))
27078 }
27079
27080 pub fn gated_rmsnorm_zv(
27085 &self,
27086 o: &CudaSlice<f32>,
27087 w: &CudaSlice<f32>,
27088 z: &cudarc::driver::CudaView<f32>,
27089 dst: &mut CudaSlice<f32>,
27090 ncols: usize,
27091 nrows: usize,
27092 eps: f32,
27093 ) -> Result<(), Box<dyn std::error::Error>> {
27094 let f = self.func("gated_rmsnorm_f32");
27095 let cfg = LaunchConfig {
27096 grid_dim: (nrows as u32, 1, 1),
27097 block_dim: (128, 1, 1),
27098 shared_mem_bytes: 0,
27099 };
27100 let (nc, e) = (ncols as i32, eps);
27101 let __s_b = self.gpu.stream();
27102 let mut b = __s_b.launch_builder(&f);
27103 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27104 unsafe {
27105 b.launch(cfg)?;
27106 }
27107 Ok(())
27108 }
27109
27110 pub fn gated_rmsnorm_f16out_zv(
27111 &self,
27112 o: &CudaSlice<f32>,
27113 w: &CudaSlice<f32>,
27114 z: &cudarc::driver::CudaView<f32>,
27115 dst: &mut CudaSlice<f32>,
27116 dst16: &mut CudaSlice<u8>,
27117 ncols: usize,
27118 nrows: usize,
27119 eps: f32,
27120 ) -> Result<(), Box<dyn std::error::Error>> {
27121 let f = self.func("gated_rmsnorm_f16out_f32");
27122 let cfg = LaunchConfig {
27124 grid_dim: (nrows as u32, 1, 1),
27125 block_dim: (128, 1, 1),
27126 shared_mem_bytes: 0,
27127 };
27128 let (nc, e) = (ncols as i32, eps);
27129 let __s_b = self.gpu.stream();
27130 let mut b = __s_b.launch_builder(&f);
27131 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27132 unsafe {
27133 b.launch(cfg)?;
27134 }
27135 Ok(())
27136 }
27137
27138 pub fn gated_rmsnorm_q8_1(
27139 &self,
27140 o: &CudaSlice<f32>,
27141 w: &CudaSlice<f32>,
27142 z: &CudaSlice<f32>,
27143 ncols: usize,
27144 nrows: usize,
27145 eps: f32,
27146 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27147 assert!(ncols % 32 == 0);
27148 let f = self.func("gated_rmsnorm_q8_1");
27149 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27150 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27151 let cfg = LaunchConfig {
27152 grid_dim: (nrows as u32, 1, 1),
27153 block_dim: (128, 1, 1),
27154 shared_mem_bytes: 0,
27155 };
27156 let (nc, ep) = (ncols as i32, eps);
27157 let __s_b = self.gpu.stream();
27158 let mut b = __s_b.launch_builder(&f);
27159 b.arg(o)
27160 .arg(w)
27161 .arg(z)
27162 .arg(&mut out_q)
27163 .arg(&mut out_d)
27164 .arg(&nc)
27165 .arg(&ep);
27166 unsafe {
27167 b.launch(cfg)?;
27168 }
27169 Ok((out_q, out_d))
27170 }
27171
27172 pub fn transpose(
27174 &self,
27175 inp: &CudaSlice<f32>,
27176 rows: usize,
27177 cols: usize,
27178 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27179 let f = self.func("transpose_f32");
27180 let mut out = self.zeros(rows * cols)?;
27181 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27182 let (r, c) = (rows as i32, cols as i32);
27183 let __s_b = self.gpu.stream();
27184 let mut b = __s_b.launch_builder(&f);
27185 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27186 unsafe {
27187 b.launch(cfg)?;
27188 }
27189 Ok(out)
27190 }
27191
27192 pub fn repeat_heads(
27194 &self,
27195 inp: &CudaSlice<f32>,
27196 out: &mut CudaSlice<f32>,
27197 head_dim: usize,
27198 n_in: usize,
27199 n_out: usize,
27200 t: usize,
27201 ) -> Result<(), Box<dyn std::error::Error>> {
27202 let f = self.func("repeat_heads_f32");
27203 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27204 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27205 let __s_b = self.gpu.stream();
27206 let mut b = __s_b.launch_builder(&f);
27207 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27208 unsafe {
27209 b.launch(cfg)?;
27210 }
27211 Ok(())
27212 }
27213
27214 pub fn q_gate_split(
27221 &self,
27222 qf: &CudaSlice<f32>,
27223 q_out: &mut CudaSlice<f32>,
27224 gate_out: &mut CudaSlice<f32>,
27225 head_dim: usize,
27226 n_head: usize,
27227 t: usize,
27228 ) -> Result<(), Box<dyn std::error::Error>> {
27229 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27230 let out_need = head_dim * n_head * t;
27231 if q_out.len() < out_need || gate_out.len() < out_need {
27232 return Err(format!(
27233 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27234 q_out.len(),
27235 gate_out.len()
27236 )
27237 .into());
27238 }
27239 let f = self.func("q_gate_split_f32");
27240 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27241 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27242 let __s_b = self.gpu.stream();
27243 let mut b = __s_b.launch_builder(&f);
27244 b.arg(qf)
27245 .arg(q_out)
27246 .arg(gate_out)
27247 .arg(&hd)
27248 .arg(&nh)
27249 .arg(&ti);
27250 unsafe {
27251 b.launch(cfg)?;
27252 }
27253 Ok(())
27254 }
27255
27256 pub fn qkv_to_gdn_repack(
27260 &self,
27261 conv_out: &CudaSlice<f32>,
27262 q_g: &mut CudaSlice<f32>,
27263 k_g: &mut CudaSlice<f32>,
27264 v_g: &mut CudaSlice<f32>,
27265 d_state: usize,
27266 num_v: usize,
27267 num_k: usize,
27268 key_dim: usize,
27269 t: usize,
27270 ) -> Result<(), Box<dyn std::error::Error>> {
27271 let f = self.func("qkv_to_gdn_repack_f32");
27272 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27273 let (ds, nv, nk, kd, ti) = (
27274 d_state as i32,
27275 num_v as i32,
27276 num_k as i32,
27277 key_dim as i32,
27278 t as i32,
27279 );
27280 let __s_b = self.gpu.stream();
27281 let mut b = __s_b.launch_builder(&f);
27282 b.arg(conv_out)
27283 .arg(q_g)
27284 .arg(k_g)
27285 .arg(v_g)
27286 .arg(&ds)
27287 .arg(&nv)
27288 .arg(&nk)
27289 .arg(&kd)
27290 .arg(&ti);
27291 unsafe {
27292 b.launch(cfg)?;
27293 }
27294 Ok(())
27295 }
27296
27297 pub fn conv_left_pad(
27300 &self,
27301 src: &CudaSlice<f32>,
27302 dst: &mut CudaSlice<f32>,
27303 conv_dim: usize,
27304 t: usize,
27305 pad: usize,
27306 ) -> Result<(), Box<dyn std::error::Error>> {
27307 let f = self.func("conv_left_pad_f32");
27308 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27309 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27310 let __s_b = self.gpu.stream();
27311 let mut b = __s_b.launch_builder(&f);
27312 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27313 unsafe {
27314 b.launch(cfg)?;
27315 }
27316 Ok(())
27317 }
27318
27319 pub fn conv_assemble_and_roll(
27323 &self,
27324 qkv_col: &CudaSlice<f32>,
27325 conv_state: &mut CudaSlice<f32>,
27326 conv_in: &mut CudaSlice<f32>,
27327 conv_dim: usize,
27328 pad: usize,
27329 ) -> Result<(), Box<dyn std::error::Error>> {
27330 let f = self.func("conv_assemble_and_roll_f32");
27331 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27332 let (cd, p) = (conv_dim as i32, pad as i32);
27333 let __s_b = self.gpu.stream();
27334 let mut b = __s_b.launch_builder(&f);
27335 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27336 unsafe {
27337 b.launch(cfg)?;
27338 }
27339 Ok(())
27340 }
27341
27342 pub fn ssm_conv1d_fused_decode(
27348 &self,
27349 qkv_col: &CudaSlice<f32>,
27350 conv_state: &mut CudaSlice<f32>,
27351 w: &CudaSlice<f32>,
27352 conv_out: &mut CudaSlice<f32>,
27353 conv_dim: usize,
27354 d_conv: usize,
27355 ) -> Result<(), Box<dyn std::error::Error>> {
27356 let f = self.func("ssm_conv1d_fused_decode_f32");
27357 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27358 let (cd, dc) = (conv_dim as i32, d_conv as i32);
27359 let __s_b = self.gpu.stream();
27360 let mut b = __s_b.launch_builder(&f);
27361 b.arg(qkv_col)
27362 .arg(conv_state)
27363 .arg(w)
27364 .arg(conv_out)
27365 .arg(&cd)
27366 .arg(&dc);
27367 unsafe {
27368 b.launch(cfg)?;
27369 }
27370 Ok(())
27371 }
27372
27373 pub fn slice_range(
27376 &self,
27377 src: &CudaSlice<f32>,
27378 start: usize,
27379 len: usize,
27380 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27381 let host = self.gpu.stream().clone_dtoh(src)?;
27382 self.gpu.stream().synchronize()?;
27383 Ok(self.htod(&host[start..start + len])?)
27384 }
27385}
27386
27387#[cfg(test)]
27388mod target_dispatch_tests {
27389 use super::legacy_quant_gemm_allowed;
27390
27391 #[test]
27392 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27393 assert!(legacy_quant_gemm_allowed(false, false, false));
27395 assert!(!legacy_quant_gemm_allowed(false, false, true));
27396 assert!(!legacy_quant_gemm_allowed(true, false, false));
27398 assert!(!legacy_quant_gemm_allowed(true, false, true));
27399 assert!(legacy_quant_gemm_allowed(true, true, false));
27401 assert!(!legacy_quant_gemm_allowed(true, true, true));
27402 }
27403
27404 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27405 #[test]
27406 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27407 assert!(!legacy_quant_gemm_allowed(
27408 cfg!(memra_portable_cuda),
27409 cfg!(memra_hopper_mma),
27410 false
27411 ));
27412 }
27413
27414 #[cfg(memra_hopper_mma)]
27415 #[test]
27416 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27417 assert!(legacy_quant_gemm_allowed(
27418 cfg!(memra_portable_cuda),
27419 cfg!(memra_hopper_mma),
27420 false
27421 ));
27422 assert!(super::portable_mma_gated() == false);
27423 }
27424}
27425
27426impl memra_kv::KvDev for Engine {
27429 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27430 Engine::zeros(self, n)
27431 }
27432 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27433 Engine::uninit(self, n)
27434 }
27435 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27436 Engine::alloc_u8(self, n)
27437 }
27438 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27439 Engine::htod_i32(self, v)
27440 }
27441 fn clone_dtod(
27442 &self,
27443 src: &CudaSlice<f32>,
27444 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27445 Engine::clone_dtod(self, src)
27446 }
27447 fn copy_into(
27448 &self,
27449 dst: &mut CudaSlice<f32>,
27450 off: usize,
27451 src: &CudaSlice<f32>,
27452 len: usize,
27453 ) -> Result<(), Box<dyn std::error::Error>> {
27454 Engine::copy_into(self, dst, off, src, len)
27455 }
27456 fn set_i32_one(
27457 &self,
27458 d: &mut CudaSlice<i32>,
27459 v: i32,
27460 ) -> Result<(), Box<dyn std::error::Error>> {
27461 Engine::set_i32_one(self, d, v)
27462 }
27463}
27464
27465#[cfg(test)]
27466mod fused_gate_bounds_tests {
27467 use super::*;
27468
27469 #[test]
27482 #[ignore = "requires a CUDA GPU"]
27483 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27484 let e = Engine::new(0).unwrap();
27485 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27486 let fused = 2 * head_dim * n_head * t;
27487 let out_n = head_dim * n_head * t;
27488
27489 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27491 let mut q = e.uninit(out_n).unwrap();
27492 let mut gate = e.uninit(out_n).unwrap();
27493 let err = e
27494 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27495 .expect_err("half-width wq must be refused, not read past")
27496 .to_string();
27497 assert!(err.contains("NO fused gate"), "{err}");
27498 assert!(err.contains(&format!("{fused}")), "{err}");
27499
27500 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27503 let wide = e.htod(&host).unwrap();
27504 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27505 .expect("full-width wq splits");
27506 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27507 for tok in 0..t {
27508 for hh in 0..n_head {
27509 for d in 0..head_dim {
27510 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27511 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27512 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27513 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27514 }
27515 }
27516 }
27517
27518 let mut small = e.uninit(out_n - 1).unwrap();
27520 assert!(
27521 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27522 .is_err()
27523 );
27524 }
27525}
27526
27527#[cfg(test)]
27531mod fused_rope_width_tests {
27532 use super::Engine;
27533
27534 #[test]
27537 fn full_width_is_accepted() {
27538 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27539 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27540 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27541 }
27542
27543 #[test]
27556 fn gemma4_official_artifact_widths_pass() {
27557 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27558 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27559 }
27560
27561 #[test]
27564 fn partial_rotary_is_refused_with_the_geometry_named() {
27565 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27567 .expect_err("partial rotary must refuse");
27568 let msg = err.to_string();
27569 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27570 assert!(msg.contains("n_rot 64"), "{msg}");
27571 assert!(msg.contains("head_dim 256"), "{msg}");
27572 assert!(
27573 msg.contains("64..256"),
27574 "names the band it would corrupt: {msg}"
27575 );
27576 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27578 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27580 }
27581}