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_trim_to_zero(&self) -> usize {
1531 use cudarc::driver::sys;
1532 let (before, _) = self.pool_reserved_used();
1533 unsafe {
1534 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1535 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1536 != sys::CUresult::CUDA_SUCCESS
1537 {
1538 return 0;
1539 }
1540 let _ = sys::cuMemPoolTrimTo(pool, 0);
1541 }
1542 let (after, _) = self.pool_reserved_used();
1543 before.saturating_sub(after)
1544 }
1545
1546 pub fn pool_reserved_used(&self) -> (usize, usize) {
1547 use cudarc::driver::sys;
1548 unsafe {
1549 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1550 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1551 != sys::CUresult::CUDA_SUCCESS
1552 {
1553 return (0, 0);
1554 }
1555 let (mut reserved, mut used) = (0u64, 0u64);
1556 if sys::cuMemPoolGetAttribute(
1557 pool,
1558 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1559 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1560 ) != sys::CUresult::CUDA_SUCCESS
1561 {
1562 return (0, 0);
1563 }
1564 if sys::cuMemPoolGetAttribute(
1565 pool,
1566 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1567 &mut used as *mut u64 as *mut core::ffi::c_void,
1568 ) != sys::CUresult::CUDA_SUCCESS
1569 {
1570 return (0, 0);
1571 }
1572 (reserved as usize, used as usize)
1573 }
1574 }
1575
1576 pub fn stream(&self) -> Arc<CudaStream> {
1579 self.gpu.stream()
1580 }
1581 pub fn gkv_on() -> bool {
1584 memra_kv::gkv_on()
1585 }
1586
1587 pub fn wkv_on() -> bool {
1599 memra_kv::wkv_on()
1600 }
1601
1602 pub fn kv_fp8_on() -> bool {
1608 memra_kv::kv_fp8_on()
1609 }
1610
1611 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1614 if head_dim == 512 && Self::gkv_on() {
1615 self.func_g(name)
1616 } else {
1617 self.func(name)
1618 }
1619 }
1620
1621 fn func_g(&self, name: &str) -> CudaFunction {
1625 let m = self.flash_g.get_or_init(|| {
1626 self.gpu
1627 .ctx
1628 .load_module(cudarc::nvrtc::Ptx::from_binary(
1629 FLASH_FATBIN_KF8VF8.to_vec(),
1630 ))
1631 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1632 });
1633 let key = format!("g:{name}");
1634 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1635 return f.clone();
1636 }
1637 let f = match m.load_function(name) {
1638 Ok(f) => f,
1639 Err(_) => self.func(name),
1640 };
1641 self.fn_cache.lock().unwrap().insert(key, f.clone());
1642 f
1643 }
1644
1645 fn func(&self, name: &str) -> CudaFunction {
1646 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1649 return f.clone();
1650 }
1651 let f = self
1652 .module
1653 .load_function(name)
1654 .or_else(|_| self.hybrid.load_function(name))
1655 .or_else(|_| self.qmatvec.load_function(name))
1656 .or_else(|_| self.flash.load_function(name))
1657 .or_else(|_| self.gemm.load_function(name))
1658 .or_else(|_| self.router.load_function(name))
1659 .or_else(|_| self.sample.load_function(name))
1660 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1661 self.fn_cache
1662 .lock()
1663 .unwrap()
1664 .insert(name.to_string(), f.clone());
1665 f
1666 }
1667
1668 pub fn scatter_trim_logits(
1671 &self,
1672 src: &CudaSlice<f32>,
1673 d2t: &CudaSlice<u32>,
1674 dst: &mut CudaSlice<f32>,
1675 d_vocab: usize,
1676 n_vocab: usize,
1677 ) -> Result<(), Box<dyn std::error::Error>> {
1678 let f1 = self.func("scatter_trim_logits_f32");
1679 let f2 = self.func("scatter_trim_logits_pass2_f32");
1680 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1681 let cfg1 = LaunchConfig {
1682 grid_dim: (256, 1, 1),
1683 block_dim: (256, 1, 1),
1684 shared_mem_bytes: 0,
1685 };
1686 let __s_b1 = self.gpu.stream();
1687 let mut b1 = __s_b1.launch_builder(&f1);
1688 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1689 unsafe {
1690 b1.launch(cfg1)?;
1691 }
1692 let cfg2 = LaunchConfig {
1693 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1694 block_dim: (256, 1, 1),
1695 shared_mem_bytes: 0,
1696 };
1697 let __s_b2 = self.gpu.stream();
1698 let mut b2 = __s_b2.launch_builder(&f2);
1699 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1700 unsafe {
1701 b2.launch(cfg2)?;
1702 }
1703 Ok(())
1704 }
1705
1706 #[allow(clippy::too_many_arguments)]
1712 pub fn filter_stats(
1713 &self,
1714 x: &CudaSlice<f32>,
1715 row_stride: usize,
1716 rows: &CudaSlice<i32>,
1717 out_th: &mut CudaSlice<f32>,
1718 out_z: &mut CudaSlice<f32>,
1719 out_max: &mut CudaSlice<f32>,
1720 n: usize,
1721 nrow: usize,
1722 temp: f32,
1723 top_k: i32,
1724 top_p: f32,
1725 min_p: f32,
1726 ) -> Result<(), Box<dyn std::error::Error>> {
1727 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1754 let coop_on =
1755 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1756 if coop_on && self.sm_count() >= 16 {
1757 let cap = self.sm_count() as usize / 16;
1758 let mut done = 0usize;
1759 while done < nrow {
1760 let chunk = cap.min(nrow - done);
1761 self.filter_stats_coop_chunk(
1762 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1763 top_p, min_p,
1764 )?;
1765 done += chunk;
1766 }
1767 return Ok(());
1768 }
1769 self.filter_stats_plain_program(
1770 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1771 )
1772 }
1773
1774 #[allow(clippy::too_many_arguments)]
1779 pub fn filter_stats_coop_chunk(
1780 &self,
1781 x: &CudaSlice<f32>,
1782 row_stride: usize,
1783 rows: &CudaSlice<i32>,
1784 row0: usize,
1785 out_th: &mut CudaSlice<f32>,
1786 out_z: &mut CudaSlice<f32>,
1787 out_max: &mut CudaSlice<f32>,
1788 n: usize,
1789 chunk: usize,
1790 temp: f32,
1791 top_k: i32,
1792 top_p: f32,
1793 min_p: f32,
1794 ) -> Result<(), Box<dyn std::error::Error>> {
1795 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1796 let f = self.func("filter_stats_coop_f32");
1797 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1798 let cfg = LaunchConfig {
1799 grid_dim: (16, chunk as u32, 1),
1800 block_dim: (512, 1, 1),
1801 shared_mem_bytes: 0,
1802 };
1803 let rows_v = rows.slice(row0..row0 + chunk);
1804 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1805 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1806 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1807 let __s_b = self.gpu.stream();
1808 let mut b = __s_b.launch_builder(&f);
1809 b.arg(x)
1810 .arg(&rs)
1811 .arg(&rows_v)
1812 .arg(&mut th_v)
1813 .arg(&mut z_v)
1814 .arg(&mut mx_v)
1815 .arg(&mut ws)
1816 .arg(&ni)
1817 .arg(&nr)
1818 .arg(&temp)
1819 .arg(&top_k)
1820 .arg(&top_p)
1821 .arg(&min_p);
1822 unsafe {
1823 b.launch_cooperative(cfg)?;
1824 }
1825 Ok(())
1826 }
1827
1828 #[allow(clippy::too_many_arguments)]
1832 pub fn filter_stats_plain_program(
1833 &self,
1834 x: &CudaSlice<f32>,
1835 row_stride: usize,
1836 rows: &CudaSlice<i32>,
1837 out_th: &mut CudaSlice<f32>,
1838 out_z: &mut CudaSlice<f32>,
1839 out_max: &mut CudaSlice<f32>,
1840 n: usize,
1841 nrow: usize,
1842 temp: f32,
1843 top_k: i32,
1844 top_p: f32,
1845 min_p: f32,
1846 ) -> Result<(), Box<dyn std::error::Error>> {
1847 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1848 let f = self.func("filter_stats_f32");
1849 let cfg = LaunchConfig {
1850 grid_dim: (nrow as u32, 1, 1),
1851 block_dim: (1024, 1, 1),
1852 shared_mem_bytes: 0,
1853 };
1854 let __s_b = self.gpu.stream();
1855 let mut b = __s_b.launch_builder(&f);
1856 b.arg(x)
1857 .arg(&rs)
1858 .arg(rows)
1859 .arg(&mut *out_th)
1860 .arg(&mut *out_z)
1861 .arg(&mut *out_max)
1862 .arg(&ni)
1863 .arg(&nr)
1864 .arg(&temp)
1865 .arg(&top_k)
1866 .arg(&top_p)
1867 .arg(&min_p);
1868 unsafe {
1869 b.launch(cfg)?;
1870 }
1871 Ok(())
1872 }
1873
1874 #[allow(clippy::too_many_arguments)]
1876 pub fn softmax_gather_filtered(
1877 &self,
1878 x: &CudaSlice<f32>,
1879 row_stride: usize,
1880 ids: &CudaSlice<u32>,
1881 rows: &CudaSlice<i32>,
1882 th: &CudaSlice<f32>,
1883 z: &CudaSlice<f32>,
1884 out: &mut CudaSlice<f32>,
1885 n: usize,
1886 npair: usize,
1887 temp: f32,
1888 ) -> Result<(), Box<dyn std::error::Error>> {
1889 let f = self.func("softmax_gather_filtered_f32");
1890 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1891 let cfg = LaunchConfig {
1892 grid_dim: (npair as u32, 1, 1),
1893 block_dim: (256, 1, 1),
1894 shared_mem_bytes: 0,
1895 };
1896 let __s_b = self.gpu.stream();
1897 let mut b = __s_b.launch_builder(&f);
1898 b.arg(x)
1899 .arg(&rs)
1900 .arg(ids)
1901 .arg(rows)
1902 .arg(th)
1903 .arg(z)
1904 .arg(&mut *out)
1905 .arg(&ni)
1906 .arg(&np)
1907 .arg(&temp);
1908 unsafe {
1909 b.launch(cfg)?;
1910 }
1911 Ok(())
1912 }
1913
1914 #[allow(clippy::too_many_arguments)]
1916 pub fn residual_sample_filtered(
1917 &self,
1918 p: &CudaSlice<f32>,
1919 q: Option<&CudaSlice<f32>>,
1920 n: usize,
1921 temp: f32,
1922 seed: u64,
1923 stream_pos: u32,
1924 p_stats: (f32, f32, f32),
1925 q_stats: (f32, f32, f32),
1926 out_tok: &mut CudaSlice<u32>,
1927 ) -> Result<(), Box<dyn std::error::Error>> {
1928 let f = self.func("residual_sample_filtered_f32");
1929 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1930 let has_q: i32 = q.is_some() as i32;
1931 let qbuf = q.unwrap_or(p);
1932 let (pm, pth, pz) = p_stats;
1933 let (qm, qth, qz) = q_stats;
1934 let cfg = LaunchConfig {
1935 grid_dim: (1, 1, 1),
1936 block_dim: (1024, 1, 1),
1937 shared_mem_bytes: 0,
1938 };
1939 let __s_b = self.gpu.stream();
1940 let mut b = __s_b.launch_builder(&f);
1941 b.arg(p)
1942 .arg(qbuf)
1943 .arg(&has_q)
1944 .arg(&ni)
1945 .arg(&temp)
1946 .arg(&slo)
1947 .arg(&shi)
1948 .arg(&stream_pos)
1949 .arg(&pm)
1950 .arg(&pth)
1951 .arg(&pz)
1952 .arg(&qm)
1953 .arg(&qth)
1954 .arg(&qz)
1955 .arg(&mut *out_tok);
1956 unsafe {
1957 b.launch(cfg)?;
1958 }
1959 Ok(())
1960 }
1961
1962 #[allow(clippy::too_many_arguments)]
1968 pub fn residual_sample_sparse_q(
1969 &self,
1970 p: &CudaSlice<f32>,
1971 cand_ids: &CudaSlice<u32>,
1972 q_probs: &CudaSlice<f32>,
1973 n_cand: usize,
1974 n: usize,
1975 temp: f32,
1976 seed: u64,
1977 stream_pos: u32,
1978 p_stats: (f32, f32, f32),
1979 out_tok: &mut CudaSlice<u32>,
1980 ) -> Result<(), Box<dyn std::error::Error>> {
1981 assert!(
1982 n_cand >= 1 && n_cand <= 32,
1983 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1984 );
1985 let f = self.func("residual_sample_sparse_q_f32");
1986 let (ni, nc) = (n as i32, n_cand as i32);
1987 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1988 let (pm, pth, pz) = p_stats;
1989 let cfg = LaunchConfig {
1990 grid_dim: (1, 1, 1),
1991 block_dim: (1024, 1, 1),
1992 shared_mem_bytes: 0,
1993 };
1994 let __s_b = self.gpu.stream();
1995 let mut b = __s_b.launch_builder(&f);
1996 b.arg(p)
1997 .arg(cand_ids)
1998 .arg(q_probs)
1999 .arg(&nc)
2000 .arg(&ni)
2001 .arg(&temp)
2002 .arg(&slo)
2003 .arg(&shi)
2004 .arg(&stream_pos)
2005 .arg(&pm)
2006 .arg(&pth)
2007 .arg(&pz)
2008 .arg(&mut *out_tok);
2009 unsafe {
2010 b.launch(cfg)?;
2011 }
2012 Ok(())
2013 }
2014
2015 #[allow(clippy::too_many_arguments)]
2017 pub fn gumbel_perturb_filtered(
2018 &self,
2019 x: &CudaSlice<f32>,
2020 y: &mut CudaSlice<f32>,
2021 n: usize,
2022 seed: u64,
2023 stream_pos: u32,
2024 temp: f32,
2025 row_max: f32,
2026 th: f32,
2027 ) -> Result<(), Box<dyn std::error::Error>> {
2028 let f = self.func("gumbel_perturb_filtered_f32");
2029 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2030 let cfg = LaunchConfig {
2031 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2032 block_dim: (256, 1, 1),
2033 shared_mem_bytes: 0,
2034 };
2035 let __s_b = self.gpu.stream();
2036 let mut b = __s_b.launch_builder(&f);
2037 b.arg(x)
2038 .arg(&mut *y)
2039 .arg(&ni)
2040 .arg(&slo)
2041 .arg(&shi)
2042 .arg(&stream_pos)
2043 .arg(&temp)
2044 .arg(&row_max)
2045 .arg(&th);
2046 unsafe {
2047 b.launch(cfg)?;
2048 }
2049 Ok(())
2050 }
2051
2052 #[allow(clippy::too_many_arguments)]
2056 pub fn penalize_logits(
2057 &self,
2058 x: &mut CudaSlice<f32>,
2059 hist: &CudaSlice<u32>,
2060 n_hist: usize,
2061 rep: f32,
2062 freq: f32,
2063 present: f32,
2064 n: usize,
2065 ) -> Result<(), Box<dyn std::error::Error>> {
2066 if n_hist == 0 {
2067 return Ok(());
2068 }
2069 let f = self.func("penalize_logits_f32");
2070 let (nh, ni) = (n_hist as i32, n as i32);
2071 let cfg = LaunchConfig {
2072 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2073 block_dim: (128, 1, 1),
2074 shared_mem_bytes: 0,
2075 };
2076 let __s_b = self.gpu.stream();
2077 let mut b = __s_b.launch_builder(&f);
2078 b.arg(&mut *x)
2079 .arg(hist)
2080 .arg(&nh)
2081 .arg(&rep)
2082 .arg(&freq)
2083 .arg(&present)
2084 .arg(&ni);
2085 unsafe {
2086 b.launch(cfg)?;
2087 }
2088 Ok(())
2089 }
2090
2091 #[allow(clippy::too_many_arguments)]
2093 pub fn penalize_logits_rows(
2094 &self,
2095 x: &mut CudaSlice<f32>,
2096 hist: &CudaSlice<u32>,
2097 n_hist: usize,
2098 rep: f32,
2099 freq: f32,
2100 present: f32,
2101 n: usize,
2102 nrow: usize,
2103 ) -> Result<(), Box<dyn std::error::Error>> {
2104 if n_hist == 0 || nrow == 0 {
2105 return Ok(());
2106 }
2107 let f = self.func("penalize_logits_rows_f32");
2108 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2109 let cfg = LaunchConfig {
2110 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2111 block_dim: (128, 1, 1),
2112 shared_mem_bytes: 0,
2113 };
2114 let __s_b = self.gpu.stream();
2115 let mut b = __s_b.launch_builder(&f);
2116 b.arg(&mut *x)
2117 .arg(hist)
2118 .arg(&nh)
2119 .arg(&rep)
2120 .arg(&freq)
2121 .arg(&present)
2122 .arg(&ni)
2123 .arg(&nr);
2124 unsafe {
2125 b.launch(cfg)?;
2126 }
2127 Ok(())
2128 }
2129
2130 #[allow(clippy::too_many_arguments)]
2136 pub fn penalize_logits_sparse_rows(
2137 &self,
2138 x: &mut CudaSlice<f32>,
2139 ids: &[u32],
2140 counts: &[u32],
2141 offsets: &[i32],
2142 rows: &[i32],
2143 reps: &[f32],
2144 freqs: &[f32],
2145 presents: &[f32],
2146 n: usize,
2147 ) -> Result<(), Box<dyn std::error::Error>> {
2148 let nrow = rows.len();
2149 if nrow == 0 {
2150 return Ok(());
2151 }
2152 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2153 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2154 let entry_count =
2155 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2156 if ids.len() != counts.len()
2157 || offsets.len() != nrow + 1
2158 || reps.len() != nrow
2159 || freqs.len() != nrow
2160 || presents.len() != nrow
2161 || offsets.first().copied() != Some(0)
2162 || offsets.last().copied() != Some(entry_count)
2163 {
2164 return Err("sparse penalty row metadata shape mismatch".into());
2165 }
2166 if counts.contains(&0) {
2167 return Err("sparse penalty counts must be positive".into());
2168 }
2169 let mut max_len = 0usize;
2170 for pair in offsets.windows(2) {
2171 if pair[0] < 0 || pair[1] < pair[0] {
2172 return Err("sparse penalty offsets must be monotonic".into());
2173 }
2174 max_len = max_len.max((pair[1] - pair[0]) as usize);
2175 }
2176 if max_len == 0 {
2177 return Ok(());
2178 }
2179
2180 let mut seen = std::collections::HashSet::with_capacity(ids.len());
2181 for (r, &row) in rows.iter().enumerate() {
2182 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2183 return Err("sparse penalty row index exceeds logits shape".into());
2184 }
2185 let begin = offsets[r] as usize;
2186 let end = offsets[r + 1] as usize;
2187 for &id in &ids[begin..end] {
2188 if id as usize >= n {
2189 return Err("sparse penalty token id exceeds logits row".into());
2190 }
2191 if !seen.insert((row, id)) {
2192 return Err("sparse penalty entries must be unique per logits row".into());
2193 }
2194 }
2195 }
2196
2197 unsafe {
2199 self.penalize_logits_sparse_rows_unchecked(
2200 x, ids, counts, offsets, rows, reps, freqs, presents, n,
2201 )
2202 }
2203 }
2204
2205 #[allow(clippy::too_many_arguments)]
2214 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2215 &self,
2216 x: &mut CudaSlice<f32>,
2217 ids: &[u32],
2218 counts: &[u32],
2219 offsets: &[i32],
2220 rows: &[i32],
2221 reps: &[f32],
2222 freqs: &[f32],
2223 presents: &[f32],
2224 n: usize,
2225 ) -> Result<(), Box<dyn std::error::Error>> {
2226 let nrow = rows.len();
2227 if nrow == 0 {
2228 return Ok(());
2229 }
2230 let max_len = offsets
2231 .windows(2)
2232 .map(|pair| (pair[1] - pair[0]) as usize)
2233 .max()
2234 .unwrap_or(0);
2235 if max_len == 0 {
2236 return Ok(());
2237 }
2238 let ids_d = self.htod_u32_v(ids)?;
2239 let counts_d = self.htod_u32_v(counts)?;
2240 let offsets_d = self.htod_i32(offsets)?;
2241 let rows_d = self.htod_i32(rows)?;
2242 let reps_d = self.htod(reps)?;
2243 let freqs_d = self.htod(freqs)?;
2244 let presents_d = self.htod(presents)?;
2245 let f = self.func("penalize_logits_sparse_rows_f32");
2246 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2247 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2248 let cfg = LaunchConfig {
2249 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2250 block_dim: (128, 1, 1),
2251 shared_mem_bytes: 0,
2252 };
2253 let __s_b = self.gpu.stream();
2254 let mut b = __s_b.launch_builder(&f);
2255 b.arg(&mut *x)
2256 .arg(&ids_d)
2257 .arg(&counts_d)
2258 .arg(&offsets_d)
2259 .arg(&rows_d)
2260 .arg(&reps_d)
2261 .arg(&freqs_d)
2262 .arg(&presents_d)
2263 .arg(&ni)
2264 .arg(&nr);
2265 unsafe {
2266 b.launch(cfg)?;
2267 }
2268 Ok(())
2269 }
2270
2271 #[allow(clippy::too_many_arguments)]
2279 pub fn penalize_logits_rows_inc(
2280 &self,
2281 x: &mut CudaSlice<f32>,
2282 hist: &CudaSlice<u32>,
2283 n_hist0: usize,
2284 rep: f32,
2285 freq: f32,
2286 present: f32,
2287 n: usize,
2288 nrow: usize,
2289 win: usize,
2290 ) -> Result<(), Box<dyn std::error::Error>> {
2291 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2292 return Ok(());
2293 }
2294 debug_assert!(
2295 hist.len() >= n_hist0 + nrow - 1,
2296 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2297 );
2298 let f = self.func("penalize_logits_rows_inc_f32");
2299 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2300 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2301 let cfg = LaunchConfig {
2302 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2303 block_dim: (128, 1, 1),
2304 shared_mem_bytes: 0,
2305 };
2306 let __s_b = self.gpu.stream();
2307 let mut b = __s_b.launch_builder(&f);
2308 b.arg(&mut *x)
2309 .arg(hist)
2310 .arg(&nh)
2311 .arg(&rep)
2312 .arg(&freq)
2313 .arg(&present)
2314 .arg(&ni)
2315 .arg(&nr)
2316 .arg(&wi);
2317 unsafe {
2318 b.launch(cfg)?;
2319 }
2320 Ok(())
2321 }
2322
2323 pub fn wpf_level() -> u32 {
2331 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2332 *ON.get_or_init(|| {
2333 std::env::var("MEMRA_WPF")
2334 .ok()
2335 .and_then(|v| v.parse().ok())
2336 .unwrap_or(1)
2337 })
2338 }
2339
2340 pub fn set_verify_exact(&self, on: bool) {
2356 self.verify_exact
2357 .store(on, std::sync::atomic::Ordering::Relaxed);
2358 }
2359 pub(crate) fn verify_exact_on(&self) -> bool {
2360 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2361 }
2362
2363 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2369 ExactScope::set(&self.verify_exact, on)
2370 }
2371
2372 pub fn qkv_append_on() -> bool {
2375 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2376 *ON.get_or_init(|| {
2377 std::env::var("MEMRA_QKV_APPEND")
2378 .map(|v| v != "0")
2379 .unwrap_or(true)
2380 })
2381 }
2382
2383 pub fn pdl_wb_on() -> bool {
2386 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2387 *ON.get_or_init(|| {
2388 std::env::var("MEMRA_PDL_WB")
2389 .map(|v| v != "0")
2390 .unwrap_or(true)
2391 })
2392 }
2393
2394 pub fn norm_ilp_on() -> bool {
2402 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2403 *ON.get_or_init(|| {
2404 std::env::var("MEMRA_NORM_ILP")
2405 .map(|v| v != "0")
2406 .unwrap_or(true)
2407 })
2408 }
2409
2410 pub fn tk_ffn_dual_on() -> bool {
2418 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2419 *ON.get_or_init(|| {
2420 std::env::var("MEMRA_TK_FFN_DUAL")
2421 .map(|v| v != "0")
2422 .unwrap_or(true)
2423 })
2424 }
2425
2426 pub fn pdl_mmvq_on() -> bool {
2430 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2431 *ON.get_or_init(|| {
2432 std::env::var("MEMRA_PDL_MMVQ")
2433 .map(|v| v != "0")
2434 .unwrap_or(true)
2435 })
2436 }
2437
2438 pub fn pdl_on() -> bool {
2439 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2440 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2441 }
2442
2443 pub fn pdl_nvfp4q8_on() -> bool {
2449 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2450 *ON.get_or_init(|| {
2451 std::env::var("MEMRA_PDL_NVFP4")
2452 .map(|v| v != "0")
2453 .unwrap_or(true)
2454 })
2455 }
2456
2457 fn q40_mr1_on() -> bool {
2463 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2464 match *Q40MR.get_or_init(|| {
2465 std::env::var("MEMRA_Q40_MR")
2466 .ok()
2467 .and_then(|v| v.parse().ok())
2468 }) {
2469 Some(v) => v == 1,
2470 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2471 }
2472 }
2473
2474 fn pdl_func_flash(
2479 &self,
2480 g: bool,
2481 name: &'static str,
2482 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2483 use cudarc::driver::sys as cu;
2484 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2491 std::sync::Mutex::new(None);
2492 static FNS: std::sync::Mutex<
2493 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2494 > = std::sync::Mutex::new(None);
2495 let ctx_key = self.ctx().cu_ctx() as usize;
2496 if let Some(&f) = FNS
2497 .lock()
2498 .unwrap()
2499 .get_or_insert_with(Default::default)
2500 .get(&(ctx_key, g, name))
2501 {
2502 return Ok(f as cu::CUfunction);
2503 }
2504 let module = {
2505 let mut mods = MODS.lock().unwrap();
2506 let map = mods.get_or_insert_with(Default::default);
2507 match map.get(&(ctx_key, g)) {
2508 Some(&m) => m,
2509 None => {
2510 let m = self.pdl_load_module_in_ctx(if g {
2511 FLASH_FATBIN_KF8VF8
2512 } else {
2513 FLASH_FATBIN
2514 })?;
2515 map.insert((ctx_key, g), m);
2516 m
2517 }
2518 }
2519 };
2520 let cname = std::ffi::CString::new(name)?;
2521 let mut f: cu::CUfunction = std::ptr::null_mut();
2522 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2523 if r != cu::CUresult::CUDA_SUCCESS {
2524 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2525 }
2526 FNS.lock()
2527 .unwrap()
2528 .get_or_insert_with(Default::default)
2529 .insert((ctx_key, g, name), f as usize);
2530 Ok(f)
2531 }
2532
2533 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2538 use cudarc::driver::sys as cu;
2539 let mut prev: cu::CUcontext = std::ptr::null_mut();
2540 unsafe {
2541 cu::cuCtxGetCurrent(&mut prev).result()?;
2542 }
2543 self.ctx().bind_to_thread()?;
2544 let mut m: cu::CUmodule = std::ptr::null_mut();
2545 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2546 let restore = if prev.is_null() {
2547 cu::CUresult::CUDA_SUCCESS
2548 } else {
2549 unsafe { cu::cuCtxSetCurrent(prev) }
2550 };
2551 if r != cu::CUresult::CUDA_SUCCESS {
2552 return Err(format!("pdl module load: {r:?}").into());
2553 }
2554 if restore != cu::CUresult::CUDA_SUCCESS {
2555 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2556 }
2557 Ok(m as usize)
2558 }
2559
2560 pub fn raw_kernel_function(
2563 &self,
2564 name: &'static str,
2565 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2566 self.pdl_func(name)
2567 }
2568
2569 fn pdl_func(
2570 &self,
2571 name: &'static str,
2572 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2573 use cudarc::driver::sys as cu;
2574 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2577 std::sync::Mutex::new(None);
2578 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2581 std::sync::Mutex::new(None);
2582 static FNS: std::sync::Mutex<
2583 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2584 > = std::sync::Mutex::new(None);
2585 let ctx_key = self.ctx().cu_ctx() as usize;
2586 if let Some(&f) = FNS
2587 .lock()
2588 .unwrap()
2589 .get_or_insert_with(Default::default)
2590 .get(&(ctx_key, name))
2591 {
2592 return Ok(f as cu::CUfunction);
2593 }
2594 let module = {
2595 let mut mods = MODULES.lock().unwrap();
2596 let map = mods.get_or_insert_with(Default::default);
2597 match map.get(&ctx_key) {
2598 Some(&m) => m,
2599 None => {
2600 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2601 map.insert(ctx_key, m);
2602 m
2603 }
2604 }
2605 };
2606 let cname = std::ffi::CString::new(name)?;
2607 let mut f: cu::CUfunction = std::ptr::null_mut();
2608 let mut r =
2609 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2610 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2611 let qmodule = {
2612 let mut mods = QMODULES.lock().unwrap();
2613 let map = mods.get_or_insert_with(Default::default);
2614 match map.get(&ctx_key) {
2615 Some(&m) => m,
2616 None => {
2617 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2618 map.insert(ctx_key, m);
2619 m
2620 }
2621 }
2622 };
2623 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2624 }
2625 if r != cu::CUresult::CUDA_SUCCESS {
2626 return Err(format!("pdl_func {name}: {r:?}").into());
2627 }
2628 FNS.lock()
2629 .unwrap()
2630 .get_or_insert_with(Default::default)
2631 .insert((ctx_key, name), f as usize);
2632 Ok(f)
2633 }
2634
2635 unsafe fn launch_pdl_flash(
2647 &self,
2648 g: bool,
2649 name: &'static str,
2650 grid: (u32, u32, u32),
2651 block: (u32, u32, u32),
2652 smem: u32,
2653 params: &mut [*mut std::ffi::c_void],
2654 ) -> Result<(), Box<dyn std::error::Error>> {
2655 use cudarc::driver::sys as cu;
2656 let f = self.pdl_func_flash(g, name)?;
2657 if smem > 0 {
2658 let r =
2660 unsafe {
2661 cu::cuFuncSetAttribute(f,
2662 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2663 smem as i32)
2664 };
2665 if r != cu::CUresult::CUDA_SUCCESS {
2666 return Err(format!("pdl smem attr {name}: {r:?}").into());
2667 }
2668 }
2669 let mut attr = cu::CUlaunchAttribute {
2670 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2671 pad: [0; 4],
2672 value: cu::CUlaunchAttributeValue {
2673 programmaticStreamSerializationAllowed: 1,
2674 },
2675 };
2676 let cfg = cu::CUlaunchConfig {
2677 gridDimX: grid.0,
2678 gridDimY: grid.1,
2679 gridDimZ: grid.2,
2680 blockDimX: block.0,
2681 blockDimY: block.1,
2682 blockDimZ: block.2,
2683 sharedMemBytes: smem,
2684 hStream: self.gpu.stream().cu_stream(),
2685 attrs: &mut attr,
2686 numAttrs: 1,
2687 };
2688 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2689 if r != cu::CUresult::CUDA_SUCCESS {
2690 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2691 }
2692 Ok(())
2693 }
2694
2695 unsafe fn launch_pdl(
2696 &self,
2697 name: &'static str,
2698 grid: (u32, u32, u32),
2699 block: (u32, u32, u32),
2700 params: &mut [*mut std::ffi::c_void],
2701 ) -> Result<(), Box<dyn std::error::Error>> {
2702 use cudarc::driver::sys as cu;
2703 let f = self.pdl_func(name)?;
2704 let mut attr = cu::CUlaunchAttribute {
2705 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2706 pad: [0; 4],
2707 value: cu::CUlaunchAttributeValue {
2708 programmaticStreamSerializationAllowed: 1,
2709 },
2710 };
2711 let cfg = cu::CUlaunchConfig {
2712 gridDimX: grid.0,
2713 gridDimY: grid.1,
2714 gridDimZ: grid.2,
2715 blockDimX: block.0,
2716 blockDimY: block.1,
2717 blockDimZ: block.2,
2718 sharedMemBytes: 0,
2719 hStream: self.gpu.stream().cu_stream(),
2720 attrs: &mut attr,
2721 numAttrs: 1,
2722 };
2723 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2724 if r != cu::CUresult::CUDA_SUCCESS {
2725 return Err(format!("launch_pdl {name}: {r:?}").into());
2726 }
2727 Ok(())
2728 }
2729
2730 pub fn prefetch_weight_l2(
2733 &self,
2734 w: &crate::model::GpuTensor,
2735 ) -> Result<(), Box<dyn std::error::Error>> {
2736 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2737 let p = rp4.as_ref().unwrap_or(bytes);
2738 self.prefetch_l2(p, p.len())?;
2739 }
2740 Ok(())
2741 }
2742
2743 pub fn gather_row_bf16(
2746 &self,
2747 table: &CudaSlice<u8>,
2748 tok: &CudaSlice<u32>,
2749 idx: usize,
2750 dst: &mut CudaSlice<f32>,
2751 ncols: usize,
2752 ) -> Result<(), Box<dyn std::error::Error>> {
2753 let f = self.func("gather_row_bf16_f32");
2754 let cfg = LaunchConfig {
2755 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2756 block_dim: (256, 1, 1),
2757 shared_mem_bytes: 0,
2758 };
2759 let (nc, ix) = (ncols as i32, idx as i32);
2760 let __s_b = self.gpu.stream();
2761 let mut b = __s_b.launch_builder(&f);
2762 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2763 unsafe {
2764 b.launch(cfg)?;
2765 }
2766 Ok(())
2767 }
2768
2769 #[allow(clippy::too_many_arguments)]
2775 pub fn dflash2_dynconv(
2776 &self,
2777 x: &CudaSlice<f32>,
2778 dyn_: &CudaSlice<f32>,
2779 base: &CudaSlice<f32>,
2780 out: &mut CudaSlice<f32>,
2781 rows: usize,
2782 hidden: usize,
2783 group_size: usize,
2784 ksize: usize,
2785 half: usize,
2786 ) -> Result<(), Box<dyn std::error::Error>> {
2787 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2788 let f = self.func("dflash2_dynconv_f32");
2789 let n = rows * hidden;
2790 let cfg = LaunchConfig {
2791 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2792 block_dim: (256, 1, 1),
2793 shared_mem_bytes: 0,
2794 };
2795 let (ri, hi, gi, ki, hf) = (
2796 rows as i32,
2797 hidden as i32,
2798 group_size as i32,
2799 ksize as i32,
2800 half as i32,
2801 );
2802 let __s_b = self.gpu.stream();
2803 let mut b = __s_b.launch_builder(&f);
2804 b.arg(x)
2805 .arg(dyn_)
2806 .arg(base)
2807 .arg(out)
2808 .arg(&ri)
2809 .arg(&hi)
2810 .arg(&gi)
2811 .arg(&ki)
2812 .arg(&hf);
2813 unsafe {
2814 b.launch(cfg)?;
2815 }
2816 Ok(())
2817 }
2818
2819 pub fn topk_rows(
2823 &self,
2824 logits: &CudaSlice<f32>,
2825 n_rows: usize,
2826 n_cols: usize,
2827 k: usize,
2828 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2829 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2830 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2831 let f = self.func("topk_rows_f32");
2832 let nth = 256usize;
2833 let mut vals = self.uninit(n_rows * k)?;
2834 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2835 let cfg = LaunchConfig {
2836 grid_dim: (n_rows as u32, 1, 1),
2837 block_dim: (nth as u32, 1, 1),
2838 shared_mem_bytes: (nth * k * 8) as u32,
2839 };
2840 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2841 let __s_b = self.gpu.stream();
2842 let mut b = __s_b.launch_builder(&f);
2843 b.arg(logits)
2844 .arg(&nr)
2845 .arg(&nc)
2846 .arg(&ki)
2847 .arg(&mut vals)
2848 .arg(&mut idxs);
2849 unsafe {
2850 b.launch(cfg)?;
2851 }
2852 Ok((vals, idxs))
2853 }
2854
2855 pub fn add_row_inplace(
2857 &self,
2858 logits: &mut CudaSlice<f32>,
2859 bias: &CudaSlice<f32>,
2860 n: usize,
2861 row_off: usize,
2862 ) -> Result<(), Box<dyn std::error::Error>> {
2863 let f = self.func("add_row_inplace_f32");
2864 let cfg = LaunchConfig {
2865 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2866 block_dim: (256, 1, 1),
2867 shared_mem_bytes: 0,
2868 };
2869 let (ni, off) = (n as i32, row_off as i64);
2870 let __s_b = self.gpu.stream();
2871 let mut b = __s_b.launch_builder(&f);
2872 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2873 unsafe {
2874 b.launch(cfg)?;
2875 }
2876 Ok(())
2877 }
2878
2879 pub fn prefetch_l2(
2881 &self,
2882 p: &CudaSlice<u8>,
2883 n: usize,
2884 ) -> Result<(), Box<dyn std::error::Error>> {
2885 let f = self.func("prefetch_l2_bytes");
2886 let lines = n.div_ceil(128);
2887 let ni = n as i64;
2888 let cfg = LaunchConfig {
2889 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2890 block_dim: (256, 1, 1),
2891 shared_mem_bytes: 0,
2892 };
2893 let __s_b = self.gpu.stream();
2894 let mut b = __s_b.launch_builder(&f);
2895 b.arg(p).arg(&ni);
2896 unsafe {
2897 b.launch(cfg)?;
2898 }
2899 Ok(())
2900 }
2901
2902 pub fn router_gemv(
2905 &self,
2906 w: &CudaSlice<f32>,
2907 x: &CudaSlice<f32>,
2908 n_embd: usize,
2909 n_experts: usize,
2910 t: usize,
2911 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2912 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2918 Ok("0") => false,
2919 Ok(_) => true,
2920 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2921 };
2922 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2932 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2933 }
2934
2935 pub fn router_gemv_form(
2938 &self,
2939 w: &CudaSlice<f32>,
2940 x: &CudaSlice<f32>,
2941 n_embd: usize,
2942 n_experts: usize,
2943 t: usize,
2944 w8: bool,
2945 batch: bool,
2946 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2947 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2948 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2949 let f = if batch {
2950 self.func("router_gemv_f32_w8_batch")
2951 } else if w8 {
2952 self.func("router_gemv_f32_w8")
2953 } else {
2954 self.func("router_gemv_f32")
2955 };
2956 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2957 let cfg = if batch {
2958 LaunchConfig {
2959 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2960 block_dim: (32, 8, 1),
2961 shared_mem_bytes: 0,
2962 }
2963 } else {
2964 LaunchConfig {
2965 grid_dim: (n_experts as u32, t as u32, 1),
2966 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2967 shared_mem_bytes: 0,
2968 }
2969 };
2970 let __s_b = self.gpu.stream();
2971 let mut b = __s_b.launch_builder(&f);
2972 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2973 unsafe {
2974 b.launch(cfg)?;
2975 }
2976 Ok(y)
2977 }
2978
2979 pub fn router_gemv_into(
2982 &self,
2983 w: &CudaSlice<f32>,
2984 x: &CudaSlice<f32>,
2985 y: &mut CudaSlice<f32>,
2986 n_embd: usize,
2987 n_experts: usize,
2988 t: usize,
2989 ) -> Result<(), Box<dyn std::error::Error>> {
2990 if y.len() < t * n_experts {
2991 return Err("router_gemv_into output too small".into());
2992 }
2993 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2994 Ok("0") => false,
2995 Ok(_) => true,
2996 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2997 };
2998 let f = if w8 {
2999 self.func("router_gemv_f32_w8")
3000 } else {
3001 self.func("router_gemv_f32")
3002 };
3003 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3004 let cfg = LaunchConfig {
3005 grid_dim: (n_experts as u32, t as u32, 1),
3006 block_dim: (32, if w8 { 8 } else { 1 }, 1),
3007 shared_mem_bytes: 0,
3008 };
3009 let __s_b = self.gpu.stream();
3010 let mut b = __s_b.launch_builder(&f);
3011 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
3012 unsafe {
3013 b.launch(cfg)?;
3014 }
3015 Ok(())
3016 }
3017
3018 pub fn rows_permute(
3020 &self,
3021 src: &CudaSlice<f32>,
3022 idx: &CudaSlice<i32>,
3023 nrows: usize,
3024 ncols: usize,
3025 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3026 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3027 let f = self.func("rows_permute_f32");
3028 let (nc, nr) = (ncols as i32, nrows as i32);
3029 let cfg = LaunchConfig {
3030 grid_dim: (nrows as u32, 1, 1),
3031 block_dim: (256, 1, 1),
3032 shared_mem_bytes: 0,
3033 };
3034 let __s_b = self.gpu.stream();
3035 let mut b = __s_b.launch_builder(&f);
3036 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3037 unsafe {
3038 b.launch(cfg)?;
3039 }
3040 Ok(dst)
3041 }
3042
3043 pub fn sigmoid_dot_rows(
3048 &self,
3049 x: &CudaSlice<f32>,
3050 w: &CudaSlice<f32>,
3051 n_embd: usize,
3052 t: usize,
3053 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3054 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3057 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3058 let gs = self.linear(x, w, t, n_embd, 1)?;
3059 let mut g = self.uninit(t)?;
3060 self.sigmoid(&gs, &mut g, t)?;
3061 return Ok(g);
3062 }
3063 let mut g = self.alloc_uninit::<f32>(t)?;
3069 let f = self.func("sigmoid_dot_rows_f32");
3070 let (ne, ti) = (n_embd as i32, t as i32);
3071 let cfg = LaunchConfig {
3072 grid_dim: (t as u32, 1, 1),
3073 block_dim: (32, 8, 1),
3074 shared_mem_bytes: 0,
3075 };
3076 let __s_b = self.gpu.stream();
3077 let mut b = __s_b.launch_builder(&f);
3078 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3079 unsafe {
3080 b.launch(cfg)?;
3081 }
3082 Ok(g)
3083 }
3084
3085 pub fn sigmoid_dot_rows_into(
3087 &self,
3088 x: &CudaSlice<f32>,
3089 w: &CudaSlice<f32>,
3090 g: &mut CudaSlice<f32>,
3091 n_embd: usize,
3092 t: usize,
3093 ) -> Result<(), Box<dyn std::error::Error>> {
3094 if g.len() < t {
3095 return Err("sigmoid_dot_rows_into output too small".into());
3096 }
3097 let f = self.func("sigmoid_dot_rows_f32");
3098 let (ne, ti) = (n_embd as i32, t as i32);
3099 let cfg = LaunchConfig {
3100 grid_dim: (t as u32, 1, 1),
3101 block_dim: (32, 8, 1),
3102 shared_mem_bytes: 0,
3103 };
3104 let __s_b = self.gpu.stream();
3105 let mut b = __s_b.launch_builder(&f);
3106 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3107 unsafe {
3108 b.launch(cfg)?;
3109 }
3110 Ok(())
3111 }
3112
3113 pub fn spec_rollback_stream(
3115 &self,
3116 len_ptrs: &CudaSlice<u64>,
3117 pos_start: &CudaSlice<i32>,
3118 acc: &CudaSlice<u32>,
3119 base: usize,
3120 n_rows: usize,
3121 ) -> Result<(), Box<dyn std::error::Error>> {
3122 let f = self.func("spec_rollback_stream");
3123 let (b, nr) = (base as i32, n_rows as i32);
3124 let cfg = LaunchConfig {
3125 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3126 block_dim: (64, 1, 1),
3127 shared_mem_bytes: 0,
3128 };
3129 let __s_bl = self.gpu.stream();
3130 let mut bl = __s_bl.launch_builder(&f);
3131 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3132 unsafe {
3133 bl.launch(cfg)?;
3134 }
3135 Ok(())
3136 }
3137
3138 pub fn plain_tok_ring(
3140 &self,
3141 vam: &CudaSlice<u32>,
3142 pos_start: &CudaSlice<i32>,
3143 base: usize,
3144 ring: &mut CudaSlice<u32>,
3145 ) -> Result<(), Box<dyn std::error::Error>> {
3146 let f = self.func("plain_tok_ring");
3147 let (b, cap) = (base as i32, ring.len() as i32);
3148 let cfg = LaunchConfig {
3149 grid_dim: (1, 1, 1),
3150 block_dim: (32, 1, 1),
3151 shared_mem_bytes: 0,
3152 };
3153 let __s_bl = self.gpu.stream();
3154 let mut bl = __s_bl.launch_builder(&f);
3155 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3156 unsafe {
3157 bl.launch(cfg)?;
3158 }
3159 Ok(())
3160 }
3161
3162 pub fn spec_ring_commit(
3164 &self,
3165 vtok: &CudaSlice<u32>,
3166 acc: &CudaSlice<u32>,
3167 brk: &CudaSlice<u32>,
3168 ring: &mut CudaSlice<u32>,
3169 pend: &mut CudaSlice<u32>,
3170 ) -> Result<(), Box<dyn std::error::Error>> {
3171 let f = self.func("spec_ring_commit");
3172 let cfg = LaunchConfig {
3173 grid_dim: (1, 1, 1),
3174 block_dim: (32, 1, 1),
3175 shared_mem_bytes: 0,
3176 };
3177 let __s_b = self.gpu.stream();
3178 let mut b = __s_b.launch_builder(&f);
3179 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3180 unsafe {
3181 b.launch(cfg)?;
3182 }
3183 Ok(())
3184 }
3185 pub fn i32_copy_add(
3186 &self,
3187 src: &CudaSlice<i32>,
3188 dst: &mut CudaSlice<i32>,
3189 delta: i32,
3190 ) -> Result<(), Box<dyn std::error::Error>> {
3191 let f = self.func("i32_copy_add");
3192 let cfg = LaunchConfig {
3193 grid_dim: (1, 1, 1),
3194 block_dim: (32, 1, 1),
3195 shared_mem_bytes: 0,
3196 };
3197 let __s_b = self.gpu.stream();
3198 let mut b = __s_b.launch_builder(&f);
3199 b.arg(src).arg(dst).arg(&delta);
3200 unsafe {
3201 b.launch(cfg)?;
3202 }
3203 Ok(())
3204 }
3205 pub fn u32_copy(
3206 &self,
3207 src: &CudaSlice<u32>,
3208 dst: &mut CudaSlice<u32>,
3209 ) -> Result<(), Box<dyn std::error::Error>> {
3210 let f = self.func("u32_copy");
3211 let cfg = LaunchConfig {
3212 grid_dim: (1, 1, 1),
3213 block_dim: (32, 1, 1),
3214 shared_mem_bytes: 0,
3215 };
3216 let __s_b = self.gpu.stream();
3217 let mut b = __s_b.launch_builder(&f);
3218 b.arg(src).arg(dst);
3219 unsafe {
3220 b.launch(cfg)?;
3221 }
3222 Ok(())
3223 }
3224
3225 pub fn spec_adapt_k(
3229 &self,
3230 acc: &CudaSlice<u32>,
3231 brk: &mut CudaSlice<u32>,
3232 floor: usize,
3233 cap: usize,
3234 ) -> Result<(), Box<dyn std::error::Error>> {
3235 let f = self.func("spec_adapt_k");
3236 let (fl, cp) = (floor as i32, cap as i32);
3237 let cfg = LaunchConfig {
3238 grid_dim: (1, 1, 1),
3239 block_dim: (32, 1, 1),
3240 shared_mem_bytes: 0,
3241 };
3242 let __s_b = self.gpu.stream();
3243 let mut b = __s_b.launch_builder(&f);
3244 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3245 unsafe {
3246 b.launch(cfg)?;
3247 }
3248 Ok(())
3249 }
3250
3251 pub fn spec_accept_greedy_dc(
3253 &self,
3254 preds: &CudaSlice<u32>,
3255 vtok: &CudaSlice<u32>,
3256 last_pred: &CudaSlice<u32>,
3257 brk: &CudaSlice<u32>,
3258 out: &mut CudaSlice<u32>,
3259 ) -> Result<(), Box<dyn std::error::Error>> {
3260 let f = self.func("spec_accept_greedy_dc");
3261 let cfg = LaunchConfig {
3262 grid_dim: (1, 1, 1),
3263 block_dim: (32, 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(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3269 unsafe {
3270 b.launch(cfg)?;
3271 }
3272 Ok(())
3273 }
3274
3275 pub fn pos_iota(
3277 &self,
3278 pos0: &CudaSlice<i32>,
3279 out: &mut CudaSlice<i32>,
3280 t: usize,
3281 ) -> Result<(), Box<dyn std::error::Error>> {
3282 let f = self.func("pos_iota_i32");
3283 let ti = t as i32;
3284 let cfg = LaunchConfig {
3285 grid_dim: (1, 1, 1),
3286 block_dim: (t.max(1) as u32, 1, 1),
3287 shared_mem_bytes: 0,
3288 };
3289 let __s_b = self.gpu.stream();
3290 let mut b = __s_b.launch_builder(&f);
3291 b.arg(pos0).arg(out).arg(&ti);
3292 unsafe {
3293 b.launch(cfg)?;
3294 }
3295 Ok(())
3296 }
3297 #[allow(clippy::too_many_arguments)]
3298 pub fn append_kv_quantized_rows_dc(
3299 &self,
3300 k_rows: &CudaSlice<f32>,
3301 v_rows: &CudaSlice<f32>,
3302 kc: &mut CudaSlice<u8>,
3303 vc: &mut CudaSlice<u8>,
3304 t0_dev: &CudaSlice<i32>,
3305 t: usize,
3306 kv_dim_k: usize,
3307 kv_dim_v: usize,
3308 k_tok_bytes: usize,
3309 v_tok_bytes: usize,
3310 g: bool,
3311 ) -> Result<(), Box<dyn std::error::Error>> {
3312 let f = if g {
3313 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3314 } else {
3315 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3316 };
3317 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3318 let cfg = LaunchConfig {
3319 grid_dim: (nblk, t as u32, 1),
3320 block_dim: (32, 1, 1),
3321 shared_mem_bytes: 0,
3322 };
3323 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3324 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3325 let __s_b = self.gpu.stream();
3326 let mut b = __s_b.launch_builder(&f);
3327 b.arg(k_rows)
3328 .arg(v_rows)
3329 .arg(kc)
3330 .arg(vc)
3331 .arg(t0_dev)
3332 .arg(&kdk)
3333 .arg(&kdv)
3334 .arg(&ktb)
3335 .arg(&vtb);
3336 unsafe {
3337 b.launch(cfg)?;
3338 }
3339 Ok(())
3340 }
3341
3342 #[allow(clippy::too_many_arguments)]
3345 pub fn append_kv_quantized_row_dc_inc(
3346 &self,
3347 k_row: &CudaSlice<f32>,
3348 v_row: &CudaSlice<f32>,
3349 kc: &mut CudaSlice<u8>,
3350 vc: &mut CudaSlice<u8>,
3351 t0_dev: &mut CudaSlice<i32>,
3352 kv_dim_k: usize,
3353 kv_dim_v: usize,
3354 k_tok_bytes: usize,
3355 v_tok_bytes: usize,
3356 g: bool,
3357 ) -> Result<(), Box<dyn std::error::Error>> {
3358 let f = if g {
3359 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3360 } else {
3361 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3362 };
3363 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3364 let cfg = LaunchConfig {
3365 grid_dim: (1, 1, 1),
3366 block_dim: (nthreads, 1, 1),
3367 shared_mem_bytes: 0,
3368 };
3369 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3370 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3371 let __s_b = self.gpu.stream();
3372 let mut b = __s_b.launch_builder(&f);
3373 b.arg(k_row)
3374 .arg(v_row)
3375 .arg(kc)
3376 .arg(vc)
3377 .arg(t0_dev)
3378 .arg(&kdk)
3379 .arg(&kdv)
3380 .arg(&ktb)
3381 .arg(&vtb);
3382 unsafe {
3383 b.launch(cfg)?;
3384 }
3385 Ok(())
3386 }
3387
3388 pub fn pack_tok_p(
3390 &self,
3391 tok: &CudaSlice<u32>,
3392 p: &CudaSlice<f32>,
3393 out: &mut CudaSlice<u32>,
3394 slot: usize,
3395 ) -> Result<(), Box<dyn std::error::Error>> {
3396 let f = self.func("pack_tok_p");
3397 let sl = slot as i32;
3398 let cfg = LaunchConfig {
3399 grid_dim: (1, 1, 1),
3400 block_dim: (32, 1, 1),
3401 shared_mem_bytes: 0,
3402 };
3403 let __s_b = self.gpu.stream();
3404 let mut b = __s_b.launch_builder(&f);
3405 b.arg(tok).arg(p).arg(out).arg(&sl);
3406 unsafe {
3407 b.launch(cfg)?;
3408 }
3409 Ok(())
3410 }
3411 pub fn tok_map_u32(
3412 &self,
3413 tok: &mut CudaSlice<u32>,
3414 map: &CudaSlice<u32>,
3415 ) -> Result<(), Box<dyn std::error::Error>> {
3416 let f = self.func("tok_map_u32");
3417 let cfg = LaunchConfig {
3418 grid_dim: (1, 1, 1),
3419 block_dim: (32, 1, 1),
3420 shared_mem_bytes: 0,
3421 };
3422 let __s_b = self.gpu.stream();
3423 let mut b = __s_b.launch_builder(&f);
3424 b.arg(tok).arg(map);
3425 unsafe {
3426 b.launch(cfg)?;
3427 }
3428 Ok(())
3429 }
3430
3431 #[allow(clippy::too_many_arguments)]
3433 pub fn spec_assemble_verify(
3434 &self,
3435 tokp: &CudaSlice<u32>,
3436 pend: &CudaSlice<u32>,
3437 d2t: Option<&CudaSlice<u32>>,
3438 vtok: &mut CudaSlice<u32>,
3439 brk: &mut CudaSlice<u32>,
3440 p_min: f32,
3441 k: usize,
3442 pmin0: bool,
3443 ) -> Result<(), Box<dyn std::error::Error>> {
3444 let f = self.func("spec_assemble_verify");
3445 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3446 let cfg = LaunchConfig {
3447 grid_dim: (1, 1, 1),
3448 block_dim: (32, 1, 1),
3449 shared_mem_bytes: 0,
3450 };
3451 let __s_b = self.gpu.stream();
3452 let mut b = __s_b.launch_builder(&f);
3453 match d2t {
3454 Some(m) => {
3455 b.arg(tokp)
3456 .arg(pend)
3457 .arg(m)
3458 .arg(vtok)
3459 .arg(brk)
3460 .arg(&p_min)
3461 .arg(&ki)
3462 .arg(&pm);
3463 unsafe {
3464 b.launch(cfg)?;
3465 }
3466 }
3467 None => {
3468 let null: u64 = 0;
3469 b.arg(tokp)
3470 .arg(pend)
3471 .arg(&null)
3472 .arg(vtok)
3473 .arg(brk)
3474 .arg(&p_min)
3475 .arg(&ki)
3476 .arg(&pm);
3477 unsafe {
3478 b.launch(cfg)?;
3479 }
3480 }
3481 }
3482 Ok(())
3483 }
3484
3485 #[allow(clippy::too_many_arguments)]
3487 pub fn ssm_conv_ring_rebuild_dc(
3488 &self,
3489 qkv_tm: &CudaSlice<f32>,
3490 ring_old: &CudaSlice<f32>,
3491 conv_state: &mut CudaSlice<f32>,
3492 conv_dim: usize,
3493 acc: &CudaSlice<u32>,
3494 base: usize,
3495 t_v: usize,
3496 d_conv: usize,
3497 ) -> Result<(), Box<dyn std::error::Error>> {
3498 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3499 let n = conv_dim * (d_conv - 1);
3500 let cfg = LaunchConfig::for_num_elems(n as u32);
3501 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3502 let __s_b = self.gpu.stream();
3503 let mut b = __s_b.launch_builder(&f);
3504 b.arg(qkv_tm)
3505 .arg(ring_old)
3506 .arg(conv_state)
3507 .arg(&cd)
3508 .arg(acc)
3509 .arg(&b0)
3510 .arg(&tv)
3511 .arg(&dc);
3512 unsafe {
3513 b.launch(cfg)?;
3514 }
3515 Ok(())
3516 }
3517 #[allow(clippy::too_many_arguments)]
3518 pub fn gdn_scan_s128_dc(
3519 &self,
3520 q: &CudaSlice<f32>,
3521 k: &CudaSlice<f32>,
3522 v: &CudaSlice<f32>,
3523 g: &CudaSlice<f32>,
3524 beta: &CudaSlice<f32>,
3525 state_in: &CudaSlice<f32>,
3526 state_out: &mut CudaSlice<f32>,
3527 o: &mut CudaSlice<f32>,
3528 n_head: usize,
3529 acc: &CudaSlice<u32>,
3530 base: usize,
3531 t_v: usize,
3532 scale: f32,
3533 ) -> Result<(), Box<dyn std::error::Error>> {
3534 let f = self.func("gdn_scan_s128_dc");
3535 const S_V: u32 = 128;
3536 const WARP: u32 = 32;
3537 const COLS_PER_BLOCK: u32 = 4;
3538 let cfg = LaunchConfig {
3539 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3540 block_dim: (WARP, COLS_PER_BLOCK, 1),
3541 shared_mem_bytes: 0,
3542 };
3543 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3544 let __s_b = self.gpu.stream();
3545 let mut b = __s_b.launch_builder(&f);
3546 b.arg(q)
3547 .arg(k)
3548 .arg(v)
3549 .arg(g)
3550 .arg(beta)
3551 .arg(state_in)
3552 .arg(state_out)
3553 .arg(o)
3554 .arg(&h)
3555 .arg(acc)
3556 .arg(&b0)
3557 .arg(&tv)
3558 .arg(&scale);
3559 unsafe {
3560 b.launch(cfg)?;
3561 }
3562 Ok(())
3563 }
3564
3565 pub fn spec_rollback_kv(
3567 &self,
3568 len_ptrs: &CudaSlice<u64>,
3569 saved: &CudaSlice<i32>,
3570 acc: &CudaSlice<u32>,
3571 base: usize,
3572 n_layer: usize,
3573 ) -> Result<(), Box<dyn std::error::Error>> {
3574 let f = self.func("spec_rollback_kv");
3575 let (b, nl) = (base as i32, n_layer as i32);
3576 let cfg = LaunchConfig {
3577 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3578 block_dim: (64, 1, 1),
3579 shared_mem_bytes: 0,
3580 };
3581 let __s_bl = self.gpu.stream();
3582 let mut bl = __s_bl.launch_builder(&f);
3583 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3584 unsafe {
3585 bl.launch(cfg)?;
3586 }
3587 Ok(())
3588 }
3589
3590 pub fn spec_fork_valid(
3592 &self,
3593 acc: &CudaSlice<u32>,
3594 optimistic_pending: u32,
3595 valid: &mut CudaSlice<u32>,
3596 ) -> Result<(), Box<dyn std::error::Error>> {
3597 let f = self.func("spec_fork_valid");
3598 let cfg = LaunchConfig {
3599 grid_dim: (1, 1, 1),
3600 block_dim: (1, 1, 1),
3601 shared_mem_bytes: 0,
3602 };
3603 let __s_bl = self.gpu.stream();
3604 let mut bl = __s_bl.launch_builder(&f);
3605 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3606 unsafe {
3607 bl.launch(cfg)?;
3608 }
3609 Ok(())
3610 }
3611
3612 pub fn spec_fork_reconcile_kv(
3614 &self,
3615 len_ptrs: &CudaSlice<u64>,
3616 saved: &CudaSlice<i32>,
3617 acc: &CudaSlice<u32>,
3618 valid: &CudaSlice<u32>,
3619 base: usize,
3620 n_layer: usize,
3621 ) -> Result<(), Box<dyn std::error::Error>> {
3622 let f = self.func("spec_fork_reconcile_kv");
3623 let (b, nl) = (base as i32, n_layer as i32);
3624 let cfg = LaunchConfig {
3625 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3626 block_dim: (64, 1, 1),
3627 shared_mem_bytes: 0,
3628 };
3629 let __s_bl = self.gpu.stream();
3630 let mut bl = __s_bl.launch_builder(&f);
3631 bl.arg(len_ptrs)
3632 .arg(saved)
3633 .arg(acc)
3634 .arg(valid)
3635 .arg(&b)
3636 .arg(&nl);
3637 unsafe {
3638 bl.launch(cfg)?;
3639 }
3640 Ok(())
3641 }
3642
3643 pub fn spec_fork_restore_f32(
3645 &self,
3646 snapshot: &CudaSlice<f32>,
3647 state: &mut CudaSlice<f32>,
3648 valid: &CudaSlice<u32>,
3649 ) -> Result<(), Box<dyn std::error::Error>> {
3650 assert_eq!(
3651 snapshot.len(),
3652 state.len(),
3653 "fork recurrent snapshot shape mismatch"
3654 );
3655 let f = self.func("spec_fork_restore_f32");
3656 let n = state.len() as i32;
3657 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3658 let cfg = LaunchConfig {
3659 grid_dim: (blocks, 1, 1),
3660 block_dim: (256, 1, 1),
3661 shared_mem_bytes: 0,
3662 };
3663 let __s_bl = self.gpu.stream();
3664 let mut bl = __s_bl.launch_builder(&f);
3665 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3666 unsafe {
3667 bl.launch(cfg)?;
3668 }
3669 Ok(())
3670 }
3671
3672 pub fn spec_seed_gather(
3675 &self,
3676 vx: &CudaSlice<f32>,
3677 fill_prev: &CudaSlice<f32>,
3678 acc: &CudaSlice<u32>,
3679 h_seed: &mut CudaSlice<f32>,
3680 base: usize,
3681 n_embd: usize,
3682 ) -> Result<(), Box<dyn std::error::Error>> {
3683 let f = self.func("spec_seed_gather");
3684 let (b, ne) = (base as i32, n_embd as i32);
3685 let cfg = LaunchConfig {
3686 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3687 block_dim: (256, 1, 1),
3688 shared_mem_bytes: 0,
3689 };
3690 let __s_bl = self.gpu.stream();
3691 let mut bl = __s_bl.launch_builder(&f);
3692 bl.arg(vx)
3693 .arg(fill_prev)
3694 .arg(acc)
3695 .arg(h_seed)
3696 .arg(&b)
3697 .arg(&ne);
3698 unsafe {
3699 bl.launch(cfg)?;
3700 }
3701 Ok(())
3702 }
3703
3704 pub fn spec_accept_greedy(
3706 &self,
3707 preds: &CudaSlice<u32>,
3708 draft: &CudaSlice<u32>,
3709 last_pred: u32,
3710 base: usize,
3711 k_round: usize,
3712 out: &mut CudaSlice<u32>,
3713 ) -> Result<(), Box<dyn std::error::Error>> {
3714 let f = self.func("spec_accept_greedy");
3715 let (b, k) = (base as i32, k_round as i32);
3716 let cfg = LaunchConfig {
3717 grid_dim: (1, 1, 1),
3718 block_dim: (32, 1, 1),
3719 shared_mem_bytes: 0,
3720 };
3721 let __s_bl = self.gpu.stream();
3722 let mut bl = __s_bl.launch_builder(&f);
3723 bl.arg(preds)
3724 .arg(draft)
3725 .arg(&last_pred)
3726 .arg(&b)
3727 .arg(&k)
3728 .arg(out);
3729 unsafe {
3730 bl.launch(cfg)?;
3731 }
3732 Ok(())
3733 }
3734
3735 pub fn gumbel_perturb(
3742 &self,
3743 x: &CudaSlice<f32>,
3744 y: &mut CudaSlice<f32>,
3745 n: usize,
3746 seed: u64,
3747 stream_pos: u32,
3748 temp: f32,
3749 ) -> Result<(), Box<dyn std::error::Error>> {
3750 let f = self.func("gumbel_perturb_f32");
3751 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3752 let cfg = LaunchConfig {
3753 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3754 block_dim: (256, 1, 1),
3755 shared_mem_bytes: 0,
3756 };
3757 let __s_b = self.gpu.stream();
3758 let mut b = __s_b.launch_builder(&f);
3759 b.arg(x)
3760 .arg(&mut *y)
3761 .arg(&ni)
3762 .arg(&slo)
3763 .arg(&shi)
3764 .arg(&stream_pos)
3765 .arg(&temp);
3766 unsafe {
3767 b.launch(cfg)?;
3768 }
3769 Ok(())
3770 }
3771
3772 pub fn mask_logits_col(
3780 &self,
3781 logits: &mut CudaSlice<f32>,
3782 mask: &CudaSlice<u32>,
3783 col: usize,
3784 n: usize,
3785 mask_words: usize,
3786 ) -> Result<(), Box<dyn std::error::Error>> {
3787 let f = self.func("mask_logits_f32");
3788 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3789 let cfg = LaunchConfig {
3790 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3791 block_dim: (256, 1, 1),
3792 shared_mem_bytes: 0,
3793 };
3794 let __s_b = self.gpu.stream();
3795 let mut b = __s_b.launch_builder(&f);
3796 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3797 unsafe {
3798 b.launch(cfg)?;
3799 }
3800 Ok(())
3801 }
3802
3803 pub fn gumbel_perturb_col(
3810 &self,
3811 x: &CudaSlice<f32>,
3812 col: usize,
3813 y: &mut CudaSlice<f32>,
3814 n: usize,
3815 seed: u64,
3816 stream_pos: u32,
3817 temp: f32,
3818 ) -> Result<(), Box<dyn std::error::Error>> {
3819 let f = self.func("gumbel_perturb_f32");
3820 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3821 let col_view = x.slice(col * n..(col + 1) * n);
3822 let cfg = LaunchConfig {
3823 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3824 block_dim: (256, 1, 1),
3825 shared_mem_bytes: 0,
3826 };
3827 let __s_b = self.gpu.stream();
3828 let mut b = __s_b.launch_builder(&f);
3829 b.arg(&col_view)
3830 .arg(&mut *y)
3831 .arg(&ni)
3832 .arg(&slo)
3833 .arg(&shi)
3834 .arg(&stream_pos)
3835 .arg(&temp);
3836 unsafe {
3837 b.launch(cfg)?;
3838 }
3839 Ok(())
3840 }
3841
3842 #[allow(clippy::too_many_arguments)]
3848 pub fn gumbel_perturb_filtered_col(
3849 &self,
3850 x: &CudaSlice<f32>,
3851 col: usize,
3852 y: &mut CudaSlice<f32>,
3853 n: usize,
3854 seed: u64,
3855 stream_pos: u32,
3856 temp: f32,
3857 stat_max: &CudaSlice<f32>,
3858 stat_th: &CudaSlice<f32>,
3859 stat_idx: usize,
3860 ) -> Result<(), Box<dyn std::error::Error>> {
3861 let f = self.func("gumbel_perturb_filtered_col_f32");
3862 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3863 let (ci, si) = (col as i32, stat_idx as i32);
3864 let cfg = LaunchConfig {
3865 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3866 block_dim: (256, 1, 1),
3867 shared_mem_bytes: 0,
3868 };
3869 let __s_b = self.gpu.stream();
3870 let mut b = __s_b.launch_builder(&f);
3871 b.arg(x)
3872 .arg(&ci)
3873 .arg(&mut *y)
3874 .arg(&ni)
3875 .arg(&slo)
3876 .arg(&shi)
3877 .arg(&stream_pos)
3878 .arg(&temp)
3879 .arg(stat_max)
3880 .arg(stat_th)
3881 .arg(&si);
3882 unsafe {
3883 b.launch(cfg)?;
3884 }
3885 Ok(())
3886 }
3887
3888 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3893 let f = self.func("memra_sctr_inc");
3894 let cfg = LaunchConfig {
3895 grid_dim: (1, 1, 1),
3896 block_dim: (1, 1, 1),
3897 shared_mem_bytes: 0,
3898 };
3899 let __s_b = self.gpu.stream();
3900 let mut b = __s_b.launch_builder(&f);
3901 b.arg(&mut *ctr);
3902 unsafe {
3903 b.launch(cfg)?;
3904 }
3905 Ok(())
3906 }
3907
3908 pub fn gumbel_perturb_ctr(
3913 &self,
3914 x: &CudaSlice<f32>,
3915 y: &mut CudaSlice<f32>,
3916 n: usize,
3917 seed: u64,
3918 ctr: &CudaSlice<u32>,
3919 temp: f32,
3920 ) -> Result<(), Box<dyn std::error::Error>> {
3921 let f = self.func("gumbel_perturb_ctr_f32");
3922 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3923 let cfg = LaunchConfig {
3924 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3925 block_dim: (256, 1, 1),
3926 shared_mem_bytes: 0,
3927 };
3928 let __s_b = self.gpu.stream();
3929 let mut b = __s_b.launch_builder(&f);
3930 b.arg(x)
3931 .arg(&mut *y)
3932 .arg(&ni)
3933 .arg(&slo)
3934 .arg(&shi)
3935 .arg(ctr)
3936 .arg(&temp);
3937 unsafe {
3938 b.launch(cfg)?;
3939 }
3940 Ok(())
3941 }
3942
3943 pub fn softmax_gather(
3947 &self,
3948 x: &CudaSlice<f32>,
3949 row_stride: usize,
3950 ids: &CudaSlice<u32>,
3951 rows: &CudaSlice<i32>,
3952 out: &mut CudaSlice<f32>,
3953 n: usize,
3954 npair: usize,
3955 temp: f32,
3956 ) -> Result<(), Box<dyn std::error::Error>> {
3957 let f = self.func("softmax_gather_f32");
3958 let (ni, rs) = (n as i32, row_stride as i64);
3959 let np = npair as i32;
3960 let cfg = LaunchConfig {
3961 grid_dim: (npair as u32, 1, 1),
3962 block_dim: (256, 1, 1),
3963 shared_mem_bytes: 0,
3964 };
3965 let __s_b = self.gpu.stream();
3966 let mut b = __s_b.launch_builder(&f);
3967 b.arg(x)
3968 .arg(&rs)
3969 .arg(ids)
3970 .arg(rows)
3971 .arg(&mut *out)
3972 .arg(&ni)
3973 .arg(&np)
3974 .arg(&temp);
3975 unsafe {
3976 b.launch(cfg)?;
3977 }
3978 Ok(())
3979 }
3980
3981 pub fn residual_sample(
3985 &self,
3986 p: &CudaSlice<f32>,
3987 q: Option<&CudaSlice<f32>>,
3988 n: usize,
3989 temp: f32,
3990 seed: u64,
3991 stream_pos: u32,
3992 out_tok: &mut CudaSlice<u32>,
3993 ) -> Result<(), Box<dyn std::error::Error>> {
3994 let f = self.func("residual_sample_f32");
3995 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3996 let nth = 1024u32;
3997 let cfg = LaunchConfig {
3998 grid_dim: (1, 1, 1),
3999 block_dim: (nth, 1, 1),
4000 shared_mem_bytes: 0,
4001 };
4002 let has_q: i32 = q.is_some() as i32;
4003 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
4005 let mut b = __s_b.launch_builder(&f);
4006 b.arg(p)
4007 .arg(qbuf)
4008 .arg(&has_q)
4009 .arg(&ni)
4010 .arg(&temp)
4011 .arg(&slo)
4012 .arg(&shi)
4013 .arg(&stream_pos)
4014 .arg(&mut *out_tok);
4015 unsafe {
4016 b.launch(cfg)?;
4017 }
4018 Ok(())
4019 }
4020
4021 pub fn with_moe_cache<R>(
4026 &self,
4027 max_block_bytes: usize,
4028 f: impl FnOnce(
4029 &mut crate::moe_cache::MoeSlotCache,
4030 &Engine,
4031 ) -> Result<R, Box<dyn std::error::Error>>,
4032 ) -> Result<R, Box<dyn std::error::Error>> {
4033 let mut guard = self.moe_cache.lock().unwrap();
4034 if guard.is_none() {
4035 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4036 }
4037 let cache = guard.as_mut().unwrap();
4038 f(cache, self)
4039 }
4040
4041 pub fn freeze_moe_cache(&self) {
4044 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4045 cache.freeze();
4046 }
4047 }
4048
4049 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4052 self.moe_cache
4053 .lock()
4054 .unwrap()
4055 .as_ref()
4056 .map(crate::moe_cache::MoeSlotCache::export_residency)
4057 }
4058
4059 pub(crate) fn moe_cache_frozen(&self) -> bool {
4060 self.moe_cache
4061 .lock()
4062 .unwrap()
4063 .as_ref()
4064 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4065 }
4066
4067 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4074 crate::cpu_experts::configured()
4075 && self.moe_cache_frozen()
4076 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4077 }
4078
4079 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4081 assert!(
4082 self.moe_cache.lock().unwrap().is_none(),
4083 "MoE cache layout configured after cache construction"
4084 );
4085 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4086 }
4087
4088 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4089 self.moe_cache_layout.lock().unwrap().clone()
4090 }
4091
4092 pub fn moe_cache_enabled() -> bool {
4094 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4095 }
4096
4097 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4100 let guard = self.moe_cache.lock().unwrap();
4101 guard
4102 .as_ref()
4103 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4104 }
4105
4106 pub fn cpu_expert_stats(
4110 &self,
4111 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4112 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4113 }
4114
4115 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4118 crate::cpu_experts::predictor_stats()
4119 }
4120
4121 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4122 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4123 }
4124
4125 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4128 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4129 }
4130
4131 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4134 let guard = self.moe_cache.lock().unwrap();
4135 guard
4136 .as_ref()
4137 .and_then(|cache| cache.pread_stats())
4138 .map(|stats| {
4139 (
4140 stats.reads,
4141 stats.bytes,
4142 stats.read_errors,
4143 stats.short_reads,
4144 stats.fallbacks,
4145 stats.buffer_waits,
4146 stats.ring_full,
4147 )
4148 })
4149 }
4150
4151 pub fn spill_config_fallbacks(&self) -> u64 {
4153 crate::spill_pread::config_fallbacks()
4154 }
4155
4156 pub fn moe_cache_reset_counters(&self) {
4158 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4159 c.reset_counters();
4160 }
4161 }
4162
4163 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4164 Ok(self.gpu.stream().clone_htod(v)?)
4165 }
4166
4167 pub fn htod_bytes_padded(
4171 &self,
4172 v: &[u8],
4173 pad: usize,
4174 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4175 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4176 {
4177 let mut view = d.slice_mut(0..v.len());
4178 self.gpu.stream().memcpy_htod(v, &mut view)?;
4179 }
4180 Ok(d)
4181 }
4182
4183 pub fn copy_into(
4185 &self,
4186 dst: &mut CudaSlice<f32>,
4187 off: usize,
4188 src: &CudaSlice<f32>,
4189 len: usize,
4190 ) -> Result<(), Box<dyn std::error::Error>> {
4191 let mut view = dst.slice_mut(off..off + len);
4192 self.gpu
4193 .stream()
4194 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4195 Ok(())
4196 }
4197
4198 pub fn copy_range_into(
4202 &self,
4203 dst: &mut CudaSlice<f32>,
4204 dst_off: usize,
4205 src: &CudaSlice<f32>,
4206 src_off: usize,
4207 len: usize,
4208 ) -> Result<(), Box<dyn std::error::Error>> {
4209 let mut view = dst.slice_mut(dst_off..dst_off + len);
4210 self.gpu
4211 .stream()
4212 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
4213 Ok(())
4214 }
4215
4216 pub fn copy_u8_into(
4219 &self,
4220 dst: &mut CudaSlice<u8>,
4221 off: usize,
4222 src: &CudaSlice<u8>,
4223 len: usize,
4224 ) -> Result<(), Box<dyn std::error::Error>> {
4225 let mut view = dst.slice_mut(off..off + len);
4226 self.gpu
4227 .stream()
4228 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4229 Ok(())
4230 }
4231
4232 pub fn copy_u8_range_into(
4234 &self,
4235 dst: &mut CudaSlice<u8>,
4236 dst_off: usize,
4237 src: &CudaSlice<u8>,
4238 src_off: usize,
4239 len: usize,
4240 ) -> Result<(), Box<dyn std::error::Error>> {
4241 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
4242 self.gpu
4243 .stream()
4244 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4245 Ok(())
4246 }
4247
4248 pub fn prepare_kv_append(
4252 &self,
4253 kv: &mut crate::cache::KvLayer,
4254 retain_from: usize,
4255 append_rows: usize,
4256 ) -> Result<usize, Box<dyn std::error::Error>> {
4257 let Some(plan) = kv
4258 .ring
4259 .as_ref()
4260 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4261 .transpose()?
4262 else {
4263 return Ok(kv.len);
4264 };
4265 match plan {
4266 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4267 crate::cache::KvRingAppend::Rebase {
4268 src_row,
4269 keep_rows,
4270 new_base,
4271 write_row,
4272 } => {
4273 if keep_rows > 0 {
4274 let k_len = keep_rows * kv.k_tok_bytes;
4275 let v_len = keep_rows * kv.v_tok_bytes;
4276 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4277 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4278 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4279 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4280 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4281 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4282 }
4283 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4284 Ok(write_row)
4285 }
4286 }
4287 }
4288
4289 pub fn htod_u8_into(
4292 &self,
4293 dst: &mut CudaSlice<u8>,
4294 off: usize,
4295 src: &[u8],
4296 ) -> Result<(), Box<dyn std::error::Error>> {
4297 let mut view = dst.slice_mut(off..off + src.len());
4298 self.gpu.stream().memcpy_htod(src, &mut view)?;
4299 Ok(())
4300 }
4301
4302 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4303 b.slice(0..len)
4304 }
4305
4306 pub fn view_u8_range<'a>(
4309 &self,
4310 b: &'a CudaSlice<u8>,
4311 start: usize,
4312 end: usize,
4313 ) -> cudarc::driver::CudaView<'a, u8> {
4314 b.slice(start..end)
4315 }
4316 pub fn view_u8<'a>(
4317 &self,
4318 b: &'a CudaSlice<u8>,
4319 len: usize,
4320 ) -> cudarc::driver::CudaView<'a, u8> {
4321 b.slice(0..len)
4322 }
4323
4324 pub fn append_kv_quantized(
4328 &self,
4329 k_row: &CudaSlice<f32>,
4330 v_row: &CudaSlice<f32>,
4331 kc: &mut CudaSlice<u8>,
4332 vc: &mut CudaSlice<u8>,
4333 t: usize,
4334 kv_dim_k: usize,
4335 kv_dim_v: usize,
4336 k_tok_bytes: usize,
4337 v_tok_bytes: usize,
4338 g: bool,
4339 ) -> Result<(), Box<dyn std::error::Error>> {
4340 let f = if g {
4341 self.func_g("append_quantize_kv_q8_0_q5_1")
4342 } else {
4343 self.func("append_quantize_kv_q8_0_q5_1")
4344 };
4345 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4346 let cfg = LaunchConfig {
4347 grid_dim: (nblk, 1, 1),
4348 block_dim: (32, 1, 1),
4349 shared_mem_bytes: 0,
4350 };
4351 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4352 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4353 let __s_b = self.gpu.stream();
4354 let mut b = __s_b.launch_builder(&f);
4355 b.arg(k_row)
4356 .arg(v_row)
4357 .arg(kc)
4358 .arg(vc)
4359 .arg(&ti)
4360 .arg(&kdk)
4361 .arg(&kdv)
4362 .arg(&ktb)
4363 .arg(&vtb);
4364 unsafe {
4365 b.launch(cfg)?;
4366 }
4367 Ok(())
4368 }
4369
4370 pub fn append_kv_quantized_dc(
4374 &self,
4375 k_row: &CudaSlice<f32>,
4376 v_row: &CudaSlice<f32>,
4377 kc: &mut CudaSlice<u8>,
4378 vc: &mut CudaSlice<u8>,
4379 t_dev: &CudaSlice<i32>,
4380 kv_dim_k: usize,
4381 kv_dim_v: usize,
4382 k_tok_bytes: usize,
4383 v_tok_bytes: usize,
4384 g: bool,
4385 ) -> Result<(), Box<dyn std::error::Error>> {
4386 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4387 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4388 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4389 if Self::pdl_on() && Self::pdl_wb_on() {
4391 use cudarc::driver::{DevicePtr, DevicePtrMut};
4392 let s = &self.gpu.stream();
4393 let (pk, _g0) = k_row.device_ptr(s);
4394 let (pv, _g1) = v_row.device_ptr(s);
4395 let (pkc, _g2) = kc.device_ptr_mut(s);
4396 let (pvc, _g3) = vc.device_ptr_mut(s);
4397 let (pt, _g4) = t_dev.device_ptr(s);
4398 let mut ps = [
4399 &pk as *const _ as *mut std::ffi::c_void,
4400 &pv as *const _ as *mut _,
4401 &pkc as *const _ as *mut _,
4402 &pvc as *const _ as *mut _,
4403 &pt as *const _ as *mut _,
4404 &kdk as *const _ as *mut _,
4405 &kdv as *const _ as *mut _,
4406 &ktb as *const _ as *mut _,
4407 &vtb as *const _ as *mut _,
4408 ];
4409 unsafe {
4410 self.launch_pdl_flash(
4411 g,
4412 "append_quantize_kv_q8_0_q5_1_dc",
4413 (nblk, 1, 1),
4414 (32, 1, 1),
4415 0,
4416 &mut ps,
4417 )?;
4418 }
4419 return Ok(());
4420 }
4421 let f = if g {
4422 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4423 } else {
4424 self.func("append_quantize_kv_q8_0_q5_1_dc")
4425 };
4426 let cfg = LaunchConfig {
4427 grid_dim: (nblk, 1, 1),
4428 block_dim: (32, 1, 1),
4429 shared_mem_bytes: 0,
4430 };
4431 let __s_b = self.gpu.stream();
4432 let mut b = __s_b.launch_builder(&f);
4433 b.arg(k_row)
4434 .arg(v_row)
4435 .arg(kc)
4436 .arg(vc)
4437 .arg(t_dev)
4438 .arg(&kdk)
4439 .arg(&kdv)
4440 .arg(&ktb)
4441 .arg(&vtb);
4442 unsafe {
4443 b.launch(cfg)?;
4444 }
4445 Ok(())
4446 }
4447
4448 #[allow(clippy::too_many_arguments)]
4455 pub fn append_kv_quantized_rows(
4456 &self,
4457 k_rows: &CudaSlice<f32>,
4458 v_rows: &CudaSlice<f32>,
4459 kc: &mut CudaSlice<u8>,
4460 vc: &mut CudaSlice<u8>,
4461 t0: usize,
4462 t: usize,
4463 kv_dim_k: usize,
4464 kv_dim_v: usize,
4465 k_tok_bytes: usize,
4466 v_tok_bytes: usize,
4467 g: bool,
4468 ) -> Result<(), Box<dyn std::error::Error>> {
4469 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4470 for i in 0..t {
4471 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4472 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4473 self.append_kv_quantized_view(
4474 &k_row,
4475 &v_row,
4476 kc,
4477 vc,
4478 t0 + i,
4479 kv_dim_k,
4480 kv_dim_v,
4481 k_tok_bytes,
4482 v_tok_bytes,
4483 g,
4484 )?;
4485 }
4486 return Ok(());
4487 }
4488 let f = if g {
4489 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4490 } else {
4491 self.func("append_quantize_kv_q8_0_q5_1_rows")
4492 };
4493 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4494 let cfg = LaunchConfig {
4495 grid_dim: (nblk, t as u32, 1),
4496 block_dim: (32, 1, 1),
4497 shared_mem_bytes: 0,
4498 };
4499 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4500 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4501 let __s_b = self.gpu.stream();
4502 let mut b = __s_b.launch_builder(&f);
4503 b.arg(k_rows)
4504 .arg(v_rows)
4505 .arg(kc)
4506 .arg(vc)
4507 .arg(&t0i)
4508 .arg(&kdk)
4509 .arg(&kdv)
4510 .arg(&ktb)
4511 .arg(&vtb);
4512 unsafe {
4513 b.launch(cfg)?;
4514 }
4515 Ok(())
4516 }
4517
4518 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4522 let f = self.func("inc_i32");
4523 let cfg = LaunchConfig {
4524 grid_dim: (1, 1, 1),
4525 block_dim: (1, 1, 1),
4526 shared_mem_bytes: 0,
4527 };
4528 let __s_b = self.gpu.stream();
4529 let mut b = __s_b.launch_builder(&f);
4530 b.arg(p);
4531 unsafe {
4532 b.launch(cfg)?;
4533 }
4534 Ok(())
4535 }
4536
4537 pub fn append_kv_quantized_view(
4540 &self,
4541 k_row: &cudarc::driver::CudaView<f32>,
4542 v_row: &cudarc::driver::CudaView<f32>,
4543 kc: &mut CudaSlice<u8>,
4544 vc: &mut CudaSlice<u8>,
4545 t: usize,
4546 kv_dim_k: usize,
4547 kv_dim_v: usize,
4548 k_tok_bytes: usize,
4549 v_tok_bytes: usize,
4550 g: bool,
4551 ) -> Result<(), Box<dyn std::error::Error>> {
4552 let stream = self.gpu.stream();
4553 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4554 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4555 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4556 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4557 let f = if g {
4558 self.func_g("append_quantize_kv_q8_0_q5_1")
4559 } else {
4560 self.func("append_quantize_kv_q8_0_q5_1")
4561 };
4562 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4563 let cfg = LaunchConfig {
4564 grid_dim: (nblk, 1, 1),
4565 block_dim: (32, 1, 1),
4566 shared_mem_bytes: 0,
4567 };
4568 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4569 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4570 let mut b = stream.launch_builder(&f);
4571 b.arg(k_row)
4572 .arg(v_row)
4573 .arg(kc)
4574 .arg(vc)
4575 .arg(&ti)
4576 .arg(&kdk)
4577 .arg(&kdv)
4578 .arg(&ktb)
4579 .arg(&vtb);
4580 unsafe {
4581 b.launch(cfg)?;
4582 }
4583 Ok(())
4584 }
4585
4586 pub fn copy_view_into(
4589 &self,
4590 dst: &mut CudaSlice<f32>,
4591 off: usize,
4592 src: &cudarc::driver::CudaView<f32>,
4593 len: usize,
4594 ) -> Result<(), Box<dyn std::error::Error>> {
4595 let mut view = dst.slice_mut(off..off + len);
4596 self.gpu
4597 .stream()
4598 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4599 Ok(())
4600 }
4601
4602 pub fn clone_dtod(
4606 &self,
4607 src: &CudaSlice<f32>,
4608 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4609 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4610 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4611 Ok(dst)
4612 }
4613
4614 pub fn dtod_copy_view(
4617 &self,
4618 src: &cudarc::driver::CudaView<f32>,
4619 dst: &mut CudaSlice<f32>,
4620 ) -> Result<(), Box<dyn std::error::Error>> {
4621 self.gpu.stream().memcpy_dtod(src, dst)?;
4622 Ok(())
4623 }
4624
4625 pub fn dtod_copy_view_i8(
4627 &self,
4628 src: &cudarc::driver::CudaView<i8>,
4629 dst: &mut CudaSlice<i8>,
4630 ) -> Result<(), Box<dyn std::error::Error>> {
4631 self.gpu.stream().memcpy_dtod(src, dst)?;
4632 Ok(())
4633 }
4634
4635 pub fn dtod_copy_into(
4637 &self,
4638 src: &CudaSlice<f32>,
4639 dst: &mut CudaSlice<f32>,
4640 offset: usize,
4641 ) -> Result<(), Box<dyn std::error::Error>> {
4642 let n = src.len();
4643 let mut dv = dst.slice_mut(offset..offset + n);
4644 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4645 Ok(())
4646 }
4647
4648 pub fn copy_batch_uniform_f32(
4654 &self,
4655 table: &CudaSlice<u64>,
4656 n: usize,
4657 words: usize,
4658 ) -> Result<(), Box<dyn std::error::Error>> {
4659 if n == 0 || words == 0 {
4660 return Ok(());
4661 }
4662 debug_assert!(
4663 table.len() >= 2 * n,
4664 "pointer table must hold n srcs + n dsts"
4665 );
4666 let f = self.func("copy_batch_uniform_f32");
4667 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4670 let (ni, wi) = (n as i32, words as i32);
4671 let cfg = LaunchConfig {
4672 grid_dim: (chunks, n as u32, 1),
4673 block_dim: (256, 1, 1),
4674 shared_mem_bytes: 0,
4675 };
4676 let __s = self.gpu.stream();
4677 let mut b = __s.launch_builder(&f);
4678 b.arg(table).arg(&ni).arg(&wi);
4679 unsafe {
4680 b.launch(cfg)?;
4681 }
4682 Ok(())
4683 }
4684
4685 pub fn htod_u64_into(
4688 &self,
4689 v: &[u64],
4690 dst: &mut CudaSlice<u64>,
4691 ) -> Result<(), Box<dyn std::error::Error>> {
4692 let mut view = dst.slice_mut(0..v.len());
4693 self.gpu.stream().memcpy_htod(v, &mut view)?;
4694 Ok(())
4695 }
4696
4697 pub fn copy_indirect_src_f32(
4702 &self,
4703 src_entry: &cudarc::driver::CudaView<u64>,
4704 dst: &mut CudaSlice<f32>,
4705 dst_off: usize,
4706 words: usize,
4707 ) -> Result<(), Box<dyn std::error::Error>> {
4708 let f = self.func("copy_indirect_src_f32");
4709 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4710 let wi = words as i32;
4711 let cfg = LaunchConfig {
4712 grid_dim: (chunks, 1, 1),
4713 block_dim: (256, 1, 1),
4714 shared_mem_bytes: 0,
4715 };
4716 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4717 let __s = self.gpu.stream();
4718 let mut b = __s.launch_builder(&f);
4719 b.arg(src_entry).arg(&mut dv).arg(&wi);
4720 unsafe {
4721 b.launch(cfg)?;
4722 }
4723 Ok(())
4724 }
4725
4726 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4728 self.alloc_uninit::<i8>(n)
4729 }
4730
4731 pub fn qmatvec(
4733 &self,
4734 w: &CudaSlice<u8>,
4735 x: &CudaSlice<f32>,
4736 m: usize,
4737 in_f: usize,
4738 out_f: usize,
4739 qtype: i32,
4740 row_bytes: usize,
4741 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4742 let f = self.func("qmatvec_f32");
4743 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4745 grid_dim: (out_f as u32, m as u32, 1),
4746 block_dim: (256, 1, 1),
4747 shared_mem_bytes: 0,
4748 };
4749 let (inf, outf, mi, qt, rb) =
4750 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4751 let __s_b = self.gpu.stream();
4752 let mut b = __s_b.launch_builder(&f);
4753 b.arg(w)
4754 .arg(x)
4755 .arg(&mut y)
4756 .arg(&inf)
4757 .arg(&outf)
4758 .arg(&mi)
4759 .arg(&qt)
4760 .arg(&rb);
4761 unsafe {
4762 b.launch(cfg)?;
4763 }
4764 Ok(y)
4765 }
4766
4767 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4769 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4770 self.keep_if_capturing(&s);
4771 Ok(s)
4772 }
4773
4774 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4778 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4779 self.keep_if_capturing(&s);
4780 Ok(s)
4781 }
4782
4783 pub fn memset_zeros_view(
4786 &self,
4787 dst: &mut cudarc::driver::CudaViewMut<f32>,
4788 ) -> Result<(), Box<dyn std::error::Error>> {
4789 self.gpu.stream().memset_zeros(dst)?;
4790 Ok(())
4791 }
4792
4793 pub fn stage_expert(
4799 &self,
4800 host_bytes: &[u8],
4801 scratch: &mut CudaSlice<u8>,
4802 off: usize,
4803 ) -> Result<(), Box<dyn std::error::Error>> {
4804 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4807 }
4808
4809 pub fn moe_router_topk(
4815 &self,
4816 logits: &CudaSlice<f32>,
4817 t: usize,
4818 n_expert: usize,
4819 n_used: usize,
4820 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4821 let f = self.func("moe_router_topk_f32");
4822 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 {
4825 grid_dim: (t as u32, 1, 1),
4826 block_dim: (n_expert as u32, 1, 1),
4827 shared_mem_bytes: 0,
4828 };
4829 let (ne, nu) = (n_expert as i32, n_used as i32);
4830 let __s_b = self.gpu.stream();
4831 let mut b = __s_b.launch_builder(&f);
4832 b.arg(logits)
4833 .arg(&mut sel_idx)
4834 .arg(&mut sel_w)
4835 .arg(&ne)
4836 .arg(&nu);
4837 unsafe {
4838 b.launch(cfg)?;
4839 }
4840 Ok((sel_idx, sel_w))
4841 }
4842
4843 pub fn moe_router_topk_scaled(
4846 &self,
4847 logits: &CudaSlice<f32>,
4848 t: usize,
4849 n_expert: usize,
4850 n_used: usize,
4851 ex_scale: &CudaSlice<f32>,
4852 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4853 let f = self.func("moe_router_topk_scaled_f32");
4858 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4859 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4860 let cfg = LaunchConfig {
4861 grid_dim: (t as u32, 1, 1),
4862 block_dim: (n_expert as u32, 1, 1),
4863 shared_mem_bytes: 0,
4864 };
4865 let (ne, nu) = (n_expert as i32, n_used as i32);
4866 let __s_b = self.gpu.stream();
4867 let mut b = __s_b.launch_builder(&f);
4868 b.arg(logits)
4869 .arg(&mut sel_idx)
4870 .arg(&mut sel_w)
4871 .arg(&ne)
4872 .arg(&nu)
4873 .arg(ex_scale);
4874 unsafe {
4875 b.launch(cfg)?;
4876 }
4877 Ok((sel_idx, sel_w))
4878 }
4879
4880 pub fn moe_router_topk_host(
4888 &self,
4889 logits: &CudaSlice<f32>,
4890 t: usize,
4891 n_expert: usize,
4892 n_used: usize,
4893 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4894 let f = self.func("moe_router_topk_f32");
4895 let n = t * n_used;
4896 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4897 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4898 let cfg = LaunchConfig {
4899 grid_dim: (t as u32, 1, 1),
4900 block_dim: (n_expert as u32, 1, 1),
4901 shared_mem_bytes: 0,
4902 };
4903 let (ne, nu) = (n_expert as i32, n_used as i32);
4904 let __s_b = self.gpu.stream();
4905 let mut b = __s_b.launch_builder(&f);
4906 b.arg(logits)
4907 .arg(&mut sel_idx)
4908 .arg(&mut sel_w)
4909 .arg(&ne)
4910 .arg(&nu);
4911 unsafe {
4912 b.launch(cfg)?;
4913 }
4914 let bytes = n * 8;
4916 let mut guard = self.router_stage.lock().unwrap();
4917 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4918 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4919 }
4920 let stage = guard.as_mut().unwrap();
4921 let (si, sw) = unsafe {
4922 (
4923 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4924 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4925 )
4926 };
4927 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()))
4931 }
4932
4933 #[allow(clippy::too_many_arguments)]
4937 pub fn moe_router_sigmoid_topk(
4938 &self,
4939 logits: &CudaSlice<f32>,
4940 t: usize,
4941 n_expert: usize,
4942 n_used: usize,
4943 active_count: usize,
4944 correction_bias: &CudaSlice<f32>,
4945 active: &CudaSlice<u8>,
4946 scaling_factor: f32,
4947 route_norm: bool,
4948 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4949 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4950 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4951 return Err(format!(
4952 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4953 )
4954 .into());
4955 }
4956 if logits.len() < t * n_expert
4957 || correction_bias.len() != n_expert
4958 || active.len() != n_expert
4959 {
4960 return Err(format!(
4961 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4962 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4963 ).into());
4964 }
4965 let f = self.func(crate::sigmoid_topk_kernel(
4966 crate::sig_expf_dev_on(),
4967 crate::topk_fast_on(),
4968 n_used,
4969 ));
4970 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4971 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4972 let threads = n_expert.div_ceil(32) * 32;
4973 let cfg = LaunchConfig {
4974 grid_dim: (t as u32, 1, 1),
4975 block_dim: (threads as u32, 1, 1),
4976 shared_mem_bytes: 0,
4977 };
4978 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4979 let __s_b = self.gpu.stream();
4980 let mut b = __s_b.launch_builder(&f);
4981 b.arg(logits)
4982 .arg(correction_bias)
4983 .arg(active)
4984 .arg(&mut sel_idx)
4985 .arg(&mut sel_w)
4986 .arg(&ne)
4987 .arg(&nu)
4988 .arg(&scaling_factor)
4989 .arg(&rn);
4990 unsafe {
4991 b.launch(cfg)?;
4992 }
4993 Ok((sel_idx, sel_w))
4994 }
4995
4996 #[allow(clippy::too_many_arguments)]
4999 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
5003 if ptr == 0 {
5004 return Err("ring_flag_raw: unarmed flag".into());
5005 }
5006 let f = self.func("memra_ring_flag");
5007 let cfg = LaunchConfig {
5008 grid_dim: (1, 1, 1),
5009 block_dim: (32, 1, 1),
5010 shared_mem_bytes: 0,
5011 };
5012 let __s_b = self.gpu.stream();
5013 let mut b = __s_b.launch_builder(&f);
5014 b.arg(&ptr).arg(&value);
5015 unsafe {
5016 b.launch(cfg)?;
5017 }
5018 Ok(())
5019 }
5020
5021 pub fn moe_sel_w_mirror(
5024 &self,
5025 sel_src: &CudaSlice<i32>,
5026 w_src: &CudaSlice<f32>,
5027 sel_dst: &mut CudaSlice<i32>,
5028 w_dst: &mut CudaSlice<f32>,
5029 n: usize,
5030 ) -> Result<(), Box<dyn std::error::Error>> {
5031 if n == 0
5032 || n > 32
5033 || sel_src.len() < n
5034 || w_src.len() < n
5035 || sel_dst.len() < n
5036 || w_dst.len() < n
5037 {
5038 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
5039 }
5040 let f = self.func("moe_sel_w_mirror");
5041 let cfg = LaunchConfig {
5042 grid_dim: (1, 1, 1),
5043 block_dim: (32, 1, 1),
5044 shared_mem_bytes: 0,
5045 };
5046 let ni = n as i32;
5047 let __s_b = self.gpu.stream();
5048 let mut b = __s_b.launch_builder(&f);
5049 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5050 unsafe {
5051 b.launch(cfg)?;
5052 }
5053 Ok(())
5054 }
5055
5056 pub fn moe_router_sigmoid_topk_into(
5057 &self,
5058 logits: &CudaSlice<f32>,
5059 t: usize,
5060 n_expert: usize,
5061 n_used: usize,
5062 active_count: usize,
5063 correction_bias: &CudaSlice<f32>,
5064 active: &CudaSlice<u8>,
5065 scaling_factor: f32,
5066 route_norm: bool,
5067 sel_idx: &mut CudaSlice<i32>,
5068 sel_w: &mut CudaSlice<f32>,
5069 ) -> Result<(), Box<dyn std::error::Error>> {
5070 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5071 if n_expert == 0
5072 || n_expert > 1024
5073 || n_used == 0
5074 || n_used > 32 || n_used > n_expert
5076 || logits.len() < t * n_expert
5077 || correction_bias.len() != n_expert
5078 || active.len() != n_expert
5079 || sel_idx.len() < t * n_used
5080 || sel_w.len() < t * n_used
5081 {
5082 return Err("sigmoid router _into geometry mismatch".into());
5083 }
5084 let f = self.func(crate::sigmoid_topk_kernel(
5085 crate::sig_expf_dev_on(),
5086 crate::topk_fast_on(),
5087 n_used,
5088 ));
5089 let threads = n_expert.div_ceil(32) * 32;
5090 let cfg = LaunchConfig {
5091 grid_dim: (t as u32, 1, 1),
5092 block_dim: (threads as u32, 1, 1),
5093 shared_mem_bytes: 0,
5094 };
5095 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5096 let __s_b = self.gpu.stream();
5097 let mut b = __s_b.launch_builder(&f);
5098 b.arg(logits)
5099 .arg(correction_bias)
5100 .arg(active)
5101 .arg(&mut *sel_idx)
5102 .arg(&mut *sel_w)
5103 .arg(&ne)
5104 .arg(&nu)
5105 .arg(&scaling_factor)
5106 .arg(&rn);
5107 unsafe {
5108 b.launch(cfg)?;
5109 }
5110 Ok(())
5111 }
5112
5113 #[allow(clippy::too_many_arguments)]
5116 pub fn moe_router_sigmoid_topk_host(
5117 &self,
5118 logits: &CudaSlice<f32>,
5119 t: usize,
5120 n_expert: usize,
5121 n_used: usize,
5122 active_count: usize,
5123 correction_bias: &CudaSlice<f32>,
5124 active: &CudaSlice<u8>,
5125 scaling_factor: f32,
5126 route_norm: bool,
5127 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5128 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5129 logits,
5130 t,
5131 n_expert,
5132 n_used,
5133 active_count,
5134 correction_bias,
5135 active,
5136 scaling_factor,
5137 route_norm,
5138 )?;
5139 let n = t * n_used;
5140 let bytes = n * 8;
5141 let mut guard = self.router_stage.lock().unwrap();
5142 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5143 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5144 }
5145 let stage = guard.as_mut().unwrap();
5146 let (si, sw) = unsafe {
5147 (
5148 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5149 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5150 )
5151 };
5152 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5153 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5154 self.gpu.stream().synchronize()?;
5155 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5156 }
5157
5158 pub fn stage_expert_async(
5162 &self,
5163 host_bytes: &[u8],
5164 scratch: &mut CudaSlice<u8>,
5165 off: usize,
5166 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5167 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5168 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5169 Ok(self.copy_stream.record_event(None)?)
5170 }
5171
5172 pub fn compute_wait(
5174 &self,
5175 ev: &cudarc::driver::CudaEvent,
5176 ) -> Result<(), Box<dyn std::error::Error>> {
5177 self.gpu.stream().wait(ev)?;
5178 Ok(())
5179 }
5180
5181 pub fn qmatvec_view(
5186 &self,
5187 w: &CudaSlice<u8>,
5188 range: std::ops::Range<usize>,
5189 x: &cudarc::driver::CudaView<f32>,
5190 m: usize,
5191 in_f: usize,
5192 out_f: usize,
5193 qtype: i32,
5194 row_bytes: usize,
5195 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5196 let f = self.func("qmatvec_f32");
5197 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
5200 grid_dim: (out_f as u32, m as u32, 1),
5201 block_dim: (256, 1, 1),
5202 shared_mem_bytes: 0,
5203 };
5204 let (inf, outf, mi, qt, rb) =
5205 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5206 let __s_b = self.gpu.stream();
5207 let mut b = __s_b.launch_builder(&f);
5208 b.arg(&wv)
5209 .arg(x)
5210 .arg(&mut y)
5211 .arg(&inf)
5212 .arg(&outf)
5213 .arg(&mi)
5214 .arg(&qt)
5215 .arg(&rb);
5216 unsafe {
5217 b.launch(cfg)?;
5218 }
5219 Ok(y)
5220 }
5221
5222 #[allow(clippy::too_many_arguments)]
5229 pub fn moe_gate_up_silu8_q8(
5233 &self,
5234 gp: WPtr8,
5235 up: WPtr8,
5236 aq: &CudaSlice<i8>,
5237 ad: &CudaSlice<f32>,
5238 in_f: usize,
5239 n_ff: usize,
5240 n_used: usize,
5241 qt_g: i32,
5242 qt_u: i32,
5243 rb_g: usize,
5244 rb_u: usize,
5245 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5246 let f = self.func("moe_gate_up_silu8_q8");
5247 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5248 let cfg = LaunchConfig {
5249 grid_dim: (n_ff as u32, n_used as u32, 1),
5250 block_dim: (32, 1, 1),
5251 shared_mem_bytes: 0,
5252 };
5253 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5254 let __s_b = self.gpu.stream();
5255 let mut b = __s_b.launch_builder(&f);
5256 b.arg(&gp)
5257 .arg(&up)
5258 .arg(aq)
5259 .arg(ad)
5260 .arg(&mut act)
5261 .arg(&inf)
5262 .arg(&nff)
5263 .arg(&qt_g)
5264 .arg(&qt_u)
5265 .arg(&rbg)
5266 .arg(&rbu);
5267 unsafe {
5268 b.launch(cfg)?;
5269 }
5270 Ok(act)
5271 }
5272
5273 #[allow(clippy::too_many_arguments)]
5274 pub fn moe_down8_fma_q8(
5275 &self,
5276 dp: WPtr8,
5277 w: F32x8,
5278 aq2: &CudaSlice<i8>,
5279 ad2: &CudaSlice<f32>,
5280 dst: &mut cudarc::driver::CudaViewMut<f32>,
5281 in_f: usize,
5282 out_f: usize,
5283 n_used: usize,
5284 qt: i32,
5285 rb: usize,
5286 ) -> Result<(), Box<dyn std::error::Error>> {
5287 let f = self.func("moe_down8_fma_q8");
5288 let cfg = LaunchConfig {
5289 grid_dim: (out_f as u32, 1, 1),
5290 block_dim: (32, 1, 1),
5291 shared_mem_bytes: 0,
5292 };
5293 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5294 let __s_b = self.gpu.stream();
5295 let mut b = __s_b.launch_builder(&f);
5296 b.arg(&dp)
5297 .arg(&w)
5298 .arg(aq2)
5299 .arg(ad2)
5300 .arg(dst)
5301 .arg(&inf)
5302 .arg(&outf)
5303 .arg(&nu)
5304 .arg(&qt)
5305 .arg(&rbi);
5306 unsafe {
5307 b.launch(cfg)?;
5308 }
5309 Ok(())
5310 }
5311
5312 pub fn qmatvec_expert_q8(
5314 &self,
5315 w: &CudaSlice<u8>,
5316 range: std::ops::Range<usize>,
5317 aq: &CudaSlice<i8>,
5318 ad: &CudaSlice<f32>,
5319 m: usize,
5320 in_f: usize,
5321 out_f: usize,
5322 qtype: i32,
5323 row_bytes: usize,
5324 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5325 let f = self.func("qmatvec_expert_q8");
5326 let wv = w.slice(range);
5327 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5328 const ROWS: u32 = 4; let cfg = LaunchConfig {
5330 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5331 block_dim: (32, ROWS, 1),
5332 shared_mem_bytes: 0,
5333 };
5334 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5335 let __s_b = self.gpu.stream();
5336 let mut b = __s_b.launch_builder(&f);
5337 b.arg(&wv)
5338 .arg(aq)
5339 .arg(ad)
5340 .arg(&mut y)
5341 .arg(&inf)
5342 .arg(&outf)
5343 .arg(&mi)
5344 .arg(&qtype)
5345 .arg(&rbi);
5346 unsafe {
5347 b.launch(cfg)?;
5348 }
5349 Ok(y)
5350 }
5351
5352 pub fn moe_gate_up_silu8(
5353 &self,
5354 gp: WPtr8,
5355 up: WPtr8,
5356 x: &cudarc::driver::CudaView<f32>,
5357 in_f: usize,
5358 n_ff: usize,
5359 n_used: usize,
5360 qt_g: i32,
5361 qt_u: i32,
5362 rb_g: usize,
5363 rb_u: usize,
5364 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5365 let f = self.func("moe_gate_up_silu8_f32");
5366 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5368 grid_dim: (n_ff as u32, n_used as u32, 1),
5369 block_dim: (256, 1, 1),
5370 shared_mem_bytes: 0,
5371 };
5372 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5373 let __s_b = self.gpu.stream();
5374 let mut b = __s_b.launch_builder(&f);
5375 b.arg(&gp)
5376 .arg(&up)
5377 .arg(x)
5378 .arg(&mut act)
5379 .arg(&inf)
5380 .arg(&nff)
5381 .arg(&qt_g)
5382 .arg(&qt_u)
5383 .arg(&rbg)
5384 .arg(&rbu);
5385 unsafe {
5386 b.launch(cfg)?;
5387 }
5388 Ok(act)
5389 }
5390
5391 #[allow(clippy::too_many_arguments)]
5397 pub fn moe_down8_fma_into(
5398 &self,
5399 dp: WPtr8,
5400 w: F32x8,
5401 act: &CudaSlice<f32>,
5402 dst: &mut cudarc::driver::CudaViewMut<f32>,
5403 in_f: usize,
5404 out_f: usize,
5405 n_used: usize,
5406 qt: i32,
5407 rb: usize,
5408 ) -> Result<(), Box<dyn std::error::Error>> {
5409 let f = self.func("moe_down8_fma_f32");
5410 let cfg = LaunchConfig {
5411 grid_dim: (out_f as u32, 1, 1),
5412 block_dim: (256, 1, 1),
5413 shared_mem_bytes: 0,
5414 };
5415 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5416 let __s_b = self.gpu.stream();
5417 let mut b = __s_b.launch_builder(&f);
5418 b.arg(&dp)
5419 .arg(&w)
5420 .arg(act)
5421 .arg(dst)
5422 .arg(&inf)
5423 .arg(&outf)
5424 .arg(&nu)
5425 .arg(&qt)
5426 .arg(&rbv);
5427 unsafe {
5428 b.launch(cfg)?;
5429 }
5430 Ok(())
5431 }
5432
5433 #[allow(clippy::too_many_arguments)]
5438 #[allow(clippy::too_many_arguments)]
5453 #[allow(clippy::too_many_arguments)]
5455 pub fn moe_pairs_matvec_q8(
5456 &self,
5457 table: &CudaSlice<u64>,
5458 proj: i32,
5459 pair_tok: &CudaSlice<i32>,
5460 pair_ex: &CudaSlice<i32>,
5461 aq: &CudaSlice<i8>,
5462 ad: &CudaSlice<f32>,
5463 in_f: usize,
5464 out_f: usize,
5465 n_expert: usize,
5466 n_pairs: usize,
5467 qtype: i32,
5468 row_bytes: usize,
5469 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5470 let f = self.func("moe_pairs_matvec_q8");
5471 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5472 const ROWS: u32 = 4;
5473 let cfg = LaunchConfig {
5474 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5475 block_dim: (32, ROWS, 1),
5476 shared_mem_bytes: 0,
5477 };
5478 let (inf, outf, ne, np, rbi) = (
5479 in_f as i32,
5480 out_f as i32,
5481 n_expert as i32,
5482 n_pairs as i32,
5483 row_bytes as i64,
5484 );
5485 let __s_b = self.gpu.stream();
5486 let mut b = __s_b.launch_builder(&f);
5487 b.arg(table)
5488 .arg(&proj)
5489 .arg(pair_tok)
5490 .arg(pair_ex)
5491 .arg(aq)
5492 .arg(ad)
5493 .arg(&mut y)
5494 .arg(&inf)
5495 .arg(&outf)
5496 .arg(&ne)
5497 .arg(&np)
5498 .arg(&qtype)
5499 .arg(&rbi);
5500 unsafe {
5501 b.launch(cfg)?;
5502 }
5503 Ok(y)
5504 }
5505
5506 #[allow(clippy::too_many_arguments)]
5508 pub fn moe_pairs_matvec_q8_em(
5509 &self,
5510 table: &CudaSlice<u64>,
5511 proj: i32,
5512 ex_ids: &CudaSlice<i32>,
5513 ex_off: &CudaSlice<i32>,
5514 ex_pairs: &CudaSlice<i32>,
5515 pair_tok: &CudaSlice<i32>,
5516 aq: &CudaSlice<i8>,
5517 ad: &CudaSlice<f32>,
5518 in_f: usize,
5519 out_f: usize,
5520 n_expert: usize,
5521 n_active: usize,
5522 n_pairs: usize,
5523 qtype: i32,
5524 row_bytes: usize,
5525 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5526 let f = self.func("moe_pairs_matvec_q8_em");
5527 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5528 const ROWS: u32 = 4;
5529 let cfg = LaunchConfig {
5530 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5531 block_dim: (32, ROWS, 1),
5532 shared_mem_bytes: 0,
5533 };
5534 let (inf, outf, ne, na, rbi) = (
5535 in_f as i32,
5536 out_f as i32,
5537 n_expert as i32,
5538 n_active as i32,
5539 row_bytes as i64,
5540 );
5541 let __s_b = self.gpu.stream();
5542 let mut b = __s_b.launch_builder(&f);
5543 b.arg(table)
5544 .arg(&proj)
5545 .arg(ex_ids)
5546 .arg(ex_off)
5547 .arg(ex_pairs)
5548 .arg(pair_tok)
5549 .arg(aq)
5550 .arg(ad)
5551 .arg(&mut y)
5552 .arg(&inf)
5553 .arg(&outf)
5554 .arg(&ne)
5555 .arg(&na)
5556 .arg(&qtype)
5557 .arg(&rbi);
5558 unsafe {
5559 b.launch(cfg)?;
5560 }
5561 Ok(y)
5562 }
5563
5564 #[allow(clippy::too_many_arguments)]
5567 pub fn moe_pairs_matvec_q8_dec(
5568 &self,
5569 table: &CudaSlice<u64>,
5570 proj: i32,
5571 ex_ids: &CudaSlice<i32>,
5572 ex_off: &CudaSlice<i32>,
5573 ex_pairs: &CudaSlice<i32>,
5574 pair_tok: &CudaSlice<i32>,
5575 aq: &CudaSlice<i8>,
5576 ad: &CudaSlice<f32>,
5577 in_f: usize,
5578 out_f: usize,
5579 n_expert: usize,
5580 n_active: usize,
5581 n_pairs: usize,
5582 qtype: i32,
5583 row_bytes: usize,
5584 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5585 let f = self.func("moe_pairs_matvec_q8_dec");
5586 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5587 const ROWS: u32 = 4;
5588 let cfg = LaunchConfig {
5589 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5590 block_dim: (32, ROWS, 1),
5591 shared_mem_bytes: 0,
5592 };
5593 let (inf, outf, ne, na, rbi) = (
5594 in_f as i32,
5595 out_f as i32,
5596 n_expert as i32,
5597 n_active as i32,
5598 row_bytes as i64,
5599 );
5600 let __s_b = self.gpu.stream();
5601 let mut b = __s_b.launch_builder(&f);
5602 b.arg(table)
5603 .arg(&proj)
5604 .arg(ex_ids)
5605 .arg(ex_off)
5606 .arg(ex_pairs)
5607 .arg(pair_tok)
5608 .arg(aq)
5609 .arg(ad)
5610 .arg(&mut y)
5611 .arg(&inf)
5612 .arg(&outf)
5613 .arg(&ne)
5614 .arg(&na)
5615 .arg(&qtype)
5616 .arg(&rbi);
5617 unsafe {
5618 b.launch(cfg)?;
5619 }
5620 Ok(y)
5621 }
5622
5623 pub fn moe_pairs_gelu_mul(
5624 &self,
5625 gate: &CudaSlice<f32>,
5626 up: &CudaSlice<f32>,
5627 n: usize,
5628 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5629 let f = self.func("moe_pairs_gelu_mul");
5630 let mut act = self.alloc_uninit::<f32>(n)?;
5631 let cfg = LaunchConfig::for_num_elems(n as u32);
5632 let nl = n as i64;
5633 let __s_b = self.gpu.stream();
5634 let mut b = __s_b.launch_builder(&f);
5635 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5636 unsafe {
5637 b.launch(cfg)?;
5638 }
5639 Ok(act)
5640 }
5641
5642 pub fn moe_pairs_silu_mul(
5643 &self,
5644 gate: &CudaSlice<f32>,
5645 up: &CudaSlice<f32>,
5646 n: usize,
5647 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5648 let f = self.func("moe_pairs_silu_mul");
5649 let mut act = self.alloc_uninit::<f32>(n)?;
5650 let cfg = LaunchConfig::for_num_elems(n as u32);
5651 let nl = n as i64;
5652 let __s_b = self.gpu.stream();
5653 let mut b = __s_b.launch_builder(&f);
5654 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5655 unsafe {
5656 b.launch(cfg)?;
5657 }
5658 Ok(act)
5659 }
5660
5661 #[allow(clippy::too_many_arguments)]
5662 pub fn moe_pairs_scatter(
5663 &self,
5664 y_down: &CudaSlice<f32>,
5665 pair_w: &CudaSlice<f32>,
5666 tok_pair_off: &CudaSlice<i32>,
5667 tok_pair_ids: &CudaSlice<i32>,
5668 moe_out: &mut CudaSlice<f32>,
5669 t: usize,
5670 n_embd: usize,
5671 ) -> Result<(), Box<dyn std::error::Error>> {
5672 let f = self.func("moe_pairs_scatter");
5673 let cfg = LaunchConfig {
5674 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5675 block_dim: (256, 1, 1),
5676 shared_mem_bytes: 0,
5677 };
5678 let ne = n_embd as i32;
5679 let __s_b = self.gpu.stream();
5680 let mut b = __s_b.launch_builder(&f);
5681 b.arg(y_down)
5682 .arg(pair_w)
5683 .arg(tok_pair_off)
5684 .arg(tok_pair_ids)
5685 .arg(moe_out)
5686 .arg(&ne);
5687 unsafe {
5688 b.launch(cfg)?;
5689 }
5690 Ok(())
5691 }
5692
5693 #[allow(clippy::too_many_arguments)]
5697 pub fn moe_gate_up_gelu8_dev_q8(
5698 &self,
5699 table: &CudaSlice<u64>,
5700 sel: &cudarc::driver::CudaView<i32>,
5701 aq: &CudaSlice<i8>,
5702 ad: &CudaSlice<f32>,
5703 in_f: usize,
5704 n_ff: usize,
5705 n_used: usize,
5706 n_expert: usize,
5707 qt_g: i32,
5708 qt_u: i32,
5709 rb_g: usize,
5710 rb_u: usize,
5711 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5712 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5713 let (inf, nff, ne, rbg, rbu) = (
5714 in_f as i32,
5715 n_ff as i32,
5716 n_expert as i32,
5717 rb_g as i64,
5718 rb_u as i64,
5719 );
5720 let f = self.func("moe_gate_up_gelu8_dev_q8");
5721 let cfg = LaunchConfig {
5722 grid_dim: (n_ff as u32, n_used as u32, 1),
5723 block_dim: (32, 1, 1),
5724 shared_mem_bytes: 0,
5725 };
5726 let __s_b = self.gpu.stream();
5727 let mut b = __s_b.launch_builder(&f);
5728 b.arg(table)
5729 .arg(sel)
5730 .arg(aq)
5731 .arg(ad)
5732 .arg(&mut act)
5733 .arg(&inf)
5734 .arg(&nff)
5735 .arg(&ne)
5736 .arg(&qt_g)
5737 .arg(&qt_u)
5738 .arg(&rbg)
5739 .arg(&rbu);
5740 unsafe {
5741 b.launch(cfg)?;
5742 }
5743 Ok(act)
5744 }
5745
5746 #[allow(clippy::too_many_arguments)]
5748 pub fn moe_gate_up_gelu8_dev_q8_rows(
5749 &self,
5750 table: &CudaSlice<u64>,
5751 sel: &CudaSlice<i32>,
5752 aq: &CudaSlice<i8>,
5753 ad: &CudaSlice<f32>,
5754 t: usize,
5755 in_f: usize,
5756 n_ff: usize,
5757 n_used: usize,
5758 n_expert: usize,
5759 qt_g: i32,
5760 qt_u: i32,
5761 rb_g: usize,
5762 rb_u: usize,
5763 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5764 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5765 let (inf, nff, ne, rbg, rbu, nu) = (
5766 in_f as i32,
5767 n_ff as i32,
5768 n_expert as i32,
5769 rb_g as i64,
5770 rb_u as i64,
5771 n_used as i32,
5772 );
5773 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5774 let cfg = LaunchConfig {
5775 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5776 block_dim: (32, 1, 1),
5777 shared_mem_bytes: 0,
5778 };
5779 let __s_b = self.gpu.stream();
5780 let mut b = __s_b.launch_builder(&f);
5781 b.arg(table)
5782 .arg(sel)
5783 .arg(aq)
5784 .arg(ad)
5785 .arg(&mut act)
5786 .arg(&inf)
5787 .arg(&nff)
5788 .arg(&ne)
5789 .arg(&qt_g)
5790 .arg(&qt_u)
5791 .arg(&rbg)
5792 .arg(&rbu)
5793 .arg(&nu);
5794 unsafe {
5795 b.launch(cfg)?;
5796 }
5797 Ok(act)
5798 }
5799
5800 #[allow(clippy::too_many_arguments)]
5802 pub fn moe_gate_up_gelu8_dev_q8_csr(
5803 &self,
5804 table: &CudaSlice<u64>,
5805 sel: &CudaSlice<i32>,
5806 aq: &CudaSlice<i8>,
5807 ad: &CudaSlice<f32>,
5808 n_pairs: usize,
5809 in_f: usize,
5810 n_ff: usize,
5811 n_used: usize,
5812 n_expert: usize,
5813 qt_g: i32,
5814 qt_u: i32,
5815 rb_g: usize,
5816 rb_u: usize,
5817 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5818 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5819 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5820 in_f as i32,
5821 n_ff as i32,
5822 n_expert as i32,
5823 rb_g as i64,
5824 rb_u as i64,
5825 n_used as i32,
5826 n_pairs as i32,
5827 );
5828 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5829 let cfg = LaunchConfig {
5830 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5831 block_dim: (32, 1, 1),
5832 shared_mem_bytes: 0,
5833 };
5834 let __s_b = self.gpu.stream();
5835 let mut b = __s_b.launch_builder(&f);
5836 b.arg(table)
5837 .arg(sel)
5838 .arg(aq)
5839 .arg(ad)
5840 .arg(&mut act)
5841 .arg(&inf)
5842 .arg(&nff)
5843 .arg(&ne)
5844 .arg(&qt_g)
5845 .arg(&qt_u)
5846 .arg(&rbg)
5847 .arg(&rbu)
5848 .arg(&nu)
5849 .arg(&npi);
5850 unsafe {
5851 b.launch(cfg)?;
5852 }
5853 Ok(act)
5854 }
5855
5856 #[allow(clippy::too_many_arguments)]
5858 pub fn moe_down8_fma_dev_q8_rows_g(
5859 &self,
5860 table: &CudaSlice<u64>,
5861 sel: &CudaSlice<i32>,
5862 w: &CudaSlice<f32>,
5863 aq2: &CudaSlice<i8>,
5864 ad2: &CudaSlice<f32>,
5865 dst: &mut CudaSlice<f32>,
5866 t: usize,
5867 in_f: usize,
5868 out_f: usize,
5869 n_used: usize,
5870 n_expert: usize,
5871 qt: i32,
5872 rb: usize,
5873 ) -> Result<(), Box<dyn std::error::Error>> {
5874 let (inf, outf, nu, ne, rbi) = (
5875 in_f as i32,
5876 out_f as i32,
5877 n_used as i32,
5878 n_expert as i32,
5879 rb as i64,
5880 );
5881 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5885 let f = self.func(if step_b1_w8 {
5886 "moe_down8_fma_dev_q8_rows_w8"
5887 } else {
5888 "moe_down8_fma_dev_q8_rows_g"
5889 });
5890 let cfg = LaunchConfig {
5891 grid_dim: (out_f as u32, 1, t as u32),
5892 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5893 shared_mem_bytes: 0,
5894 };
5895 let __s_b = self.gpu.stream();
5896 let mut b = __s_b.launch_builder(&f);
5897 b.arg(table)
5898 .arg(sel)
5899 .arg(w)
5900 .arg(aq2)
5901 .arg(ad2)
5902 .arg(dst)
5903 .arg(&inf)
5904 .arg(&outf)
5905 .arg(&nu)
5906 .arg(&ne)
5907 .arg(&qt)
5908 .arg(&rbi);
5909 unsafe {
5910 b.launch(cfg)?;
5911 }
5912 Ok(())
5913 }
5914
5915 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5919 let (out_f, in_f) = (2048usize, 2816usize);
5920 let nblk = in_f / 32;
5921 let mut seed = 0x9E3779B97F4A7C15u64;
5922 let mut rng = move || {
5923 seed = seed
5924 .wrapping_mul(6364136223846793005)
5925 .wrapping_add(1442695040888963407);
5926 (seed >> 33) as u8
5927 };
5928 let mut w = vec![0u8; out_f * nblk * 18];
5929 for b in w.iter_mut() {
5930 *b = rng();
5931 }
5932 for r in 0..out_f {
5933 for g in 0..nblk {
5934 let off = (r * nblk + g) * 18;
5935 w[off] = 0x00;
5936 w[off + 1] = 0x2C; }
5938 }
5939 let qplane = out_f * nblk * 16;
5940 let mut wrp = vec![0u8; w.len()];
5941 for r in 0..out_f {
5942 for g in 0..nblk {
5943 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5944 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5945 .copy_from_slice(&src[0..2]);
5946 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5947 }
5948 }
5949 let w_d = self.htod_bytes(&w)?;
5950 let wrp_d = self.htod_bytes(&wrp)?;
5951 let mut aq = vec![0i8; m * in_f];
5952 for v in aq.iter_mut() {
5953 *v = rng() as i8;
5954 }
5955 let aq_d = self.htod_i8(&aq)?;
5956 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5957 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5958 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5959 const RPB: u32 = 4;
5960 let cfg = LaunchConfig {
5961 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5962 block_dim: (32, RPB, 1),
5963 shared_mem_bytes: 0,
5964 };
5965 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5966 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5967 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5968 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5969 {
5970 let __s_b = self.gpu.stream();
5971 let mut b = __s_b.launch_builder(&fb);
5972 b.arg(&w_d)
5973 .arg(&aq_d)
5974 .arg(&ad_d)
5975 .arg(&mut y0)
5976 .arg(&inf)
5977 .arg(&outf)
5978 .arg(&mi)
5979 .arg(&rb);
5980 unsafe {
5981 b.launch(cfg)?;
5982 }
5983 let __s_b = self.gpu.stream();
5984 let mut b = __s_b.launch_builder(&fr);
5985 b.arg(&wrp_d)
5986 .arg(&aq_d)
5987 .arg(&ad_d)
5988 .arg(&mut y1)
5989 .arg(&inf)
5990 .arg(&outf)
5991 .arg(&mi)
5992 .arg(&qp);
5993 unsafe {
5994 b.launch(cfg)?;
5995 }
5996 }
5997 self.gpu.stream().synchronize()?;
5998 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5999 let nd = h0
6000 .iter()
6001 .zip(&h1)
6002 .filter(|(a, b)| a.to_bits() != b.to_bits())
6003 .count();
6004 if nd != 0 {
6005 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
6006 }
6007 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
6008 self.gpu.stream().synchronize()?;
6009 let t0 = std::time::Instant::now();
6010 for _ in 0..500 {
6011 if rp {
6012 let __s_b = self.gpu.stream();
6013 let mut b = __s_b.launch_builder(&fr);
6014 b.arg(&wrp_d)
6015 .arg(&aq_d)
6016 .arg(&ad_d)
6017 .arg(&mut y1)
6018 .arg(&inf)
6019 .arg(&outf)
6020 .arg(&mi)
6021 .arg(&qp);
6022 unsafe {
6023 b.launch(cfg)?;
6024 }
6025 } else {
6026 let __s_b = self.gpu.stream();
6027 let mut b = __s_b.launch_builder(&fb);
6028 b.arg(&w_d)
6029 .arg(&aq_d)
6030 .arg(&ad_d)
6031 .arg(&mut y0)
6032 .arg(&inf)
6033 .arg(&outf)
6034 .arg(&mi)
6035 .arg(&rb);
6036 unsafe {
6037 b.launch(cfg)?;
6038 }
6039 }
6040 }
6041 self.gpu.stream().synchronize()?;
6042 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6043 };
6044 let _ = time(false)?;
6045 let _ = time(true)?; Ok((time(false)?, time(true)?))
6047 }
6048
6049 pub fn build_q4_rp4(
6054 &self,
6055 t: &mut crate::model::GpuTensor,
6056 ) -> Result<(), Box<dyn std::error::Error>> {
6057 use crate::model::GpuTensor;
6058 let GpuTensor::Quant {
6059 bytes,
6060 qtype,
6061 row_bytes,
6062 ne,
6063 rp4,
6064 ..
6065 } = t
6066 else {
6067 return Ok(());
6068 };
6069 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6070 return Ok(());
6071 }
6072 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6073 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6074 return Ok(());
6075 }
6076 let nblk = in_f / 32;
6077 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6078 let f = self.func("q4_0_split_rp_build");
6079 let n = (out_f * nblk) as i32;
6080 let cfg = LaunchConfig {
6081 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6082 block_dim: (256, 1, 1),
6083 shared_mem_bytes: 0,
6084 };
6085 let (of, nb) = (out_f as i32, nblk as i32);
6086 let _ = n;
6087 let __s_b = self.gpu.stream();
6088 let mut b = __s_b.launch_builder(&f);
6089 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6090 unsafe {
6091 b.launch(cfg)?;
6092 }
6093 *rp4 = Some(dst);
6094 Ok(())
6095 }
6096
6097 pub fn build_q8_rp4(
6102 &self,
6103 t: &mut crate::model::GpuTensor,
6104 ) -> Result<(), Box<dyn std::error::Error>> {
6105 use crate::model::GpuTensor;
6106 let GpuTensor::Quant {
6107 bytes,
6108 qtype,
6109 row_bytes,
6110 ne,
6111 rp4,
6112 ..
6113 } = t
6114 else {
6115 return Ok(());
6116 };
6117 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6118 return Ok(());
6119 }
6120 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6121 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6122 return Ok(());
6123 }
6124 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6125 Ok(())
6126 }
6127
6128 pub fn build_q8_rp4_raw(
6131 &self,
6132 bytes: &CudaSlice<u8>,
6133 in_f: usize,
6134 out_f: usize,
6135 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6136 assert!(in_f % 32 == 0);
6137 let nblk = in_f / 32;
6138 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6139 let f = self.func("q8_0_split_rp_build");
6140 let cfg = LaunchConfig {
6141 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6142 block_dim: (256, 1, 1),
6143 shared_mem_bytes: 0,
6144 };
6145 let (of, nb) = (out_f as i32, nblk as i32);
6146 let __s_b = self.gpu.stream();
6147 let mut b = __s_b.launch_builder(&f);
6148 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6149 unsafe {
6150 b.launch(cfg)?;
6151 }
6152 Ok(dst)
6153 }
6154
6155 pub fn build_q4k_rp4(
6163 &self,
6164 t: &mut crate::model::GpuTensor,
6165 ) -> Result<(), Box<dyn std::error::Error>> {
6166 use crate::model::GpuTensor;
6167 let GpuTensor::Quant {
6168 bytes,
6169 qtype,
6170 row_bytes,
6171 ne,
6172 rp4,
6173 ..
6174 } = t
6175 else {
6176 return Ok(());
6177 };
6178 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6179 return Ok(());
6180 }
6181 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6182 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6183 return Ok(());
6184 }
6185 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6186 Ok(())
6187 }
6188
6189 pub fn build_q6k_rp4(
6190 &self,
6191 t: &mut crate::model::GpuTensor,
6192 ) -> Result<(), Box<dyn std::error::Error>> {
6193 use crate::model::GpuTensor;
6194 let GpuTensor::Quant {
6195 bytes,
6196 qtype,
6197 row_bytes,
6198 ne,
6199 rp4,
6200 ..
6201 } = t
6202 else {
6203 return Ok(());
6204 };
6205 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6206 return Ok(());
6207 }
6208 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6209 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6210 return Ok(());
6211 }
6212 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6213 Ok(())
6214 }
6215
6216 pub fn build_kq_rp4_raw(
6218 &self,
6219 bytes: &CudaSlice<u8>,
6220 in_f: usize,
6221 out_f: usize,
6222 qtype: i32,
6223 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6224 assert!(in_f % 256 == 0);
6225 let nsbk = in_f / 256;
6226 let (sb_bytes, kname) = match qtype {
6227 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6228 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6229 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6230 };
6231 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6232 let f = self.func(kname);
6233 let cfg = LaunchConfig {
6234 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6235 block_dim: (256, 1, 1),
6236 shared_mem_bytes: 0,
6237 };
6238 let (of, nb) = (out_f as i32, nsbk as i32);
6239 let __s_b = self.gpu.stream();
6240 let mut b = __s_b.launch_builder(&f);
6241 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6242 unsafe {
6243 b.launch(cfg)?;
6244 }
6245 Ok(dst)
6246 }
6247
6248 pub fn kqrp_enabled() -> bool {
6252 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6253 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6254 Ok("0") => false,
6255 Ok(_) => true,
6256 Err(_) => cfg!(memra_hopper_mma),
6257 })
6258 }
6259
6260 pub fn build_q4_rp_swap(
6266 &self,
6267 t: &mut crate::model::GpuTensor,
6268 ) -> Result<bool, Box<dyn std::error::Error>> {
6269 use crate::model::GpuTensor;
6270 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6280 return Ok(false);
6281 }
6282 self.build_q4_rp4(t)?;
6283 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6285 return Ok(false);
6286 };
6287 match rp4.take() {
6288 Some(split) => {
6289 *bytes = split; *rp = true;
6291 Ok(true)
6292 }
6293 None => Ok(false),
6294 }
6295 }
6296
6297 pub fn q4rp_enabled() -> bool {
6299 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6300 *ON.get_or_init(|| {
6301 std::env::var("MEMRA_Q4RP")
6302 .map(|v| v != "0")
6303 .unwrap_or(true)
6304 })
6305 }
6306
6307 pub fn copy_rows_strided(
6310 &self,
6311 src: &CudaSlice<f32>,
6312 dst: &mut CudaSlice<f32>,
6313 row_elems: usize,
6314 n_rows: usize,
6315 src_stride: usize,
6316 src_off: usize,
6317 ) -> Result<(), Box<dyn std::error::Error>> {
6318 let f = self.func("copy_rows_strided_f32");
6319 let cfg = LaunchConfig {
6320 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6321 block_dim: (256, 1, 1),
6322 shared_mem_bytes: 0,
6323 };
6324 let (re, nr) = (row_elems as i32, n_rows as i32);
6325 let (st, off) = (src_stride as i64, src_off as i64);
6326 let __s_b = self.gpu.stream();
6327 let mut b = __s_b.launch_builder(&f);
6328 b.arg(src)
6329 .arg(&mut *dst)
6330 .arg(&re)
6331 .arg(&nr)
6332 .arg(&st)
6333 .arg(&off);
6334 unsafe {
6335 b.launch(cfg)?;
6336 }
6337 Ok(())
6338 }
6339
6340 pub fn place_rows_strided(
6346 &self,
6347 src: &CudaSlice<f32>,
6348 dst: &mut CudaSlice<f32>,
6349 row_elems: usize,
6350 n_rows: usize,
6351 dst_stride: usize,
6352 dst_off: usize,
6353 ) -> Result<(), Box<dyn std::error::Error>> {
6354 if row_elems == 0 || n_rows == 0 {
6355 return Err("strided row placement requires nonzero rows and row width".into());
6356 }
6357 let src_len = n_rows
6358 .checked_mul(row_elems)
6359 .ok_or("strided row placement source size overflow")?;
6360 let dst_len = n_rows
6361 .checked_sub(1)
6362 .and_then(|rows| rows.checked_mul(dst_stride))
6363 .and_then(|base| base.checked_add(dst_off))
6364 .and_then(|base| base.checked_add(row_elems))
6365 .ok_or("strided row placement destination size overflow")?;
6366 let row_end = dst_off
6367 .checked_add(row_elems)
6368 .ok_or("strided row placement row size overflow")?;
6369 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6370 return Err(format!(
6371 "strided row placement geometry mismatch: src={} need_src={src_len} \
6372 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6373 dst_stride={dst_stride} dst_off={dst_off}",
6374 src.len(),
6375 dst.len(),
6376 )
6377 .into());
6378 }
6379 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6380 return Err("strided row placement exceeds CUDA kernel geometry".into());
6381 }
6382 let f = self.func("place_rows_strided_f32");
6383 let cfg = LaunchConfig {
6384 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6385 block_dim: (256, 1, 1),
6386 shared_mem_bytes: 0,
6387 };
6388 let (re, nr) = (row_elems as i32, n_rows as i32);
6389 let (st, off) = (dst_stride as i64, dst_off as i64);
6390 let __s_b = self.gpu.stream();
6391 let mut b = __s_b.launch_builder(&f);
6392 b.arg(src)
6393 .arg(&mut *dst)
6394 .arg(&re)
6395 .arg(&nr)
6396 .arg(&st)
6397 .arg(&off);
6398 unsafe {
6399 b.launch(cfg)?;
6400 }
6401 Ok(())
6402 }
6403
6404 pub fn u32_set_k(
6406 &self,
6407 dst: &mut CudaSlice<u32>,
6408 v: u32,
6409 idx: usize,
6410 ) -> Result<(), Box<dyn std::error::Error>> {
6411 let f = self.func("u32_set_k");
6412 let cfg = LaunchConfig {
6413 grid_dim: (1, 1, 1),
6414 block_dim: (1, 1, 1),
6415 shared_mem_bytes: 0,
6416 };
6417 let ii = idx as i32;
6418 let __s_b = self.gpu.stream();
6419 let mut b = __s_b.launch_builder(&f);
6420 b.arg(dst).arg(&v).arg(&ii);
6421 unsafe {
6422 b.launch(cfg)?;
6423 }
6424 Ok(())
6425 }
6426
6427 pub fn i32_add_k(
6429 &self,
6430 d: &mut CudaSlice<i32>,
6431 v: i32,
6432 ) -> Result<(), Box<dyn std::error::Error>> {
6433 let f = self.func("i32_add_k");
6434 let cfg = LaunchConfig {
6435 grid_dim: (1, 1, 1),
6436 block_dim: (32, 1, 1),
6437 shared_mem_bytes: 0,
6438 };
6439 let __s_b = self.gpu.stream();
6440 let mut b = __s_b.launch_builder(&f);
6441 b.arg(d).arg(&v);
6442 unsafe {
6443 b.launch(cfg)?;
6444 }
6445 Ok(())
6446 }
6447
6448 pub fn i32_iota_from(
6450 &self,
6451 ctr: &CudaSlice<i32>,
6452 dst: &mut CudaSlice<i32>,
6453 n: usize,
6454 ) -> Result<(), Box<dyn std::error::Error>> {
6455 let f = self.func("i32_iota_from");
6456 let cfg = LaunchConfig::for_num_elems(n as u32);
6457 let ni = n as i32;
6458 let __s_b = self.gpu.stream();
6459 let mut b = __s_b.launch_builder(&f);
6460 b.arg(ctr).arg(dst).arg(&ni);
6461 unsafe {
6462 b.launch(cfg)?;
6463 }
6464 Ok(())
6465 }
6466
6467 pub fn u32_map_k(
6469 &self,
6470 buf: &mut CudaSlice<u32>,
6471 map: &CudaSlice<u32>,
6472 idx: usize,
6473 ) -> Result<(), Box<dyn std::error::Error>> {
6474 let f = self.func("u32_map_k");
6475 let cfg = LaunchConfig {
6476 grid_dim: (1, 1, 1),
6477 block_dim: (1, 1, 1),
6478 shared_mem_bytes: 0,
6479 };
6480 let ii = idx as i32;
6481 let __s_b = self.gpu.stream();
6482 let mut b = __s_b.launch_builder(&f);
6483 b.arg(buf).arg(map).arg(&ii);
6484 unsafe {
6485 b.launch(cfg)?;
6486 }
6487 Ok(())
6488 }
6489
6490 #[allow(clippy::too_many_arguments)]
6492 pub fn u32_pack2(
6493 &self,
6494 a: &CudaSlice<u32>,
6495 off_a: usize,
6496 n1: usize,
6497 b_in: &CudaSlice<u32>,
6498 n2: usize,
6499 out: &mut CudaSlice<u32>,
6500 ) -> Result<(), Box<dyn std::error::Error>> {
6501 let f = self.func("u32_pack2");
6502 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6503 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6504 let __s_b = self.gpu.stream();
6505 let mut b = __s_b.launch_builder(&f);
6506 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6507 unsafe {
6508 b.launch(cfg)?;
6509 }
6510 Ok(())
6511 }
6512
6513 pub fn moe_w_exscale(
6515 &self,
6516 w: &mut CudaSlice<f32>,
6517 sel: &CudaSlice<i32>,
6518 s: &CudaSlice<f32>,
6519 n: usize,
6520 ) -> Result<(), Box<dyn std::error::Error>> {
6521 let f = self.func("moe_w_exscale");
6522 let cfg = LaunchConfig::for_num_elems(n as u32);
6523 let ni = n as i32;
6524 let __s_b = self.gpu.stream();
6525 let mut b = __s_b.launch_builder(&f);
6526 b.arg(w).arg(sel).arg(s).arg(&ni);
6527 unsafe {
6528 b.launch(cfg)?;
6529 }
6530 Ok(())
6531 }
6532
6533 pub fn moe_w_scale_by_expert(
6536 &self,
6537 w: &mut CudaSlice<f32>,
6538 sel: &CudaSlice<i32>,
6539 macros: &CudaSlice<f32>,
6540 n_expert: usize,
6541 n: usize,
6542 ) -> Result<(), Box<dyn std::error::Error>> {
6543 let f = self.func("moe_w_scale_by_expert");
6544 let cfg = LaunchConfig {
6545 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6546 block_dim: (64, 1, 1),
6547 shared_mem_bytes: 0,
6548 };
6549 let (ne, nn) = (n_expert as i32, n as i32);
6550 let __s_b = self.gpu.stream();
6551 let mut b = __s_b.launch_builder(&f);
6552 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6553 unsafe {
6554 b.launch(cfg)?;
6555 }
6556 Ok(())
6557 }
6558
6559 pub fn moe_gate_up_silu8_dev_q8(
6560 &self,
6561 table: &CudaSlice<u64>,
6562 sel: &cudarc::driver::CudaView<i32>,
6563 aq: &CudaSlice<i8>,
6564 ad: &CudaSlice<f32>,
6565 in_f: usize,
6566 n_ff: usize,
6567 n_used: usize,
6568 n_expert: usize,
6569 qt_g: i32,
6570 qt_u: i32,
6571 rb_g: usize,
6572 rb_u: usize,
6573 macros: &CudaSlice<f32>,
6574 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6575 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6576 let (mode, wpb) = GU.get_or_init(|| {
6577 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6578 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6579 .ok()
6580 .and_then(|v| v.parse().ok())
6581 .unwrap_or(4u32)
6582 .clamp(1, 16);
6583 (mode, wpb)
6584 });
6585 let (mode, wpb) = (mode.as_str(), *wpb);
6586 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6587 let (inf, nff, ne, rbg, rbu) = (
6588 in_f as i32,
6589 n_ff as i32,
6590 n_expert as i32,
6591 rb_g as i64,
6592 rb_u as i64,
6593 );
6594 let (f, cfg) = match mode {
6595 "1" | "2" | "4" => {
6596 let rpw: u32 = mode.parse().unwrap();
6597 let f = self.func(match rpw {
6598 1 => "moe_gate_up_silu8_dev_q8_r1",
6599 2 => "moe_gate_up_silu8_dev_q8_r2",
6600 _ => "moe_gate_up_silu8_dev_q8_r4",
6601 });
6602 let rows_per_block = (rpw * wpb) as usize;
6603 let gx = n_ff.div_ceil(rows_per_block) as u32;
6604 (
6605 f,
6606 LaunchConfig {
6607 grid_dim: (gx, n_used as u32, 1),
6608 block_dim: (32, wpb, 1),
6609 shared_mem_bytes: 0,
6610 },
6611 )
6612 }
6613 "j8" if n_used <= 32 => (
6614 self.func("moe_gate_up_silu8_dev_q8_j8"),
6615 LaunchConfig {
6616 grid_dim: (n_ff as u32, 1, 1),
6617 block_dim: (32, n_used as u32, 1),
6618 shared_mem_bytes: 0,
6619 },
6620 ),
6621 "vsm2" => {
6623 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6624 let sh = (rb_g + rb_u) as u32;
6625 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6626 f.set_attribute(
6627 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6628 sh as i32,
6629 )?;
6630 (
6631 f,
6632 LaunchConfig {
6633 grid_dim: (n_ff as u32, n_used as u32, 1),
6634 block_dim: (32, 1, 1),
6635 shared_mem_bytes: sh,
6636 },
6637 )
6638 }
6639 "vsm" => {
6640 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6641 let sh = (rb_g + rb_u) as u32;
6642 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6643 f.set_attribute(
6644 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6645 sh as i32,
6646 )?;
6647 (
6648 f,
6649 LaunchConfig {
6650 grid_dim: (n_ff as u32, n_used as u32, 1),
6651 block_dim: (32, 1, 1),
6652 shared_mem_bytes: sh,
6653 },
6654 )
6655 }
6656 "sg" => (
6657 self.func("moe_gate_up_silu8_dev_q8_sg"),
6658 LaunchConfig {
6659 grid_dim: (n_ff as u32, n_used as u32, 1),
6660 block_dim: (32, 1, 1),
6661 shared_mem_bytes: 0,
6662 },
6663 ),
6664 "j8sg" if n_used <= 32 => (
6665 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6666 LaunchConfig {
6667 grid_dim: (n_ff as u32, 1, 1),
6668 block_dim: (32, n_used as u32, 1),
6669 shared_mem_bytes: 0,
6670 },
6671 ),
6672 "u64" if in_f == 2048 => (
6673 self.func("moe_gate_up_silu8_dev_q8_u64"),
6674 LaunchConfig {
6675 grid_dim: (n_ff as u32, n_used as u32, 1),
6676 block_dim: (32, 1, 1),
6677 shared_mem_bytes: 0,
6678 },
6679 ),
6680 "gs4" if in_f == 2048 => (
6681 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6682 LaunchConfig {
6683 grid_dim: (n_ff as u32, n_used as u32, 1),
6684 block_dim: (32, 4, 1),
6685 shared_mem_bytes: 0,
6686 },
6687 ),
6688 "v" | "" => (
6690 self.func("moe_gate_up_silu8_dev_q8_v"),
6691 LaunchConfig {
6692 grid_dim: (n_ff as u32, n_used as u32, 1),
6693 block_dim: (32, 1, 1),
6694 shared_mem_bytes: 0,
6695 },
6696 ),
6697 "s2" => (
6698 self.func("moe_gate_up_silu8_dev_q8_s2"),
6699 LaunchConfig {
6700 grid_dim: (n_ff as u32, n_used as u32, 1),
6701 block_dim: (32, 2, 1),
6702 shared_mem_bytes: 0,
6703 },
6704 ),
6705 "s2z" => {
6706 let rz = wpb.min(16); (
6708 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6709 LaunchConfig {
6710 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6711 block_dim: (32, 2, rz),
6712 shared_mem_bytes: 0,
6713 },
6714 )
6715 }
6716 _ => (
6717 self.func("moe_gate_up_silu8_dev_q8"),
6718 LaunchConfig {
6719 grid_dim: (n_ff as u32, n_used as u32, 1),
6720 block_dim: (32, 1, 1),
6721 shared_mem_bytes: 0,
6722 },
6723 ),
6724 };
6725 let __s_b = self.gpu.stream();
6726 let mut b = __s_b.launch_builder(&f);
6727 b.arg(table)
6728 .arg(sel)
6729 .arg(aq)
6730 .arg(ad)
6731 .arg(&mut act)
6732 .arg(&inf)
6733 .arg(&nff)
6734 .arg(&ne)
6735 .arg(&qt_g)
6736 .arg(&qt_u)
6737 .arg(&rbg)
6738 .arg(&rbu)
6739 .arg(macros);
6740 unsafe {
6741 b.launch(cfg)?;
6742 }
6743 Ok(act)
6744 }
6745
6746 #[allow(clippy::too_many_arguments)]
6747 pub fn moe_down8_fma_dev_q8(
6748 &self,
6749 table: &CudaSlice<u64>,
6750 sel: &cudarc::driver::CudaView<i32>,
6751 w: &cudarc::driver::CudaView<f32>,
6752 aq2: &CudaSlice<i8>,
6753 ad2: &CudaSlice<f32>,
6754 dst: &mut cudarc::driver::CudaViewMut<f32>,
6755 in_f: usize,
6756 out_f: usize,
6757 n_used: usize,
6758 n_expert: usize,
6759 qt: i32,
6760 rb: usize,
6761 ) -> Result<(), Box<dyn std::error::Error>> {
6762 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6763 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6764 let (inf, outf, nu, ne, rbi) = (
6765 in_f as i32,
6766 out_f as i32,
6767 n_used as i32,
6768 n_expert as i32,
6769 rb as i64,
6770 );
6771 let (f, cfg) = match mode.as_str() {
6774 m @ ("1" | "2" | "4") if n_used <= 8 => {
6775 let rpw: usize = m.parse().unwrap();
6776 let f = self.func(match rpw {
6777 1 => "moe_down8_fma_dev_q8_w8r1",
6778 2 => "moe_down8_fma_dev_q8_w8r2",
6779 _ => "moe_down8_fma_dev_q8_w8r4",
6780 });
6781 (
6782 f,
6783 LaunchConfig {
6784 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6785 block_dim: (32, n_used as u32, 1),
6786 shared_mem_bytes: 0,
6787 },
6788 )
6789 }
6790 "h2" if in_f == 512 => (
6791 self.func("moe_down8_fma_dev_q8_h2"),
6792 LaunchConfig {
6793 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6794 block_dim: (32, 1, 1),
6795 shared_mem_bytes: 0,
6796 },
6797 ),
6798 "" if in_f == 704 && n_used <= 8 => (
6801 self.func("moe_down8_fma_dev_q8_w8r2"),
6802 LaunchConfig {
6803 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6804 block_dim: (32, n_used as u32, 1),
6805 shared_mem_bytes: 0,
6806 },
6807 ),
6808 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6812 self.func("moe_down8_fma_dev_q8_w8h2v"),
6813 LaunchConfig {
6814 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6815 block_dim: (32, n_used as u32, 1),
6816 shared_mem_bytes: 0,
6817 },
6818 ),
6819 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6820 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6821 LaunchConfig {
6822 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6823 block_dim: (32, n_used as u32, 1),
6824 shared_mem_bytes: 0,
6825 },
6826 ),
6827 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6828 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6829 LaunchConfig {
6830 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6831 block_dim: (32, n_used as u32, 1),
6832 shared_mem_bytes: 0,
6833 },
6834 ),
6835 "w8h2" if in_f == 512 && n_used <= 8 => (
6836 self.func("moe_down8_fma_dev_q8_w8h2"),
6837 LaunchConfig {
6838 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6839 block_dim: (32, n_used as u32, 1),
6840 shared_mem_bytes: 0,
6841 },
6842 ),
6843 _ => (
6844 self.func("moe_down8_fma_dev_q8"),
6845 LaunchConfig {
6846 grid_dim: (out_f as u32, 1, 1),
6847 block_dim: (32, 1, 1),
6848 shared_mem_bytes: 0,
6849 },
6850 ),
6851 };
6852 let __s_b = self.gpu.stream();
6853 let mut b = __s_b.launch_builder(&f);
6854 b.arg(table)
6855 .arg(sel)
6856 .arg(w)
6857 .arg(aq2)
6858 .arg(ad2)
6859 .arg(dst)
6860 .arg(&inf)
6861 .arg(&outf)
6862 .arg(&nu)
6863 .arg(&ne)
6864 .arg(&qt)
6865 .arg(&rbi);
6866 unsafe {
6867 b.launch(cfg)?;
6868 }
6869 Ok(())
6870 }
6871
6872 #[allow(clippy::too_many_arguments)]
6879 pub fn moe_gate_up_silu8_dev_q8_rows(
6880 &self,
6881 table: &CudaSlice<u64>,
6882 sel: &CudaSlice<i32>,
6883 aq: &CudaSlice<i8>,
6884 ad: &CudaSlice<f32>,
6885 t: usize,
6886 in_f: usize,
6887 n_ff: usize,
6888 n_used: usize,
6889 n_expert: usize,
6890 qt_g: i32,
6891 qt_u: i32,
6892 rb_g: usize,
6893 rb_u: usize,
6894 macros: &CudaSlice<f32>,
6895 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6896 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6897 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6898 let cfg = LaunchConfig {
6899 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6900 block_dim: (32, 1, 1),
6901 shared_mem_bytes: 0,
6902 };
6903 let (inf, nff, ne, nu, rbg, rbu) = (
6904 in_f as i32,
6905 n_ff as i32,
6906 n_expert as i32,
6907 n_used as i32,
6908 rb_g as i64,
6909 rb_u as i64,
6910 );
6911 let __s_b = self.gpu.stream();
6912 let mut b = __s_b.launch_builder(&f);
6913 b.arg(table)
6914 .arg(sel)
6915 .arg(aq)
6916 .arg(ad)
6917 .arg(&mut act)
6918 .arg(&inf)
6919 .arg(&nff)
6920 .arg(&ne)
6921 .arg(&qt_g)
6922 .arg(&qt_u)
6923 .arg(&rbg)
6924 .arg(&rbu)
6925 .arg(&nu)
6926 .arg(macros);
6927 unsafe {
6928 b.launch(cfg)?;
6929 }
6930 Ok(act)
6931 }
6932
6933 #[allow(clippy::too_many_arguments)]
6938 pub fn moe_down8_fma_dev_q8_rows(
6939 &self,
6940 table: &CudaSlice<u64>,
6941 sel: &CudaSlice<i32>,
6942 w: &CudaSlice<f32>,
6943 aq2: &CudaSlice<i8>,
6944 ad2: &CudaSlice<f32>,
6945 dst: &mut CudaSlice<f32>,
6946 t: usize,
6947 in_f: usize,
6948 out_f: usize,
6949 n_used: usize,
6950 n_expert: usize,
6951 qt: i32,
6952 rb: usize,
6953 ) -> Result<(), Box<dyn std::error::Error>> {
6954 assert!(
6955 in_f == 512 && n_used <= 8,
6956 "down rows twin is w8h2v shape-gated"
6957 );
6958 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6959 let cfg = LaunchConfig {
6960 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6961 block_dim: (32, n_used as u32, 1),
6962 shared_mem_bytes: 0,
6963 };
6964 let (inf, outf, nu, ne, rbi) = (
6965 in_f as i32,
6966 out_f as i32,
6967 n_used as i32,
6968 n_expert as i32,
6969 rb as i64,
6970 );
6971 let __s_b = self.gpu.stream();
6972 let mut b = __s_b.launch_builder(&f);
6973 b.arg(table)
6974 .arg(sel)
6975 .arg(w)
6976 .arg(aq2)
6977 .arg(ad2)
6978 .arg(dst)
6979 .arg(&inf)
6980 .arg(&outf)
6981 .arg(&nu)
6982 .arg(&ne)
6983 .arg(&qt)
6984 .arg(&rbi);
6985 unsafe {
6986 b.launch(cfg)?;
6987 }
6988 Ok(())
6989 }
6990
6991 #[allow(clippy::too_many_arguments)]
6995 pub fn moe_gate_up_silu8_dev_q8_csr(
6996 &self,
6997 table: &CudaSlice<u64>,
6998 sel: &CudaSlice<i32>,
6999 aq: &CudaSlice<i8>,
7000 ad: &CudaSlice<f32>,
7001 n_pairs: usize,
7002 in_f: usize,
7003 n_ff: usize,
7004 n_used: usize,
7005 n_expert: usize,
7006 qt_g: i32,
7007 qt_u: i32,
7008 rb_g: usize,
7009 rb_u: usize,
7010 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7011 let f = if qt_g == crate::QT_NVFP4 {
7014 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
7015 } else {
7016 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
7017 };
7018 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7019 let cfg = LaunchConfig {
7020 grid_dim: (n_ff as u32, n_pairs as u32, 1),
7021 block_dim: (32, 1, 1),
7022 shared_mem_bytes: 0,
7023 };
7024 let (inf, nff, ne, nu, npi, rbg, rbu) = (
7025 in_f as i32,
7026 n_ff as i32,
7027 n_expert as i32,
7028 n_used as i32,
7029 n_pairs as i32,
7030 rb_g as i64,
7031 rb_u as i64,
7032 );
7033 let __s_b = self.gpu.stream();
7034 let mut b = __s_b.launch_builder(&f);
7035 b.arg(table)
7036 .arg(sel)
7037 .arg(aq)
7038 .arg(ad)
7039 .arg(&mut act)
7040 .arg(&inf)
7041 .arg(&nff)
7042 .arg(&ne)
7043 .arg(&qt_g)
7044 .arg(&qt_u)
7045 .arg(&rbg)
7046 .arg(&rbu)
7047 .arg(&nu)
7048 .arg(&npi);
7049 unsafe {
7050 b.launch(cfg)?;
7051 }
7052 Ok(act)
7053 }
7054
7055 #[allow(clippy::too_many_arguments)]
7059 pub fn moe_down8_fma_dev_q8_variant(
7060 &self,
7061 variant: &str,
7062 table: &CudaSlice<u64>,
7063 sel: &cudarc::driver::CudaView<i32>,
7064 w: &cudarc::driver::CudaView<f32>,
7065 aq2: &CudaSlice<i8>,
7066 ad2: &CudaSlice<f32>,
7067 dst: &mut cudarc::driver::CudaViewMut<f32>,
7068 in_f: usize,
7069 out_f: usize,
7070 n_used: usize,
7071 n_expert: usize,
7072 qt: i32,
7073 rb: usize,
7074 ) -> Result<(), Box<dyn std::error::Error>> {
7075 let (inf, outf, nu, ne, rbi) = (
7076 in_f as i32,
7077 out_f as i32,
7078 n_used as i32,
7079 n_expert as i32,
7080 rb as i64,
7081 );
7082 let (f, cfg) = match variant {
7083 "w8h2" | "w8h2v" => (
7084 self.func(if variant == "w8h2" {
7085 "moe_down8_fma_dev_q8_w8h2"
7086 } else {
7087 "moe_down8_fma_dev_q8_w8h2v"
7088 }),
7089 LaunchConfig {
7090 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7091 block_dim: (32, n_used as u32, 1),
7092 shared_mem_bytes: 0,
7093 },
7094 ),
7095 "w8h2r2" | "w8h2r2v" => (
7096 self.func(if variant == "w8h2r2" {
7097 "moe_down8_fma_dev_q8_w8h2r2"
7098 } else {
7099 "moe_down8_fma_dev_q8_w8h2r2v"
7100 }),
7101 LaunchConfig {
7102 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7103 block_dim: (32, n_used as u32, 1),
7104 shared_mem_bytes: 0,
7105 },
7106 ),
7107 _ => (
7108 self.func("moe_down8_fma_dev_q8"),
7109 LaunchConfig {
7110 grid_dim: (out_f as u32, 1, 1),
7111 block_dim: (32, 1, 1),
7112 shared_mem_bytes: 0,
7113 },
7114 ),
7115 };
7116 let __s_b = self.gpu.stream();
7117 let mut b = __s_b.launch_builder(&f);
7118 b.arg(table)
7119 .arg(sel)
7120 .arg(w)
7121 .arg(aq2)
7122 .arg(ad2)
7123 .arg(dst)
7124 .arg(&inf)
7125 .arg(&outf)
7126 .arg(&nu)
7127 .arg(&ne)
7128 .arg(&qt)
7129 .arg(&rbi);
7130 unsafe {
7131 b.launch(cfg)?;
7132 }
7133 Ok(())
7134 }
7135
7136 #[allow(clippy::too_many_arguments)]
7138 pub fn moe_gate_up_silu8_dev_q8_variant(
7139 &self,
7140 variant: &str,
7141 table: &CudaSlice<u64>,
7142 sel: &cudarc::driver::CudaView<i32>,
7143 aq: &CudaSlice<i8>,
7144 ad: &CudaSlice<f32>,
7145 in_f: usize,
7146 n_ff: usize,
7147 n_used: usize,
7148 n_expert: usize,
7149 qt_g: i32,
7150 qt_u: i32,
7151 rb_g: usize,
7152 rb_u: usize,
7153 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7154 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7155 let (inf, nff, ne, rbg, rbu) = (
7156 in_f as i32,
7157 n_ff as i32,
7158 n_expert as i32,
7159 rb_g as i64,
7160 rb_u as i64,
7161 );
7162 let f = self.func(if variant == "v" {
7163 "moe_gate_up_silu8_dev_q8_v"
7164 } else {
7165 "moe_gate_up_silu8_dev_q8"
7166 });
7167 let cfg = LaunchConfig {
7168 grid_dim: (n_ff as u32, n_used as u32, 1),
7169 block_dim: (32, 1, 1),
7170 shared_mem_bytes: 0,
7171 };
7172 let __s_b = self.gpu.stream();
7173 let mut b = __s_b.launch_builder(&f);
7174 b.arg(table)
7175 .arg(sel)
7176 .arg(aq)
7177 .arg(ad)
7178 .arg(&mut act)
7179 .arg(&inf)
7180 .arg(&nff)
7181 .arg(&ne)
7182 .arg(&qt_g)
7183 .arg(&qt_u)
7184 .arg(&rbg)
7185 .arg(&rbu);
7186 unsafe {
7187 b.launch(cfg)?;
7188 }
7189 Ok(act)
7190 }
7191
7192 pub fn moe_gate_up_silu8_dev(
7193 &self,
7194 table: &CudaSlice<u64>,
7195 sel: &cudarc::driver::CudaView<i32>,
7196 x: &cudarc::driver::CudaView<f32>,
7197 in_f: usize,
7198 n_ff: usize,
7199 n_used: usize,
7200 n_expert: usize,
7201 qt_g: i32,
7202 qt_u: i32,
7203 rb_g: usize,
7204 rb_u: usize,
7205 macros: &CudaSlice<f32>,
7206 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7207 let f = self.func("moe_gate_up_silu8_dev");
7208 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
7210 grid_dim: (n_ff as u32, n_used as u32, 1),
7211 block_dim: (256, 1, 1),
7212 shared_mem_bytes: 0,
7213 };
7214 let (inf, nff, ne, rbg, rbu) = (
7215 in_f as i32,
7216 n_ff as i32,
7217 n_expert as i32,
7218 rb_g as i64,
7219 rb_u as i64,
7220 );
7221 let __s_b = self.gpu.stream();
7222 let mut b = __s_b.launch_builder(&f);
7223 b.arg(table)
7224 .arg(sel)
7225 .arg(x)
7226 .arg(&mut act)
7227 .arg(&inf)
7228 .arg(&nff)
7229 .arg(&ne)
7230 .arg(&qt_g)
7231 .arg(&qt_u)
7232 .arg(&rbg)
7233 .arg(&rbu)
7234 .arg(macros);
7235 unsafe {
7236 b.launch(cfg)?;
7237 }
7238 Ok(act)
7239 }
7240
7241 #[allow(clippy::too_many_arguments)]
7244 pub fn moe_down8_fma_dev(
7245 &self,
7246 table: &CudaSlice<u64>,
7247 sel: &cudarc::driver::CudaView<i32>,
7248 w: &cudarc::driver::CudaView<f32>,
7249 act: &CudaSlice<f32>,
7250 dst: &mut cudarc::driver::CudaViewMut<f32>,
7251 in_f: usize,
7252 out_f: usize,
7253 n_used: usize,
7254 n_expert: usize,
7255 qt: i32,
7256 rb: usize,
7257 ) -> Result<(), Box<dyn std::error::Error>> {
7258 let f = self.func("moe_down8_fma_dev");
7259 let cfg = LaunchConfig {
7260 grid_dim: (out_f as u32, 1, 1),
7261 block_dim: (256, 1, 1),
7262 shared_mem_bytes: 0,
7263 };
7264 let (inf, outf, nu, ne, rbv) = (
7265 in_f as i32,
7266 out_f as i32,
7267 n_used as i32,
7268 n_expert as i32,
7269 rb as i64,
7270 );
7271 let __s_b = self.gpu.stream();
7272 let mut b = __s_b.launch_builder(&f);
7273 b.arg(table)
7274 .arg(sel)
7275 .arg(w)
7276 .arg(act)
7277 .arg(dst)
7278 .arg(&inf)
7279 .arg(&outf)
7280 .arg(&nu)
7281 .arg(&ne)
7282 .arg(&qt)
7283 .arg(&rbv);
7284 unsafe {
7285 b.launch(cfg)?;
7286 }
7287 Ok(())
7288 }
7289
7290 pub fn axpy_into(
7292 &self,
7293 src: &CudaSlice<f32>,
7294 alpha: f32,
7295 dst: &mut cudarc::driver::CudaViewMut<f32>,
7296 n: usize,
7297 ) -> Result<(), Box<dyn std::error::Error>> {
7298 let f = self.func("axpy_f32");
7299 let cfg = LaunchConfig::for_num_elems(n as u32);
7300 let (a, ni) = (alpha, n as i32);
7301 let __s_b = self.gpu.stream();
7302 let mut b = __s_b.launch_builder(&f);
7303 b.arg(src).arg(dst).arg(&a).arg(&ni);
7304 unsafe {
7305 b.launch(cfg)?;
7306 }
7307 Ok(())
7308 }
7309
7310 pub fn axpy_host_into(
7312 &self,
7313 src: &cudarc::driver::CudaView<'_, f32>,
7314 alpha: f32,
7315 dst: &mut cudarc::driver::CudaViewMut<f32>,
7316 n: usize,
7317 ) -> Result<(), Box<dyn std::error::Error>> {
7318 let f = self.func("axpy_host_f32");
7319 let cfg = LaunchConfig::for_num_elems(n as u32);
7320 let (a, ni) = (alpha, n as i32);
7321 let __s_b = self.gpu.stream();
7322 let mut b = __s_b.launch_builder(&f);
7323 b.arg(src).arg(dst).arg(&a).arg(&ni);
7324 unsafe {
7325 b.launch(cfg)?;
7326 }
7327 Ok(())
7328 }
7329
7330 pub fn add_scaled_rows(
7332 &self,
7333 src: &CudaSlice<f32>,
7334 scale: &CudaSlice<f32>,
7335 dst: &mut CudaSlice<f32>,
7336 ncols: usize,
7337 nrows: usize,
7338 ) -> Result<(), Box<dyn std::error::Error>> {
7339 let f = self.func("add_scaled_rows_f32");
7340 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7341 let (nc, nr) = (ncols as i32, nrows as i32);
7342 let __s_b = self.gpu.stream();
7343 let mut b = __s_b.launch_builder(&f);
7344 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7345 unsafe {
7346 b.launch(cfg)?;
7347 }
7348 Ok(())
7349 }
7350
7351 pub fn gather_rows(
7355 &self,
7356 src: &CudaSlice<f32>,
7357 idx: &CudaSlice<i32>,
7358 dst: &mut CudaSlice<f32>,
7359 ncols: usize,
7360 m_e: usize,
7361 ) -> Result<(), Box<dyn std::error::Error>> {
7362 let f = self.func("gather_rows_f32");
7363 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7364 let (nc, me) = (ncols as i32, m_e as i32);
7365 let __s_b = self.gpu.stream();
7366 let mut b = __s_b.launch_builder(&f);
7367 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7368 unsafe {
7369 b.launch(cfg)?;
7370 }
7371 Ok(())
7372 }
7373
7374 pub fn scatter_slot(
7379 &self,
7380 src: &CudaSlice<f32>,
7381 tok_idx: &CudaSlice<i32>,
7382 slot_idx: &CudaSlice<i32>,
7383 weight: &CudaSlice<f32>,
7384 dst: &mut CudaSlice<f32>,
7385 wbuf: &mut CudaSlice<f32>,
7386 ncols: usize,
7387 n_used: usize,
7388 m_e: usize,
7389 ) -> Result<(), Box<dyn std::error::Error>> {
7390 let f = self.func("scatter_add_slot_f32");
7391 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7392 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7393 let __s_b = self.gpu.stream();
7394 let mut b = __s_b.launch_builder(&f);
7395 b.arg(src)
7396 .arg(tok_idx)
7397 .arg(slot_idx)
7398 .arg(weight)
7399 .arg(dst)
7400 .arg(wbuf)
7401 .arg(&nc)
7402 .arg(&nu)
7403 .arg(&me);
7404 unsafe {
7405 b.launch(cfg)?;
7406 }
7407 Ok(())
7408 }
7409
7410 pub fn reduce_slots(
7414 &self,
7415 slots: &CudaSlice<f32>,
7416 wbuf: &CudaSlice<f32>,
7417 dst: &mut CudaSlice<f32>,
7418 ncols: usize,
7419 n_used: usize,
7420 t: usize,
7421 ) -> Result<(), Box<dyn std::error::Error>> {
7422 let f = self.func("reduce_slots_f32");
7423 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7424 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7425 let __s_b = self.gpu.stream();
7426 let mut b = __s_b.launch_builder(&f);
7427 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7428 unsafe {
7429 b.launch(cfg)?;
7430 }
7431 Ok(())
7432 }
7433
7434 pub fn reduce_slots_host(
7439 &self,
7440 slots: &CudaSlice<f32>,
7441 wbuf: &CudaSlice<f32>,
7442 dst: &mut CudaSlice<f32>,
7443 ncols: usize,
7444 n_used: usize,
7445 t: usize,
7446 ) -> Result<(), Box<dyn std::error::Error>> {
7447 let f = self.func("reduce_slots_host_f32");
7448 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7449 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7450 let __s_b = self.gpu.stream();
7451 let mut b = __s_b.launch_builder(&f);
7452 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7453 unsafe {
7454 b.launch(cfg)?;
7455 }
7456 Ok(())
7457 }
7458
7459 pub fn quantize_q8_1_view(
7466 &self,
7467 x: &cudarc::driver::CudaView<f32>,
7468 m: usize,
7469 in_f: usize,
7470 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7471 let f = self.func("quantize_q8_1");
7472 let nblk = in_f / 32;
7473 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7474 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7475 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7476 let (inf, mi) = (in_f as i32, m as i32);
7477 let __s_b = self.gpu.stream();
7478 let mut b = __s_b.launch_builder(&f);
7479 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7480 unsafe {
7481 b.launch(cfg)?;
7482 }
7483 Ok((q, d))
7484 }
7485
7486 pub fn quantize_q8_1(
7487 &self,
7488 x: &CudaSlice<f32>,
7489 m: usize,
7490 in_f: usize,
7491 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7492 let nblk = in_f / 32;
7493 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);
7497 let (inf, mi) = (in_f as i32, m as i32);
7498 if Self::pdl_on() && Self::pdl_wb_on() {
7499 {
7500 use cudarc::driver::{DevicePtr, DevicePtrMut};
7501 let s = &self.gpu.stream();
7502 let (px, _g0) = x.device_ptr(s);
7503 let (pq, _g1) = q.device_ptr_mut(s);
7504 let (pd, _g2) = d.device_ptr_mut(s);
7505 let mut ps = [
7506 &px as *const _ as *mut std::ffi::c_void,
7507 &pq as *const _ as *mut _,
7508 &pd as *const _ as *mut _,
7509 &inf as *const _ as *mut _,
7510 &mi as *const _ as *mut _,
7511 ];
7512 unsafe {
7513 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7514 }
7515 }
7516 return Ok((q, d));
7517 }
7518 let f = self.func("quantize_q8_1");
7519 let __s_b = self.gpu.stream();
7520 let mut b = __s_b.launch_builder(&f);
7521 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7522 unsafe {
7523 b.launch(cfg)?;
7524 }
7525 Ok((q, d))
7526 }
7527
7528 pub fn quantize_fp4_act(
7532 &self,
7533 x: &CudaSlice<f32>,
7534 m: usize,
7535 in_f: usize,
7536 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7537 let f = self.func("quantize_fp4_act");
7538 let nb16 = in_f / 16;
7539 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);
7542 let (inf, mi) = (in_f as i32, m as i32);
7543 let __s_b = self.gpu.stream();
7544 let mut b = __s_b.launch_builder(&f);
7545 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7546 unsafe {
7547 b.launch(cfg)?;
7548 }
7549 Ok((aq4, ad4))
7550 }
7551
7552 pub fn qmatvec_gemm_nvfp4_fp4(
7557 &self,
7558 bytes: &CudaSlice<u8>,
7559 x: &CudaSlice<f32>,
7560 m: usize,
7561 in_f: usize,
7562 out_f: usize,
7563 row_bytes: usize,
7564 scale: f32,
7565 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7566 assert!(
7567 in_f % 64 == 0,
7568 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7569 );
7570 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7571 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7572 if scale != 1.0 {
7573 self.scale_inplace(&mut y, scale, m * out_f)?;
7574 }
7575 Ok(y)
7576 }
7577
7578 fn fp4_gemm_launch(
7581 &self,
7582 bytes: &CudaSlice<u8>,
7583 aq4: &CudaSlice<u32>,
7584 ad4: &CudaSlice<u8>,
7585 m: usize,
7586 in_f: usize,
7587 out_f: usize,
7588 row_bytes: usize,
7589 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7590 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7591 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7593 const BN: u32 = 256;
7594 let cfg = LaunchConfig {
7595 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7596 block_dim: (32, 4, 1),
7597 shared_mem_bytes: 0,
7598 };
7599 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7600 let __s_b = self.gpu.stream();
7601 let mut b = __s_b.launch_builder(&f);
7602 b.arg(bytes)
7603 .arg(aq4)
7604 .arg(ad4)
7605 .arg(&mut y)
7606 .arg(&inf)
7607 .arg(&outf)
7608 .arg(&mi)
7609 .arg(&rb);
7610 unsafe {
7611 b.launch(cfg)?;
7612 }
7613 Ok(y)
7614 }
7615
7616 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7618 &self,
7619 bytes: &CudaSlice<u8>,
7620 x: &CudaSlice<f32>,
7621 m: usize,
7622 in_f: usize,
7623 out_f: usize,
7624 row_bytes: usize,
7625 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7626 assert!(
7627 in_f % 64 == 0,
7628 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7629 );
7630 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7631 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7632 }
7633
7634 pub fn qmatvec_q8_0_fast(
7636 &self,
7637 w: &CudaSlice<u8>,
7638 x: &CudaSlice<f32>,
7639 m: usize,
7640 in_f: usize,
7641 out_f: usize,
7642 row_bytes: usize,
7643 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7644 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7645 let f = self.func("qmatvec_q8_0_dp4a");
7646 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7648 grid_dim: (out_f as u32, m as u32, 1),
7649 block_dim: (128, 1, 1),
7650 shared_mem_bytes: 0,
7651 };
7652 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7653 let __s_b = self.gpu.stream();
7654 let mut b = __s_b.launch_builder(&f);
7655 b.arg(w)
7656 .arg(&aq)
7657 .arg(&ad)
7658 .arg(&mut y)
7659 .arg(&inf)
7660 .arg(&outf)
7661 .arg(&mi)
7662 .arg(&rb);
7663 unsafe {
7664 b.launch(cfg)?;
7665 }
7666 Ok(y)
7667 }
7668
7669 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7672 &self,
7673 w: &CudaSlice<u8>,
7674 x: &CudaSlice<f32>,
7675 m: usize,
7676 in_f: usize,
7677 out_f: usize,
7678 row_bytes: usize,
7679 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7680 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7681 let f = self.func("qmatvec_q4_K_dp4a");
7682 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7684 grid_dim: (out_f as u32, m as u32, 1),
7685 block_dim: (128, 1, 1),
7686 shared_mem_bytes: 0,
7687 };
7688 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7689 let __s_b = self.gpu.stream();
7690 let mut b = __s_b.launch_builder(&f);
7691 b.arg(w)
7692 .arg(&aq)
7693 .arg(&ad)
7694 .arg(&mut y)
7695 .arg(&inf)
7696 .arg(&outf)
7697 .arg(&mi)
7698 .arg(&rb);
7699 unsafe {
7700 b.launch(cfg)?;
7701 }
7702 Ok(y)
7703 }
7704
7705 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7708 &self,
7709 w: &CudaSlice<u8>,
7710 x: &CudaSlice<f32>,
7711 m: usize,
7712 in_f: usize,
7713 out_f: usize,
7714 row_bytes: usize,
7715 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7716 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7717 let f = self.func("qmatvec_q6_K_dp4a");
7718 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7720 grid_dim: (out_f as u32, m as u32, 1),
7721 block_dim: (128, 1, 1),
7722 shared_mem_bytes: 0,
7723 };
7724 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7725 let __s_b = self.gpu.stream();
7726 let mut b = __s_b.launch_builder(&f);
7727 b.arg(w)
7728 .arg(&aq)
7729 .arg(&ad)
7730 .arg(&mut y)
7731 .arg(&inf)
7732 .arg(&outf)
7733 .arg(&mi)
7734 .arg(&rb);
7735 unsafe {
7736 b.launch(cfg)?;
7737 }
7738 Ok(y)
7739 }
7740
7741 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
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 self.qmatvec_dp4a_named(
7753 "qmatvec_q5_K_dp4a",
7754 &w.slice(0..w.len()),
7755 x,
7756 m,
7757 in_f,
7758 out_f,
7759 row_bytes,
7760 )
7761 }
7762 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7765 &self,
7766 w: &CudaSlice<u8>,
7767 x: &CudaSlice<f32>,
7768 m: usize,
7769 in_f: usize,
7770 out_f: usize,
7771 row_bytes: usize,
7772 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7773 self.qmatvec_dp4a_named(
7774 "qmatvec_q3_K_dp4a",
7775 &w.slice(0..w.len()),
7776 x,
7777 m,
7778 in_f,
7779 out_f,
7780 row_bytes,
7781 )
7782 }
7783 pub fn qmatvec_nvfp4_fast_rp(
7785 &self,
7786 w: &CudaSlice<u8>,
7787 x: &CudaSlice<f32>,
7788 m: usize,
7789 in_f: usize,
7790 out_f: usize,
7791 row_bytes: usize,
7792 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7793 assert!(
7794 in_f % 64 == 0,
7795 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7796 );
7797 self.qmatvec_dp4a_named(
7798 "qmatvec_nvfp4_dp4a_rp",
7799 &w.slice(0..w.len()),
7800 x,
7801 m,
7802 in_f,
7803 out_f,
7804 row_bytes,
7805 )
7806 }
7807 pub fn qmatvec_nvfp4_fast(
7809 &self,
7810 w: &cudarc::driver::CudaView<'_, u8>,
7811 x: &CudaSlice<f32>,
7812 m: usize,
7813 in_f: usize,
7814 out_f: usize,
7815 row_bytes: usize,
7816 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7817 assert!(
7820 in_f % 64 == 0,
7821 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7822 );
7823 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7824 }
7825 pub fn qmatvec_nvfp4_fast_v2(
7828 &self,
7829 w: &cudarc::driver::CudaView<'_, u8>,
7830 x: &CudaSlice<f32>,
7831 m: usize,
7832 in_f: usize,
7833 out_f: usize,
7834 row_bytes: usize,
7835 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7836 assert!(
7837 in_f % 64 == 0,
7838 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7839 );
7840 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7841 }
7842 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7845 &self,
7846 w: &CudaSlice<u8>,
7847 x: &CudaSlice<f32>,
7848 m: usize,
7849 in_f: usize,
7850 out_f: usize,
7851 row_bytes: usize,
7852 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7853 self.qmatvec_dp4a_named(
7854 "qmatvec_iq4_XS_dp4a",
7855 &w.slice(0..w.len()),
7856 x,
7857 m,
7858 in_f,
7859 out_f,
7860 row_bytes,
7861 )
7862 }
7863
7864 fn qmatvec_dp4a_named(
7866 &self,
7867 name: &str,
7868 w: &cudarc::driver::CudaView<'_, u8>,
7869 x: &CudaSlice<f32>,
7870 m: usize,
7871 in_f: usize,
7872 out_f: usize,
7873 row_bytes: usize,
7874 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7875 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7876 let f = self.func(name);
7877 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7879 grid_dim: (out_f as u32, m as u32, 1),
7880 block_dim: (128, 1, 1),
7881 shared_mem_bytes: 0,
7882 };
7883 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7884 let __s_b = self.gpu.stream();
7885 let mut b = __s_b.launch_builder(&f);
7886 b.arg(w)
7887 .arg(&aq)
7888 .arg(&ad)
7889 .arg(&mut y)
7890 .arg(&inf)
7891 .arg(&outf)
7892 .arg(&mi)
7893 .arg(&rb);
7894 unsafe {
7895 b.launch(cfg)?;
7896 }
7897 Ok(y)
7898 }
7899
7900 #[allow(clippy::too_many_arguments)]
7906 pub fn qmatvec_nvfp4_fast_prequant_into(
7907 &self,
7908 w: &CudaSlice<u8>,
7909 aq: &CudaSlice<i8>,
7910 ad: &CudaSlice<f32>,
7911 y: &mut CudaSlice<f32>,
7912 m: usize,
7913 in_f: usize,
7914 out_f: usize,
7915 row_bytes: usize,
7916 ) -> Result<(), Box<dyn std::error::Error>> {
7917 assert!(
7918 in_f % 64 == 0,
7919 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7920 );
7921 if y.len() < m * out_f {
7922 return Err(format!(
7923 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7924 y.len()
7925 )
7926 .into());
7927 }
7928 let f = self.func("qmatvec_nvfp4_dp4a");
7929 let cfg = LaunchConfig {
7930 grid_dim: (out_f as u32, m as u32, 1),
7931 block_dim: (128, 1, 1),
7932 shared_mem_bytes: 0,
7933 };
7934 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7935 let __s_b = self.gpu.stream();
7936 let mut b = __s_b.launch_builder(&f);
7937 b.arg(w)
7938 .arg(aq)
7939 .arg(ad)
7940 .arg(y)
7941 .arg(&inf)
7942 .arg(&outf)
7943 .arg(&mi)
7944 .arg(&rb);
7945 unsafe {
7946 b.launch(cfg)?;
7947 }
7948 Ok(())
7949 }
7950
7951 #[allow(clippy::too_many_arguments)]
7954 pub fn matvec_f32_qkv_into(
7955 &self,
7956 wq: &CudaSlice<f32>,
7957 wk: &CudaSlice<f32>,
7958 wv: &CudaSlice<f32>,
7959 wg: &CudaSlice<f32>,
7960 x: &CudaSlice<f32>,
7961 yq: &mut CudaSlice<f32>,
7962 yk: &mut CudaSlice<f32>,
7963 yv: &mut CudaSlice<f32>,
7964 yg: &mut CudaSlice<f32>,
7965 in_f: usize,
7966 out_q: usize,
7967 out_kv: usize,
7968 out_g: usize,
7969 ) -> Result<(), Box<dyn std::error::Error>> {
7970 if in_f % 4 != 0
7971 || wq.len() != out_q * in_f
7972 || wk.len() != out_kv * in_f
7973 || wv.len() != out_kv * in_f
7974 || wg.len() < out_g * in_f
7975 || x.len() < in_f
7976 || yq.len() < out_q
7977 || yk.len() < out_kv
7978 || yv.len() < out_kv
7979 || (out_g > 0 && yg.len() < out_g)
7980 {
7981 return Err(format!(
7982 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7983 wq={} wk={} wv={} wg={}",
7984 wq.len(),
7985 wk.len(),
7986 wv.len(),
7987 wg.len()
7988 )
7989 .into());
7990 }
7991 let f = self.func("matvec_f32_qkv");
7992 let cfg = LaunchConfig {
7993 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7994 block_dim: (128, 1, 1),
7995 shared_mem_bytes: 0,
7996 };
7997 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7998 let __s_b = self.gpu.stream();
7999 let mut b = __s_b.launch_builder(&f);
8000 b.arg(wq)
8001 .arg(wk)
8002 .arg(wv)
8003 .arg(wg)
8004 .arg(x)
8005 .arg(yq)
8006 .arg(yk)
8007 .arg(yv)
8008 .arg(yg)
8009 .arg(&inf)
8010 .arg(&oq)
8011 .arg(&okv)
8012 .arg(&og);
8013 unsafe {
8014 b.launch(cfg)?;
8015 }
8016 Ok(())
8017 }
8018
8019 #[allow(clippy::too_many_arguments)]
8025 #[allow(clippy::too_many_arguments)]
8028 pub fn qmatvec_nvfp4_sel_gu_into(
8029 &self,
8030 gate_bank: &CudaSlice<u8>,
8031 up_bank: &CudaSlice<u8>,
8032 sel: &CudaSlice<i32>,
8033 aq: &CudaSlice<i8>,
8034 ad: &CudaSlice<f32>,
8035 yg: &mut CudaSlice<f32>,
8036 yu: &mut CudaSlice<f32>,
8037 n_sel: usize,
8038 in_f: usize,
8039 out_f: usize,
8040 row_bytes: usize,
8041 expert_stride: usize,
8042 ) -> Result<(), Box<dyn std::error::Error>> {
8043 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8044 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8045 return Err("NVFP4 gu sel geometry".into());
8046 }
8047 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
8050 let rpw = *RPW.get_or_init(|| {
8051 std::env::var("MEMRA_SEL_GU_RPW")
8052 .ok()
8053 .and_then(|v| v.parse().ok())
8054 .filter(|r| *r == 2 || *r == 4)
8055 .unwrap_or(1)
8056 });
8057 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
8058 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8061 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
8062 let f = self.func(match (wpr, rpw) {
8063 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
8064 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
8065 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
8066 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
8067 });
8068 let cfg = LaunchConfig {
8069 grid_dim: if wpr {
8070 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
8071 } else if rpw == 1 {
8072 ((2 * out_f) as u32, n_sel as u32, 1)
8073 } else {
8074 ((out_f / rpw) as u32, n_sel as u32, 1)
8075 },
8076 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
8077 shared_mem_bytes: 0,
8078 };
8079 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8080 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8081 let (ars, adrs) = (0i64, 0i64);
8082 let __s_b = self.gpu.stream();
8083 let mut b = __s_b.launch_builder(&f);
8084 b.arg(gate_bank)
8085 .arg(up_bank)
8086 .arg(sel)
8087 .arg(aq)
8088 .arg(ad)
8089 .arg(yg)
8090 .arg(yu)
8091 .arg(&inf)
8092 .arg(&outf)
8093 .arg(&ns)
8094 .arg(&rb)
8095 .arg(&es)
8096 .arg(&ars)
8097 .arg(&adrs);
8098 unsafe {
8099 b.launch(cfg)?;
8100 }
8101 Ok(())
8102 }
8103
8104 #[allow(clippy::too_many_arguments)]
8112 pub fn qmatvec_nvfp4_sel_down8_into(
8113 &self,
8114 bank: &CudaSlice<u8>,
8115 sel: &CudaSlice<i32>,
8116 aq: &CudaSlice<i8>,
8117 ad: &CudaSlice<f32>,
8118 route_w: &CudaSlice<f32>,
8119 md: &CudaSlice<f32>,
8120 dst: &mut CudaSlice<f32>,
8121 n_sel: usize,
8122 in_f: usize,
8123 out_f: usize,
8124 row_bytes: usize,
8125 expert_stride: usize,
8126 act_row_stride: usize,
8127 ad_row_stride: usize,
8128 ) -> Result<(), Box<dyn std::error::Error>> {
8129 if in_f % 64 != 0
8130 || n_sel == 0
8131 || n_sel > 8
8132 || (in_f >> 5) > 32
8133 || dst.len() < out_f
8134 || sel.len() < n_sel
8135 || route_w.len() < n_sel
8136 {
8137 return Err(format!(
8138 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
8139 dst.len()
8140 )
8141 .into());
8142 }
8143 if !crate::tp::nvfp4_bank_v2_on() {
8144 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
8145 }
8146 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
8147 let cfg = LaunchConfig {
8148 grid_dim: (out_f as u32, 1, 1),
8149 block_dim: (32, n_sel as u32, 1),
8150 shared_mem_bytes: 0,
8151 };
8152 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8153 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8154 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8155 let __s_b = self.gpu.stream();
8156 let mut b = __s_b.launch_builder(&f);
8157 b.arg(bank)
8158 .arg(sel)
8159 .arg(aq)
8160 .arg(ad)
8161 .arg(route_w)
8162 .arg(md)
8163 .arg(dst)
8164 .arg(&inf)
8165 .arg(&outf)
8166 .arg(&ns)
8167 .arg(&rb)
8168 .arg(&es)
8169 .arg(&ars)
8170 .arg(&adrs);
8171 unsafe {
8172 b.launch(cfg)?;
8173 }
8174 Ok(())
8175 }
8176
8177 #[allow(clippy::too_many_arguments)]
8181 pub fn qmatvec_nvfp4_sel_down8_rows_into(
8182 &self,
8183 bank: &CudaSlice<u8>,
8184 sel: &CudaSlice<i32>,
8185 aq: &CudaSlice<i8>,
8186 ad: &CudaSlice<f32>,
8187 route_w: &CudaSlice<f32>,
8188 md: &CudaSlice<f32>,
8189 dst: &mut CudaSlice<f32>,
8190 t: usize,
8191 n_sel_col: usize,
8192 in_f: usize,
8193 out_f: usize,
8194 row_bytes: usize,
8195 expert_stride: usize,
8196 act_row_stride: usize,
8197 ad_row_stride: usize,
8198 ) -> Result<(), Box<dyn std::error::Error>> {
8199 let n_sel = t * n_sel_col;
8200 if in_f % 64 != 0
8201 || n_sel_col == 0
8202 || n_sel_col > 8
8203 || t == 0
8204 || t > 64
8205 || (in_f >> 5) > 32
8206 || dst.len() < t * out_f
8207 || sel.len() < n_sel
8208 || route_w.len() < n_sel
8209 {
8210 return Err(format!(
8211 "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
8212 dst.len()
8213 )
8214 .into());
8215 }
8216 if !crate::tp::nvfp4_bank_v2_on() {
8217 return Err(
8218 "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8219 );
8220 }
8221 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8222 let cfg = LaunchConfig {
8223 grid_dim: (out_f as u32, t as u32, 1),
8224 block_dim: (32, n_sel_col as u32, 1),
8225 shared_mem_bytes: 0,
8226 };
8227 let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8228 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8229 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8230 let __s_b = self.gpu.stream();
8231 let mut b = __s_b.launch_builder(&f);
8232 b.arg(bank)
8233 .arg(sel)
8234 .arg(aq)
8235 .arg(ad)
8236 .arg(route_w)
8237 .arg(md)
8238 .arg(dst)
8239 .arg(&inf)
8240 .arg(&outf)
8241 .arg(&nsc)
8242 .arg(&rb)
8243 .arg(&es)
8244 .arg(&ars)
8245 .arg(&adrs);
8246 unsafe {
8247 b.launch(cfg)?;
8248 }
8249 Ok(())
8250 }
8251
8252 #[allow(clippy::too_many_arguments)]
8255 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8256 &self,
8257 gate_bank: &CudaSlice<u8>,
8258 up_bank: &CudaSlice<u8>,
8259 sel: &CudaSlice<i32>,
8260 aq: &CudaSlice<i8>,
8261 ad: &CudaSlice<f32>,
8262 yg: &mut CudaSlice<f32>,
8263 yu: &mut CudaSlice<f32>,
8264 n_sel: usize,
8265 in_f: usize,
8266 out_f: usize,
8267 row_bytes: usize,
8268 expert_stride: usize,
8269 owner: usize,
8270 ) -> Result<(), Box<dyn std::error::Error>> {
8271 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8272 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8273 return Err("NVFP4 gu ep geometry".into());
8274 }
8275 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8276 let cfg = LaunchConfig {
8277 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8278 block_dim: (128, 1, 1),
8279 shared_mem_bytes: 0,
8280 };
8281 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8282 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8283 let (ars, adrs) = (0i64, 0i64);
8284 let __s_b = self.gpu.stream();
8285 let mut b = __s_b.launch_builder(&f);
8286 b.arg(gate_bank)
8287 .arg(up_bank)
8288 .arg(sel)
8289 .arg(aq)
8290 .arg(ad)
8291 .arg(yg)
8292 .arg(yu)
8293 .arg(&inf)
8294 .arg(&outf)
8295 .arg(&ns)
8296 .arg(&rb)
8297 .arg(&es)
8298 .arg(&ars)
8299 .arg(&adrs)
8300 .arg(&own);
8301 unsafe {
8302 b.launch(cfg)?;
8303 }
8304 Ok(())
8305 }
8306
8307 #[allow(clippy::too_many_arguments)]
8309 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8310 &self,
8311 gate: &CudaSlice<f32>,
8312 up: &CudaSlice<f32>,
8313 gmac: &CudaSlice<f32>,
8314 umac: &CudaSlice<f32>,
8315 sel: &CudaSlice<i32>,
8316 limit: Option<f32>,
8317 out_q: &mut CudaSlice<i8>,
8318 out_d: &mut CudaSlice<f32>,
8319 n_per: usize,
8320 n_sel: usize,
8321 owner: usize,
8322 ) -> Result<(), Box<dyn std::error::Error>> {
8323 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8324 return Err("NVFP4 silu ep geometry".into());
8325 }
8326 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8327 let warps = n_sel * n_per / 32;
8328 let cfg = LaunchConfig {
8329 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8330 block_dim: (128, 1, 1),
8331 shared_mem_bytes: 0,
8332 };
8333 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8334 let (lim, has) = match limit {
8335 Some(l) => (l, 1i32),
8336 None => (0.0f32, 0i32),
8337 };
8338 let __s_b = self.gpu.stream();
8339 let mut b = __s_b.launch_builder(&f);
8340 b.arg(gate)
8341 .arg(up)
8342 .arg(gmac)
8343 .arg(umac)
8344 .arg(sel)
8345 .arg(&lim)
8346 .arg(&has)
8347 .arg(out_q)
8348 .arg(out_d)
8349 .arg(&np)
8350 .arg(&ns)
8351 .arg(&own);
8352 unsafe {
8353 b.launch(cfg)?;
8354 }
8355 Ok(())
8356 }
8357
8358 #[allow(clippy::too_many_arguments)]
8360 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8361 &self,
8362 bank: &CudaSlice<u8>,
8363 sel: &CudaSlice<i32>,
8364 aq: &CudaSlice<i8>,
8365 ad: &CudaSlice<f32>,
8366 route_w: &CudaSlice<f32>,
8367 md: &CudaSlice<f32>,
8368 dst: &mut CudaSlice<f32>,
8369 n_sel: usize,
8370 in_f: usize,
8371 out_f: usize,
8372 row_bytes: usize,
8373 expert_stride: usize,
8374 act_row_stride: usize,
8375 ad_row_stride: usize,
8376 owner: usize,
8377 ) -> Result<(), Box<dyn std::error::Error>> {
8378 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8379 return Err("NVFP4 down8 ep geometry".into());
8380 }
8381 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8382 let cfg = LaunchConfig {
8383 grid_dim: (out_f as u32, 1, 1),
8384 block_dim: (32, n_sel as u32, 1),
8385 shared_mem_bytes: 0,
8386 };
8387 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8388 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8389 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8390 let __s_b = self.gpu.stream();
8391 let mut b = __s_b.launch_builder(&f);
8392 b.arg(bank)
8393 .arg(sel)
8394 .arg(aq)
8395 .arg(ad)
8396 .arg(route_w)
8397 .arg(md)
8398 .arg(dst)
8399 .arg(&inf)
8400 .arg(&outf)
8401 .arg(&ns)
8402 .arg(&rb)
8403 .arg(&es)
8404 .arg(&ars)
8405 .arg(&adrs)
8406 .arg(&own);
8407 unsafe {
8408 b.launch(cfg)?;
8409 }
8410 Ok(())
8411 }
8412
8413 pub fn qmatvec_nvfp4_sel_into(
8414 &self,
8415 bank: &CudaSlice<u8>,
8416 sel: &CudaSlice<i32>,
8417 aq: &CudaSlice<i8>,
8418 ad: &CudaSlice<f32>,
8419 y: &mut CudaSlice<f32>,
8420 n_sel: usize,
8421 in_f: usize,
8422 out_f: usize,
8423 row_bytes: usize,
8424 expert_stride: usize,
8425 act_row_stride: usize,
8426 ad_row_stride: usize,
8427 ) -> Result<(), Box<dyn std::error::Error>> {
8428 assert!(
8429 in_f % 64 == 0,
8430 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8431 );
8432 if y.len() < n_sel * out_f || sel.len() < n_sel {
8433 return Err(format!(
8434 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8435 y.len(),
8436 sel.len()
8437 )
8438 .into());
8439 }
8440 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8447 let mode = *MR.get_or_init(|| {
8448 if crate::tp::nvfp4_bank_v2_on() {
8449 3
8450 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8451 2
8452 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8453 1
8454 } else {
8455 0
8456 }
8457 });
8458 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8459 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8463 let v2s = mode == 3
8464 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8465 && row_bytes % 16 == 0
8466 && in_f <= 4096;
8467 let f = match (mode, v2s) {
8468 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8469 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8470 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8471 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8472 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8473 };
8474 let nsb = in_f >> 5;
8479 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8480 32
8481 } else if mode == 1 {
8482 512
8483 } else {
8484 128
8485 };
8486 let cfg = LaunchConfig {
8487 grid_dim: (
8488 if v2s {
8489 (out_f as u32).div_ceil(8)
8490 } else {
8491 match mode {
8492 2 => (out_f as u32).div_ceil(16),
8493 1 => (out_f as u32).div_ceil(4),
8494 _ => out_f as u32,
8495 }
8496 },
8497 n_sel as u32,
8498 1,
8499 ),
8500 block_dim: (fit_block, 1, 1),
8501 shared_mem_bytes: 0,
8502 };
8503 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8504 let (rb, es, ars, adrs) = (
8505 row_bytes as i64,
8506 expert_stride as i64,
8507 act_row_stride as i64,
8508 ad_row_stride as i64,
8509 );
8510 let __s_b = self.gpu.stream();
8511 let mut b = __s_b.launch_builder(&f);
8512 b.arg(bank)
8513 .arg(sel)
8514 .arg(aq)
8515 .arg(ad)
8516 .arg(y)
8517 .arg(&inf)
8518 .arg(&outf)
8519 .arg(&ns)
8520 .arg(&rb)
8521 .arg(&es)
8522 .arg(&ars)
8523 .arg(&adrs);
8524 unsafe {
8525 b.launch(cfg)?;
8526 }
8527 Ok(())
8528 }
8529
8530 #[allow(clippy::too_many_arguments)]
8535 pub fn silu_mul_scaled_q8_1_sel_into(
8536 &self,
8537 gate: &CudaSlice<f32>,
8538 up: &CudaSlice<f32>,
8539 gmac: &CudaSlice<f32>,
8540 umac: &CudaSlice<f32>,
8541 sel: &CudaSlice<i32>,
8542 limit: Option<f32>,
8543 out_q: &mut CudaSlice<i8>,
8544 out_d: &mut CudaSlice<f32>,
8545 n_per: usize,
8546 n_sel: usize,
8547 ) -> Result<(), Box<dyn std::error::Error>> {
8548 let n = n_per * n_sel;
8549 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8550 return Err(format!(
8551 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8552 out_q.len(),
8553 out_d.len()
8554 )
8555 .into());
8556 }
8557 if let Some(limit) = limit {
8558 if limit <= 1e-6 {
8559 return Err(format!(
8560 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8561 )
8562 .into());
8563 }
8564 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8565 let cfg = LaunchConfig::for_num_elems(n as u32);
8566 let (np, ns) = (n_per as i32, n_sel as i32);
8567 let __s_b = self.gpu.stream();
8568 let mut b = __s_b.launch_builder(&f);
8569 b.arg(gate)
8570 .arg(up)
8571 .arg(gmac)
8572 .arg(umac)
8573 .arg(sel)
8574 .arg(&limit)
8575 .arg(out_q)
8576 .arg(out_d)
8577 .arg(&np)
8578 .arg(&ns);
8579 unsafe {
8580 b.launch(cfg)?;
8581 }
8582 return Ok(());
8583 }
8584 let f = self.func("silu_mul_scaled_q8_1_sel");
8585 let cfg = LaunchConfig::for_num_elems(n as u32);
8586 let (np, ns) = (n_per as i32, n_sel as i32);
8587 let __s_b = self.gpu.stream();
8588 let mut b = __s_b.launch_builder(&f);
8589 b.arg(gate)
8590 .arg(up)
8591 .arg(gmac)
8592 .arg(umac)
8593 .arg(sel)
8594 .arg(out_q)
8595 .arg(out_d)
8596 .arg(&np)
8597 .arg(&ns);
8598 unsafe {
8599 b.launch(cfg)?;
8600 }
8601 Ok(())
8602 }
8603
8604 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8605 Ok(self.gpu.stream().clone_htod(v)?)
8606 }
8607 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8608 Ok(self.gpu.stream().clone_htod(v)?)
8609 }
8610 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8612 Ok(self.gpu.stream().clone_htod(v)?)
8613 }
8614 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8615 Ok(self.gpu.stream().clone_htod(v)?)
8616 }
8617 pub fn dtoh_view(
8619 &self,
8620 d: &cudarc::driver::CudaView<f32>,
8621 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8622 let v = self.gpu.stream().clone_dtoh(d)?;
8623 self.gpu.stream().synchronize()?;
8624 Ok(v)
8625 }
8626 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8627 let v = self.gpu.stream().clone_dtoh(d)?;
8628 self.gpu.stream().synchronize()?;
8629 Ok(v)
8630 }
8631 pub fn dtoh_pair(
8635 &self,
8636 a: &CudaSlice<f32>,
8637 b: &CudaSlice<f32>,
8638 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8639 let av = self.gpu.stream().clone_dtoh(a)?;
8640 let bv = self.gpu.stream().clone_dtoh(b)?;
8641 self.gpu.stream().synchronize()?;
8642 Ok((av, bv))
8643 }
8644 pub fn dtoh_pair_views(
8647 &self,
8648 a: &cudarc::driver::CudaView<f32>,
8649 b: &cudarc::driver::CudaView<f32>,
8650 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8651 let av = self.gpu.stream().clone_dtoh(a)?;
8652 let bv = self.gpu.stream().clone_dtoh(b)?;
8653 self.gpu.stream().synchronize()?;
8654 Ok((av, bv))
8655 }
8656 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8658 let v = self.gpu.stream().clone_dtoh(d)?;
8659 self.gpu.stream().synchronize()?;
8660 Ok(v)
8661 }
8662 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8664 let v = self.gpu.stream().clone_dtoh(d)?;
8665 self.gpu.stream().synchronize()?;
8666 Ok(v)
8667 }
8668 pub fn dtoh_u8_view(
8669 &self,
8670 d: &cudarc::driver::CudaView<u8>,
8671 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8672 let v = self.gpu.stream().clone_dtoh(d)?;
8673 self.gpu.stream().synchronize()?;
8674 Ok(v)
8675 }
8676 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8677 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8678 self.keep_if_capturing(&s);
8679 Ok(s)
8680 }
8681
8682 pub fn prob_of_token_device(
8691 &self,
8692 logits: &CudaSlice<f32>,
8693 tok: &CudaSlice<u32>,
8694 n_vocab: usize,
8695 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8696 let nb = ARGMAX_NB;
8697 let mut part = self.alloc_uninit::<f32>(nb)?;
8698 let mut p = self.alloc_uninit::<f32>(1)?;
8699 let f1 = self.func("prob_of_token_partial_f32");
8700 let cfg1 = LaunchConfig {
8701 grid_dim: (nb as u32, 1, 1),
8702 block_dim: (256, 1, 1),
8703 shared_mem_bytes: 0,
8704 };
8705 let nv = n_vocab as i32;
8706 let __s_b1 = self.gpu.stream();
8707 let mut b1 = __s_b1.launch_builder(&f1);
8708 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8709 unsafe {
8710 b1.launch(cfg1)?;
8711 }
8712 let f2 = self.func("prob_of_token_final_f32");
8713 let cfg2 = LaunchConfig {
8714 grid_dim: (1, 1, 1),
8715 block_dim: (256, 1, 1),
8716 shared_mem_bytes: 0,
8717 };
8718 let nbi = nb as i32;
8719 let __s_b2 = self.gpu.stream();
8720 let mut b2 = __s_b2.launch_builder(&f2);
8721 b2.arg(&part).arg(&mut p).arg(&nbi);
8722 unsafe {
8723 b2.launch(cfg2)?;
8724 }
8725 Ok(p)
8726 }
8727
8728 pub fn prob_of_token_device_col(
8735 &self,
8736 logits: &CudaSlice<f32>,
8737 tok_all: &CudaSlice<u32>,
8738 tok_idx: usize,
8739 p_out: &mut CudaSlice<f32>,
8740 p_idx: usize,
8741 n_vocab: usize,
8742 ) -> Result<(), Box<dyn std::error::Error>> {
8743 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8744 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8745 let nb = ARGMAX_NB;
8746 let mut part = self.alloc_uninit::<f32>(nb)?;
8747 let f1 = self.func("prob_of_token_partial_f32");
8748 let cfg1 = LaunchConfig {
8749 grid_dim: (nb as u32, 1, 1),
8750 block_dim: (256, 1, 1),
8751 shared_mem_bytes: 0,
8752 };
8753 let nv = n_vocab as i32;
8754 let __s_b1 = self.gpu.stream();
8755 let mut b1 = __s_b1.launch_builder(&f1);
8756 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8757 unsafe {
8758 b1.launch(cfg1)?;
8759 }
8760 let f2 = self.func("prob_of_token_final_f32");
8761 let cfg2 = LaunchConfig {
8762 grid_dim: (1, 1, 1),
8763 block_dim: (256, 1, 1),
8764 shared_mem_bytes: 0,
8765 };
8766 let nbi = nb as i32;
8767 let __s_b2 = self.gpu.stream();
8768 let mut b2 = __s_b2.launch_builder(&f2);
8769 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8770 unsafe {
8771 b2.launch(cfg2)?;
8772 }
8773 Ok(())
8774 }
8775
8776 pub fn prob_of_token_device_into(
8777 &self,
8778 logits: &CudaSlice<f32>,
8779 tok: &CudaSlice<u32>,
8780 p_out: &mut CudaSlice<f32>,
8781 n_vocab: usize,
8782 ) -> Result<(), Box<dyn std::error::Error>> {
8783 let nb = ARGMAX_NB;
8784 let mut part = self.alloc_uninit::<f32>(nb)?;
8785 let f1 = self.func("prob_of_token_partial_f32");
8786 let cfg1 = LaunchConfig {
8787 grid_dim: (nb as u32, 1, 1),
8788 block_dim: (256, 1, 1),
8789 shared_mem_bytes: 0,
8790 };
8791 let nv = n_vocab as i32;
8792 let __s_b1 = self.gpu.stream();
8793 let mut b1 = __s_b1.launch_builder(&f1);
8794 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8795 unsafe {
8796 b1.launch(cfg1)?;
8797 }
8798 let f2 = self.func("prob_of_token_final_f32");
8799 let cfg2 = LaunchConfig {
8800 grid_dim: (1, 1, 1),
8801 block_dim: (256, 1, 1),
8802 shared_mem_bytes: 0,
8803 };
8804 let nbi = nb as i32;
8805 let __s_b2 = self.gpu.stream();
8806 let mut b2 = __s_b2.launch_builder(&f2);
8807 b2.arg(&part).arg(p_out).arg(&nbi);
8808 unsafe {
8809 b2.launch(cfg2)?;
8810 }
8811 Ok(())
8812 }
8813
8814 pub fn u32_hist_append(
8817 &self,
8818 tok: &CudaSlice<u32>,
8819 hist: &mut CudaSlice<u32>,
8820 idx: &mut CudaSlice<i32>,
8821 ) -> Result<(), Box<dyn std::error::Error>> {
8822 let f = self.func("u32_hist_append");
8823 let cfg = LaunchConfig {
8824 grid_dim: (1, 1, 1),
8825 block_dim: (32, 1, 1),
8826 shared_mem_bytes: 0,
8827 };
8828 let __s_b = self.gpu.stream();
8829 let mut b = __s_b.launch_builder(&f);
8830 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8831 unsafe {
8832 b.launch(cfg)?;
8833 }
8834 Ok(())
8835 }
8836
8837 pub fn argmax_token_device(
8838 &self,
8839 logits: &CudaSlice<f32>,
8840 n_vocab: usize,
8841 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8842 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8843 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8844 Ok(tok)
8845 }
8846 pub fn argmax_token_device_into(
8853 &self,
8854 logits: &CudaSlice<f32>,
8855 tok: &mut CudaSlice<u32>,
8856 n_vocab: usize,
8857 ) -> Result<(), Box<dyn std::error::Error>> {
8858 let nb = ARGMAX_NB;
8859 let f1 = self.func("argmax_partial_f32");
8860 let f2 = self.func("argmax_final_f32");
8861 let mut guard = self.argmax_partials.lock().unwrap();
8862 if guard.is_none() {
8863 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8866 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8867 *guard = Some((pv, pi));
8868 }
8869 let (part_v, part_i) = guard.as_mut().unwrap();
8870 let nv = n_vocab as i32;
8871 let nbi = nb as i32;
8872 let cfg1 = LaunchConfig {
8874 grid_dim: (nb as u32, 1, 1),
8875 block_dim: (256, 1, 1),
8876 shared_mem_bytes: 0,
8877 };
8878 let __s_b1 = self.gpu.stream();
8879 let mut b1 = __s_b1.launch_builder(&f1);
8880 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8881 unsafe {
8882 b1.launch(cfg1)?;
8883 }
8884 let cfg2 = LaunchConfig {
8886 grid_dim: (1, 1, 1),
8887 block_dim: (256, 1, 1),
8888 shared_mem_bytes: 0,
8889 };
8890 let __s_b2 = self.gpu.stream();
8891 let mut b2 = __s_b2.launch_builder(&f2);
8892 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8893 unsafe {
8894 b2.launch(cfg2)?;
8895 }
8896 Ok(())
8897 }
8898 pub fn argmax_token_device_col(
8904 &self,
8905 logits: &CudaSlice<f32>,
8906 col: usize,
8907 n_vocab: usize,
8908 toks: &mut CudaSlice<u32>,
8909 out_idx: usize,
8910 ) -> Result<(), Box<dyn std::error::Error>> {
8911 let nb = ARGMAX_NB;
8912 let f1 = self.func("argmax_partial_f32");
8913 let f2 = self.func("argmax_final_f32");
8914 let mut guard = self.argmax_partials.lock().unwrap();
8915 if guard.is_none() {
8916 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8917 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8918 *guard = Some((pv, pi));
8919 }
8920 let (part_v, part_i) = guard.as_mut().unwrap();
8921 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8922 let nv = n_vocab as i32;
8923 let nbi = nb as i32;
8924 let cfg1 = LaunchConfig {
8925 grid_dim: (nb as u32, 1, 1),
8926 block_dim: (256, 1, 1),
8927 shared_mem_bytes: 0,
8928 };
8929 let __s_b1 = self.gpu.stream();
8930 let mut b1 = __s_b1.launch_builder(&f1);
8931 b1.arg(&col_view)
8932 .arg(&mut *part_v)
8933 .arg(&mut *part_i)
8934 .arg(&nv);
8935 unsafe {
8936 b1.launch(cfg1)?;
8937 }
8938 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8939 let cfg2 = LaunchConfig {
8940 grid_dim: (1, 1, 1),
8941 block_dim: (256, 1, 1),
8942 shared_mem_bytes: 0,
8943 };
8944 let __s_b2 = self.gpu.stream();
8945 let mut b2 = __s_b2.launch_builder(&f2);
8946 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8947 unsafe {
8948 b2.launch(cfg2)?;
8949 }
8950 Ok(())
8951 }
8952 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8954 Ok(self.gpu.stream().clone_htod(v)?)
8955 }
8956 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8957 let v = self.gpu.stream().clone_dtoh(d)?;
8958 self.gpu.stream().synchronize()?;
8959 Ok(v)
8960 }
8961
8962 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8963 let v = self.gpu.stream().clone_dtoh(d)?;
8964 self.gpu.stream().synchronize()?;
8965 Ok(v)
8966 }
8967 pub fn htod_u32_into(
8971 &self,
8972 dst: &mut CudaSlice<u32>,
8973 src: &[u32],
8974 ) -> Result<(), Box<dyn std::error::Error>> {
8975 let mut view = dst.slice_mut(0..src.len());
8976 self.gpu.stream().memcpy_htod(src, &mut view)?;
8977 Ok(())
8978 }
8979
8980 pub fn htod_i32_into(
8983 &self,
8984 dst: &mut CudaSlice<i32>,
8985 src: &[i32],
8986 ) -> Result<(), Box<dyn std::error::Error>> {
8987 let mut view = dst.slice_mut(0..src.len());
8988 self.gpu.stream().memcpy_htod(src, &mut view)?;
8989 Ok(())
8990 }
8991
8992 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8993 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8994 self.keep_if_capturing(&s);
8995 Ok(s)
8996 }
8997 pub fn embed_gather_device_into(
9000 &self,
9001 embd: &CudaSlice<u8>,
9002 token_d: &CudaSlice<u32>,
9003 x_out: &mut CudaSlice<f32>,
9004 n_embd: usize,
9005 qtype: i32,
9006 row_bytes: usize,
9007 ) -> Result<(), Box<dyn std::error::Error>> {
9008 let f = self.func("embed_gather_u32");
9009 let cfg = LaunchConfig {
9010 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9011 block_dim: (256, 1, 1),
9012 shared_mem_bytes: 0,
9013 };
9014 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9015 let __s_b = self.gpu.stream();
9016 let mut b = __s_b.launch_builder(&f);
9017 b.arg(embd)
9018 .arg(token_d)
9019 .arg(x_out)
9020 .arg(&ne)
9021 .arg(&qt)
9022 .arg(&rb);
9023 unsafe {
9024 b.launch(cfg)?;
9025 }
9026 Ok(())
9027 }
9028 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
9030 let v = self.gpu.stream().clone_dtoh(d)?;
9031 self.gpu.stream().synchronize()?;
9032 Ok(v[0])
9033 }
9034 pub fn i32_set_k(
9041 &self,
9042 dst: &mut CudaSlice<i32>,
9043 v: i32,
9044 ) -> Result<(), Box<dyn std::error::Error>> {
9045 let f = self.func("i32_set_k");
9046 let cfg = LaunchConfig {
9047 grid_dim: (1, 1, 1),
9048 block_dim: (1, 1, 1),
9049 shared_mem_bytes: 0,
9050 };
9051 let idx = 0i32;
9052 let __s_b = self.gpu.stream();
9053 let mut b = __s_b.launch_builder(&f);
9054 b.arg(dst).arg(&v).arg(&idx);
9055 unsafe {
9056 b.launch(cfg)?;
9057 }
9058 Ok(())
9059 }
9060
9061 pub fn set_i32_one(
9062 &self,
9063 d: &mut CudaSlice<i32>,
9064 v: i32,
9065 ) -> Result<(), Box<dyn std::error::Error>> {
9066 self.gpu.stream().memcpy_htod(&[v], d)?;
9067 Ok(())
9068 }
9069 pub fn set_u32_one(
9072 &self,
9073 d: &mut CudaSlice<u32>,
9074 v: u32,
9075 ) -> Result<(), Box<dyn std::error::Error>> {
9076 self.gpu.stream().memcpy_htod(&[v], d)?;
9077 Ok(())
9078 }
9079 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9081 let v = self.gpu.stream().clone_dtoh(d)?;
9082 self.gpu.stream().synchronize()?;
9083 Ok(v[0])
9084 }
9085 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9087 Ok(self.gpu.stream().clone_htod(bytes)?)
9088 }
9089 pub fn embed_gather_device(
9093 &self,
9094 embd: &CudaSlice<u8>,
9095 token_d: &CudaSlice<u32>,
9096 n_embd: usize,
9097 qtype: i32,
9098 row_bytes: usize,
9099 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9100 let f = self.func("embed_gather_u32");
9101 let mut x = self.alloc_uninit::<f32>(n_embd)?;
9102 let cfg = LaunchConfig {
9103 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9104 block_dim: (256, 1, 1),
9105 shared_mem_bytes: 0,
9106 };
9107 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9108 let __s_b = self.gpu.stream();
9109 let mut b = __s_b.launch_builder(&f);
9110 b.arg(embd)
9111 .arg(token_d)
9112 .arg(&mut x)
9113 .arg(&ne)
9114 .arg(&qt)
9115 .arg(&rb);
9116 unsafe {
9117 b.launch(cfg)?;
9118 }
9119 Ok(x)
9120 }
9121
9122 pub fn embed_gather_device_t(
9126 &self,
9127 embd: &CudaSlice<u8>,
9128 tokens: &[u32],
9129 n_embd: usize,
9130 qtype: i32,
9131 row_bytes: usize,
9132 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9133 let t = tokens.len();
9134 let tok_d = self.gpu.stream().clone_htod(tokens)?;
9135 let f = self.func("embed_gather_u32_t");
9136 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9137 let cfg = LaunchConfig {
9138 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9139 block_dim: (256, 1, 1),
9140 shared_mem_bytes: 0,
9141 };
9142 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9143 let __s_b = self.gpu.stream();
9144 let mut b = __s_b.launch_builder(&f);
9145 b.arg(embd)
9146 .arg(&tok_d)
9147 .arg(&mut x)
9148 .arg(&ne)
9149 .arg(&qt)
9150 .arg(&rb)
9151 .arg(&ti);
9152 unsafe {
9153 b.launch(cfg)?;
9154 }
9155 Ok(x)
9156 }
9157
9158 pub fn embed_gather_device_tv(
9163 &self,
9164 embd: &CudaSlice<u8>,
9165 tok_v: &cudarc::driver::CudaView<u32>,
9166 t: usize,
9167 n_embd: usize,
9168 qtype: i32,
9169 row_bytes: usize,
9170 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9171 let f = self.func("embed_gather_u32_t");
9172 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9173 let cfg = LaunchConfig {
9174 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9175 block_dim: (256, 1, 1),
9176 shared_mem_bytes: 0,
9177 };
9178 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9179 let __s_b = self.gpu.stream();
9180 let mut b = __s_b.launch_builder(&f);
9181 b.arg(embd)
9182 .arg(tok_v)
9183 .arg(&mut x)
9184 .arg(&ne)
9185 .arg(&qt)
9186 .arg(&rb)
9187 .arg(&ti);
9188 unsafe {
9189 b.launch(cfg)?;
9190 }
9191 Ok(x)
9192 }
9193
9194 pub fn embed_gather_device_td(
9195 &self,
9196 embd: &CudaSlice<u8>,
9197 tok_d: &CudaSlice<u32>,
9198 t: usize,
9199 n_embd: usize,
9200 qtype: i32,
9201 row_bytes: usize,
9202 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9203 let f = self.func("embed_gather_u32_t");
9204 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9205 let cfg = LaunchConfig {
9206 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9207 block_dim: (256, 1, 1),
9208 shared_mem_bytes: 0,
9209 };
9210 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9211 let __s_b = self.gpu.stream();
9212 let mut b = __s_b.launch_builder(&f);
9213 b.arg(embd)
9214 .arg(tok_d)
9215 .arg(&mut x)
9216 .arg(&ne)
9217 .arg(&qt)
9218 .arg(&rb)
9219 .arg(&ti);
9220 unsafe {
9221 b.launch(cfg)?;
9222 }
9223 Ok(x)
9224 }
9225
9226 #[inline]
9232 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9234 if self
9235 .capture_keep_on
9236 .load(std::sync::atomic::Ordering::Relaxed)
9237 {
9238 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9239 }
9240 }
9241
9242 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9243 &self,
9244 n: usize,
9245 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9246 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9247 {
9251 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9252 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9253 use cudarc::driver::DevicePtrMut;
9255 let n_bytes = s.len() * std::mem::size_of::<T>();
9256 let stream = self.gpu.stream();
9257 let (p_, _g) = s.device_ptr_mut(&stream);
9258 unsafe {
9259 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9260 .result()?;
9261 }
9262 }
9263 }
9264 self.keep_if_capturing(&s);
9265 Ok(s)
9266 }
9267
9268 pub fn uninit_q8_pair(
9273 &self,
9274 n: usize,
9275 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9276 Ok((
9277 self.alloc_uninit::<i8>(n)?,
9278 self.alloc_uninit::<f32>(n / 32)?,
9279 ))
9280 }
9281
9282 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9283 self.alloc_uninit::<f32>(n)
9284 }
9285
9286 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9288 self.alloc_uninit::<i8>(n)
9289 }
9290
9291 #[allow(clippy::too_many_arguments)]
9295 pub fn rms_norm3(
9296 &self,
9297 x: &CudaSlice<f32>,
9298 w0: &CudaSlice<f32>,
9299 w1: &CudaSlice<f32>,
9300 w2: &CudaSlice<f32>,
9301 d0: &mut CudaSlice<f32>,
9302 d1: &mut CudaSlice<f32>,
9303 d2: &mut CudaSlice<f32>,
9304 ncols: usize,
9305 nrows: usize,
9306 eps: f32,
9307 ) -> Result<(), Box<dyn std::error::Error>> {
9308 let f = self.func("rms_norm3_f32");
9309 let cfg = LaunchConfig {
9310 grid_dim: (nrows as u32, 1, 1),
9311 block_dim: (rms_block(), 1, 1),
9312 shared_mem_bytes: 0,
9313 };
9314 let (nc, e) = (ncols as i32, eps);
9315 let __s_b = self.gpu.stream();
9316 let mut b = __s_b.launch_builder(&f);
9317 b.arg(x)
9318 .arg(w0)
9319 .arg(w1)
9320 .arg(w2)
9321 .arg(d0)
9322 .arg(d1)
9323 .arg(d2)
9324 .arg(&nc)
9325 .arg(&e);
9326 unsafe {
9327 b.launch(cfg)?;
9328 }
9329 Ok(())
9330 }
9331
9332 #[allow(clippy::too_many_arguments)]
9334 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9337 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9338 *WARP_ON.get_or_init(|| {
9339 std::env::var("MEMRA_QKVNORM_W")
9340 .map(|v| v != "0")
9341 .unwrap_or(true)
9342 }) && ncols % 4 == 0
9343 && rows >= 64
9344 }
9345
9346 #[allow(clippy::too_many_arguments)]
9349 pub fn rms_norm_qkv_w4b(
9350 &self,
9351 q: &CudaSlice<f32>,
9352 k: &CudaSlice<f32>,
9353 v: &CudaSlice<f32>,
9354 wq: &CudaSlice<f32>,
9355 wk: &CudaSlice<f32>,
9356 wv: &CudaSlice<f32>,
9357 dq: &mut CudaSlice<f32>,
9358 dk: &mut CudaSlice<f32>,
9359 dv: &mut CudaSlice<f32>,
9360 dvb: &mut CudaSlice<u8>,
9361 ncols: usize,
9362 rq: usize,
9363 rk: usize,
9364 eps: f32,
9365 vf16: bool,
9366 ) -> Result<(), Box<dyn std::error::Error>> {
9367 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9368 let f = self.func("rms_norm_qkv_w4b_f32");
9369 let rows = (rq + 2 * rk) as u32;
9370 let cfg = LaunchConfig {
9371 grid_dim: (rows.div_ceil(8), 1, 1),
9372 block_dim: (256, 1, 1),
9373 shared_mem_bytes: 0,
9374 };
9375 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9376 let vf = vf16 as i32;
9377 let __s_b = self.gpu.stream();
9378 let mut b = __s_b.launch_builder(&f);
9379 b.arg(q)
9380 .arg(k)
9381 .arg(v)
9382 .arg(wq)
9383 .arg(wk)
9384 .arg(wv)
9385 .arg(dq)
9386 .arg(dk)
9387 .arg(dv)
9388 .arg(&mut *dvb)
9389 .arg(&nc)
9390 .arg(&rqi)
9391 .arg(&rki)
9392 .arg(&rvi)
9393 .arg(&e)
9394 .arg(&vf);
9395 unsafe {
9396 b.launch(cfg)?;
9397 }
9398 Ok(())
9399 }
9400
9401 pub fn rms_norm_qkv(
9402 &self,
9403 q: &CudaSlice<f32>,
9404 k: &CudaSlice<f32>,
9405 v: &CudaSlice<f32>,
9406 wq: &CudaSlice<f32>,
9407 wk: &CudaSlice<f32>,
9408 wv: &CudaSlice<f32>,
9409 dq: &mut CudaSlice<f32>,
9410 dk: &mut CudaSlice<f32>,
9411 dv: &mut CudaSlice<f32>,
9412 ncols: usize,
9413 rq: usize,
9414 rk: usize,
9415 eps: f32,
9416 ) -> Result<(), Box<dyn std::error::Error>> {
9417 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9421 let warp_on = *WARP_ON.get_or_init(|| {
9422 std::env::var("MEMRA_QKVNORM_W")
9423 .map(|v| v != "0")
9424 .unwrap_or(true)
9425 });
9426 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9429 let f = self.func("rms_norm_qkv_w4_f32");
9430 let rows = (rq + 2 * rk) as u32;
9431 let cfg = LaunchConfig {
9432 grid_dim: (rows.div_ceil(8), 1, 1),
9433 block_dim: (256, 1, 1),
9434 shared_mem_bytes: 0,
9435 };
9436 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9437 let __s_b = self.gpu.stream();
9438 let mut b = __s_b.launch_builder(&f);
9439 b.arg(q)
9440 .arg(k)
9441 .arg(v)
9442 .arg(wq)
9443 .arg(wk)
9444 .arg(wv)
9445 .arg(dq)
9446 .arg(dk)
9447 .arg(dv)
9448 .arg(&nc)
9449 .arg(&rqi)
9450 .arg(&rki)
9451 .arg(&rvi)
9452 .arg(&e);
9453 unsafe {
9454 b.launch(cfg)?;
9455 }
9456 return Ok(());
9457 }
9458 let f = self.func("rms_norm_qkv_f32");
9459 let grid = (rq + 2 * rk) as u32;
9460 let cfg = LaunchConfig {
9461 grid_dim: (grid, 1, 1),
9462 block_dim: (rms_block(), 1, 1),
9463 shared_mem_bytes: 0,
9464 };
9465 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9466 let __s_b = self.gpu.stream();
9467 let mut b = __s_b.launch_builder(&f);
9468 b.arg(q)
9469 .arg(k)
9470 .arg(v)
9471 .arg(wq)
9472 .arg(wk)
9473 .arg(wv)
9474 .arg(dq)
9475 .arg(dk)
9476 .arg(dv)
9477 .arg(&nc)
9478 .arg(&rqi)
9479 .arg(&rki)
9480 .arg(&e);
9481 unsafe {
9482 b.launch(cfg)?;
9483 }
9484 Ok(())
9485 }
9486
9487 #[allow(clippy::too_many_arguments)]
9489 pub fn rms_norm2x(
9490 &self,
9491 a: &CudaSlice<f32>,
9492 bb: &CudaSlice<f32>,
9493 wa: &CudaSlice<f32>,
9494 wb: &CudaSlice<f32>,
9495 da: &mut CudaSlice<f32>,
9496 db: &mut CudaSlice<f32>,
9497 ncols: usize,
9498 nrows: usize,
9499 eps: f32,
9500 ) -> Result<(), Box<dyn std::error::Error>> {
9501 let f = self.func("rms_norm2x_f32");
9502 let cfg = LaunchConfig {
9503 grid_dim: (2 * nrows as u32, 1, 1),
9504 block_dim: (rms_block(), 1, 1),
9505 shared_mem_bytes: 0,
9506 };
9507 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9508 let __s_b = self.gpu.stream();
9509 let mut b = __s_b.launch_builder(&f);
9510 b.arg(a)
9511 .arg(bb)
9512 .arg(wa)
9513 .arg(wb)
9514 .arg(da)
9515 .arg(db)
9516 .arg(&nc)
9517 .arg(&nr)
9518 .arg(&e);
9519 unsafe {
9520 b.launch(cfg)?;
9521 }
9522 Ok(())
9523 }
9524
9525 pub fn softcap(
9527 &self,
9528 y: &mut CudaSlice<f32>,
9529 cap: f32,
9530 n: usize,
9531 ) -> Result<(), Box<dyn std::error::Error>> {
9532 let f = self.func("softcap_f32");
9533 let cfg = LaunchConfig::for_num_elems(n as u32);
9534 let ni = n as i32;
9535 let __s_b = self.gpu.stream();
9536 let mut b = __s_b.launch_builder(&f);
9537 b.arg(y).arg(&cap).arg(&ni);
9538 unsafe {
9539 b.launch(cfg)?;
9540 }
9541 Ok(())
9542 }
9543
9544 pub fn mask_ids_rows(
9547 &self,
9548 y: &mut CudaSlice<f32>,
9549 ids: &CudaSlice<i32>,
9550 n_ids: usize,
9551 n_vocab: usize,
9552 t: usize,
9553 ) -> Result<(), Box<dyn std::error::Error>> {
9554 let f = self.func("mask_ids_rows_f32");
9555 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9556 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9557 let __s_b = self.gpu.stream();
9558 let mut b = __s_b.launch_builder(&f);
9559 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9560 unsafe {
9561 b.launch(cfg)?;
9562 }
9563 Ok(())
9564 }
9565
9566 #[allow(clippy::too_many_arguments)]
9568 pub fn add_scale_rms_norm(
9569 &self,
9570 a: &CudaSlice<f32>,
9571 b_in: &CudaSlice<f32>,
9572 c: f32,
9573 w: &CudaSlice<f32>,
9574 res: &mut CudaSlice<f32>,
9575 dst: &mut CudaSlice<f32>,
9576 ncols: usize,
9577 nrows: usize,
9578 eps: f32,
9579 ) -> Result<(), Box<dyn std::error::Error>> {
9580 let f = self.func("add_scale_rms_norm_f32");
9581 let cfg = LaunchConfig {
9582 grid_dim: (nrows as u32, 1, 1),
9583 block_dim: (rms_block(), 1, 1),
9584 shared_mem_bytes: 0,
9585 };
9586 let (nc, e2) = (ncols as i32, eps);
9587 let __s_b = self.gpu.stream();
9588 let mut b = __s_b.launch_builder(&f);
9589 b.arg(a)
9590 .arg(b_in)
9591 .arg(&c)
9592 .arg(w)
9593 .arg(res)
9594 .arg(dst)
9595 .arg(&nc)
9596 .arg(&e2);
9597 unsafe {
9598 b.launch(cfg)?;
9599 }
9600 Ok(())
9601 }
9602
9603 #[allow(clippy::too_many_arguments)]
9606 pub fn add_scale_rms_norm_q8_1(
9607 &self,
9608 a: &CudaSlice<f32>,
9609 b_in: &CudaSlice<f32>,
9610 c: f32,
9611 w: &CudaSlice<f32>,
9612 res: &mut CudaSlice<f32>,
9613 ncols: usize,
9614 nrows: usize,
9615 eps: f32,
9616 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9617 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9618 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9619 let (nc, e2) = (ncols as i32, eps);
9620 if Self::pdl_on() && Self::pdl_wb_on() {
9621 {
9622 use cudarc::driver::{DevicePtr, DevicePtrMut};
9623 let s = &self.gpu.stream();
9624 let (pa, _g0) = a.device_ptr(s);
9625 let (pb, _g1) = b_in.device_ptr(s);
9626 let (pw, _g2) = w.device_ptr(s);
9627 let (pr, _g3) = res.device_ptr_mut(s);
9628 let (pq, _g4) = out_q.device_ptr_mut(s);
9629 let (pd, _g5) = out_d.device_ptr_mut(s);
9630 let mut ps = [
9631 &pa as *const _ as *mut std::ffi::c_void,
9632 &pb as *const _ as *mut _,
9633 &c as *const _ as *mut _,
9634 &pw as *const _ as *mut _,
9635 &pr as *const _ as *mut _,
9636 &pq as *const _ as *mut _,
9637 &pd as *const _ as *mut _,
9638 &nc as *const _ as *mut _,
9639 &e2 as *const _ as *mut _,
9640 ];
9641 unsafe {
9642 self.launch_pdl(
9643 "add_scale_rms_norm_q8_1",
9644 (nrows as u32, 1, 1),
9645 (rms_block(), 1, 1),
9646 &mut ps,
9647 )?;
9648 }
9649 }
9650 return Ok((out_q, out_d));
9651 }
9652 let f = self.func("add_scale_rms_norm_q8_1");
9653 let cfg = LaunchConfig {
9654 grid_dim: (nrows as u32, 1, 1),
9655 block_dim: (rms_block(), 1, 1),
9656 shared_mem_bytes: 0,
9657 };
9658 let __s_b = self.gpu.stream();
9659 let mut b = __s_b.launch_builder(&f);
9660 b.arg(a)
9661 .arg(b_in)
9662 .arg(&c)
9663 .arg(w)
9664 .arg(res)
9665 .arg(&mut out_q)
9666 .arg(&mut out_d)
9667 .arg(&nc)
9668 .arg(&e2);
9669 unsafe {
9670 b.launch(cfg)?;
9671 }
9672 Ok((out_q, out_d))
9673 }
9674
9675 #[allow(clippy::too_many_arguments)]
9677 pub fn add_scale_rms_norm_q8_1_into(
9678 &self,
9679 a: &CudaSlice<f32>,
9680 b_in: &CudaSlice<f32>,
9681 c: f32,
9682 w: &CudaSlice<f32>,
9683 res: &mut CudaSlice<f32>,
9684 ncols: usize,
9685 nrows: usize,
9686 eps: f32,
9687 out_q: &mut CudaSlice<i8>,
9688 out_d: &mut CudaSlice<f32>,
9689 ) -> Result<(), Box<dyn std::error::Error>> {
9690 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9691 let (nc, e2) = (ncols as i32, eps);
9692 if Self::pdl_on() && Self::pdl_wb_on() {
9693 use cudarc::driver::{DevicePtr, DevicePtrMut};
9694 let s = &self.gpu.stream();
9695 let (pa, _g0) = a.device_ptr(s);
9696 let (pb, _g1) = b_in.device_ptr(s);
9697 let (pw, _g2) = w.device_ptr(s);
9698 let (pr, _g3) = res.device_ptr_mut(s);
9699 let (pq, _g4) = out_q.device_ptr_mut(s);
9700 let (pd, _g5) = out_d.device_ptr_mut(s);
9701 let mut ps = [
9702 &pa as *const _ as *mut std::ffi::c_void,
9703 &pb as *const _ as *mut _,
9704 &c as *const _ as *mut _,
9705 &pw as *const _ as *mut _,
9706 &pr as *const _ as *mut _,
9707 &pq as *const _ as *mut _,
9708 &pd as *const _ as *mut _,
9709 &nc as *const _ as *mut _,
9710 &e2 as *const _ as *mut _,
9711 ];
9712 unsafe {
9713 self.launch_pdl(
9714 "add_scale_rms_norm_q8_1",
9715 (nrows as u32, 1, 1),
9716 (rms_block(), 1, 1),
9717 &mut ps,
9718 )?;
9719 }
9720 return Ok(());
9721 }
9722 let f = self.func("add_scale_rms_norm_q8_1");
9723 let cfg = LaunchConfig {
9724 grid_dim: (nrows as u32, 1, 1),
9725 block_dim: (rms_block(), 1, 1),
9726 shared_mem_bytes: 0,
9727 };
9728 let __s_b = self.gpu.stream();
9729 let mut b = __s_b.launch_builder(&f);
9730 b.arg(a)
9731 .arg(b_in)
9732 .arg(&c)
9733 .arg(w)
9734 .arg(res)
9735 .arg(&mut *out_q)
9736 .arg(&mut *out_d)
9737 .arg(&nc)
9738 .arg(&e2);
9739 unsafe {
9740 b.launch(cfg)?;
9741 }
9742 Ok(())
9743 }
9744
9745 #[allow(clippy::too_many_arguments)]
9748 pub fn rms_pre_add_scale_rms_norm_q8_1(
9749 &self,
9750 a: &CudaSlice<f32>,
9751 wa: &CudaSlice<f32>,
9752 b_in: &CudaSlice<f32>,
9753 c: f32,
9754 w: &CudaSlice<f32>,
9755 res: &mut CudaSlice<f32>,
9756 ncols: usize,
9757 nrows: usize,
9758 eps: f32,
9759 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9760 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9761 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9762 let (nc, e2) = (ncols as i32, eps);
9763 if Self::pdl_on() {
9764 {
9765 use cudarc::driver::{DevicePtr, DevicePtrMut};
9766 let s = &self.gpu.stream();
9767 let (pa, _g0) = a.device_ptr(s);
9768 let (pwa, _g1) = wa.device_ptr(s);
9769 let (pb, _g2) = b_in.device_ptr(s);
9770 let (pw, _g3) = w.device_ptr(s);
9771 let (pr, _g4) = res.device_ptr_mut(s);
9772 let (pq, _g5) = out_q.device_ptr_mut(s);
9773 let (pd, _g6) = out_d.device_ptr_mut(s);
9774 let mut ps = [
9775 &pa as *const _ as *mut std::ffi::c_void,
9776 &pwa as *const _ as *mut _,
9777 &pb as *const _ as *mut _,
9778 &c as *const _ as *mut _,
9779 &pw as *const _ as *mut _,
9780 &pr as *const _ as *mut _,
9781 &pq as *const _ as *mut _,
9782 &pd as *const _ as *mut _,
9783 &nc as *const _ as *mut _,
9784 &e2 as *const _ as *mut _,
9785 ];
9786 unsafe {
9787 self.launch_pdl(
9788 "rms_pre_add_scale_rms_norm_q8_1",
9789 (nrows as u32, 1, 1),
9790 (rms_block(), 1, 1),
9791 &mut ps,
9792 )?;
9793 }
9794 }
9795 return Ok((out_q, out_d));
9796 }
9797 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9798 let cfg = LaunchConfig {
9799 grid_dim: (nrows as u32, 1, 1),
9800 block_dim: (rms_block(), 1, 1),
9801 shared_mem_bytes: 0,
9802 };
9803 let __s_b = self.gpu.stream();
9804 let mut b = __s_b.launch_builder(&f);
9805 b.arg(a)
9806 .arg(wa)
9807 .arg(b_in)
9808 .arg(&c)
9809 .arg(w)
9810 .arg(res)
9811 .arg(&mut out_q)
9812 .arg(&mut out_d)
9813 .arg(&nc)
9814 .arg(&e2);
9815 unsafe {
9816 b.launch(cfg)?;
9817 }
9818 Ok((out_q, out_d))
9819 }
9820
9821 pub fn gelu_tanh_mul_q8_1(
9824 &self,
9825 gate: &CudaSlice<f32>,
9826 up: &cudarc::driver::CudaView<f32>,
9827 act: &mut CudaSlice<f32>,
9828 ncols: usize,
9829 nrows: usize,
9830 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9831 debug_assert!(ncols % 128 == 0);
9832 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9833 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9834 let nc = ncols as i32;
9835 if Self::pdl_on() {
9836 {
9837 use cudarc::driver::{DevicePtr, DevicePtrMut};
9838 let s = &self.gpu.stream();
9839 let (pg, _g0) = gate.device_ptr(s);
9840 let (pu, _g1) = up.device_ptr(s);
9841 let (pact, _g2) = act.device_ptr_mut(s);
9842 let (pq, _g3) = out_q.device_ptr_mut(s);
9843 let (pd, _g4) = out_d.device_ptr_mut(s);
9844 let mut ps = [
9845 &pg as *const _ as *mut std::ffi::c_void,
9846 &pu as *const _ as *mut _,
9847 &pact as *const _ as *mut _,
9848 &pq as *const _ as *mut _,
9849 &pd as *const _ as *mut _,
9850 &nc as *const _ as *mut _,
9851 ];
9852 unsafe {
9853 self.launch_pdl(
9854 "gelu_tanh_mul_q8_1",
9855 (nrows as u32, 1, 1),
9856 (rms_block(), 1, 1),
9857 &mut ps,
9858 )?;
9859 }
9860 }
9861 return Ok((out_q, out_d));
9862 }
9863 let f = self.func("gelu_tanh_mul_q8_1");
9864 let cfg = LaunchConfig {
9865 grid_dim: (nrows as u32, 1, 1),
9866 block_dim: (rms_block(), 1, 1),
9867 shared_mem_bytes: 0,
9868 };
9869 let __s_b = self.gpu.stream();
9870 let mut b = __s_b.launch_builder(&f);
9871 b.arg(gate)
9872 .arg(up)
9873 .arg(act)
9874 .arg(&mut out_q)
9875 .arg(&mut out_d)
9876 .arg(&nc);
9877 unsafe {
9878 b.launch(cfg)?;
9879 }
9880 Ok((out_q, out_d))
9881 }
9882
9883 #[allow(clippy::too_many_arguments)]
9885 pub fn gelu_tanh_mul_q8_1_into(
9886 &self,
9887 gate: &CudaSlice<f32>,
9888 up: &cudarc::driver::CudaView<f32>,
9889 act: &mut CudaSlice<f32>,
9890 ncols: usize,
9891 nrows: usize,
9892 out_q: &mut CudaSlice<i8>,
9893 out_d: &mut CudaSlice<f32>,
9894 ) -> Result<(), Box<dyn std::error::Error>> {
9895 debug_assert!(ncols % 128 == 0);
9896 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9897 let nc = ncols as i32;
9898 if Self::pdl_on() {
9899 use cudarc::driver::{DevicePtr, DevicePtrMut};
9900 let s = &self.gpu.stream();
9901 let (pg, _g0) = gate.device_ptr(s);
9902 let (pu, _g1) = up.device_ptr(s);
9903 let (pact, _g2) = act.device_ptr_mut(s);
9904 let (pq, _g3) = out_q.device_ptr_mut(s);
9905 let (pd, _g4) = out_d.device_ptr_mut(s);
9906 let mut ps = [
9907 &pg as *const _ as *mut std::ffi::c_void,
9908 &pu as *const _ as *mut _,
9909 &pact as *const _ as *mut _,
9910 &pq as *const _ as *mut _,
9911 &pd as *const _ as *mut _,
9912 &nc as *const _ as *mut _,
9913 ];
9914 unsafe {
9915 self.launch_pdl(
9916 "gelu_tanh_mul_q8_1",
9917 (nrows as u32, 1, 1),
9918 (rms_block(), 1, 1),
9919 &mut ps,
9920 )?;
9921 }
9922 return Ok(());
9923 }
9924 let f = self.func("gelu_tanh_mul_q8_1");
9925 let cfg = LaunchConfig {
9926 grid_dim: (nrows as u32, 1, 1),
9927 block_dim: (rms_block(), 1, 1),
9928 shared_mem_bytes: 0,
9929 };
9930 let __s_b = self.gpu.stream();
9931 let mut b = __s_b.launch_builder(&f);
9932 b.arg(gate)
9933 .arg(up)
9934 .arg(&mut *act)
9935 .arg(&mut *out_q)
9936 .arg(&mut *out_d)
9937 .arg(&nc);
9938 unsafe {
9939 b.launch(cfg)?;
9940 }
9941 Ok(())
9942 }
9943
9944 #[allow(clippy::too_many_arguments)]
9946 pub fn add_rms_norm3_q8z(
9947 &self,
9948 a: &CudaSlice<f32>,
9949 b_in: &CudaSlice<f32>,
9950 w0: &CudaSlice<f32>,
9951 w1: &CudaSlice<f32>,
9952 w2: &CudaSlice<f32>,
9953 res: &mut CudaSlice<f32>,
9954 out1: &mut CudaSlice<f32>,
9955 ncols: usize,
9956 nrows: usize,
9957 eps: f32,
9958 ) -> Result<
9959 (
9960 (CudaSlice<i8>, CudaSlice<f32>),
9961 (CudaSlice<i8>, CudaSlice<f32>),
9962 ),
9963 Box<dyn std::error::Error>,
9964 > {
9965 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9966 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9967 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9968 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9969 let f = self.func("add_rms_norm3_q8z_f32");
9970 let cfg = LaunchConfig {
9971 grid_dim: (nrows as u32, 1, 1),
9972 block_dim: (rms_block(), 1, 1),
9973 shared_mem_bytes: 0,
9974 };
9975 let (nc, e2) = (ncols as i32, eps);
9976 let __s_b = self.gpu.stream();
9977 let mut b = __s_b.launch_builder(&f);
9978 b.arg(a)
9979 .arg(b_in)
9980 .arg(w0)
9981 .arg(w1)
9982 .arg(w2)
9983 .arg(res)
9984 .arg(&mut q0)
9985 .arg(&mut d0)
9986 .arg(out1)
9987 .arg(&mut q2)
9988 .arg(&mut d2)
9989 .arg(&nc)
9990 .arg(&e2);
9991 unsafe {
9992 b.launch(cfg)?;
9993 }
9994 Ok(((q0, d0), (q2, d2)))
9995 }
9996
9997 #[allow(clippy::too_many_arguments)]
9999 pub fn add_rms_norm3(
10000 &self,
10001 a: &CudaSlice<f32>,
10002 b_in: &CudaSlice<f32>,
10003 w0: &CudaSlice<f32>,
10004 w1: &CudaSlice<f32>,
10005 w2: &CudaSlice<f32>,
10006 res: &mut CudaSlice<f32>,
10007 d0: &mut CudaSlice<f32>,
10008 d1: &mut CudaSlice<f32>,
10009 d2: &mut CudaSlice<f32>,
10010 ncols: usize,
10011 nrows: usize,
10012 eps: f32,
10013 ) -> Result<(), Box<dyn std::error::Error>> {
10014 let f = self.func("add_rms_norm3_f32");
10015 let cfg = LaunchConfig {
10016 grid_dim: (nrows as u32, 1, 1),
10017 block_dim: (rms_block(), 1, 1),
10018 shared_mem_bytes: 0,
10019 };
10020 let (nc, e2) = (ncols as i32, eps);
10021 let __s_b = self.gpu.stream();
10022 let mut b = __s_b.launch_builder(&f);
10023 b.arg(a)
10024 .arg(b_in)
10025 .arg(w0)
10026 .arg(w1)
10027 .arg(w2)
10028 .arg(res)
10029 .arg(d0)
10030 .arg(d1)
10031 .arg(d2)
10032 .arg(&nc)
10033 .arg(&e2);
10034 unsafe {
10035 b.launch(cfg)?;
10036 }
10037 Ok(())
10038 }
10039
10040 pub fn add_scale(
10042 &self,
10043 a: &CudaSlice<f32>,
10044 b_in: &CudaSlice<f32>,
10045 c: f32,
10046 dst: &mut CudaSlice<f32>,
10047 n: usize,
10048 ) -> Result<(), Box<dyn std::error::Error>> {
10049 let f = self.func("add_scale_f32");
10050 let cfg = LaunchConfig::for_num_elems(n as u32);
10051 let ni = n as i32;
10052 let __s_b = self.gpu.stream();
10053 let mut b = __s_b.launch_builder(&f);
10054 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10055 unsafe {
10056 b.launch(cfg)?;
10057 }
10058 Ok(())
10059 }
10060
10061 pub fn layer_norm_bias(
10063 &self,
10064 x: &CudaSlice<f32>,
10065 w: &CudaSlice<f32>,
10066 b: &CudaSlice<f32>,
10067 dst: &mut CudaSlice<f32>,
10068 ncols: usize,
10069 nrows: usize,
10070 eps: f32,
10071 ) -> Result<(), Box<dyn std::error::Error>> {
10072 let f = self.func("layer_norm_bias_f32");
10073 let (nc, e) = (ncols as i32, eps);
10074 let cfg = LaunchConfig {
10075 grid_dim: (nrows as u32, 1, 1),
10076 block_dim: (256, 1, 1),
10077 shared_mem_bytes: 0,
10078 };
10079 let __s_b = self.gpu.stream();
10080 let mut lb = __s_b.launch_builder(&f);
10081 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10082 unsafe {
10083 lb.launch(cfg)?;
10084 }
10085 Ok(())
10086 }
10087
10088 pub fn gelu_tanh(
10090 &self,
10091 x: &CudaSlice<f32>,
10092 dst: &mut CudaSlice<f32>,
10093 n: usize,
10094 ) -> Result<(), Box<dyn std::error::Error>> {
10095 let f = self.func("gelu_tanh_f32");
10096 let ni = n as i64;
10097 let cfg = LaunchConfig {
10098 grid_dim: (n.div_ceil(256) as u32, 1, 1),
10099 block_dim: (256, 1, 1),
10100 shared_mem_bytes: 0,
10101 };
10102 let __s_b = self.gpu.stream();
10103 let mut lb = __s_b.launch_builder(&f);
10104 lb.arg(x).arg(&mut *dst).arg(&ni);
10105 unsafe {
10106 lb.launch(cfg)?;
10107 }
10108 Ok(())
10109 }
10110
10111 pub fn row_softmax(
10113 &self,
10114 x: &mut CudaSlice<f32>,
10115 ncols: usize,
10116 nrows: usize,
10117 ) -> Result<(), Box<dyn std::error::Error>> {
10118 let f = self.func("row_softmax_f32");
10119 let nc = ncols as i32;
10120 let cfg = LaunchConfig {
10121 grid_dim: (nrows as u32, 1, 1),
10122 block_dim: (256, 1, 1),
10123 shared_mem_bytes: 0,
10124 };
10125 let __s_b = self.gpu.stream();
10126 let mut lb = __s_b.launch_builder(&f);
10127 lb.arg(&mut *x).arg(&nc);
10128 unsafe {
10129 lb.launch(cfg)?;
10130 }
10131 Ok(())
10132 }
10133
10134 pub fn rms_norm(
10135 &self,
10136 x: &CudaSlice<f32>,
10137 w: &CudaSlice<f32>,
10138 dst: &mut CudaSlice<f32>,
10139 ncols: usize,
10140 nrows: usize,
10141 eps: f32,
10142 ) -> Result<(), Box<dyn std::error::Error>> {
10143 let (nc, e) = (ncols as i32, eps);
10144 let kname = if Self::norm_ilp_on() {
10145 "rms_norm_f32_v2"
10146 } else {
10147 "rms_norm_f32"
10148 };
10149 if Self::pdl_on() && Self::pdl_wb_on() {
10150 use cudarc::driver::{DevicePtr, DevicePtrMut};
10151 let s = &self.gpu.stream();
10152 let (px, _g0) = x.device_ptr(s);
10153 let (pw, _g1) = w.device_ptr(s);
10154 let (pd, _g2) = dst.device_ptr_mut(s);
10155 let mut ps = [
10156 &px as *const _ as *mut std::ffi::c_void,
10157 &pw as *const _ as *mut _,
10158 &pd as *const _ as *mut _,
10159 &nc as *const _ as *mut _,
10160 &e as *const _ as *mut _,
10161 ];
10162 unsafe {
10163 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10164 }
10165 return Ok(());
10166 }
10167 let f = self.func(kname);
10168 let cfg = LaunchConfig {
10169 grid_dim: (nrows as u32, 1, 1),
10170 block_dim: (rms_block(), 1, 1),
10171 shared_mem_bytes: 0,
10172 };
10173 let __s_b = self.gpu.stream();
10174 let mut b = __s_b.launch_builder(&f);
10175 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10176 unsafe {
10177 b.launch(cfg)?;
10178 }
10179 Ok(())
10180 }
10181
10182 pub fn rms_norm_decode(
10190 &self,
10191 x: &CudaSlice<f32>,
10192 w: &CudaSlice<f32>,
10193 dst: &mut CudaSlice<f32>,
10194 ncols: usize,
10195 nrows: usize,
10196 eps: f32,
10197 ) -> Result<(), Box<dyn std::error::Error>> {
10198 let f = self.func(if Self::norm_ilp_on() {
10199 "rms_norm_f32_v2"
10200 } else {
10201 "rms_norm_f32"
10202 });
10203 let cfg = LaunchConfig {
10204 grid_dim: (nrows as u32, 1, 1),
10205 block_dim: (1024, 1, 1),
10206 shared_mem_bytes: 0,
10207 };
10208 let (nc, e) = (ncols as i32, eps);
10209 let __s_b = self.gpu.stream();
10210 let mut b = __s_b.launch_builder(&f);
10211 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10212 unsafe {
10213 b.launch(cfg)?;
10214 }
10215 Ok(())
10216 }
10217
10218 pub fn rms_norm_q8_1(
10222 &self,
10223 x: &CudaSlice<f32>,
10224 w: &CudaSlice<f32>,
10225 ncols: usize,
10226 nrows: usize,
10227 eps: f32,
10228 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10229 let nblk = ncols / 32;
10230 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10231 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10232 let (nc, e) = (ncols as i32, eps);
10233 if Self::pdl_on() {
10234 {
10235 use cudarc::driver::{DevicePtr, DevicePtrMut};
10236 let s = &self.gpu.stream();
10237 let (px, _g0) = x.device_ptr(s);
10238 let (pw, _g1) = w.device_ptr(s);
10239 let (pq, _g2) = q.device_ptr_mut(s);
10240 let (pd, _g3) = d.device_ptr_mut(s);
10241 let mut ps = [
10242 &px as *const _ as *mut std::ffi::c_void,
10243 &pw as *const _ as *mut _,
10244 &pq as *const _ as *mut _,
10245 &pd as *const _ as *mut _,
10246 &nc as *const _ as *mut _,
10247 &e as *const _ as *mut _,
10248 ];
10249 unsafe {
10250 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10251 }
10252 }
10253 return Ok((q, d));
10254 }
10255 let f = self.func("rms_norm_q8_1");
10256 let cfg = LaunchConfig {
10259 grid_dim: (nrows as u32, 1, 1),
10260 block_dim: (1024, 1, 1),
10261 shared_mem_bytes: 0,
10262 };
10263 let __s_b = self.gpu.stream();
10264 let mut b = __s_b.launch_builder(&f);
10265 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10266 unsafe {
10267 b.launch(cfg)?;
10268 }
10269 Ok((q, d))
10270 }
10271
10272 pub fn rms_norm_q8_1_into(
10275 &self,
10276 x: &CudaSlice<f32>,
10277 w: &CudaSlice<f32>,
10278 ncols: usize,
10279 nrows: usize,
10280 eps: f32,
10281 q: &mut CudaSlice<i8>,
10282 d: &mut CudaSlice<f32>,
10283 ) -> Result<(), Box<dyn std::error::Error>> {
10284 let nblk = ncols / 32;
10285 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10286 let (nc, e) = (ncols as i32, eps);
10287 if Self::pdl_on() {
10288 use cudarc::driver::{DevicePtr, DevicePtrMut};
10289 let s = &self.gpu.stream();
10290 let (px, _g0) = x.device_ptr(s);
10291 let (pw, _g1) = w.device_ptr(s);
10292 let (pq, _g2) = q.device_ptr_mut(s);
10293 let (pd, _g3) = d.device_ptr_mut(s);
10294 let mut ps = [
10295 &px as *const _ as *mut std::ffi::c_void,
10296 &pw as *const _ as *mut _,
10297 &pq as *const _ as *mut _,
10298 &pd as *const _ as *mut _,
10299 &nc as *const _ as *mut _,
10300 &e as *const _ as *mut _,
10301 ];
10302 unsafe {
10303 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10304 }
10305 return Ok(());
10306 }
10307 let f = self.func("rms_norm_q8_1");
10308 let cfg = LaunchConfig {
10309 grid_dim: (nrows as u32, 1, 1),
10310 block_dim: (1024, 1, 1),
10311 shared_mem_bytes: 0,
10312 };
10313 let __s_b = self.gpu.stream();
10314 let mut b = __s_b.launch_builder(&f);
10315 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10316 unsafe {
10317 b.launch(cfg)?;
10318 }
10319 Ok(())
10320 }
10321
10322 pub fn quantize_q8_1_into(
10324 &self,
10325 x: &CudaSlice<f32>,
10326 m: usize,
10327 in_f: usize,
10328 q: &mut CudaSlice<i8>,
10329 d: &mut CudaSlice<f32>,
10330 ) -> Result<(), Box<dyn std::error::Error>> {
10331 let nblk = in_f / 32;
10332 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10333 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10334 let (inf, mi) = (in_f as i32, m as i32);
10335 if Self::pdl_on() && Self::pdl_wb_on() {
10336 use cudarc::driver::{DevicePtr, DevicePtrMut};
10337 let s = &self.gpu.stream();
10338 let (px, _g0) = x.device_ptr(s);
10339 let (pq, _g1) = q.device_ptr_mut(s);
10340 let (pd, _g2) = d.device_ptr_mut(s);
10341 let mut ps = [
10342 &px as *const _ as *mut std::ffi::c_void,
10343 &pq as *const _ as *mut _,
10344 &pd as *const _ as *mut _,
10345 &inf as *const _ as *mut _,
10346 &mi as *const _ as *mut _,
10347 ];
10348 unsafe {
10349 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10350 }
10351 return Ok(());
10352 }
10353 let f = self.func("quantize_q8_1");
10354 let __s_b = self.gpu.stream();
10355 let mut b = __s_b.launch_builder(&f);
10356 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10357 unsafe {
10358 b.launch(cfg)?;
10359 }
10360 Ok(())
10361 }
10362
10363 pub fn add_rms_norm_q8_1(
10367 &self,
10368 a: &CudaSlice<f32>,
10369 b_in: &CudaSlice<f32>,
10370 w: &CudaSlice<f32>,
10371 res: &mut CudaSlice<f32>,
10372 ncols: usize,
10373 nrows: usize,
10374 eps: f32,
10375 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10376 let nblk = ncols / 32;
10377 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10378 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10379 let f = self.func("add_rms_norm_q8_1");
10380 let cfg = LaunchConfig {
10382 grid_dim: (nrows as u32, 1, 1),
10383 block_dim: (1024, 1, 1),
10384 shared_mem_bytes: 0,
10385 };
10386 let (nc, e) = (ncols as i32, eps);
10387 let __s_bld = self.gpu.stream();
10388 let mut bld = __s_bld.launch_builder(&f);
10389 bld.arg(a)
10390 .arg(b_in)
10391 .arg(w)
10392 .arg(res)
10393 .arg(&mut q)
10394 .arg(&mut d)
10395 .arg(&nc)
10396 .arg(&e);
10397 unsafe {
10398 bld.launch(cfg)?;
10399 }
10400 Ok((q, d))
10401 }
10402
10403 #[allow(clippy::too_many_arguments)]
10409 pub fn join_add_rms_norm_raw(
10410 &self,
10411 a0_raw: u64,
10412 a1_raw: u64,
10413 x: &CudaSlice<f32>,
10414 w: &CudaSlice<f32>,
10415 res: &mut CudaSlice<f32>,
10416 dst: &mut CudaSlice<f32>,
10417 ncols: usize,
10418 eps: f32,
10419 ) -> Result<(), Box<dyn std::error::Error>> {
10420 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10421 return Err("join_add_rms_norm geometry".into());
10422 }
10423 let f = self.func("join_add_rms_norm_f32");
10424 let cfg = LaunchConfig {
10425 grid_dim: (1, 1, 1),
10426 block_dim: (rms_block(), 1, 1),
10427 shared_mem_bytes: 0,
10428 };
10429 let (nc, e) = (ncols as i32, eps);
10430 let __s_b = self.gpu.stream();
10431 let mut b = __s_b.launch_builder(&f);
10432 b.arg(&a0_raw)
10433 .arg(&a1_raw)
10434 .arg(x)
10435 .arg(w)
10436 .arg(&mut *res)
10437 .arg(&mut *dst)
10438 .arg(&nc)
10439 .arg(&e);
10440 unsafe {
10441 b.launch(cfg)?;
10442 }
10443 Ok(())
10444 }
10445
10446 pub fn add_rms_norm(
10447 &self,
10448 a: &CudaSlice<f32>,
10449 b: &CudaSlice<f32>,
10450 w: &CudaSlice<f32>,
10451 res: &mut CudaSlice<f32>,
10452 dst: &mut CudaSlice<f32>,
10453 ncols: usize,
10454 nrows: usize,
10455 eps: f32,
10456 ) -> Result<(), Box<dyn std::error::Error>> {
10457 let (nc, e) = (ncols as i32, eps);
10458 let kname = if Self::norm_ilp_on() {
10459 "add_rms_norm_f32_v2"
10460 } else {
10461 "add_rms_norm_f32"
10462 };
10463 if Self::pdl_on() && Self::pdl_wb_on() {
10464 use cudarc::driver::{DevicePtr, DevicePtrMut};
10465 let s = &self.gpu.stream();
10466 let (pa, _g0) = a.device_ptr(s);
10467 let (pb, _g1) = b.device_ptr(s);
10468 let (pw, _g2) = w.device_ptr(s);
10469 let (pr, _g3) = res.device_ptr_mut(s);
10470 let (pd, _g4) = dst.device_ptr_mut(s);
10471 let mut ps = [
10472 &pa as *const _ as *mut std::ffi::c_void,
10473 &pb as *const _ as *mut _,
10474 &pw as *const _ as *mut _,
10475 &pr as *const _ as *mut _,
10476 &pd as *const _ as *mut _,
10477 &nc as *const _ as *mut _,
10478 &e as *const _ as *mut _,
10479 ];
10480 unsafe {
10481 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10482 }
10483 return Ok(());
10484 }
10485 let f = self.func(kname);
10486 let cfg = LaunchConfig {
10487 grid_dim: (nrows as u32, 1, 1),
10488 block_dim: (rms_block(), 1, 1),
10489 shared_mem_bytes: 0,
10490 };
10491 let __s_b2 = self.gpu.stream();
10492 let mut b2 = __s_b2.launch_builder(&f);
10493 b2.arg(a)
10494 .arg(b)
10495 .arg(w)
10496 .arg(&mut *res)
10497 .arg(&mut *dst)
10498 .arg(&nc)
10499 .arg(&e);
10500 unsafe {
10501 b2.launch(cfg)?;
10502 }
10503 Ok(())
10504 }
10505
10506 #[allow(clippy::too_many_arguments)]
10509 pub fn rms_pre_add_rms_norm(
10510 &self,
10511 a: &CudaSlice<f32>,
10512 wa: &CudaSlice<f32>,
10513 b: &CudaSlice<f32>,
10514 w: &CudaSlice<f32>,
10515 res: &mut CudaSlice<f32>,
10516 dst: &mut CudaSlice<f32>,
10517 ncols: usize,
10518 nrows: usize,
10519 eps: f32,
10520 ) -> Result<(), Box<dyn std::error::Error>> {
10521 let f = self.func("rms_pre_add_rms_norm_f32");
10522 let cfg = LaunchConfig {
10523 grid_dim: (nrows as u32, 1, 1),
10524 block_dim: (rms_block(), 1, 1),
10525 shared_mem_bytes: 0,
10526 };
10527 let (nc, e) = (ncols as i32, eps);
10528 let __s_b2 = self.gpu.stream();
10529 let mut b2 = __s_b2.launch_builder(&f);
10530 b2.arg(a)
10531 .arg(wa)
10532 .arg(b)
10533 .arg(w)
10534 .arg(&mut *res)
10535 .arg(&mut *dst)
10536 .arg(&nc)
10537 .arg(&e);
10538 unsafe {
10539 b2.launch(cfg)?;
10540 }
10541 Ok(())
10542 }
10543
10544 #[allow(clippy::too_many_arguments)]
10546 pub fn rms_pre_add_rms_norm_q8z(
10547 &self,
10548 a: &CudaSlice<f32>,
10549 wa: &CudaSlice<f32>,
10550 b: &CudaSlice<f32>,
10551 w: &CudaSlice<f32>,
10552 res: &mut CudaSlice<f32>,
10553 dst: &mut CudaSlice<f32>,
10554 ncols: usize,
10555 nrows: usize,
10556 eps: f32,
10557 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10558 debug_assert!(ncols % 128 == 0);
10559 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10560 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10561 let (nc, e) = (ncols as i32, eps);
10562 if Self::pdl_on() {
10563 {
10564 use cudarc::driver::{DevicePtr, DevicePtrMut};
10565 let s = &self.gpu.stream();
10566 let (pa, _g0) = a.device_ptr(s);
10567 let (pwa, _g1) = wa.device_ptr(s);
10568 let (pb, _g2) = b.device_ptr(s);
10569 let (pw, _g3) = w.device_ptr(s);
10570 let (pr, _g4) = res.device_ptr_mut(s);
10571 let (pdst, _g5) = dst.device_ptr_mut(s);
10572 let (pq, _g6) = out_q.device_ptr_mut(s);
10573 let (pd, _g7) = out_d.device_ptr_mut(s);
10574 let mut ps = [
10575 &pa as *const _ as *mut std::ffi::c_void,
10576 &pwa as *const _ as *mut _,
10577 &pb as *const _ as *mut _,
10578 &pw as *const _ as *mut _,
10579 &pr as *const _ as *mut _,
10580 &pdst as *const _ as *mut _,
10581 &pq as *const _ as *mut _,
10582 &pd as *const _ as *mut _,
10583 &nc as *const _ as *mut _,
10584 &e as *const _ as *mut _,
10585 ];
10586 unsafe {
10587 self.launch_pdl(
10588 "rms_pre_add_rms_norm_q8z_f32",
10589 (nrows as u32, 1, 1),
10590 (rms_block(), 1, 1),
10591 &mut ps,
10592 )?;
10593 }
10594 }
10595 return Ok((out_q, out_d));
10596 }
10597 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10598 let cfg = LaunchConfig {
10599 grid_dim: (nrows as u32, 1, 1),
10600 block_dim: (rms_block(), 1, 1),
10601 shared_mem_bytes: 0,
10602 };
10603 let __s_b2 = self.gpu.stream();
10604 let mut b2 = __s_b2.launch_builder(&f);
10605 b2.arg(a)
10606 .arg(wa)
10607 .arg(b)
10608 .arg(w)
10609 .arg(&mut *res)
10610 .arg(&mut *dst)
10611 .arg(&mut out_q)
10612 .arg(&mut out_d)
10613 .arg(&nc)
10614 .arg(&e);
10615 unsafe {
10616 b2.launch(cfg)?;
10617 }
10618 Ok((out_q, out_d))
10619 }
10620
10621 #[allow(clippy::too_many_arguments)]
10625 pub fn rms_pre_add_rms_norm_q8z_into(
10626 &self,
10627 a: &CudaSlice<f32>,
10628 wa: &CudaSlice<f32>,
10629 b: &CudaSlice<f32>,
10630 w: &CudaSlice<f32>,
10631 res: &mut CudaSlice<f32>,
10632 dst: &mut CudaSlice<f32>,
10633 ncols: usize,
10634 nrows: usize,
10635 eps: f32,
10636 out_q: &mut CudaSlice<i8>,
10637 out_d: &mut CudaSlice<f32>,
10638 ) -> Result<(), Box<dyn std::error::Error>> {
10639 debug_assert!(ncols % 128 == 0);
10640 let (nc, e) = (ncols as i32, eps);
10641 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10642 let cfg = LaunchConfig {
10643 grid_dim: (nrows as u32, 1, 1),
10644 block_dim: (rms_block(), 1, 1),
10645 shared_mem_bytes: 0,
10646 };
10647 let __s_b = self.gpu.stream();
10648 let mut b2 = __s_b.launch_builder(&f);
10649 b2.arg(a)
10650 .arg(wa)
10651 .arg(b)
10652 .arg(w)
10653 .arg(&mut *res)
10654 .arg(&mut *dst)
10655 .arg(&mut *out_q)
10656 .arg(&mut *out_d)
10657 .arg(&nc)
10658 .arg(&e);
10659 unsafe {
10660 b2.launch(cfg)?;
10661 }
10662 Ok(())
10663 }
10664
10665 #[allow(clippy::too_many_arguments)]
10668 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10669 &self,
10670 a: &CudaSlice<f32>,
10671 wa: &CudaSlice<f32>,
10672 b_in: &CudaSlice<f32>,
10673 c: f32,
10674 w: &CudaSlice<f32>,
10675 res: &mut CudaSlice<f32>,
10676 ncols: usize,
10677 nrows: usize,
10678 eps: f32,
10679 out_q: &mut CudaSlice<i8>,
10680 out_d: &mut CudaSlice<f32>,
10681 ) -> Result<(), Box<dyn std::error::Error>> {
10682 debug_assert!(ncols % 128 == 0);
10683 let (nc, e2) = (ncols as i32, eps);
10684 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10685 let cfg = LaunchConfig {
10686 grid_dim: (nrows as u32, 1, 1),
10687 block_dim: (rms_block(), 1, 1),
10688 shared_mem_bytes: 0,
10689 };
10690 let __s_b = self.gpu.stream();
10691 let mut b2 = __s_b.launch_builder(&f);
10692 b2.arg(a)
10693 .arg(wa)
10694 .arg(b_in)
10695 .arg(&c)
10696 .arg(w)
10697 .arg(&mut *res)
10698 .arg(&mut *out_q)
10699 .arg(&mut *out_d)
10700 .arg(&nc)
10701 .arg(&e2);
10702 unsafe {
10703 b2.launch(cfg)?;
10704 }
10705 Ok(())
10706 }
10707
10708 pub fn g4_pnfold_on() -> bool {
10716 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10717 *ON.get_or_init(|| {
10718 std::env::var("MEMRA_G4_PNFOLD")
10719 .map(|v| v != "0")
10720 .unwrap_or(true)
10721 })
10722 }
10723
10724 pub fn build_q4_out_concat3(
10728 &self,
10729 w0: &crate::model::GpuTensor,
10730 w1: &crate::model::GpuTensor,
10731 w2: &crate::model::GpuTensor,
10732 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10733 use crate::model::GpuTensor;
10734 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10735 match w {
10736 GpuTensor::Quant {
10737 qtype,
10738 row_bytes,
10739 rp,
10740 ..
10741 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10742 _ => None,
10743 }
10744 };
10745 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10746 else {
10747 return Ok(None);
10748 };
10749 if rb0 != rb1
10750 || rb0 != rb2
10751 || w0.in_features() != w1.in_features()
10752 || w0.in_features() != w2.in_features()
10753 {
10754 return Ok(None);
10755 }
10756 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10757 match w {
10758 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10759 _ => unreachable!(),
10760 }
10761 }
10762 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10763 let total = rb0 * (o0 + o1 + o2);
10764 let mut cat = self.alloc_u8(total)?;
10765 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10766 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10767 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10768 Ok(Some(GpuTensor::Quant {
10769 bytes: cat,
10770 qtype: QT_Q4_0,
10771 row_bytes: rb0,
10772 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10773 scale: 1.0,
10774 rp: false,
10775 #[cfg(memra_cutlass)]
10776 cutlass: None,
10777 fp8: None,
10778 blk: None,
10779 rp4: None,
10780 f16: None,
10781 }))
10782 }
10783
10784 fn full_width_rope_only(
10802 kernel: &str,
10803 n_rot: usize,
10804 head_dim: usize,
10805 ) -> Result<(), Box<dyn std::error::Error>> {
10806 if n_rot == head_dim {
10807 return Ok(());
10808 }
10809 Err(format!(
10810 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10811 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10812 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10813 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10814 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10815 )
10816 .into())
10817 }
10818
10819 #[allow(clippy::too_many_arguments)]
10823 pub fn rms_norm_qkv_rope_cat(
10824 &self,
10825 qkv: &CudaSlice<f32>,
10826 wq: &CudaSlice<f32>,
10827 wk: &CudaSlice<f32>,
10828 wv: &CudaSlice<f32>,
10829 q: &mut CudaSlice<f32>,
10830 k: &mut CudaSlice<f32>,
10831 v: &mut CudaSlice<f32>,
10832 head_dim: usize,
10833 n_rot: usize,
10834 rq: usize,
10835 rk: usize,
10836 pos: &CudaSlice<i32>,
10837 nh_q: usize,
10838 nh_k: usize,
10839 base: f32,
10840 freq_scale: f32,
10841 ff: Option<&CudaSlice<f32>>,
10842 eps: f32,
10843 ) -> Result<(), Box<dyn std::error::Error>> {
10844 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10845 let rows = rq + rk + rk;
10846 let theta_scale = base.powf(-2.0 / head_dim as f32);
10847 let (nc, rqi, rki, nhq, nhk) = (
10848 head_dim as i32,
10849 rq as i32,
10850 rk as i32,
10851 nh_q as i32,
10852 nh_k as i32,
10853 );
10854 if Self::pdl_on() {
10855 use cudarc::driver::{DevicePtr, DevicePtrMut};
10856 let s = &self.gpu.stream();
10857 let (pqkv, _g0) = qkv.device_ptr(s);
10858 let (pwq, _g1) = wq.device_ptr(s);
10859 let (pwk, _g2) = wk.device_ptr(s);
10860 let (pwv, _g3) = wv.device_ptr(s);
10861 let (pq, _g4) = q.device_ptr_mut(s);
10862 let (pk, _g5) = k.device_ptr_mut(s);
10863 let (pv, _g6) = v.device_ptr_mut(s);
10864 let (ppos, _g7) = pos.device_ptr(s);
10865 let (pff, _g8) = match ff {
10866 Some(t) => {
10867 let (p, g) = t.device_ptr(s);
10868 (p, Some(g))
10869 }
10870 None => (0, None),
10871 };
10872 let mut ps = [
10873 &pqkv as *const _ as *mut std::ffi::c_void,
10874 &pwq as *const _ as *mut _,
10875 &pwk as *const _ as *mut _,
10876 &pwv as *const _ as *mut _,
10877 &pq as *const _ as *mut _,
10878 &pk as *const _ as *mut _,
10879 &pv as *const _ as *mut _,
10880 &nc as *const _ as *mut _,
10881 &rqi as *const _ as *mut _,
10882 &rki as *const _ as *mut _,
10883 &ppos as *const _ as *mut _,
10884 &nhq as *const _ as *mut _,
10885 &nhk as *const _ as *mut _,
10886 &theta_scale as *const _ as *mut _,
10887 &freq_scale as *const _ as *mut _,
10888 &pff as *const _ as *mut _,
10889 &eps as *const _ as *mut _,
10890 ];
10891 unsafe {
10892 self.launch_pdl(
10893 "rms_norm_qkv_rope_cat_f32",
10894 (rows as u32, 1, 1),
10895 (rms_block(), 1, 1),
10896 &mut ps,
10897 )?;
10898 }
10899 return Ok(());
10900 }
10901 let f = self.func("rms_norm_qkv_rope_cat_f32");
10902 let cfg = LaunchConfig {
10903 grid_dim: (rows as u32, 1, 1),
10904 block_dim: (rms_block(), 1, 1),
10905 shared_mem_bytes: 0,
10906 };
10907 let __s_b = self.gpu.stream();
10908 let mut b = __s_b.launch_builder(&f);
10909 match ff {
10910 Some(t) => {
10911 b.arg(qkv)
10912 .arg(wq)
10913 .arg(wk)
10914 .arg(wv)
10915 .arg(&mut *q)
10916 .arg(&mut *k)
10917 .arg(&mut *v)
10918 .arg(&nc)
10919 .arg(&rqi)
10920 .arg(&rki)
10921 .arg(pos)
10922 .arg(&nhq)
10923 .arg(&nhk)
10924 .arg(&theta_scale)
10925 .arg(&freq_scale)
10926 .arg(t)
10927 .arg(&eps);
10928 unsafe {
10929 b.launch(cfg)?;
10930 }
10931 }
10932 None => {
10933 let null: u64 = 0;
10934 b.arg(qkv)
10935 .arg(wq)
10936 .arg(wk)
10937 .arg(wv)
10938 .arg(&mut *q)
10939 .arg(&mut *k)
10940 .arg(&mut *v)
10941 .arg(&nc)
10942 .arg(&rqi)
10943 .arg(&rki)
10944 .arg(pos)
10945 .arg(&nhq)
10946 .arg(&nhk)
10947 .arg(&theta_scale)
10948 .arg(&freq_scale)
10949 .arg(&null)
10950 .arg(&eps);
10951 unsafe {
10952 b.launch(cfg)?;
10953 }
10954 }
10955 }
10956 Ok(())
10957 }
10958
10959 #[allow(clippy::too_many_arguments)]
10963 pub fn rms_norm_qkv_rope(
10964 &self,
10965 q0: &CudaSlice<f32>,
10966 k0: &CudaSlice<f32>,
10967 v0: &CudaSlice<f32>,
10968 wq: &CudaSlice<f32>,
10969 wk: &CudaSlice<f32>,
10970 wv: &CudaSlice<f32>,
10971 q: &mut CudaSlice<f32>,
10972 k: &mut CudaSlice<f32>,
10973 v: &mut CudaSlice<f32>,
10974 head_dim: usize,
10975 n_rot: usize,
10976 rq: usize,
10977 rk: usize,
10978 pos: &CudaSlice<i32>,
10979 nh_q: usize,
10980 nh_k: usize,
10981 base: f32,
10982 freq_scale: f32,
10983 ff: Option<&CudaSlice<f32>>,
10984 eps: f32,
10985 ) -> Result<(), Box<dyn std::error::Error>> {
10986 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10987 let f = self.func("rms_norm_qkv_rope_f32");
10988 let rows = rq + rk + rk; let cfg = LaunchConfig {
10990 grid_dim: (rows as u32, 1, 1),
10991 block_dim: (rms_block(), 1, 1),
10992 shared_mem_bytes: 0,
10993 };
10994 let theta_scale = base.powf(-2.0 / head_dim as f32);
10995 let (nc, rqi, rki, nhq, nhk) = (
10996 head_dim as i32,
10997 rq as i32,
10998 rk as i32,
10999 nh_q as i32,
11000 nh_k as i32,
11001 );
11002 let __s_b = self.gpu.stream();
11003 let mut b = __s_b.launch_builder(&f);
11004 match ff {
11005 Some(t) => {
11006 b.arg(q0)
11007 .arg(k0)
11008 .arg(v0)
11009 .arg(wq)
11010 .arg(wk)
11011 .arg(wv)
11012 .arg(&mut *q)
11013 .arg(&mut *k)
11014 .arg(&mut *v)
11015 .arg(&nc)
11016 .arg(&rqi)
11017 .arg(&rki)
11018 .arg(pos)
11019 .arg(&nhq)
11020 .arg(&nhk)
11021 .arg(&theta_scale)
11022 .arg(&freq_scale)
11023 .arg(t)
11024 .arg(&eps);
11025 unsafe {
11026 b.launch(cfg)?;
11027 }
11028 }
11029 None => {
11030 let null: u64 = 0;
11031 b.arg(q0)
11032 .arg(k0)
11033 .arg(v0)
11034 .arg(wq)
11035 .arg(wk)
11036 .arg(wv)
11037 .arg(&mut *q)
11038 .arg(&mut *k)
11039 .arg(&mut *v)
11040 .arg(&nc)
11041 .arg(&rqi)
11042 .arg(&rki)
11043 .arg(pos)
11044 .arg(&nhq)
11045 .arg(&nhk)
11046 .arg(&theta_scale)
11047 .arg(&freq_scale)
11048 .arg(&null)
11049 .arg(&eps);
11050 unsafe {
11051 b.launch(cfg)?;
11052 }
11053 }
11054 }
11055 Ok(())
11056 }
11057
11058 #[allow(clippy::too_many_arguments)]
11064 pub fn rms_norm_qkv_rope_append_dc(
11065 &self,
11066 q0: &CudaSlice<f32>,
11067 k0: &CudaSlice<f32>,
11068 v0: &CudaSlice<f32>,
11069 wq: &CudaSlice<f32>,
11070 wk: &CudaSlice<f32>,
11071 wv: &CudaSlice<f32>,
11072 q: &mut CudaSlice<f32>,
11073 k: &mut CudaSlice<f32>,
11074 v: &mut CudaSlice<f32>,
11075 head_dim: usize,
11076 n_rot: usize,
11077 rq: usize,
11078 rk: usize,
11079 pos: &CudaSlice<i32>,
11080 nh_q: usize,
11081 nh_k: usize,
11082 base: f32,
11083 freq_scale: f32,
11084 ff: Option<&CudaSlice<f32>>,
11085 eps: f32,
11086 kc: &mut CudaSlice<u8>,
11087 vc: &mut CudaSlice<u8>,
11088 t_dev: &CudaSlice<i32>,
11089 k_tok_bytes: usize,
11090 v_tok_bytes: usize,
11091 g: bool,
11092 ) -> Result<(), Box<dyn std::error::Error>> {
11093 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11094 let rows = rq + rk + rk;
11095 let theta_scale = base.powf(-2.0 / head_dim as f32);
11096 let (nc, rqi, rki, nhq, nhk) = (
11097 head_dim as i32,
11098 rq as i32,
11099 rk as i32,
11100 nh_q as i32,
11101 nh_k as i32,
11102 );
11103 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11104 if Self::pdl_on() && Self::pdl_wb_on() {
11105 use cudarc::driver::{DevicePtr, DevicePtrMut};
11106 let s = &self.gpu.stream();
11107 let (p0, _a0) = q0.device_ptr(s);
11108 let (p1, _a1) = k0.device_ptr(s);
11109 let (p2, _a2) = v0.device_ptr(s);
11110 let (pwq, _a3) = wq.device_ptr(s);
11111 let (pwk, _a4) = wk.device_ptr(s);
11112 let (pwv, _a5) = wv.device_ptr(s);
11113 let (pq, _a6) = q.device_ptr_mut(s);
11114 let (pk, _a7) = k.device_ptr_mut(s);
11115 let (pv, _a8) = v.device_ptr_mut(s);
11116 let (pp, _a9) = pos.device_ptr(s);
11117 let pff: u64 = match ff {
11118 Some(t) => {
11119 let (p, _gg) = t.device_ptr(s);
11120 p as u64
11121 }
11122 None => 0,
11123 };
11124 let (pkc, _a10) = kc.device_ptr_mut(s);
11125 let (pvc, _a11) = vc.device_ptr_mut(s);
11126 let (pt, _a12) = t_dev.device_ptr(s);
11127 let mut ps = [
11128 &p0 as *const _ as *mut std::ffi::c_void,
11129 &p1 as *const _ as *mut _,
11130 &p2 as *const _ as *mut _,
11131 &pwq as *const _ as *mut _,
11132 &pwk as *const _ as *mut _,
11133 &pwv as *const _ as *mut _,
11134 &pq as *const _ as *mut _,
11135 &pk as *const _ as *mut _,
11136 &pv as *const _ as *mut _,
11137 &nc as *const _ as *mut _,
11138 &rqi as *const _ as *mut _,
11139 &rki as *const _ as *mut _,
11140 &pp as *const _ as *mut _,
11141 &nhq as *const _ as *mut _,
11142 &nhk as *const _ as *mut _,
11143 &theta_scale as *const _ as *mut _,
11144 &freq_scale as *const _ as *mut _,
11145 &pff as *const _ as *mut _,
11146 &eps as *const _ as *mut _,
11147 &pkc as *const _ as *mut _,
11148 &pvc as *const _ as *mut _,
11149 &pt as *const _ as *mut _,
11150 &ktb as *const _ as *mut _,
11151 &vtb as *const _ as *mut _,
11152 ];
11153 unsafe {
11154 self.launch_pdl_flash(
11155 g,
11156 "rms_norm_qkv_rope_append_dc_f32",
11157 (rows as u32, 1, 1),
11158 (rms_block(), 1, 1),
11159 0,
11160 &mut ps,
11161 )?;
11162 }
11163 return Ok(());
11164 }
11165 let f = if g {
11166 self.func_g("rms_norm_qkv_rope_append_dc_f32")
11167 } else {
11168 self.func("rms_norm_qkv_rope_append_dc_f32")
11169 };
11170 let cfg = LaunchConfig {
11171 grid_dim: (rows as u32, 1, 1),
11172 block_dim: (rms_block(), 1, 1),
11173 shared_mem_bytes: 0,
11174 };
11175 let __s_b = self.gpu.stream();
11176 let mut b = __s_b.launch_builder(&f);
11177 match ff {
11178 Some(t) => {
11179 b.arg(q0)
11180 .arg(k0)
11181 .arg(v0)
11182 .arg(wq)
11183 .arg(wk)
11184 .arg(wv)
11185 .arg(&mut *q)
11186 .arg(&mut *k)
11187 .arg(&mut *v)
11188 .arg(&nc)
11189 .arg(&rqi)
11190 .arg(&rki)
11191 .arg(pos)
11192 .arg(&nhq)
11193 .arg(&nhk)
11194 .arg(&theta_scale)
11195 .arg(&freq_scale)
11196 .arg(t)
11197 .arg(&eps)
11198 .arg(&mut *kc)
11199 .arg(&mut *vc)
11200 .arg(t_dev)
11201 .arg(&ktb)
11202 .arg(&vtb);
11203 unsafe {
11204 b.launch(cfg)?;
11205 }
11206 }
11207 None => {
11208 let null: u64 = 0;
11209 b.arg(q0)
11210 .arg(k0)
11211 .arg(v0)
11212 .arg(wq)
11213 .arg(wk)
11214 .arg(wv)
11215 .arg(&mut *q)
11216 .arg(&mut *k)
11217 .arg(&mut *v)
11218 .arg(&nc)
11219 .arg(&rqi)
11220 .arg(&rki)
11221 .arg(pos)
11222 .arg(&nhq)
11223 .arg(&nhk)
11224 .arg(&theta_scale)
11225 .arg(&freq_scale)
11226 .arg(&null)
11227 .arg(&eps)
11228 .arg(&mut *kc)
11229 .arg(&mut *vc)
11230 .arg(t_dev)
11231 .arg(&ktb)
11232 .arg(&vtb);
11233 unsafe {
11234 b.launch(cfg)?;
11235 }
11236 }
11237 }
11238 Ok(())
11239 }
11240
11241 #[allow(clippy::too_many_arguments)]
11249 pub fn rms_norm_qkv_rope_append(
11250 &self,
11251 q0: &CudaSlice<f32>,
11252 k0: &CudaSlice<f32>,
11253 v0: &CudaSlice<f32>,
11254 wq: &CudaSlice<f32>,
11255 wk: &CudaSlice<f32>,
11256 wv: &CudaSlice<f32>,
11257 q: &mut CudaSlice<f32>,
11258 k: &mut CudaSlice<f32>,
11259 v: &mut CudaSlice<f32>,
11260 head_dim: usize,
11261 n_rot: usize,
11262 rq: usize,
11263 rk: usize,
11264 pos: &CudaSlice<i32>,
11265 nh_q: usize,
11266 nh_k: usize,
11267 base: f32,
11268 freq_scale: f32,
11269 ff: Option<&CudaSlice<f32>>,
11270 eps: f32,
11271 kc: &mut CudaSlice<u8>,
11272 vc: &mut CudaSlice<u8>,
11273 t: usize,
11274 k_tok_bytes: usize,
11275 v_tok_bytes: usize,
11276 g: bool,
11277 ) -> Result<(), Box<dyn std::error::Error>> {
11278 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11279 let rows = rq + rk + rk;
11280 let theta_scale = base.powf(-2.0 / head_dim as f32);
11281 let (nc, rqi, rki, nhq, nhk) = (
11282 head_dim as i32,
11283 rq as i32,
11284 rk as i32,
11285 nh_q as i32,
11286 nh_k as i32,
11287 );
11288 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11289 let ti = t as i32;
11290 if Self::pdl_on() && Self::pdl_wb_on() {
11291 use cudarc::driver::{DevicePtr, DevicePtrMut};
11292 let s = &self.gpu.stream();
11293 let (p0, _a0) = q0.device_ptr(s);
11294 let (p1, _a1) = k0.device_ptr(s);
11295 let (p2, _a2) = v0.device_ptr(s);
11296 let (pwq, _a3) = wq.device_ptr(s);
11297 let (pwk, _a4) = wk.device_ptr(s);
11298 let (pwv, _a5) = wv.device_ptr(s);
11299 let (pq, _a6) = q.device_ptr_mut(s);
11300 let (pk, _a7) = k.device_ptr_mut(s);
11301 let (pv, _a8) = v.device_ptr_mut(s);
11302 let (pp, _a9) = pos.device_ptr(s);
11303 let pff: u64 = match ff {
11304 Some(t) => {
11305 let (p, _gg) = t.device_ptr(s);
11306 p as u64
11307 }
11308 None => 0,
11309 };
11310 let (pkc, _a10) = kc.device_ptr_mut(s);
11311 let (pvc, _a11) = vc.device_ptr_mut(s);
11312 let mut ps = [
11313 &p0 as *const _ as *mut std::ffi::c_void,
11314 &p1 as *const _ as *mut _,
11315 &p2 as *const _ as *mut _,
11316 &pwq as *const _ as *mut _,
11317 &pwk as *const _ as *mut _,
11318 &pwv as *const _ as *mut _,
11319 &pq as *const _ as *mut _,
11320 &pk as *const _ as *mut _,
11321 &pv as *const _ as *mut _,
11322 &nc as *const _ as *mut _,
11323 &rqi as *const _ as *mut _,
11324 &rki as *const _ as *mut _,
11325 &pp as *const _ as *mut _,
11326 &nhq as *const _ as *mut _,
11327 &nhk as *const _ as *mut _,
11328 &theta_scale as *const _ as *mut _,
11329 &freq_scale as *const _ as *mut _,
11330 &pff as *const _ as *mut _,
11331 &eps as *const _ as *mut _,
11332 &pkc as *const _ as *mut _,
11333 &pvc as *const _ as *mut _,
11334 &ti as *const _ as *mut _,
11335 &ktb as *const _ as *mut _,
11336 &vtb as *const _ as *mut _,
11337 ];
11338 unsafe {
11339 self.launch_pdl_flash(
11340 g,
11341 "rms_norm_qkv_rope_append_f32",
11342 (rows as u32, 1, 1),
11343 (rms_block(), 1, 1),
11344 0,
11345 &mut ps,
11346 )?;
11347 }
11348 return Ok(());
11349 }
11350 let f = if g {
11351 self.func_g("rms_norm_qkv_rope_append_f32")
11352 } else {
11353 self.func("rms_norm_qkv_rope_append_f32")
11354 };
11355 let cfg = LaunchConfig {
11356 grid_dim: (rows as u32, 1, 1),
11357 block_dim: (rms_block(), 1, 1),
11358 shared_mem_bytes: 0,
11359 };
11360 let __s_b = self.gpu.stream();
11361 let mut b = __s_b.launch_builder(&f);
11362 let null: u64 = 0;
11363 b.arg(q0)
11364 .arg(k0)
11365 .arg(v0)
11366 .arg(wq)
11367 .arg(wk)
11368 .arg(wv)
11369 .arg(&mut *q)
11370 .arg(&mut *k)
11371 .arg(&mut *v)
11372 .arg(&nc)
11373 .arg(&rqi)
11374 .arg(&rki)
11375 .arg(pos)
11376 .arg(&nhq)
11377 .arg(&nhk)
11378 .arg(&theta_scale)
11379 .arg(&freq_scale);
11380 match ff {
11381 Some(t) => {
11382 b.arg(t);
11383 }
11384 None => {
11385 b.arg(&null);
11386 }
11387 }
11388 b.arg(&eps)
11389 .arg(&mut *kc)
11390 .arg(&mut *vc)
11391 .arg(&ti)
11392 .arg(&ktb)
11393 .arg(&vtb);
11394 unsafe {
11395 b.launch(cfg)?;
11396 }
11397 Ok(())
11398 }
11399
11400 pub fn add_q8_1(
11401 &self,
11402 a: &CudaSlice<f32>,
11403 b: &CudaSlice<f32>,
11404 res: &mut CudaSlice<f32>,
11405 ncols: usize,
11406 nrows: usize,
11407 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11408 debug_assert!(ncols % 128 == 0);
11409 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11410 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11411 let f = self.func("add_q8_1_f32");
11412 let cfg = LaunchConfig {
11413 grid_dim: (nrows as u32, 1, 1),
11414 block_dim: (rms_block(), 1, 1),
11415 shared_mem_bytes: 0,
11416 };
11417 let nc = ncols as i32;
11418 let __s_b2 = self.gpu.stream();
11419 let mut b2 = __s_b2.launch_builder(&f);
11420 b2.arg(a)
11421 .arg(b)
11422 .arg(&mut *res)
11423 .arg(&mut out_q)
11424 .arg(&mut out_d)
11425 .arg(&nc);
11426 unsafe {
11427 b2.launch(cfg)?;
11428 }
11429 Ok((out_q, out_d))
11430 }
11431
11432 pub fn rms_pre_add_q8_1(
11436 &self,
11437 a: &CudaSlice<f32>,
11438 wa: &CudaSlice<f32>,
11439 b: &CudaSlice<f32>,
11440 res: &mut CudaSlice<f32>,
11441 ncols: usize,
11442 nrows: usize,
11443 eps: f32,
11444 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11445 debug_assert!(ncols % 128 == 0);
11446 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11447 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11448 let f = self.func("rms_pre_add_q8_1_f32");
11449 let cfg = LaunchConfig {
11450 grid_dim: (nrows as u32, 1, 1),
11451 block_dim: (rms_block(), 1, 1),
11452 shared_mem_bytes: 0,
11453 };
11454 let (nc, ep) = (ncols as i32, eps);
11455 let __s_b2 = self.gpu.stream();
11456 let mut b2 = __s_b2.launch_builder(&f);
11457 b2.arg(a)
11458 .arg(wa)
11459 .arg(b)
11460 .arg(&mut *res)
11461 .arg(&mut out_q)
11462 .arg(&mut out_d)
11463 .arg(&nc)
11464 .arg(&ep);
11465 unsafe {
11466 b2.launch(cfg)?;
11467 }
11468 Ok((out_q, out_d))
11469 }
11470
11471 pub fn l2_v2_on(ncols: usize) -> bool {
11475 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11476 }
11477
11478 pub fn l2_norm_pp(
11479 &self,
11480 x: &CudaSlice<f32>,
11481 dst: &mut CudaSlice<f32>,
11482 dst16: Option<&mut CudaSlice<u8>>,
11483 ncols: usize,
11484 nrows: usize,
11485 eps: f32,
11486 ) -> Result<(), Box<dyn std::error::Error>> {
11487 if Self::l2_v2_on(ncols) {
11488 let f = self.func("l2_norm_pp_v2_f32");
11489 let rows_per_block = 8u32; let cfg = LaunchConfig {
11491 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11492 block_dim: (256, 1, 1),
11493 shared_mem_bytes: 0,
11494 };
11495 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11496 let d16: u64 = match dst16 {
11498 Some(d) => self.addr_u8(d),
11499 None => 0,
11500 };
11501 let __s_b = self.gpu.stream();
11502 let mut b = __s_b.launch_builder(&f);
11503 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11504 unsafe {
11505 b.launch(cfg)?;
11506 }
11507 return Ok(());
11508 }
11509 self.l2_norm(x, dst, ncols, nrows, eps)
11510 }
11511
11512 pub fn l2_norm(
11513 &self,
11514 x: &CudaSlice<f32>,
11515 dst: &mut CudaSlice<f32>,
11516 ncols: usize,
11517 nrows: usize,
11518 eps: f32,
11519 ) -> Result<(), Box<dyn std::error::Error>> {
11520 let f = self.func("l2_norm_f32");
11521 let cfg = LaunchConfig {
11522 grid_dim: (nrows as u32, 1, 1),
11523 block_dim: (256, 1, 1),
11524 shared_mem_bytes: 0,
11525 };
11526 let (nc, e) = (ncols as i32, eps);
11527 let __s_b = self.gpu.stream();
11528 let mut b = __s_b.launch_builder(&f);
11529 b.arg(x).arg(dst).arg(&nc).arg(&e);
11530 unsafe {
11531 b.launch(cfg)?;
11532 }
11533 Ok(())
11534 }
11535
11536 pub fn l2_norm_decode(
11542 &self,
11543 x: &CudaSlice<f32>,
11544 dst: &mut CudaSlice<f32>,
11545 ncols: usize,
11546 nrows: usize,
11547 eps: f32,
11548 ) -> Result<(), Box<dyn std::error::Error>> {
11549 let f = self.func("l2_norm_f32");
11550 let cfg = LaunchConfig {
11551 grid_dim: (nrows as u32, 1, 1),
11552 block_dim: (32, 1, 1),
11553 shared_mem_bytes: 0,
11554 };
11555 let (nc, e) = (ncols as i32, eps);
11556 let __s_b = self.gpu.stream();
11557 let mut b = __s_b.launch_builder(&f);
11558 b.arg(x).arg(dst).arg(&nc).arg(&e);
11559 unsafe {
11560 b.launch(cfg)?;
11561 }
11562 Ok(())
11563 }
11564
11565 pub fn rope_neox(
11567 &self,
11568 x: &mut CudaSlice<f32>,
11569 pos: &CudaSlice<i32>,
11570 head_dim: usize,
11571 n_dims: usize,
11572 n_heads: usize,
11573 n_tokens: usize,
11574 freq_base: f32,
11575 freq_scale: f32,
11576 ) -> Result<(), Box<dyn std::error::Error>> {
11577 let f = self.func("rope_neox_f32");
11578 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11579 let grid = (n_heads * n_tokens) as u32;
11580 let cfg = LaunchConfig {
11581 grid_dim: (grid, 1, 1),
11582 block_dim: ((head_dim / 2) as u32, 1, 1),
11583 shared_mem_bytes: 0,
11584 };
11585 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11586 let __s_b = self.gpu.stream();
11587 let mut b = __s_b.launch_builder(&f);
11588 b.arg(x)
11589 .arg(pos)
11590 .arg(&hd)
11591 .arg(&nd)
11592 .arg(&nh)
11593 .arg(&theta_scale)
11594 .arg(&freq_scale);
11595 unsafe {
11596 b.launch(cfg)?;
11597 }
11598 Ok(())
11599 }
11600
11601 pub fn rope_neox_ff(
11603 &self,
11604 x: &mut CudaSlice<f32>,
11605 pos: &CudaSlice<i32>,
11606 head_dim: usize,
11607 n_dims: usize,
11608 n_heads: usize,
11609 n_tokens: usize,
11610 freq_base: f32,
11611 freq_scale: f32,
11612 ff: &CudaSlice<f32>,
11613 ) -> Result<(), Box<dyn std::error::Error>> {
11614 let f = self.func("rope_neox_ff_f32");
11615 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11616 let grid = (n_heads * 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, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11623 let __s_b = self.gpu.stream();
11624 let mut b = __s_b.launch_builder(&f);
11625 b.arg(x)
11626 .arg(pos)
11627 .arg(&hd)
11628 .arg(&nd)
11629 .arg(&nh)
11630 .arg(&theta_scale)
11631 .arg(&freq_scale)
11632 .arg(ff);
11633 unsafe {
11634 b.launch(cfg)?;
11635 }
11636 Ok(())
11637 }
11638
11639 #[allow(clippy::too_many_arguments)]
11641 pub fn rope_neox2(
11642 &self,
11643 q: &mut CudaSlice<f32>,
11644 k: &mut CudaSlice<f32>,
11645 pos: &CudaSlice<i32>,
11646 head_dim: usize,
11647 n_dims: usize,
11648 nh_q: usize,
11649 nh_k: usize,
11650 n_tokens: usize,
11651 freq_base: f32,
11652 freq_scale: f32,
11653 ff: Option<&CudaSlice<f32>>,
11654 ) -> Result<(), Box<dyn std::error::Error>> {
11655 let f = self.func("rope_neox2_f32");
11656 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11657 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11658 let cfg = LaunchConfig {
11659 grid_dim: (grid, 1, 1),
11660 block_dim: ((head_dim / 2) as u32, 1, 1),
11661 shared_mem_bytes: 0,
11662 };
11663 let (hd, nd, nq, nk, nt) = (
11664 head_dim as i32,
11665 n_dims as i32,
11666 nh_q as i32,
11667 nh_k as i32,
11668 n_tokens as i32,
11669 );
11670 let __s_b = self.gpu.stream();
11671 let mut b = __s_b.launch_builder(&f);
11672 b.arg(q)
11673 .arg(k)
11674 .arg(pos)
11675 .arg(&hd)
11676 .arg(&nd)
11677 .arg(&nq)
11678 .arg(&nk)
11679 .arg(&nt)
11680 .arg(&theta_scale)
11681 .arg(&freq_scale);
11682 match ff {
11683 Some(ffv) => {
11684 b.arg(ffv);
11685 unsafe {
11686 b.launch(cfg)?;
11687 }
11688 }
11689 None => {
11690 let null: u64 = 0;
11691 b.arg(&null);
11692 unsafe {
11693 b.launch(cfg)?;
11694 }
11695 }
11696 }
11697 Ok(())
11698 }
11699
11700 pub fn gelu_tanh_mul(
11702 &self,
11703 gate: &CudaSlice<f32>,
11704 up: &CudaSlice<f32>,
11705 dst: &mut CudaSlice<f32>,
11706 n: usize,
11707 ) -> Result<(), Box<dyn std::error::Error>> {
11708 let f = self.func("gelu_tanh_mul_f32");
11709 let cfg = LaunchConfig::for_num_elems(n as u32);
11710 let ni = n as i32;
11711 let __s_b = self.gpu.stream();
11712 let mut b = __s_b.launch_builder(&f);
11713 b.arg(gate).arg(up).arg(dst).arg(&ni);
11714 unsafe {
11715 b.launch(cfg)?;
11716 }
11717 Ok(())
11718 }
11719
11720 pub fn silu_mul(
11721 &self,
11722 gate: &CudaSlice<f32>,
11723 up: &CudaSlice<f32>,
11724 dst: &mut CudaSlice<f32>,
11725 n: usize,
11726 ) -> Result<(), Box<dyn std::error::Error>> {
11727 let f = self.func("silu_mul_f32");
11728 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11730 let ni = n as i32;
11731 let __s_b = self.gpu.stream();
11732 let mut b = __s_b.launch_builder(&f);
11733 b.arg(gate).arg(up).arg(dst).arg(&ni);
11734 unsafe {
11735 b.launch(cfg)?;
11736 }
11737 Ok(())
11738 }
11739
11740 pub fn silu_mul_host_expf(
11742 &self,
11743 gate: &CudaSlice<f32>,
11744 up: &CudaSlice<f32>,
11745 dst: &mut CudaSlice<f32>,
11746 n: usize,
11747 ) -> Result<(), Box<dyn std::error::Error>> {
11748 let f = self.func("silu_mul_host_expf_f32");
11749 let cfg = LaunchConfig::for_num_elems(n as u32);
11750 let ni = n as i32;
11751 let __s_b = self.gpu.stream();
11752 let mut b = __s_b.launch_builder(&f);
11753 b.arg(gate).arg(up).arg(dst).arg(&ni);
11754 unsafe {
11755 b.launch(cfg)?;
11756 }
11757 Ok(())
11758 }
11759
11760 pub fn silu_clamped_mul_host_expf(
11762 &self,
11763 gate: &CudaSlice<f32>,
11764 up: &CudaSlice<f32>,
11765 limit: f32,
11766 dst: &mut CudaSlice<f32>,
11767 n: usize,
11768 ) -> Result<(), Box<dyn std::error::Error>> {
11769 if !limit.is_finite() || limit <= 0.0 {
11770 return Err(
11771 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11772 );
11773 }
11774 let f = self.func("silu_clamped_mul_host_expf_f32");
11775 let cfg = LaunchConfig::for_num_elems(n as u32);
11776 let ni = n as i32;
11777 let __s_b = self.gpu.stream();
11778 let mut b = __s_b.launch_builder(&f);
11779 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11780 unsafe {
11781 b.launch(cfg)?;
11782 }
11783 Ok(())
11784 }
11785
11786 pub fn silu_mul_f16out(
11789 &self,
11790 gate: &CudaSlice<f32>,
11791 up: &CudaSlice<f32>,
11792 dst: &mut CudaSlice<f32>,
11793 dst16: &mut CudaSlice<u8>,
11794 n: usize,
11795 ) -> Result<(), Box<dyn std::error::Error>> {
11796 let f = self.func("silu_mul_f16out_f32");
11797 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11798 let ni = n as i32;
11799 let __s_b = self.gpu.stream();
11800 let mut b = __s_b.launch_builder(&f);
11801 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11802 unsafe {
11803 b.launch(cfg)?;
11804 }
11805 Ok(())
11806 }
11807
11808 pub fn silu_mul_scaled(
11815 &self,
11816 gate: &CudaSlice<f32>,
11817 up: &CudaSlice<f32>,
11818 gs: f32,
11819 us: f32,
11820 dst: &mut CudaSlice<f32>,
11821 n: usize,
11822 ) -> Result<(), Box<dyn std::error::Error>> {
11823 let f = self.func("silu_mul_scaled_f32");
11824 let cfg = LaunchConfig::for_num_elems(n as u32);
11825 let ni = n as i32;
11826 let (gsf, usf) = (gs, us);
11827 let __s_b = self.gpu.stream();
11828 let mut b = __s_b.launch_builder(&f);
11829 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11830 unsafe {
11831 b.launch(cfg)?;
11832 }
11833 Ok(())
11834 }
11835
11836 #[allow(clippy::too_many_arguments)]
11840 pub fn swigluoai_mul_scaled(
11841 &self,
11842 gate: &CudaSlice<f32>,
11843 up: &CudaSlice<f32>,
11844 gs: f32,
11845 us: f32,
11846 alpha: f32,
11847 limit: f32,
11848 dst: &mut CudaSlice<f32>,
11849 n: usize,
11850 ) -> Result<(), Box<dyn std::error::Error>> {
11851 let f = self.func("swigluoai_mul_scaled_f32");
11852 let cfg = LaunchConfig::for_num_elems(n as u32);
11853 let ni = n as i32;
11854 let __s_b = self.gpu.stream();
11855 let mut b = __s_b.launch_builder(&f);
11856 b.arg(gate)
11857 .arg(up)
11858 .arg(&gs)
11859 .arg(&us)
11860 .arg(&alpha)
11861 .arg(&limit)
11862 .arg(dst)
11863 .arg(&ni);
11864 unsafe {
11865 b.launch(cfg)?;
11866 }
11867 Ok(())
11868 }
11869
11870 pub fn silu_mul_scaled_q8_1(
11878 &self,
11879 gate: &CudaSlice<f32>,
11880 up: &CudaSlice<f32>,
11881 gs: f32,
11882 us: f32,
11883 n: usize,
11884 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11885 let f = self.func("silu_mul_scaled_q8_1");
11886 let nblk = n / 32;
11887 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);
11891 let (gsf, usf, ni) = (gs, us, n as i32);
11892 let __s_b = self.gpu.stream();
11893 let mut b = __s_b.launch_builder(&f);
11894 b.arg(gate)
11895 .arg(up)
11896 .arg(&gsf)
11897 .arg(&usf)
11898 .arg(&mut aq)
11899 .arg(&mut ad)
11900 .arg(&ni);
11901 unsafe {
11902 b.launch(cfg)?;
11903 }
11904 Ok((aq, ad))
11905 }
11906
11907 pub fn add(
11908 &self,
11909 a: &CudaSlice<f32>,
11910 b_in: &CudaSlice<f32>,
11911 dst: &mut CudaSlice<f32>,
11912 n: usize,
11913 ) -> Result<(), Box<dyn std::error::Error>> {
11914 let f = self.func("add_f32");
11915 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11917 let ni = n as i32;
11918 let __s_bld = self.gpu.stream();
11919 let mut bld = __s_bld.launch_builder(&f);
11920 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11921 unsafe {
11922 bld.launch(cfg)?;
11923 }
11924 Ok(())
11925 }
11926
11927 pub fn mul(
11928 &self,
11929 a: &CudaSlice<f32>,
11930 b_in: &CudaSlice<f32>,
11931 dst: &mut CudaSlice<f32>,
11932 n: usize,
11933 ) -> Result<(), Box<dyn std::error::Error>> {
11934 let f = self.func("mul_f32");
11935 let cfg = LaunchConfig::for_num_elems(n as u32);
11936 let ni = n as i32;
11937 let __s_bld = self.gpu.stream();
11938 let mut bld = __s_bld.launch_builder(&f);
11939 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11940 unsafe {
11941 bld.launch(cfg)?;
11942 }
11943 Ok(())
11944 }
11945
11946 pub fn matmul(
11949 &self,
11950 w: &crate::model::GpuTensor,
11951 x: &CudaSlice<f32>,
11952 m: usize,
11953 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11954 use crate::model::GpuTensor;
11955 let in_f = w.in_features();
11956 let out_f = w.out_features();
11957 #[allow(non_snake_case)]
11965 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11968 usize::MAX
11969 } else {
11970 16usize
11971 };
11972
11973 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11998 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11999 return Ok(y);
12000 }
12001 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
12008 return Ok(y);
12009 }
12010 if let Some(y) = self.try_f16_gemm(w, x, m)? {
12013 return Ok(y);
12014 }
12015 }
12016 if let GpuTensor::Quant { qtype, .. } = w {
12031 if *qtype == QT_F8_E4M3_BLK {
12032 if m >= GEMM_M_THRESHOLD {
12033 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
12034 return Ok(y);
12035 }
12036 }
12037 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12038 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12039 return Ok(y);
12040 }
12041 }
12042 }
12043 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12044 return self.qmatvec_mmq(w, x, m);
12045 }
12046 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12047 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12048 return self.qmatvec_gemm(w, &aq, &ad, m);
12049 }
12050 if m >= GEMM_M_THRESHOLD {
12053 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12054 return Ok(y);
12055 }
12056 }
12057 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12061 if m == 1 && fast {
12066 if let GpuTensor::Quant {
12067 bytes,
12068 qtype,
12069 row_bytes,
12070 rp,
12071 rp4,
12072 scale,
12073 ..
12074 } = w
12075 {
12076 if self.mmvq_supports(*qtype) {
12077 let (bytes, rp) = match rp4 {
12081 Some(m4) => (m4, true),
12082 None => (bytes, *rp),
12083 };
12084 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12085 return self.qmatvec_mmvq(
12086 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12087 );
12088 }
12089 }
12090 }
12091 if (2..=16).contains(&m)
12107 && fast
12108 && std::env::var("MEMRA_NO_BATCHED").is_err()
12109 && (m <= 4 || Self::b8_enabled())
12110 {
12111 let m_ok = m <= 8
12121 || matches!(w, GpuTensor::Quant { qtype, .. }
12122 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12123 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12124 if m_ok {
12125 if let GpuTensor::Quant {
12126 bytes,
12127 qtype,
12128 row_bytes,
12129 rp,
12130 rp4,
12131 ..
12132 } = w
12133 {
12134 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12135 let (bytes, rp) = match rp4 {
12136 Some(m4) => (m4, true),
12137 None => (bytes, *rp),
12138 };
12139 let mcols = Self::batched_mcols(m);
12140 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12141 let mut y = self.qmatvec_mmvq_batched(
12142 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12143 )?;
12144 if let GpuTensor::Quant { scale, .. } = w {
12145 if *scale != 1.0 {
12146 self.scale_inplace(&mut y, *scale, m * out_f)?;
12147 }
12148 }
12149 return Ok(y);
12150 }
12151 }
12152 }
12153 }
12154 if fast {
12160 if let GpuTensor::Quant {
12161 bytes,
12162 qtype,
12163 row_bytes,
12164 scale,
12165 ..
12166 } = w
12167 {
12168 if *qtype == QT_F8_E4M3 {
12169 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12170 return self.qmatvec_mmvq(
12171 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12172 );
12173 }
12174 }
12175 }
12176 let mut y = match w {
12177 GpuTensor::Quant {
12178 bytes,
12179 qtype,
12180 row_bytes,
12181 ..
12182 } if fast && *qtype == QT_Q8_0 => {
12183 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12184 }
12185 GpuTensor::Quant {
12186 bytes,
12187 qtype,
12188 row_bytes,
12189 ..
12190 } if fast && *qtype == QT_Q4_K => {
12191 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12192 }
12193 GpuTensor::Quant {
12194 bytes,
12195 qtype,
12196 row_bytes,
12197 ..
12198 } if fast && *qtype == QT_Q6_K => {
12199 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12200 }
12201 GpuTensor::Quant {
12202 bytes,
12203 qtype,
12204 row_bytes,
12205 ..
12206 } if fast && *qtype == QT_Q5_K => {
12207 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12208 }
12209 GpuTensor::Quant {
12210 bytes,
12211 qtype,
12212 row_bytes,
12213 ..
12214 } if fast && *qtype == QT_Q3_K => {
12215 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12216 }
12217 GpuTensor::Quant {
12218 bytes,
12219 qtype,
12220 row_bytes,
12221 rp,
12222 ..
12223 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12224 if *rp {
12225 "qmatvec_nvfp4_dp4a_rp"
12226 } else {
12227 "qmatvec_nvfp4_dp4a"
12228 },
12229 &bytes.slice(0..bytes.len()),
12230 x,
12231 m,
12232 in_f,
12233 out_f,
12234 *row_bytes,
12235 )?,
12236 GpuTensor::Quant {
12240 bytes,
12241 qtype,
12242 row_bytes,
12243 ..
12244 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12245 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12246 }
12247 GpuTensor::Quant {
12252 bytes,
12253 qtype,
12254 row_bytes,
12255 rp,
12256 ..
12257 } =>
12258 {
12261 self.qmatvec(
12262 bytes,
12263 x,
12264 m,
12265 in_f,
12266 out_f,
12267 if *rp && *qtype == QT_NVFP4 {
12268 QT_NVFP4_RP
12269 } else {
12270 *qtype
12271 },
12272 *row_bytes,
12273 )?
12274 }
12275 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12276 GpuTensor::FloatBf16 { data, .. } => {
12279 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12286 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12287 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12288 y
12289 } else {
12290 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12291 }
12292 }
12293 };
12294 if let GpuTensor::Quant { scale, .. } = w {
12296 if *scale != 1.0 {
12297 self.scale_inplace(&mut y, *scale, m * out_f)?;
12298 }
12299 }
12300 Ok(y)
12301 }
12302
12303 pub fn stage_a_raw_needed() -> bool {
12313 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12314 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12315 }
12316
12317 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12320 use crate::model::GpuTensor;
12321 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12322 return false;
12323 }
12324 match w {
12325 GpuTensor::Quant { qtype, .. } => {
12332 matches!(
12333 *qtype,
12334 QT_Q8_0
12335 | QT_Q4_K
12336 | QT_Q6_K
12337 | QT_Q5_K
12338 | QT_Q3_K
12339 | QT_NVFP4
12340 | QT_F8_E4M3
12341 | QT_F8_E4M3_BLK
12342 | QT_Q4_0
12343 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12344 }
12345 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12346 }
12347 }
12348
12349 pub fn matmul_pre(
12354 &self,
12355 w: &crate::model::GpuTensor,
12356 aq: &CudaSlice<i8>,
12357 ad: &CudaSlice<f32>,
12358 x_fallback: &CudaSlice<f32>,
12359 m: usize,
12360 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12361 use crate::model::GpuTensor;
12362 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12368 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12371 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12372 return Ok(y);
12373 }
12374 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12377 return Ok(y);
12378 }
12379 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12381 return Ok(y);
12382 }
12383 }
12384 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12390 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12391 return Ok(y);
12392 }
12393 }
12394 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12395 return Ok(y);
12396 }
12397 if m >= 16
12402 && w.out_features() >= 128
12403 && self.mmq_supports(w)
12404 && !self.verify_exact_on()
12405 && x_raw_ok
12406 {
12407 return self.qmatvec_mmq(w, x_fallback, m);
12408 }
12409 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12412 if let Some(y) =
12413 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12414 {
12415 return Ok(y);
12416 }
12417 }
12418 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12421 return self.qmatvec_gemm(w, aq, ad, m);
12422 }
12423 if !self.uses_q8_1_fast(w) {
12442 if !x_raw_ok {
12443 return Err(format!(
12444 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12445 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12446 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12447 activation (see Engine::rms_norm_decode, which is bit-identical to \
12448 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12449 x_fallback.len(),
12450 m,
12451 w.in_features(),
12452 m * w.in_features()
12453 )
12454 .into());
12455 }
12456 return self.matmul(w, x_fallback, m);
12457 }
12458 let in_f = w.in_features();
12459 let out_f = w.out_features();
12460 let (bytes, qtype, row_bytes, scale, rp) = match w {
12461 GpuTensor::Quant {
12462 bytes,
12463 qtype,
12464 row_bytes,
12465 scale,
12466 rp,
12467 ..
12468 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12469 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12470 };
12471 let (mbytes, mrp) = match w {
12474 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12475 _ => (bytes, rp),
12476 };
12477 if m == 1 && self.mmvq_supports(qtype) {
12481 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12482 }
12483 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12496 && std::env::var("MEMRA_NO_BATCHED").is_err()
12497 && (m <= 4 || Self::b8_enabled())
12498 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12502 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12503 {
12504 let mcols = Self::batched_mcols(m);
12505 return self.qmatvec_mmvq_batched(
12506 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12507 );
12508 }
12509 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12515 let (b2, r2) = if qtype == QT_Q4_0 {
12516 (mbytes, mrp)
12517 } else {
12518 (bytes, rp)
12519 };
12520 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12521 }
12522 let name = match qtype {
12523 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12524 QT_Q4_K => "qmatvec_q4_K_dp4a",
12525 QT_Q6_K => "qmatvec_q6_K_dp4a",
12526 QT_Q5_K => "qmatvec_q5_K_dp4a",
12527 QT_Q3_K => "qmatvec_q3_K_dp4a",
12528 QT_NVFP4 => {
12529 if rp {
12530 "qmatvec_nvfp4_dp4a_rp"
12531 } else {
12532 "qmatvec_nvfp4_dp4a"
12533 }
12534 }
12535 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12536 _ => unreachable!(),
12537 };
12538 let f = self.func(name);
12539 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12541 grid_dim: (out_f as u32, m as u32, 1),
12542 block_dim: (128, 1, 1),
12543 shared_mem_bytes: 0,
12544 };
12545 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12546 let __s_b = self.gpu.stream();
12547 let mut b = __s_b.launch_builder(&f);
12548 b.arg(bytes)
12549 .arg(aq)
12550 .arg(ad)
12551 .arg(&mut y)
12552 .arg(&inf)
12553 .arg(&outf)
12554 .arg(&mi)
12555 .arg(&rb);
12556 unsafe {
12557 b.launch(cfg)?;
12558 }
12559 if scale != 1.0 {
12560 self.scale_inplace(&mut y, scale, m * out_f)?;
12561 }
12562 Ok(y)
12563 }
12564
12565 pub fn matmul_decode_exact(
12573 &self,
12574 w: &crate::model::GpuTensor,
12575 x: &CudaSlice<f32>,
12576 m: usize,
12577 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12578 use crate::model::GpuTensor;
12579 if let GpuTensor::Float { data, .. } = w {
12587 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12588 }
12589 if let GpuTensor::FloatBf16 { data, .. } = w {
12592 let (in_f, out_f) = (w.in_features(), w.out_features());
12593 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12596 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12597 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12598 return Ok(y);
12599 }
12600 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12601 }
12602 if !self.uses_q8_1_fast(w) {
12603 return self.matmul(w, x, m);
12604 }
12605 let in_f = w.in_features();
12606 let out_f = w.out_features();
12607 let (bytes, qtype, row_bytes, scale, rp) = match w {
12608 GpuTensor::Quant {
12609 bytes,
12610 qtype,
12611 row_bytes,
12612 scale,
12613 rp,
12614 ..
12615 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12616 _ => return self.matmul(w, x, m),
12617 };
12618 let (bytes, rp) = match w {
12621 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12622 _ => (bytes, rp),
12623 };
12624 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12625 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12629 return Ok(y);
12630 }
12631 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12640 && std::env::var("MEMRA_NO_BATCHED").is_err()
12641 && (m <= 4 || Self::b8_enabled())
12642 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12645 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12646 {
12647 let mcols = Self::batched_mcols(m);
12648 return self.qmatvec_mmvq_batched(
12649 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12650 );
12651 }
12652 if self.mmvq_supports(qtype) {
12653 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12656 }
12657 self.matmul_pre(w, &aq, &ad, x, m)
12660 }
12661
12662 pub fn matmul_decode_exact_pre(
12672 &self,
12673 w: &crate::model::GpuTensor,
12674 aq: &CudaSlice<i8>,
12675 ad: &CudaSlice<f32>,
12676 m: usize,
12677 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12678 use crate::model::GpuTensor;
12679 debug_assert!(
12680 self.uses_q8_1_fast(w),
12681 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12682 );
12683 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12685 return Ok(y);
12686 }
12687 let in_f = w.in_features();
12688 let out_f = w.out_features();
12689 let (bytes, qtype, row_bytes, scale, rp) = match w {
12690 GpuTensor::Quant {
12691 bytes,
12692 qtype,
12693 row_bytes,
12694 scale,
12695 rp,
12696 ..
12697 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12698 _ => {
12699 return Err(
12700 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12701 );
12702 }
12703 };
12704 let (bytes, rp) = match w {
12706 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12707 _ => (bytes, rp),
12708 };
12709 if (2..=16).contains(&m)
12711 && self.batched_supports(qtype)
12712 && self.mmvq_supports(qtype)
12713 && std::env::var("MEMRA_NO_BATCHED").is_err()
12714 && (m <= 4 || Self::b8_enabled())
12715 && (m <= 8
12716 || qtype == QT_Q4_0
12717 || qtype == QT_Q6_K
12718 || qtype == QT_F8_E4M3
12719 || qtype == QT_NVFP4
12720 || qtype == QT_Q4_K
12721 || qtype == QT_Q5_K
12722 || qtype == QT_Q8_0)
12723 {
12724 let mcols = Self::batched_mcols(m);
12725 return self.qmatvec_mmvq_batched(
12726 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12727 );
12728 }
12729 if self.mmvq_supports(qtype) {
12730 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12731 }
12732 let x0 = self.zeros(0)?;
12735 self.matmul_pre(w, aq, ad, &x0, m)
12736 }
12737
12738 pub fn matmul_decode_exact_dual_pre(
12747 &self,
12748 w0: &crate::model::GpuTensor,
12749 w1: &crate::model::GpuTensor,
12750 aq: &CudaSlice<i8>,
12751 ad: &CudaSlice<f32>,
12752 m: usize,
12753 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12754 {
12755 use crate::model::GpuTensor;
12756 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12757 let on = *ON.get_or_init(|| {
12758 std::env::var("MEMRA_SPEC_DUAL_T")
12759 .map(|v| v != "0")
12760 .unwrap_or(true)
12761 });
12762 if !on
12763 || !(2..=7).contains(&m)
12764 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12765 || !self.uses_q8_1_fast(w0)
12766 || !self.uses_q8_1_fast(w1)
12767 {
12768 return Ok(None);
12769 }
12770 if !self.mmvq_supports(QT_NVFP4) {
12775 return Ok(None);
12776 }
12777 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12778 if w1.in_features() != in_f || w1.out_features() != out_f {
12779 return Ok(None);
12780 }
12781 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12782 (
12783 GpuTensor::Quant {
12784 bytes: b0,
12785 qtype: q0,
12786 row_bytes: rb0,
12787 scale: s0,
12788 rp: rp0,
12789 rp4: None,
12790 ..
12791 },
12792 GpuTensor::Quant {
12793 bytes: b1,
12794 qtype: q1,
12795 row_bytes: rb1,
12796 scale: s1,
12797 rp: rp1,
12798 rp4: None,
12799 ..
12800 },
12801 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12802 (b0, b1, *rb0, *s0, *s1, *rp0)
12803 }
12804 _ => return Ok(None),
12805 };
12806 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12809 {
12810 return Ok(None);
12811 }
12812 let (y0, y1) =
12813 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12814 Ok(Some(((y0, s0), (y1, s1))))
12815 }
12816
12817 pub fn matmul_decode_exact_group4_pre(
12829 &self,
12830 ws: [&crate::model::GpuTensor; 4],
12831 aq: &CudaSlice<i8>,
12832 ad: &CudaSlice<f32>,
12833 m: usize,
12834 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12835 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12836 let on = *ON.get_or_init(|| {
12837 std::env::var("MEMRA_TK_GDN_GROUP")
12838 .map(|v| v != "0")
12839 .unwrap_or(true)
12840 });
12841 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12842 }
12843
12844 pub fn matmul_decode_exact_group3_pre(
12849 &self,
12850 ws: [&crate::model::GpuTensor; 3],
12851 aq: &CudaSlice<i8>,
12852 ad: &CudaSlice<f32>,
12853 m: usize,
12854 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12855 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12856 let on = *ON.get_or_init(|| {
12857 std::env::var("MEMRA_TK_FA_GROUP")
12858 .map(|v| v != "0")
12859 .unwrap_or(true)
12860 });
12861 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12862 }
12863
12864 fn matmul_decode_exact_group_pre(
12868 &self,
12869 ws: &[&crate::model::GpuTensor],
12870 aq: &CudaSlice<i8>,
12871 ad: &CudaSlice<f32>,
12872 m: usize,
12873 on: bool,
12874 tag: &'static str,
12875 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12876 use crate::model::GpuTensor;
12877 if !on
12878 || !(2..=16).contains(&m)
12879 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12880 || (m > 4 && !Self::b8_enabled())
12881 || !self.mmvq_supports(QT_NVFP4)
12882 || !self.batched_supports(QT_NVFP4)
12883 {
12884 return Ok(None);
12885 }
12886 let in_f = ws[0].in_features();
12887 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12888 for w in ws {
12889 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12890 return Ok(None);
12891 }
12892 match w {
12893 GpuTensor::Quant {
12894 bytes,
12895 qtype,
12896 scale,
12897 rp: true,
12898 rp4: None,
12899 ..
12900 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12901 parts.push((bytes, w.out_features(), *scale));
12902 }
12903 _ => return Ok(None),
12904 }
12905 }
12906 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12908 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12909 let mcols = if (5..=7).contains(&m) && b567 {
12910 m
12911 } else {
12912 Self::batched_mcols(m)
12913 };
12914 let kname: &'static str = match mcols {
12915 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12916 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12917 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12918 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12919 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12920 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12921 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12922 _ => return Ok(None),
12923 };
12924 if std::env::var("MEMRA_DEBUG").is_ok() {
12927 use std::sync::Mutex;
12928 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12929 let mut seen = SEEN.lock().unwrap();
12930 if !seen.contains(&tag) {
12931 seen.push(tag);
12932 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12933 }
12934 }
12935 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12938 let three = parts.len() == 3;
12939 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12940 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12941 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12942 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12945 let cfg = LaunchConfig {
12946 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12947 block_dim: (32, ROWS_PER_BLOCK, 1),
12948 shared_mem_bytes: 0,
12949 };
12950 let (inf, mi) = (in_f as i32, m as i32);
12951 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12952 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12953 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12954 let s3 = if three { 1.0f32 } else { parts[3].2 };
12955 let w3 = if three { parts[0].0 } else { parts[3].0 };
12956 let f = self.func(kname);
12957 let __s_b = self.gpu.stream();
12958 let mut b = __s_b.launch_builder(&f);
12959 b.arg(parts[0].0)
12960 .arg(parts[1].0)
12961 .arg(parts[2].0)
12962 .arg(w3)
12963 .arg(aq)
12964 .arg(ad)
12965 .arg(&mut y0)
12966 .arg(&mut y1)
12967 .arg(&mut y2)
12968 .arg(&mut y3)
12969 .arg(&inf)
12970 .arg(&n0)
12971 .arg(&n1)
12972 .arg(&n2)
12973 .arg(&n3)
12974 .arg(&mi)
12975 .arg(&s0)
12976 .arg(&s1)
12977 .arg(&s2)
12978 .arg(&s3);
12979 unsafe {
12980 b.launch(cfg)?;
12981 }
12982 Ok(Some(if three {
12983 vec![y0, y1, y2]
12984 } else {
12985 vec![y0, y1, y2, y3]
12986 }))
12987 }
12988
12989 pub fn matmul_decode_exact_dual(
13005 &self,
13006 w0: &crate::model::GpuTensor,
13007 w1: &crate::model::GpuTensor,
13008 x: &CudaSlice<f32>,
13009 m: usize,
13010 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13011 use crate::model::GpuTensor;
13012 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13013 let on = *ON.get_or_init(|| {
13014 std::env::var("MEMRA_SPEC_DUAL_T")
13015 .map(|v| v != "0")
13016 .unwrap_or(true)
13017 });
13018 if !on
13019 || !(2..=4).contains(&m)
13020 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13021 || !self.uses_q8_1_fast(w0)
13022 || !self.uses_q8_1_fast(w1)
13023 {
13024 return Ok(None);
13025 }
13026 if !self.mmvq_supports(QT_NVFP4) {
13031 return Ok(None);
13032 }
13033 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13034 if w1.in_features() != in_f || w1.out_features() != out_f {
13035 return Ok(None);
13036 }
13037 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
13038 (
13039 GpuTensor::Quant {
13040 bytes: b0,
13041 qtype: q0,
13042 row_bytes: rb0,
13043 scale: s0,
13044 rp: rp0,
13045 rp4: None,
13046 ..
13047 },
13048 GpuTensor::Quant {
13049 bytes: b1,
13050 qtype: q1,
13051 row_bytes: rb1,
13052 scale: s1,
13053 rp: rp1,
13054 rp4: None,
13055 ..
13056 },
13057 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13058 (b0, b1, *rb0, *s0, *s1, *rp0)
13059 }
13060 _ => return Ok(None),
13061 };
13062 if std::env::var("MEMRA_DEBUG").is_ok() {
13065 static ONCE: std::sync::Once = std::sync::Once::new();
13066 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13067 }
13068 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13069 let (y0, y1) =
13070 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13071 let mut y0 = y0;
13072 let mut y1 = y1;
13073 if s0 != 1.0 {
13074 self.scale_inplace(&mut y0, s0, m * out_f)?;
13075 }
13076 if s1 != 1.0 {
13077 self.scale_inplace(&mut y1, s1, m * out_f)?;
13078 }
13079 Ok(Some((y0, y1)))
13080 }
13081
13082 #[allow(clippy::too_many_arguments)]
13087 pub fn qmatvec_batched_dual_raw(
13088 &self,
13089 b0: &CudaSlice<u8>,
13090 b1: &CudaSlice<u8>,
13091 aq: &CudaSlice<i8>,
13092 ad: &CudaSlice<f32>,
13093 m: usize,
13094 in_f: usize,
13095 out_f: usize,
13096 row_bytes: usize,
13097 rp: bool,
13098 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13099 const ROWS_PER_BLOCK: u32 = 4;
13100 let mcols = Self::batched_mcols(m);
13101 let tiny_rp1 = rp
13104 && mcols == 4
13105 && out_f <= 128
13106 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13107 let (name, rows_per_block) = if tiny_rp1 {
13108 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13109 } else {
13110 match (mcols, rp, m) {
13111 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13112 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13113 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13114 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13115 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13116 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13117 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13118 _ => {
13119 return Err(
13120 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13121 );
13122 }
13123 }
13124 };
13125 let f = self.func(name);
13126 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13127 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13128 let cfg = LaunchConfig {
13129 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13130 block_dim: (32, ROWS_PER_BLOCK, 1),
13131 shared_mem_bytes: 0,
13132 };
13133 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13134 let __s_b = self.gpu.stream();
13135 let mut b = __s_b.launch_builder(&f);
13136 b.arg(b0)
13137 .arg(b1)
13138 .arg(aq)
13139 .arg(ad)
13140 .arg(&mut y0)
13141 .arg(&mut y1)
13142 .arg(&inf)
13143 .arg(&outf)
13144 .arg(&mi)
13145 .arg(&rb);
13146 unsafe {
13147 b.launch(cfg)?;
13148 }
13149 Ok((y0, y1))
13150 }
13151
13152 pub fn matmul_pre_dual_noscale(
13164 &self,
13165 w0: &crate::model::GpuTensor,
13166 w1: &crate::model::GpuTensor,
13167 aq: &CudaSlice<i8>,
13168 ad: &CudaSlice<f32>,
13169 m: usize,
13170 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13171 {
13172 use crate::model::GpuTensor;
13173 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13174 return Ok(None);
13175 }
13176 if !self.mmvq_supports(QT_NVFP4) {
13186 return Ok(None);
13187 }
13188 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13189 if w1.in_features() != in_f || w1.out_features() != out_f {
13190 return Ok(None);
13191 }
13192 let no_mirror =
13205 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13206 if self.q8_ffn_fuse2_on()
13207 && no_mirror(w0)
13208 && no_mirror(w1)
13209 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13210 {
13211 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13212 return Ok(Some(((y0, 1.0), (y1, 1.0))));
13213 }
13214 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13224 let (y0, y1) =
13225 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13226 return Ok(Some(((y0, p0.3), (y1, p1.3))));
13227 }
13228 let (b0, q0, rb0, s0, rp0) = match w0 {
13229 GpuTensor::Quant {
13230 bytes,
13231 qtype,
13232 row_bytes,
13233 scale,
13234 rp,
13235 ..
13236 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13237 _ => return Ok(None),
13238 };
13239 let (b1, q1, rb1, s1, rp1) = match w1 {
13240 GpuTensor::Quant {
13241 bytes,
13242 qtype,
13243 row_bytes,
13244 scale,
13245 rp,
13246 ..
13247 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13248 _ => return Ok(None),
13249 };
13250 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13251 return Ok(None);
13252 }
13253 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13255 let rows_per_block = ROWS_PER_BLOCK * RPW;
13256 let f = self.func(if rp0 {
13257 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13258 } else {
13259 "qmatvec_nvfp4_mmvq_dual_mr2"
13260 });
13261 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13262 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13263 let cfg = LaunchConfig {
13264 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13265 block_dim: (32, ROWS_PER_BLOCK, 1),
13266 shared_mem_bytes: 0,
13267 };
13268 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13269 let one = 1.0f32;
13272 let __s_b = self.gpu.stream();
13273 let mut b = __s_b.launch_builder(&f);
13274 b.arg(b0)
13275 .arg(b1)
13276 .arg(aq)
13277 .arg(ad)
13278 .arg(&mut y0)
13279 .arg(&mut y1)
13280 .arg(&inf)
13281 .arg(&outf)
13282 .arg(&mi)
13283 .arg(&rb)
13284 .arg(&one)
13285 .arg(&one);
13286 unsafe {
13287 b.launch(cfg)?;
13288 }
13289 Ok(Some(((y0, s0), (y1, s1))))
13290 }
13291
13292 #[allow(clippy::too_many_arguments)]
13300 pub fn matmul_nvfp4_fused3(
13301 &self,
13302 w0: &crate::model::GpuTensor,
13303 w1: &crate::model::GpuTensor,
13304 w2: &crate::model::GpuTensor,
13305 aq: &CudaSlice<i8>,
13306 ad: &CudaSlice<f32>,
13307 m: usize,
13308 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13309 {
13310 use crate::model::GpuTensor;
13311 if !self.mmvq_supports(QT_NVFP4)
13318 || !self.uses_q8_1_fast(w0)
13319 || !self.uses_q8_1_fast(w1)
13320 || !self.uses_q8_1_fast(w2)
13321 {
13322 return Ok(None);
13323 }
13324 if (9..=16).contains(&m) {
13327 return Ok(
13328 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13329 Some(mut ys) => {
13330 let y2 = ys.pop().unwrap();
13331 let y1 = ys.pop().unwrap();
13332 let y0 = ys.pop().unwrap();
13333 Some((y0, y1, y2))
13334 }
13335 None => None,
13336 },
13337 );
13338 }
13339 if !(1..=8).contains(&m) {
13340 return Ok(None);
13341 }
13342 if m > 1 {
13343 let in_f = w0.in_features();
13344 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13345 || !self.batched_supports(QT_NVFP4)
13346 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13347 || (m > 4 && !Self::b8_enabled())
13348 || in_f % 512 != 0
13349 || in_f / 64 > 272
13350 {
13351 return Ok(None);
13352 }
13353 }
13354 let unpack = |w: &crate::model::GpuTensor| match w {
13355 GpuTensor::Quant {
13356 bytes,
13357 qtype,
13358 scale,
13359 rp,
13360 ..
13361 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13362 _ => None,
13363 };
13364 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13365 return Ok(None);
13366 };
13367 let in_f = w0.in_features();
13368 if w1.in_features() != in_f || w2.in_features() != in_f {
13369 return Ok(None);
13370 }
13371 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13372 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13374 let rows_pb = ROWS_PER_BLOCK * RPW;
13375 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13376 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13377 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13378 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13379 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13380 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13383 if m > 1 {
13384 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13386 return Ok(None);
13387 }
13388 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13389 let cfg = LaunchConfig {
13390 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13391 block_dim: (32, ROWS_PER_BLOCK, 1),
13392 shared_mem_bytes: 0,
13393 };
13394 let __s_b = self.gpu.stream();
13395 let mut b = __s_b.launch_builder(&f);
13396 b.arg(b0)
13397 .arg(b1)
13398 .arg(b2)
13399 .arg(aq)
13400 .arg(ad)
13401 .arg(&mut y0)
13402 .arg(&mut y1)
13403 .arg(&mut y2)
13404 .arg(&inf)
13405 .arg(&oi0)
13406 .arg(&oi1)
13407 .arg(&oi2)
13408 .arg(&mi);
13409 unsafe {
13410 b.launch(cfg)?;
13411 }
13412 return Ok(Some((y0, y1, y2)));
13413 }
13414 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13415 let cfg = LaunchConfig {
13416 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13417 block_dim: (32, ROWS_PER_BLOCK, 1),
13418 shared_mem_bytes: 0,
13419 };
13420 let __s_b = self.gpu.stream();
13421 let mut b = __s_b.launch_builder(&f);
13422 b.arg(b0)
13423 .arg(b1)
13424 .arg(b2)
13425 .arg(aq)
13426 .arg(ad)
13427 .arg(&mut y0)
13428 .arg(&mut y1)
13429 .arg(&mut y2)
13430 .arg(&inf)
13431 .arg(&oi0)
13432 .arg(&oi1)
13433 .arg(&oi2)
13434 .arg(&mi)
13435 .arg(&p0.1)
13436 .arg(&p1.1)
13437 .arg(&p2.1);
13438 unsafe {
13439 b.launch(cfg)?;
13440 }
13441 Ok(Some((y0, y1, y2)))
13442 }
13443
13444 pub fn matmul_nvfp4_fused2(
13453 &self,
13454 w0: &crate::model::GpuTensor,
13455 w1: &crate::model::GpuTensor,
13456 aq: &CudaSlice<i8>,
13457 ad: &CudaSlice<f32>,
13458 m: usize,
13459 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13460 use crate::model::GpuTensor;
13461 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13462 let off =
13463 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13464 if off
13467 || m != 1
13468 || !self.mmvq_supports(QT_NVFP4)
13469 || !self.uses_q8_1_fast(w0)
13470 || !self.uses_q8_1_fast(w1)
13471 {
13472 return Ok(None);
13473 }
13474 let unpack = |w: &crate::model::GpuTensor| match w {
13475 GpuTensor::Quant {
13476 bytes,
13477 qtype,
13478 scale,
13479 rp,
13480 ..
13481 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13482 _ => None,
13483 };
13484 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13485 return Ok(None);
13486 };
13487 let in_f = w0.in_features();
13488 if w1.in_features() != in_f {
13489 return Ok(None);
13490 }
13491 let (o0, o1) = (w0.out_features(), w1.out_features());
13492 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13494 let rows_pb = ROWS_PER_BLOCK * RPW;
13495 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13496 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13497 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13498 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13499 let cfg = LaunchConfig {
13500 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13501 block_dim: (32, ROWS_PER_BLOCK, 1),
13502 shared_mem_bytes: 0,
13503 };
13504 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13505 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13508 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13511 {
13512 use cudarc::driver::{DevicePtr, DevicePtrMut};
13513 let s = &self.gpu.stream();
13514 let (pw0, _g0) = b0.device_ptr(s);
13515 let (pw1, _g1) = b1.device_ptr(s);
13516 let (paq, _g2) = aq.device_ptr(s);
13517 let (pad, _g3) = ad.device_ptr(s);
13518 let (py0, _g4) = y0.device_ptr_mut(s);
13519 let (py1, _g5) = y1.device_ptr_mut(s);
13520 let (s0, s1) = (p0.1, p1.1);
13521 let mut ps = [
13522 &pw0 as *const _ as *mut std::ffi::c_void,
13523 &pw1 as *const _ as *mut _,
13524 &paq as *const _ as *mut _,
13525 &pad as *const _ as *mut _,
13526 &py0 as *const _ as *mut _,
13527 &py1 as *const _ as *mut _,
13528 &inf as *const _ as *mut _,
13529 &oi0 as *const _ as *mut _,
13530 &oi1 as *const _ as *mut _,
13531 &mi as *const _ as *mut _,
13532 &s0 as *const _ as *mut _,
13533 &s1 as *const _ as *mut _,
13534 ];
13535 unsafe {
13536 self.launch_pdl(
13537 "qmatvec_nvfp4_mmvq_fused2_rp",
13538 cfg.grid_dim,
13539 cfg.block_dim,
13540 &mut ps,
13541 )?;
13542 }
13543 }
13544 return Ok(Some((y0, y1)));
13545 }
13546 let __s_b = self.gpu.stream();
13547 let mut b = __s_b.launch_builder(&f);
13548 b.arg(b0)
13549 .arg(b1)
13550 .arg(aq)
13551 .arg(ad)
13552 .arg(&mut y0)
13553 .arg(&mut y1)
13554 .arg(&inf)
13555 .arg(&oi0)
13556 .arg(&oi1)
13557 .arg(&mi)
13558 .arg(&p0.1)
13559 .arg(&p1.1);
13560 unsafe {
13561 b.launch(cfg)?;
13562 }
13563 Ok(Some((y0, y1)))
13564 }
13565
13566 pub fn matmul_nvfp4_fused2_into(
13571 &self,
13572 w0: &crate::model::GpuTensor,
13573 w1: &crate::model::GpuTensor,
13574 aq: &CudaSlice<i8>,
13575 ad: &CudaSlice<f32>,
13576 y0: &mut CudaSlice<f32>,
13577 y1: &mut CudaSlice<f32>,
13578 ) -> Result<bool, Box<dyn std::error::Error>> {
13579 use crate::model::GpuTensor;
13580 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13581 let off =
13582 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13583 if off
13584 || !self.mmvq_supports(QT_NVFP4)
13585 || !self.uses_q8_1_fast(w0)
13586 || !self.uses_q8_1_fast(w1)
13587 {
13588 return Ok(false);
13589 }
13590 let unpack = |w: &crate::model::GpuTensor| match w {
13591 GpuTensor::Quant {
13592 bytes,
13593 qtype,
13594 scale,
13595 rp,
13596 ..
13597 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13598 _ => None,
13599 };
13600 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13601 return Ok(false);
13602 };
13603 let in_f = w0.in_features();
13604 if w1.in_features() != in_f {
13605 return Ok(false);
13606 }
13607 let (o0, o1) = (w0.out_features(), w1.out_features());
13608 if y0.len() < o0 || y1.len() < o1 {
13609 return Ok(false);
13610 }
13611 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13613 let rows_pb = ROWS_PER_BLOCK * RPW;
13614 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13615 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13616 let cfg = LaunchConfig {
13617 grid_dim: (nb(o0) + nb(o1), 1, 1),
13618 block_dim: (32, ROWS_PER_BLOCK, 1),
13619 shared_mem_bytes: 0,
13620 };
13621 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13622 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13625 let __s_b = self.gpu.stream();
13626 let mut b = __s_b.launch_builder(&f);
13627 b.arg(b0)
13628 .arg(b1)
13629 .arg(aq)
13630 .arg(ad)
13631 .arg(&mut *y0)
13632 .arg(&mut *y1)
13633 .arg(&inf)
13634 .arg(&oi0)
13635 .arg(&oi1)
13636 .arg(&mi)
13637 .arg(&p0.1)
13638 .arg(&p1.1);
13639 unsafe {
13640 b.launch(cfg)?;
13641 }
13642 Ok(true)
13643 }
13644
13645 #[allow(clippy::type_complexity)]
13650 pub fn matmul_nvfp4_fused4(
13651 &self,
13652 w0: &crate::model::GpuTensor,
13653 w1: &crate::model::GpuTensor,
13654 w2: &crate::model::GpuTensor,
13655 w3: &crate::model::GpuTensor,
13656 aq: &CudaSlice<i8>,
13657 ad: &CudaSlice<f32>,
13658 m: usize,
13659 ) -> Result<
13660 Option<(
13661 CudaSlice<f32>,
13662 CudaSlice<f32>,
13663 CudaSlice<f32>,
13664 CudaSlice<f32>,
13665 )>,
13666 Box<dyn std::error::Error>,
13667 > {
13668 use crate::model::GpuTensor;
13669 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13676 || !self.mmvq_supports(QT_NVFP4)
13677 || !self.uses_q8_1_fast(w0)
13678 || !self.uses_q8_1_fast(w1)
13679 || !self.uses_q8_1_fast(w2)
13680 || !self.uses_q8_1_fast(w3)
13681 {
13682 return Ok(None);
13683 }
13684 if (9..=16).contains(&m) {
13689 return Ok(
13690 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13691 Some(mut ys) => {
13692 let y3 = ys.pop().unwrap();
13693 let y2 = ys.pop().unwrap();
13694 let y1 = ys.pop().unwrap();
13695 let y0 = ys.pop().unwrap();
13696 Some((y0, y1, y2, y3))
13697 }
13698 None => None,
13699 },
13700 );
13701 }
13702 if !(1..=8).contains(&m) {
13703 return Ok(None);
13704 }
13705 if m > 1 {
13706 let in_f = w0.in_features();
13709 if !self.batched_supports(QT_NVFP4)
13710 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13711 || (m > 4 && !Self::b8_enabled())
13712 || in_f % 512 != 0
13713 || in_f / 64 > 272
13714 {
13715 return Ok(None);
13716 }
13717 }
13718 let unpack = |w: &crate::model::GpuTensor| match w {
13719 GpuTensor::Quant {
13720 bytes,
13721 qtype,
13722 scale,
13723 rp,
13724 ..
13725 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13726 _ => None,
13727 };
13728 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13729 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13730 else {
13731 return Ok(None);
13732 };
13733 let in_f = w0.in_features();
13734 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13735 return Ok(None);
13736 }
13737 let (o0, o1, o2, o3) = (
13738 w0.out_features(),
13739 w1.out_features(),
13740 w2.out_features(),
13741 w3.out_features(),
13742 );
13743 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13745 let rows_pb = ROWS_PER_BLOCK * RPW;
13746 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13747 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13748 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13749 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13750 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13751 let (inf, oi0, oi1, oi2, oi3, mi) = (
13752 in_f as i32,
13753 o0 as i32,
13754 o1 as i32,
13755 o2 as i32,
13756 o3 as i32,
13757 m as i32,
13758 );
13759 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13762 if m > 1 {
13763 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13766 return Ok(None);
13767 }
13768 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13769 let cfg = LaunchConfig {
13770 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13771 block_dim: (32, ROWS_PER_BLOCK, 1),
13772 shared_mem_bytes: 0,
13773 };
13774 let __s_b = self.gpu.stream();
13775 let mut b = __s_b.launch_builder(&f);
13776 b.arg(b0)
13777 .arg(b1)
13778 .arg(b2)
13779 .arg(b3)
13780 .arg(aq)
13781 .arg(ad)
13782 .arg(&mut y0)
13783 .arg(&mut y1)
13784 .arg(&mut y2)
13785 .arg(&mut y3)
13786 .arg(&inf)
13787 .arg(&oi0)
13788 .arg(&oi1)
13789 .arg(&oi2)
13790 .arg(&oi3)
13791 .arg(&mi);
13792 unsafe {
13793 b.launch(cfg)?;
13794 }
13795 return Ok(Some((y0, y1, y2, y3)));
13796 }
13797 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13798 let cfg = LaunchConfig {
13799 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13800 block_dim: (32, ROWS_PER_BLOCK, 1),
13801 shared_mem_bytes: 0,
13802 };
13803 let __s_b = self.gpu.stream();
13804 let mut b = __s_b.launch_builder(&f);
13805 b.arg(b0)
13806 .arg(b1)
13807 .arg(b2)
13808 .arg(b3)
13809 .arg(aq)
13810 .arg(ad)
13811 .arg(&mut y0)
13812 .arg(&mut y1)
13813 .arg(&mut y2)
13814 .arg(&mut y3)
13815 .arg(&inf)
13816 .arg(&oi0)
13817 .arg(&oi1)
13818 .arg(&oi2)
13819 .arg(&oi3)
13820 .arg(&mi)
13821 .arg(&p0.1)
13822 .arg(&p1.1)
13823 .arg(&p2.1)
13824 .arg(&p3.1);
13825 unsafe {
13826 b.launch(cfg)?;
13827 }
13828 Ok(Some((y0, y1, y2, y3)))
13829 }
13830
13831 pub fn matmul_q8_fused2(
13839 &self,
13840 w0: &crate::model::GpuTensor,
13841 w1: &crate::model::GpuTensor,
13842 aq: &CudaSlice<i8>,
13843 ad: &CudaSlice<f32>,
13844 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13845 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13851 return Ok(Some(self.e4m3_fused2_core(
13852 p0.0,
13853 p1.0,
13854 aq,
13855 ad,
13856 w0.in_features(),
13857 p0.1,
13858 p1.1,
13859 p0.2,
13860 p0.3,
13861 p1.3,
13862 )?));
13863 }
13864 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13865 return Ok(None);
13866 };
13867 Ok(Some(self.q8_fused2_core(
13868 p0.0,
13869 p1.0,
13870 aq,
13871 ad,
13872 w0.in_features(),
13873 p0.1,
13874 p1.1,
13875 p0.2,
13876 )?))
13877 }
13878
13879 #[allow(clippy::too_many_arguments)]
13880 fn q8_fused2_core(
13881 &self,
13882 b0: &CudaSlice<u8>,
13883 b1: &CudaSlice<u8>,
13884 aq: &CudaSlice<i8>,
13885 ad: &CudaSlice<f32>,
13886 in_f: usize,
13887 out0: usize,
13888 out1: usize,
13889 row_bytes: usize,
13890 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13891 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13893 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13894 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13895 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13896 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13897 let cfg = LaunchConfig {
13898 grid_dim: (nb0 + nb1, 1, 1),
13899 block_dim: (32, ROWS_PER_BLOCK, 1),
13900 shared_mem_bytes: 0,
13901 };
13902 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13903 let __s_b = self.gpu.stream();
13904 let mut b = __s_b.launch_builder(&f);
13905 b.arg(b0)
13906 .arg(b1)
13907 .arg(aq)
13908 .arg(ad)
13909 .arg(&mut y0)
13910 .arg(&mut y1)
13911 .arg(&inf)
13912 .arg(&o0)
13913 .arg(&o1)
13914 .arg(&rbl);
13915 unsafe {
13916 b.launch(cfg)?;
13917 }
13918 Ok((y0, y1))
13919 }
13920
13921 pub fn matmul_q8_fused2_x(
13927 &self,
13928 w0: &crate::model::GpuTensor,
13929 w1: &crate::model::GpuTensor,
13930 x: &CudaSlice<f32>,
13931 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13932 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13933 return Ok(None);
13934 }
13935 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13936 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13937 return Ok(Some(self.e4m3_fused2_core(
13938 p0.0,
13939 p1.0,
13940 &aq,
13941 &ad,
13942 w0.in_features(),
13943 p0.1,
13944 p1.1,
13945 p0.2,
13946 p0.3,
13947 p1.3,
13948 )?));
13949 }
13950 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13951 return Ok(None);
13952 };
13953 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13954 Ok(Some(self.q8_fused2_core(
13955 p0.0,
13956 p1.0,
13957 &aq,
13958 &ad,
13959 w0.in_features(),
13960 p0.1,
13961 p1.1,
13962 p0.2,
13963 )?))
13964 }
13965
13966 #[allow(clippy::too_many_arguments)]
13969 pub fn qmatvec_q8_fused2_raw(
13970 &self,
13971 b0: &CudaSlice<u8>,
13972 b1: &CudaSlice<u8>,
13973 x: &CudaSlice<f32>,
13974 in_f: usize,
13975 out0: usize,
13976 out1: usize,
13977 row_bytes: usize,
13978 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13979 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13980 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13981 }
13982
13983 pub fn matmul_q4_fused3(
13989 &self,
13990 w0: &crate::model::GpuTensor,
13991 w1: &crate::model::GpuTensor,
13992 w2: &crate::model::GpuTensor,
13993 aq: &CudaSlice<i8>,
13994 ad: &CudaSlice<f32>,
13995 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13996 {
13997 use crate::model::GpuTensor;
13998 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13999 match w {
14000 GpuTensor::Quant {
14001 qtype, row_bytes, ..
14002 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14003 _ => None,
14004 }
14005 };
14006 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14007 return Ok(None);
14008 };
14009 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14010 return Ok(None);
14011 }
14012 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14016 match w {
14017 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14018 Some(m) => (m, true),
14019 None => (bytes, *rp),
14020 },
14021 _ => unreachable!(),
14022 }
14023 }
14024 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14025 if rp0 != rp1 || rp1 != rp2 {
14026 return Ok(None);
14027 }
14028 let rp = rp0;
14029 let rpb: u32 = 4;
14030 let mr1 = rp && Self::q40_mr1_on();
14034 let nb = |o: usize| {
14035 if mr1 {
14036 (o as u32).div_ceil(rpb)
14037 } else {
14038 (o as u32).div_ceil(2).div_ceil(rpb)
14039 }
14040 };
14041 let grid = nb(o0) + nb(o1) + nb(o2);
14042 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14043 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14044 let mut y2 = self.alloc_uninit::<f32>(o2)?;
14045 let f = self.func(if mr1 {
14046 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14047 } else if rp {
14048 "qmatvec_q4_0_mmvq_fused3_rp"
14049 } else {
14050 "qmatvec_q4_0_mmvq_fused3"
14051 });
14052 let cfg = LaunchConfig {
14053 grid_dim: (grid, 1, 1),
14054 block_dim: (32, rpb, 1),
14055 shared_mem_bytes: 0,
14056 };
14057 let inf = w0.in_features() as i32;
14058 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14059 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14060 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14063 {
14064 use cudarc::driver::{DevicePtr, DevicePtrMut};
14065 let s = &self.gpu.stream();
14066 let (p0, _g0) = b0.device_ptr(s);
14067 let (p1, _g1) = b1.device_ptr(s);
14068 let (p2, _g2) = b2.device_ptr(s);
14069 let (paq, _g3) = aq.device_ptr(s);
14070 let (pad, _g4) = ad.device_ptr(s);
14071 let (py0, _g5) = y0.device_ptr_mut(s);
14072 let (py1, _g6) = y1.device_ptr_mut(s);
14073 let (py2, _g7) = y2.device_ptr_mut(s);
14074 let mut ps = [
14075 &p0 as *const _ as *mut std::ffi::c_void,
14076 &p1 as *const _ as *mut _,
14077 &p2 as *const _ as *mut _,
14078 &paq as *const _ as *mut _,
14079 &pad as *const _ as *mut _,
14080 &py0 as *const _ as *mut _,
14081 &py1 as *const _ as *mut _,
14082 &py2 as *const _ as *mut _,
14083 &inf as *const _ as *mut _,
14084 &oo0 as *const _ as *mut _,
14085 &oo1 as *const _ as *mut _,
14086 &oo2 as *const _ as *mut _,
14087 &r0 as *const _ as *mut _,
14088 &r1 as *const _ as *mut _,
14089 &r2 as *const _ as *mut _,
14090 ];
14091 unsafe {
14092 self.launch_pdl(
14093 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14094 (grid, 1, 1),
14095 (32, rpb, 1),
14096 &mut ps,
14097 )?;
14098 }
14099 }
14100 return Ok(Some((y0, y1, y2)));
14101 }
14102 let __s_b = self.gpu.stream();
14103 let mut b = __s_b.launch_builder(&f);
14104 b.arg(b0)
14105 .arg(b1)
14106 .arg(b2)
14107 .arg(aq)
14108 .arg(ad)
14109 .arg(&mut y0)
14110 .arg(&mut y1)
14111 .arg(&mut y2)
14112 .arg(&inf)
14113 .arg(&oo0)
14114 .arg(&oo1)
14115 .arg(&oo2)
14116 .arg(&r0)
14117 .arg(&r1)
14118 .arg(&r2);
14119 unsafe {
14120 b.launch(cfg)?;
14121 }
14122 Ok(Some((y0, y1, y2)))
14123 }
14124
14125 #[allow(clippy::too_many_arguments)]
14128 pub fn matmul_q4_fused3_into(
14129 &self,
14130 w0: &crate::model::GpuTensor,
14131 w1: &crate::model::GpuTensor,
14132 w2: &crate::model::GpuTensor,
14133 aq: &CudaSlice<i8>,
14134 ad: &CudaSlice<f32>,
14135 y0: &mut CudaSlice<f32>,
14136 y1: &mut CudaSlice<f32>,
14137 y2: &mut CudaSlice<f32>,
14138 ) -> Result<bool, Box<dyn std::error::Error>> {
14139 use crate::model::GpuTensor;
14140 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14141 match w {
14142 GpuTensor::Quant {
14143 qtype, row_bytes, ..
14144 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14145 _ => None,
14146 }
14147 };
14148 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14149 return Ok(false);
14150 };
14151 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14152 return Ok(false);
14153 }
14154 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14155 match w {
14156 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14157 Some(m) => (m, true),
14158 None => (bytes, *rp),
14159 },
14160 _ => unreachable!(),
14161 }
14162 }
14163 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14164 if rp0 != rp1 || rp1 != rp2 {
14165 return Ok(false);
14166 }
14167 let rp = rp0;
14168 let rpb: u32 = 4;
14169 let mr1 = rp && Self::q40_mr1_on();
14170 let nb = |o: usize| {
14171 if mr1 {
14172 (o as u32).div_ceil(rpb)
14173 } else {
14174 (o as u32).div_ceil(2).div_ceil(rpb)
14175 }
14176 };
14177 let grid = nb(o0) + nb(o1) + nb(o2);
14178 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14179 let f = self.func(if mr1 {
14180 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14181 } else if rp {
14182 "qmatvec_q4_0_mmvq_fused3_rp"
14183 } else {
14184 "qmatvec_q4_0_mmvq_fused3"
14185 });
14186 let cfg = LaunchConfig {
14187 grid_dim: (grid, 1, 1),
14188 block_dim: (32, rpb, 1),
14189 shared_mem_bytes: 0,
14190 };
14191 let inf = w0.in_features() as i32;
14192 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14193 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14194 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14196 use cudarc::driver::{DevicePtr, DevicePtrMut};
14197 let s = &self.gpu.stream();
14198 let (p0, _g0) = b0.device_ptr(s);
14199 let (p1, _g1) = b1.device_ptr(s);
14200 let (p2, _g2) = b2.device_ptr(s);
14201 let (paq, _g3) = aq.device_ptr(s);
14202 let (pad, _g4) = ad.device_ptr(s);
14203 let (py0, _g5) = y0.device_ptr_mut(s);
14204 let (py1, _g6) = y1.device_ptr_mut(s);
14205 let (py2, _g7) = y2.device_ptr_mut(s);
14206 let mut ps = [
14207 &p0 as *const _ as *mut std::ffi::c_void,
14208 &p1 as *const _ as *mut _,
14209 &p2 as *const _ as *mut _,
14210 &paq as *const _ as *mut _,
14211 &pad as *const _ as *mut _,
14212 &py0 as *const _ as *mut _,
14213 &py1 as *const _ as *mut _,
14214 &py2 as *const _ as *mut _,
14215 &inf as *const _ as *mut _,
14216 &oo0 as *const _ as *mut _,
14217 &oo1 as *const _ as *mut _,
14218 &oo2 as *const _ as *mut _,
14219 &r0 as *const _ as *mut _,
14220 &r1 as *const _ as *mut _,
14221 &r2 as *const _ as *mut _,
14222 ];
14223 unsafe {
14224 self.launch_pdl(
14225 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14226 (grid, 1, 1),
14227 (32, rpb, 1),
14228 &mut ps,
14229 )?;
14230 }
14231 return Ok(true);
14232 }
14233 let __s_b = self.gpu.stream();
14234 let mut b = __s_b.launch_builder(&f);
14235 b.arg(b0)
14236 .arg(b1)
14237 .arg(b2)
14238 .arg(aq)
14239 .arg(ad)
14240 .arg(&mut *y0)
14241 .arg(&mut *y1)
14242 .arg(&mut *y2)
14243 .arg(&inf)
14244 .arg(&oo0)
14245 .arg(&oo1)
14246 .arg(&oo2)
14247 .arg(&r0)
14248 .arg(&r1)
14249 .arg(&r2);
14250 unsafe {
14251 b.launch(cfg)?;
14252 }
14253 Ok(true)
14254 }
14255
14256 pub fn matmul_q4_fused2(
14258 &self,
14259 w0: &crate::model::GpuTensor,
14260 w1: &crate::model::GpuTensor,
14261 aq: &CudaSlice<i8>,
14262 ad: &CudaSlice<f32>,
14263 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14264 use crate::model::GpuTensor;
14265 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14266 match w {
14267 GpuTensor::Quant {
14268 qtype, row_bytes, ..
14269 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14270 _ => None,
14271 }
14272 };
14273 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14274 return Ok(None);
14275 };
14276 if w0.in_features() != w1.in_features() {
14277 return Ok(None);
14278 }
14279 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14281 match w {
14282 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14283 Some(m) => (m, true),
14284 None => (bytes, *rp),
14285 },
14286 _ => unreachable!(),
14287 }
14288 }
14289 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14290 if rp0 != rp1 {
14291 return Ok(None);
14292 }
14293 let rp = rp0;
14294 let rpb: u32 = 4;
14295 let mr1 = rp && Self::q40_mr1_on();
14297 let nb = |o: usize| {
14298 if mr1 {
14299 (o as u32).div_ceil(rpb)
14300 } else {
14301 (o as u32).div_ceil(2).div_ceil(rpb)
14302 }
14303 };
14304 let grid = nb(o0) + nb(o1);
14305 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14306 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14307 let f = self.func(if mr1 {
14308 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14309 } else if rp {
14310 "qmatvec_q4_0_mmvq_fused2_rp"
14311 } else {
14312 "qmatvec_q4_0_mmvq_fused2"
14313 });
14314 let cfg = LaunchConfig {
14315 grid_dim: (grid, 1, 1),
14316 block_dim: (32, rpb, 1),
14317 shared_mem_bytes: 0,
14318 };
14319 let inf = w0.in_features() as i32;
14320 let (oo0, oo1) = (o0 as i32, o1 as i32);
14321 let (r0, r1) = (rb0 as i64, rb1 as i64);
14322 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14324 {
14325 use cudarc::driver::{DevicePtr, DevicePtrMut};
14326 let s = &self.gpu.stream();
14327 let (p0, _g0) = b0.device_ptr(s);
14328 let (p1, _g1) = b1.device_ptr(s);
14329 let (paq, _g2) = aq.device_ptr(s);
14330 let (pad, _g3) = ad.device_ptr(s);
14331 let (py0, _g4) = y0.device_ptr_mut(s);
14332 let (py1, _g5) = y1.device_ptr_mut(s);
14333 let mut ps = [
14334 &p0 as *const _ as *mut std::ffi::c_void,
14335 &p1 as *const _ as *mut _,
14336 &paq as *const _ as *mut _,
14337 &pad as *const _ as *mut _,
14338 &py0 as *const _ as *mut _,
14339 &py1 as *const _ as *mut _,
14340 &inf as *const _ as *mut _,
14341 &oo0 as *const _ as *mut _,
14342 &oo1 as *const _ as *mut _,
14343 &r0 as *const _ as *mut _,
14344 &r1 as *const _ as *mut _,
14345 ];
14346 unsafe {
14347 self.launch_pdl(
14348 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14349 (grid, 1, 1),
14350 (32, rpb, 1),
14351 &mut ps,
14352 )?;
14353 }
14354 }
14355 return Ok(Some((y0, y1)));
14356 }
14357 let __s_b = self.gpu.stream();
14358 let mut b = __s_b.launch_builder(&f);
14359 b.arg(b0)
14360 .arg(b1)
14361 .arg(aq)
14362 .arg(ad)
14363 .arg(&mut y0)
14364 .arg(&mut y1)
14365 .arg(&inf)
14366 .arg(&oo0)
14367 .arg(&oo1)
14368 .arg(&r0)
14369 .arg(&r1);
14370 unsafe {
14371 b.launch(cfg)?;
14372 }
14373 Ok(Some((y0, y1)))
14374 }
14375
14376 pub fn matmul_q4_fused2_into(
14378 &self,
14379 w0: &crate::model::GpuTensor,
14380 w1: &crate::model::GpuTensor,
14381 aq: &CudaSlice<i8>,
14382 ad: &CudaSlice<f32>,
14383 y0: &mut CudaSlice<f32>,
14384 y1: &mut CudaSlice<f32>,
14385 ) -> Result<bool, Box<dyn std::error::Error>> {
14386 use crate::model::GpuTensor;
14387 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14388 match w {
14389 GpuTensor::Quant {
14390 qtype, row_bytes, ..
14391 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14392 _ => None,
14393 }
14394 };
14395 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14396 return Ok(false);
14397 };
14398 if w0.in_features() != w1.in_features() {
14399 return Ok(false);
14400 }
14401 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14402 match w {
14403 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14404 Some(m) => (m, true),
14405 None => (bytes, *rp),
14406 },
14407 _ => unreachable!(),
14408 }
14409 }
14410 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14411 if rp0 != rp1 {
14412 return Ok(false);
14413 }
14414 let rp = rp0;
14415 let rpb: u32 = 4;
14416 let mr1 = rp && Self::q40_mr1_on();
14417 let nb = |o: usize| {
14418 if mr1 {
14419 (o as u32).div_ceil(rpb)
14420 } else {
14421 (o as u32).div_ceil(2).div_ceil(rpb)
14422 }
14423 };
14424 let grid = nb(o0) + nb(o1);
14425 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14426 let f = self.func(if mr1 {
14427 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14428 } else if rp {
14429 "qmatvec_q4_0_mmvq_fused2_rp"
14430 } else {
14431 "qmatvec_q4_0_mmvq_fused2"
14432 });
14433 let cfg = LaunchConfig {
14434 grid_dim: (grid, 1, 1),
14435 block_dim: (32, rpb, 1),
14436 shared_mem_bytes: 0,
14437 };
14438 let inf = w0.in_features() as i32;
14439 let (oo0, oo1) = (o0 as i32, o1 as i32);
14440 let (r0, r1) = (rb0 as i64, rb1 as i64);
14441 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14443 use cudarc::driver::{DevicePtr, DevicePtrMut};
14444 let s = &self.gpu.stream();
14445 let (p0, _g0) = b0.device_ptr(s);
14446 let (p1, _g1) = b1.device_ptr(s);
14447 let (paq, _g2) = aq.device_ptr(s);
14448 let (pad, _g3) = ad.device_ptr(s);
14449 let (py0, _g4) = y0.device_ptr_mut(s);
14450 let (py1, _g5) = y1.device_ptr_mut(s);
14451 let mut ps = [
14452 &p0 as *const _ as *mut std::ffi::c_void,
14453 &p1 as *const _ as *mut _,
14454 &paq as *const _ as *mut _,
14455 &pad as *const _ as *mut _,
14456 &py0 as *const _ as *mut _,
14457 &py1 as *const _ as *mut _,
14458 &inf as *const _ as *mut _,
14459 &oo0 as *const _ as *mut _,
14460 &oo1 as *const _ as *mut _,
14461 &r0 as *const _ as *mut _,
14462 &r1 as *const _ as *mut _,
14463 ];
14464 unsafe {
14465 self.launch_pdl(
14466 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14467 (grid, 1, 1),
14468 (32, rpb, 1),
14469 &mut ps,
14470 )?;
14471 }
14472 return Ok(true);
14473 }
14474 let __s_b = self.gpu.stream();
14475 let mut b = __s_b.launch_builder(&f);
14476 b.arg(b0)
14477 .arg(b1)
14478 .arg(aq)
14479 .arg(ad)
14480 .arg(&mut *y0)
14481 .arg(&mut *y1)
14482 .arg(&inf)
14483 .arg(&oo0)
14484 .arg(&oo1)
14485 .arg(&r0)
14486 .arg(&r1);
14487 unsafe {
14488 b.launch(cfg)?;
14489 }
14490 Ok(true)
14491 }
14492
14493 pub fn matmul_q4_fused2_batched(
14498 &self,
14499 w0: &crate::model::GpuTensor,
14500 w1: &crate::model::GpuTensor,
14501 aq: &CudaSlice<i8>,
14502 ad: &CudaSlice<f32>,
14503 m: usize,
14504 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14505 use crate::model::GpuTensor;
14506 if m < 2 || m > 8 {
14507 return Ok(None);
14508 }
14509 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14510 match w {
14511 GpuTensor::Quant {
14512 qtype, row_bytes, ..
14513 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14514 _ => None,
14515 }
14516 };
14517 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14518 return Ok(None);
14519 };
14520 if w0.in_features() != w1.in_features() {
14521 return Ok(None);
14522 }
14523 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14524 match w {
14525 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14526 Some(mr) => (mr, true),
14527 None => (bytes, *rp),
14528 },
14529 _ => unreachable!(),
14530 }
14531 }
14532 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14533 if !rp0 || !rp1 {
14534 return Ok(None);
14535 }
14536 let mcols = Self::batched_mcols(m);
14537 let rpb: u32 = 4;
14538 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14539 let grid = nb(o0) + nb(o1);
14540 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14541 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14542 let f = self.func(match mcols {
14543 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14544 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14545 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14546 });
14547 let cfg = LaunchConfig {
14548 grid_dim: (grid, 1, 1),
14549 block_dim: (32, rpb, 1),
14550 shared_mem_bytes: 0,
14551 };
14552 let inf = w0.in_features() as i32;
14553 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14554 let rb = rb0 as i64;
14555 let __s_b = self.gpu.stream();
14556 let mut b = __s_b.launch_builder(&f);
14557 b.arg(b0)
14558 .arg(b1)
14559 .arg(aq)
14560 .arg(ad)
14561 .arg(&mut y0)
14562 .arg(&mut y1)
14563 .arg(&inf)
14564 .arg(&oo0)
14565 .arg(&oo1)
14566 .arg(&mi)
14567 .arg(&rb);
14568 unsafe {
14569 b.launch(cfg)?;
14570 }
14571 Ok(Some((y0, y1)))
14572 }
14573
14574 #[allow(clippy::too_many_arguments)]
14577 pub fn matmul_q4_fused3_batched(
14578 &self,
14579 w0: &crate::model::GpuTensor,
14580 w1: &crate::model::GpuTensor,
14581 w2: &crate::model::GpuTensor,
14582 aq: &CudaSlice<i8>,
14583 ad: &CudaSlice<f32>,
14584 m: usize,
14585 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14586 {
14587 use crate::model::GpuTensor;
14588 if m < 2 || m > 8 {
14589 return Ok(None);
14590 }
14591 let q4 = |w: &GpuTensor| -> Option<usize> {
14592 match w {
14593 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14594 _ => None,
14595 }
14596 };
14597 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14598 return Ok(None);
14599 };
14600 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14601 return Ok(None);
14602 }
14603 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14604 match w {
14605 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14606 Some(mr) => (mr, true),
14607 None => (bytes, *rp),
14608 },
14609 _ => unreachable!(),
14610 }
14611 }
14612 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14613 if !rp0 || !rp1 || !rp2 {
14614 return Ok(None);
14615 }
14616 let mcols = Self::batched_mcols(m);
14617 let rpb: u32 = 4;
14618 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14619 let grid = nb(o0) + nb(o1) + nb(o2);
14620 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14621 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14622 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14623 let f = self.func(match mcols {
14624 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14625 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14626 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14627 });
14628 let cfg = LaunchConfig {
14629 grid_dim: (grid, 1, 1),
14630 block_dim: (32, rpb, 1),
14631 shared_mem_bytes: 0,
14632 };
14633 let inf = w0.in_features() as i32;
14634 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14635 let rb = 0i64;
14636 let __s_b = self.gpu.stream();
14637 let mut b = __s_b.launch_builder(&f);
14638 b.arg(b0)
14639 .arg(b1)
14640 .arg(b2)
14641 .arg(aq)
14642 .arg(ad)
14643 .arg(&mut y0)
14644 .arg(&mut y1)
14645 .arg(&mut y2)
14646 .arg(&inf)
14647 .arg(&oo0)
14648 .arg(&oo1)
14649 .arg(&oo2)
14650 .arg(&mi)
14651 .arg(&rb);
14652 unsafe {
14653 b.launch(cfg)?;
14654 }
14655 Ok(Some((y0, y1, y2)))
14656 }
14657
14658 pub fn matmul_q8_fused3(
14659 &self,
14660 w0: &crate::model::GpuTensor,
14661 w1: &crate::model::GpuTensor,
14662 w2: &crate::model::GpuTensor,
14663 aq: &CudaSlice<i8>,
14664 ad: &CudaSlice<f32>,
14665 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14666 {
14667 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14670 return Ok(Some(self.e4m3_fused3_core(
14671 p0.0,
14672 p1.0,
14673 p2.0,
14674 aq,
14675 ad,
14676 w0.in_features(),
14677 p0.1,
14678 p1.1,
14679 p2.1,
14680 p0.2,
14681 p0.3,
14682 p1.3,
14683 p2.3,
14684 )?));
14685 }
14686 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14687 return Ok(None);
14688 };
14689 Ok(Some(self.q8_fused3_core(
14690 p0.0,
14691 p1.0,
14692 p2.0,
14693 aq,
14694 ad,
14695 w0.in_features(),
14696 p0.1,
14697 p1.1,
14698 p2.1,
14699 p0.2,
14700 )?))
14701 }
14702
14703 #[allow(clippy::too_many_arguments)]
14704 fn q8_fused3_core(
14705 &self,
14706 b0: &CudaSlice<u8>,
14707 b1: &CudaSlice<u8>,
14708 b2: &CudaSlice<u8>,
14709 aq: &CudaSlice<i8>,
14710 ad: &CudaSlice<f32>,
14711 in_f: usize,
14712 out0: usize,
14713 out1: usize,
14714 out2: usize,
14715 row_bytes: usize,
14716 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14717 const ROWS_PER_BLOCK: u32 = 4;
14718 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14719 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14720 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14721 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14722 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14723 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14724 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14725 let cfg = LaunchConfig {
14726 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14727 block_dim: (32, ROWS_PER_BLOCK, 1),
14728 shared_mem_bytes: 0,
14729 };
14730 let (inf, o0, o1, o2, rbl) = (
14731 in_f as i32,
14732 out0 as i32,
14733 out1 as i32,
14734 out2 as i32,
14735 row_bytes as i64,
14736 );
14737 let __s_b = self.gpu.stream();
14738 let mut b = __s_b.launch_builder(&f);
14739 b.arg(b0)
14740 .arg(b1)
14741 .arg(b2)
14742 .arg(aq)
14743 .arg(ad)
14744 .arg(&mut y0)
14745 .arg(&mut y1)
14746 .arg(&mut y2)
14747 .arg(&inf)
14748 .arg(&o0)
14749 .arg(&o1)
14750 .arg(&o2)
14751 .arg(&rbl);
14752 unsafe {
14753 b.launch(cfg)?;
14754 }
14755 Ok((y0, y1, y2))
14756 }
14757
14758 #[allow(clippy::too_many_arguments)]
14760 pub fn qmatvec_q8_fused3_raw(
14761 &self,
14762 b0: &CudaSlice<u8>,
14763 b1: &CudaSlice<u8>,
14764 b2: &CudaSlice<u8>,
14765 x: &CudaSlice<f32>,
14766 in_f: usize,
14767 out0: usize,
14768 out1: usize,
14769 out2: usize,
14770 row_bytes: usize,
14771 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14772 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14773 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14774 }
14775
14776 pub fn matmul_q8_fused2_t(
14787 &self,
14788 w0: &crate::model::GpuTensor,
14789 w1: &crate::model::GpuTensor,
14790 aq: &CudaSlice<i8>,
14791 ad: &CudaSlice<f32>,
14792 m: usize,
14793 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14794 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14798 return Ok(None);
14799 }
14800 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14803 if m > 4 && !Self::b8_enabled() {
14804 return Ok(None);
14805 }
14806 return Ok(Some(self.e4m3_fused2_t_core(
14807 p0.0,
14808 p1.0,
14809 aq,
14810 ad,
14811 m,
14812 w0.in_features(),
14813 p0.1,
14814 p1.1,
14815 p0.2,
14816 p0.3,
14817 p1.3,
14818 )?));
14819 }
14820 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14821 return Ok(None);
14822 };
14823 Ok(Some(self.q8_fused2_t_core(
14824 p0.0,
14825 p1.0,
14826 aq,
14827 ad,
14828 m,
14829 w0.in_features(),
14830 p0.1,
14831 p1.1,
14832 p0.2,
14833 )?))
14834 }
14835
14836 #[allow(clippy::too_many_arguments)]
14837 fn q8_fused2_t_core(
14838 &self,
14839 b0: &CudaSlice<u8>,
14840 b1: &CudaSlice<u8>,
14841 aq: &CudaSlice<i8>,
14842 ad: &CudaSlice<f32>,
14843 m: usize,
14844 in_f: usize,
14845 out0: usize,
14846 out1: usize,
14847 row_bytes: usize,
14848 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14849 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14851 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14852 let f = self.func(match Self::batched_mcols(m) {
14853 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14854 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14855 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14857 });
14858 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14859 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14860 let cfg = LaunchConfig {
14861 grid_dim: (nb0 + nb1, 1, 1),
14862 block_dim: (32, ROWS_PER_BLOCK, 1),
14863 shared_mem_bytes: 0,
14864 };
14865 let (inf, o0, o1, mi, rbl) = (
14866 in_f as i32,
14867 out0 as i32,
14868 out1 as i32,
14869 m as i32,
14870 row_bytes as i64,
14871 );
14872 let __s_b = self.gpu.stream();
14873 let mut b = __s_b.launch_builder(&f);
14874 b.arg(b0)
14875 .arg(b1)
14876 .arg(aq)
14877 .arg(ad)
14878 .arg(&mut y0)
14879 .arg(&mut y1)
14880 .arg(&inf)
14881 .arg(&o0)
14882 .arg(&o1)
14883 .arg(&mi)
14884 .arg(&rbl);
14885 unsafe {
14886 b.launch(cfg)?;
14887 }
14888 Ok((y0, y1))
14889 }
14890
14891 #[allow(clippy::too_many_arguments)]
14894 pub fn qmatvec_q8_fused2_t_raw(
14895 &self,
14896 b0: &CudaSlice<u8>,
14897 b1: &CudaSlice<u8>,
14898 x: &CudaSlice<f32>,
14899 m: usize,
14900 in_f: usize,
14901 out0: usize,
14902 out1: usize,
14903 row_bytes: usize,
14904 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14905 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14906 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14907 }
14908
14909 #[allow(clippy::too_many_arguments)]
14912 pub fn matmul_q8_fused3_t(
14913 &self,
14914 w0: &crate::model::GpuTensor,
14915 w1: &crate::model::GpuTensor,
14916 w2: &crate::model::GpuTensor,
14917 aq: &CudaSlice<i8>,
14918 ad: &CudaSlice<f32>,
14919 m: usize,
14920 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14921 {
14922 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14923 return Ok(None);
14924 }
14925 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14926 return Ok(Some(self.e4m3_fused3_t_core(
14927 p0.0,
14928 p1.0,
14929 p2.0,
14930 aq,
14931 ad,
14932 m,
14933 w0.in_features(),
14934 p0.1,
14935 p1.1,
14936 p2.1,
14937 p0.2,
14938 p0.3,
14939 p1.3,
14940 p2.3,
14941 )?));
14942 }
14943 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14944 return Ok(None);
14945 };
14946 Ok(Some(self.q8_fused3_t_core(
14947 p0.0,
14948 p1.0,
14949 p2.0,
14950 aq,
14951 ad,
14952 m,
14953 w0.in_features(),
14954 p0.1,
14955 p1.1,
14956 p2.1,
14957 p0.2,
14958 )?))
14959 }
14960
14961 #[allow(clippy::too_many_arguments)]
14962 fn q8_fused3_t_core(
14963 &self,
14964 b0: &CudaSlice<u8>,
14965 b1: &CudaSlice<u8>,
14966 b2: &CudaSlice<u8>,
14967 aq: &CudaSlice<i8>,
14968 ad: &CudaSlice<f32>,
14969 m: usize,
14970 in_f: usize,
14971 out0: usize,
14972 out1: usize,
14973 out2: usize,
14974 row_bytes: usize,
14975 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14976 const ROWS_PER_BLOCK: u32 = 4;
14977 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14978 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14979 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14980 let f = self.func(if Self::batched_mcols(m) == 2 {
14981 "qmatvec_q8_0_mmvq_fused3_b2"
14982 } else {
14983 "qmatvec_q8_0_mmvq_fused3_b4"
14984 });
14985 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14986 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14987 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14988 let cfg = LaunchConfig {
14989 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14990 block_dim: (32, ROWS_PER_BLOCK, 1),
14991 shared_mem_bytes: 0,
14992 };
14993 let (inf, o0, o1, o2, mi, rbl) = (
14994 in_f as i32,
14995 out0 as i32,
14996 out1 as i32,
14997 out2 as i32,
14998 m as i32,
14999 row_bytes as i64,
15000 );
15001 let __s_b = self.gpu.stream();
15002 let mut b = __s_b.launch_builder(&f);
15003 b.arg(b0)
15004 .arg(b1)
15005 .arg(b2)
15006 .arg(aq)
15007 .arg(ad)
15008 .arg(&mut y0)
15009 .arg(&mut y1)
15010 .arg(&mut y2)
15011 .arg(&inf)
15012 .arg(&o0)
15013 .arg(&o1)
15014 .arg(&o2)
15015 .arg(&mi)
15016 .arg(&rbl);
15017 unsafe {
15018 b.launch(cfg)?;
15019 }
15020 Ok((y0, y1, y2))
15021 }
15022
15023 #[allow(clippy::too_many_arguments)]
15025 pub fn qmatvec_q8_fused3_t_raw(
15026 &self,
15027 b0: &CudaSlice<u8>,
15028 b1: &CudaSlice<u8>,
15029 b2: &CudaSlice<u8>,
15030 x: &CudaSlice<f32>,
15031 m: usize,
15032 in_f: usize,
15033 out0: usize,
15034 out1: usize,
15035 out2: usize,
15036 row_bytes: usize,
15037 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15038 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15039 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
15040 }
15041
15042 pub fn q8_ffn_fuse2_on(&self) -> bool {
15046 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15047 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15048 }
15049
15050 #[allow(clippy::type_complexity)]
15056 fn q8_fused_params<'w, const N: usize>(
15057 &self,
15058 ws: &[&'w crate::model::GpuTensor; N],
15059 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15060 use crate::model::GpuTensor;
15061 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15062 return None;
15063 }
15064 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15065 return None;
15066 }
15067 let in_f = ws[0].in_features();
15068 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15069 for (i, w) in ws.iter().enumerate() {
15070 match w {
15071 GpuTensor::Quant {
15072 bytes,
15073 qtype,
15074 row_bytes,
15075 scale,
15076 ..
15077 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15078 out[i] = Some((bytes, w.out_features(), *row_bytes))
15079 }
15080 _ => return None,
15081 }
15082 }
15083 Some(out.map(|o| o.unwrap()))
15084 }
15085
15086 pub fn e4m3_dual_on(&self) -> bool {
15089 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15090 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15091 }
15092
15093 #[allow(clippy::type_complexity)]
15105 fn e4m3_fused_params<'w, const N: usize>(
15106 &self,
15107 ws: &[&'w crate::model::GpuTensor; N],
15108 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15109 use crate::model::GpuTensor;
15110 if !self.e4m3_dual_on() {
15111 return None;
15112 }
15113 let in_f = ws[0].in_features();
15114 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15115 for (i, w) in ws.iter().enumerate() {
15116 match w {
15117 GpuTensor::Quant {
15118 bytes,
15119 qtype,
15120 row_bytes,
15121 scale,
15122 rp,
15123 rp4,
15124 ..
15125 } if *qtype == QT_F8_E4M3
15126 && w.in_features() == in_f
15127 && *row_bytes == in_f
15128 && !*rp
15129 && rp4.is_none() =>
15130 {
15131 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15132 }
15133 _ => return None,
15134 }
15135 }
15136 Some(out.map(|o| o.unwrap()))
15137 }
15138
15139 #[allow(clippy::too_many_arguments)]
15143 fn e4m3_fused2_core(
15144 &self,
15145 b0: &CudaSlice<u8>,
15146 b1: &CudaSlice<u8>,
15147 aq: &CudaSlice<i8>,
15148 ad: &CudaSlice<f32>,
15149 in_f: usize,
15150 out0: usize,
15151 out1: usize,
15152 row_bytes: usize,
15153 ws0: f32,
15154 ws1: f32,
15155 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15156 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15158 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15159 let f = self.func("qmatvec_e4m3_mmvq_fused2");
15160 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15161 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15162 let cfg = LaunchConfig {
15163 grid_dim: (nb0 + nb1, 1, 1),
15164 block_dim: (32, ROWS_PER_BLOCK, 1),
15165 shared_mem_bytes: 0,
15166 };
15167 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15168 let __s_b = self.gpu.stream();
15169 let mut b = __s_b.launch_builder(&f);
15170 b.arg(b0)
15171 .arg(b1)
15172 .arg(aq)
15173 .arg(ad)
15174 .arg(&mut y0)
15175 .arg(&mut y1)
15176 .arg(&inf)
15177 .arg(&o0)
15178 .arg(&o1)
15179 .arg(&rbl)
15180 .arg(&ws0)
15181 .arg(&ws1);
15182 unsafe {
15183 b.launch(cfg)?;
15184 }
15185 Ok((y0, y1))
15186 }
15187
15188 #[allow(clippy::too_many_arguments)]
15190 fn e4m3_fused3_core(
15191 &self,
15192 b0: &CudaSlice<u8>,
15193 b1: &CudaSlice<u8>,
15194 b2: &CudaSlice<u8>,
15195 aq: &CudaSlice<i8>,
15196 ad: &CudaSlice<f32>,
15197 in_f: usize,
15198 out0: usize,
15199 out1: usize,
15200 out2: usize,
15201 row_bytes: usize,
15202 ws0: f32,
15203 ws1: f32,
15204 ws2: f32,
15205 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15206 const ROWS_PER_BLOCK: u32 = 4;
15207 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15208 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15209 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15210 let f = self.func("qmatvec_e4m3_mmvq_fused3");
15211 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15212 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15213 let mut y2 = self.alloc_uninit::<f32>(out2)?;
15214 let cfg = LaunchConfig {
15215 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15216 block_dim: (32, ROWS_PER_BLOCK, 1),
15217 shared_mem_bytes: 0,
15218 };
15219 let (inf, o0, o1, o2, rbl) = (
15220 in_f as i32,
15221 out0 as i32,
15222 out1 as i32,
15223 out2 as i32,
15224 row_bytes as i64,
15225 );
15226 let __s_b = self.gpu.stream();
15227 let mut b = __s_b.launch_builder(&f);
15228 b.arg(b0)
15229 .arg(b1)
15230 .arg(b2)
15231 .arg(aq)
15232 .arg(ad)
15233 .arg(&mut y0)
15234 .arg(&mut y1)
15235 .arg(&mut y2)
15236 .arg(&inf)
15237 .arg(&o0)
15238 .arg(&o1)
15239 .arg(&o2)
15240 .arg(&rbl)
15241 .arg(&ws0)
15242 .arg(&ws1)
15243 .arg(&ws2);
15244 unsafe {
15245 b.launch(cfg)?;
15246 }
15247 Ok((y0, y1, y2))
15248 }
15249
15250 #[allow(clippy::too_many_arguments)]
15254 fn e4m3_fused2_t_core(
15255 &self,
15256 b0: &CudaSlice<u8>,
15257 b1: &CudaSlice<u8>,
15258 aq: &CudaSlice<i8>,
15259 ad: &CudaSlice<f32>,
15260 m: usize,
15261 in_f: usize,
15262 out0: usize,
15263 out1: usize,
15264 row_bytes: usize,
15265 ws0: f32,
15266 ws1: f32,
15267 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15268 const ROWS_PER_BLOCK: u32 = 4;
15269 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15270 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15271 let f = self.func(match Self::batched_mcols(m) {
15272 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15273 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15274 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15275 });
15276 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15277 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15278 let cfg = LaunchConfig {
15279 grid_dim: (nb0 + nb1, 1, 1),
15280 block_dim: (32, ROWS_PER_BLOCK, 1),
15281 shared_mem_bytes: 0,
15282 };
15283 let (inf, o0, o1, mi, rbl) = (
15284 in_f as i32,
15285 out0 as i32,
15286 out1 as i32,
15287 m as i32,
15288 row_bytes as i64,
15289 );
15290 let __s_b = self.gpu.stream();
15291 let mut b = __s_b.launch_builder(&f);
15292 b.arg(b0)
15293 .arg(b1)
15294 .arg(aq)
15295 .arg(ad)
15296 .arg(&mut y0)
15297 .arg(&mut y1)
15298 .arg(&inf)
15299 .arg(&o0)
15300 .arg(&o1)
15301 .arg(&mi)
15302 .arg(&rbl);
15303 unsafe {
15304 b.launch(cfg)?;
15305 }
15306 if ws0 != 1.0 {
15307 self.scale_inplace(&mut y0, ws0, m * out0)?;
15308 }
15309 if ws1 != 1.0 {
15310 self.scale_inplace(&mut y1, ws1, m * out1)?;
15311 }
15312 Ok((y0, y1))
15313 }
15314
15315 #[allow(clippy::too_many_arguments)]
15317 fn e4m3_fused3_t_core(
15318 &self,
15319 b0: &CudaSlice<u8>,
15320 b1: &CudaSlice<u8>,
15321 b2: &CudaSlice<u8>,
15322 aq: &CudaSlice<i8>,
15323 ad: &CudaSlice<f32>,
15324 m: usize,
15325 in_f: usize,
15326 out0: usize,
15327 out1: usize,
15328 out2: usize,
15329 row_bytes: usize,
15330 ws0: f32,
15331 ws1: f32,
15332 ws2: f32,
15333 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15334 const ROWS_PER_BLOCK: u32 = 4;
15335 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15336 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15337 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15338 let f = self.func(if Self::batched_mcols(m) == 2 {
15339 "qmatvec_e4m3_mmvq_fused3_b2"
15340 } else {
15341 "qmatvec_e4m3_mmvq_fused3_b4"
15342 });
15343 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15344 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15345 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15346 let cfg = LaunchConfig {
15347 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15348 block_dim: (32, ROWS_PER_BLOCK, 1),
15349 shared_mem_bytes: 0,
15350 };
15351 let (inf, o0, o1, o2, mi, rbl) = (
15352 in_f as i32,
15353 out0 as i32,
15354 out1 as i32,
15355 out2 as i32,
15356 m as i32,
15357 row_bytes as i64,
15358 );
15359 let __s_b = self.gpu.stream();
15360 let mut b = __s_b.launch_builder(&f);
15361 b.arg(b0)
15362 .arg(b1)
15363 .arg(b2)
15364 .arg(aq)
15365 .arg(ad)
15366 .arg(&mut y0)
15367 .arg(&mut y1)
15368 .arg(&mut y2)
15369 .arg(&inf)
15370 .arg(&o0)
15371 .arg(&o1)
15372 .arg(&o2)
15373 .arg(&mi)
15374 .arg(&rbl);
15375 unsafe {
15376 b.launch(cfg)?;
15377 }
15378 if ws0 != 1.0 {
15379 self.scale_inplace(&mut y0, ws0, m * out0)?;
15380 }
15381 if ws1 != 1.0 {
15382 self.scale_inplace(&mut y1, ws1, m * out1)?;
15383 }
15384 if ws2 != 1.0 {
15385 self.scale_inplace(&mut y2, ws2, m * out2)?;
15386 }
15387 Ok((y0, y1, y2))
15388 }
15389
15390 pub fn qmatvec_e4m3_blk_mmvq(
15400 &self,
15401 bytes: &CudaSlice<u8>,
15402 aq: &CudaSlice<i8>,
15403 ad: &CudaSlice<f32>,
15404 scales: &CudaSlice<f32>,
15405 m: usize,
15406 in_f: usize,
15407 out_f: usize,
15408 row_bytes: usize,
15409 scale_cols: usize,
15410 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15411 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15413 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15414 )?;
15415 Ok(y)
15416 }
15417
15418 #[allow(clippy::too_many_arguments)]
15420 pub fn qmatvec_e4m3_blk_mmvq_into(
15421 &self,
15422 bytes: &CudaSlice<u8>,
15423 aq: &CudaSlice<i8>,
15424 ad: &CudaSlice<f32>,
15425 scales: &CudaSlice<f32>,
15426 m: usize,
15427 in_f: usize,
15428 out_f: usize,
15429 row_bytes: usize,
15430 scale_cols: usize,
15431 y: &mut CudaSlice<f32>,
15432 ) -> Result<(), Box<dyn std::error::Error>> {
15433 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15435 let cfg = LaunchConfig {
15436 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15437 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15440 let (inf, outf, mi, rb, sc) = (
15441 in_f as i32,
15442 out_f as i32,
15443 m as i32,
15444 row_bytes as i64,
15445 scale_cols as i32,
15446 );
15447 let __s_b = self.gpu.stream();
15448 let mut b = __s_b.launch_builder(&f);
15449 b.arg(bytes)
15450 .arg(aq)
15451 .arg(ad)
15452 .arg(scales)
15453 .arg(&mut *y)
15454 .arg(&inf)
15455 .arg(&outf)
15456 .arg(&mi)
15457 .arg(&rb)
15458 .arg(&sc);
15459 unsafe {
15460 b.launch(cfg)?;
15461 }
15462 Ok(())
15463 }
15464
15465 #[allow(clippy::too_many_arguments)]
15471 pub fn qmatvec_e4m3_blk_mmvq_batched(
15472 &self,
15473 bytes: &CudaSlice<u8>,
15474 aq: &CudaSlice<i8>,
15475 ad: &CudaSlice<f32>,
15476 scales: &CudaSlice<f32>,
15477 m: usize,
15478 in_f: usize,
15479 out_f: usize,
15480 row_bytes: usize,
15481 scale_cols: usize,
15482 mcols: usize,
15483 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15484 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15486 let name = match mcols {
15487 2 => "qmatvec_e4m3_blk_mmvq_b2",
15488 4 => "qmatvec_e4m3_blk_mmvq_b4",
15489 8 => "qmatvec_e4m3_blk_mmvq_b8",
15490 16 => "qmatvec_e4m3_blk_mmvq_b16",
15491 _ => {
15492 return Err(
15493 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15494 );
15495 }
15496 };
15497 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15498 let f = self.func(name);
15499 let cfg = LaunchConfig {
15500 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15501 block_dim: (32, ROWS_PER_BLOCK, 1),
15502 shared_mem_bytes: 0,
15503 };
15504 let (inf, outf, mi, rb, sc) = (
15505 in_f as i32,
15506 out_f as i32,
15507 m as i32,
15508 row_bytes as i64,
15509 scale_cols as i32,
15510 );
15511 let __s_b = self.gpu.stream();
15512 let mut b = __s_b.launch_builder(&f);
15513 b.arg(bytes)
15514 .arg(aq)
15515 .arg(ad)
15516 .arg(scales)
15517 .arg(&mut y)
15518 .arg(&inf)
15519 .arg(&outf)
15520 .arg(&mi)
15521 .arg(&rb)
15522 .arg(&sc);
15523 unsafe {
15524 b.launch(cfg)?;
15525 }
15526 Ok(y)
15527 }
15528
15529 #[allow(clippy::too_many_arguments)]
15532 pub fn qmatvec_e4m3_blk_batched_raw(
15533 &self,
15534 bytes: &CudaSlice<u8>,
15535 x: &CudaSlice<f32>,
15536 scales: &CudaSlice<f32>,
15537 m: usize,
15538 in_f: usize,
15539 out_f: usize,
15540 row_bytes: usize,
15541 scale_cols: usize,
15542 mcols: usize,
15543 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15544 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15545 self.qmatvec_e4m3_blk_mmvq_batched(
15546 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15547 )
15548 }
15549
15550 #[allow(clippy::too_many_arguments)]
15553 pub fn qmatvec_e4m3_blk_mmvq_raw(
15554 &self,
15555 bytes: &CudaSlice<u8>,
15556 x: &CudaSlice<f32>,
15557 scales: &CudaSlice<f32>,
15558 m: usize,
15559 in_f: usize,
15560 out_f: usize,
15561 row_bytes: usize,
15562 scale_cols: usize,
15563 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15564 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15565 self.qmatvec_e4m3_blk_mmvq(
15566 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15567 )
15568 }
15569
15570 #[allow(clippy::too_many_arguments)]
15573 pub fn qmatvec_e4m3_fused2_raw(
15574 &self,
15575 b0: &CudaSlice<u8>,
15576 b1: &CudaSlice<u8>,
15577 x: &CudaSlice<f32>,
15578 in_f: usize,
15579 out0: usize,
15580 out1: usize,
15581 row_bytes: usize,
15582 ws0: f32,
15583 ws1: f32,
15584 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15585 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15586 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15587 }
15588
15589 #[allow(clippy::too_many_arguments)]
15590 pub fn qmatvec_e4m3_fused3_raw(
15591 &self,
15592 b0: &CudaSlice<u8>,
15593 b1: &CudaSlice<u8>,
15594 b2: &CudaSlice<u8>,
15595 x: &CudaSlice<f32>,
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, 1, in_f)?;
15606 self.e4m3_fused3_core(
15607 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15608 )
15609 }
15610
15611 #[allow(clippy::too_many_arguments)]
15612 pub fn qmatvec_e4m3_fused2_t_raw(
15613 &self,
15614 b0: &CudaSlice<u8>,
15615 b1: &CudaSlice<u8>,
15616 x: &CudaSlice<f32>,
15617 m: usize,
15618 in_f: usize,
15619 out0: usize,
15620 out1: usize,
15621 row_bytes: usize,
15622 ws0: f32,
15623 ws1: f32,
15624 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15625 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15626 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15627 }
15628
15629 #[allow(clippy::too_many_arguments)]
15630 pub fn qmatvec_e4m3_fused3_t_raw(
15631 &self,
15632 b0: &CudaSlice<u8>,
15633 b1: &CudaSlice<u8>,
15634 b2: &CudaSlice<u8>,
15635 x: &CudaSlice<f32>,
15636 m: usize,
15637 in_f: usize,
15638 out0: usize,
15639 out1: usize,
15640 out2: usize,
15641 row_bytes: usize,
15642 ws0: f32,
15643 ws1: f32,
15644 ws2: f32,
15645 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15646 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15647 self.e4m3_fused3_t_core(
15648 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15649 )
15650 }
15651
15652 fn try_e4m3_blk_pre(
15663 &self,
15664 w: &crate::model::GpuTensor,
15665 aq: &CudaSlice<i8>,
15666 ad: &CudaSlice<f32>,
15667 m: usize,
15668 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15669 use crate::model::GpuTensor;
15670 if let GpuTensor::Quant {
15671 bytes,
15672 qtype,
15673 row_bytes,
15674 blk: Some(g),
15675 ..
15676 } = w
15677 {
15678 if *qtype == QT_F8_E4M3_BLK {
15679 if (2..=16).contains(&m)
15685 && std::env::var("MEMRA_NO_BATCHED").is_err()
15686 && (m <= 4 || Self::b8_enabled())
15687 {
15688 let mcols = Self::batched_mcols(m);
15689 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15690 bytes,
15691 aq,
15692 ad,
15693 &g.scales,
15694 m,
15695 w.in_features(),
15696 w.out_features(),
15697 *row_bytes,
15698 g.cols,
15699 mcols,
15700 )?));
15701 }
15702 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15703 bytes,
15704 aq,
15705 ad,
15706 &g.scales,
15707 m,
15708 w.in_features(),
15709 w.out_features(),
15710 *row_bytes,
15711 g.cols,
15712 )?));
15713 }
15714 }
15715 Ok(None)
15716 }
15717
15718 fn try_e4m3_blk_prefill(
15765 &self,
15766 w: &crate::model::GpuTensor,
15767 x: &CudaSlice<f32>,
15768 m: usize,
15769 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15770 use crate::model::GpuTensor;
15771 let GpuTensor::Quant {
15772 bytes,
15773 qtype,
15774 blk: Some(g),
15775 ..
15776 } = w
15777 else {
15778 return Ok(None);
15779 };
15780 if *qtype != QT_F8_E4M3_BLK {
15781 return Ok(None);
15782 }
15783 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15788 return Ok(Some(y));
15789 }
15790 let (in_f, out_f) = (w.in_features(), w.out_features());
15791 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15792 let tmp = GpuTensor::Quant {
15793 bytes: slab,
15794 qtype: QT_Q8_0,
15795 row_bytes: in_f / 32 * 34,
15796 ne: vec![in_f as u64, out_f as u64],
15797 scale: 1.0,
15798 rp: false,
15799 #[cfg(memra_cutlass)]
15800 cutlass: None,
15801 fp8: None,
15802 blk: None,
15803 f16: None,
15804 rp4: None,
15805 };
15806 Ok(Some(self.matmul(&tmp, x, m)?))
15808 }
15809
15810 pub fn matmul_pre_noscale(
15811 &self,
15812 w: &crate::model::GpuTensor,
15813 aq: &CudaSlice<i8>,
15814 ad: &CudaSlice<f32>,
15815 m: usize,
15816 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15817 use crate::model::GpuTensor;
15818 if m == 1 {
15822 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15823 return Ok(Some((y, 1.0)));
15824 }
15825 }
15826 if m != 1 || !self.uses_q8_1_fast(w) {
15828 return Ok(None);
15829 }
15830 let in_f = w.in_features();
15831 let out_f = w.out_features();
15832 let (bytes, qtype, row_bytes, scale, rp) = match w {
15833 GpuTensor::Quant {
15834 bytes,
15835 qtype,
15836 row_bytes,
15837 scale,
15838 rp,
15839 ..
15840 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15841 _ => return Ok(None),
15842 };
15843 if self.mmvq_supports(qtype) {
15845 let (mbytes, mrp) = match w {
15847 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15848 _ => (bytes, rp),
15849 };
15850 let y = self.qmatvec_mmvq(
15851 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15852 )?;
15853 return Ok(Some((y, scale)));
15854 }
15855 let name = match qtype {
15857 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15858 QT_Q4_K => "qmatvec_q4_K_dp4a",
15859 QT_Q6_K => "qmatvec_q6_K_dp4a",
15860 QT_Q5_K => "qmatvec_q5_K_dp4a",
15861 QT_Q3_K => "qmatvec_q3_K_dp4a",
15862 QT_NVFP4 => {
15863 if rp {
15864 "qmatvec_nvfp4_dp4a_rp"
15865 } else {
15866 "qmatvec_nvfp4_dp4a"
15867 }
15868 }
15869 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15870 _ => return Ok(None),
15871 };
15872 let f = self.func(name);
15873 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15874 let cfg = LaunchConfig {
15875 grid_dim: (out_f as u32, m as u32, 1),
15876 block_dim: (128, 1, 1),
15877 shared_mem_bytes: 0,
15878 };
15879 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15880 let __s_b = self.gpu.stream();
15881 let mut b = __s_b.launch_builder(&f);
15882 b.arg(bytes)
15883 .arg(aq)
15884 .arg(ad)
15885 .arg(&mut y)
15886 .arg(&inf)
15887 .arg(&outf)
15888 .arg(&mi)
15889 .arg(&rb);
15890 unsafe {
15891 b.launch(cfg)?;
15892 }
15893 Ok(Some((y, scale)))
15894 }
15895
15896 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15899 if qtype == QT_F8_E4M3 {
15904 return true;
15905 }
15906 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15907 return false;
15908 }
15909 matches!(
15910 qtype,
15911 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15912 )
15913 }
15914
15915 pub fn qmatvec_mmvq(
15920 &self,
15921 bytes: &CudaSlice<u8>,
15922 aq: &CudaSlice<i8>,
15923 ad: &CudaSlice<f32>,
15924 m: usize,
15925 in_f: usize,
15926 out_f: usize,
15927 qtype: i32,
15928 row_bytes: usize,
15929 scale: f32,
15930 rp: bool,
15931 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15932 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15934 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15935 )?;
15936 Ok(y)
15937 }
15938
15939 #[allow(clippy::too_many_arguments)]
15941 pub fn qmatvec_mmvq_into(
15942 &self,
15943 bytes: &CudaSlice<u8>,
15944 aq: &CudaSlice<i8>,
15945 ad: &CudaSlice<f32>,
15946 m: usize,
15947 in_f: usize,
15948 out_f: usize,
15949 qtype: i32,
15950 row_bytes: usize,
15951 scale: f32,
15952 rp: bool,
15953 y: &mut CudaSlice<f32>,
15954 ) -> Result<(), Box<dyn std::error::Error>> {
15955 debug_assert!(y.len() >= m * out_f);
15956 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15962 && rp
15963 && m == 1
15964 && out_f >= 64
15965 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15966 && {
15967 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15968 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15969 }
15970 {
15971 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15972 let cfg = LaunchConfig {
15973 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15974 block_dim: (32, 2, 1),
15975 shared_mem_bytes: 0,
15976 };
15977 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15978 let __s_b = self.gpu.stream();
15979 let mut b = __s_b.launch_builder(&f);
15980 b.arg(bytes)
15981 .arg(aq)
15982 .arg(ad)
15983 .arg(&mut *y)
15984 .arg(&inf)
15985 .arg(&outf)
15986 .arg(&mi)
15987 .arg(&rb);
15988 unsafe {
15989 b.launch(cfg)?;
15990 }
15991 if scale != 1.0 {
15992 self.scale_inplace(y, scale, out_f)?;
15993 }
15994 return Ok(());
15995 }
15996 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
16005 2
16006 } else {
16007 1
16008 };
16009 if m == 1 && qtype == QT_Q4_0 {
16014 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16015 mr = *Q40MR.get_or_init(|| {
16018 std::env::var("MEMRA_Q40_MR")
16019 .ok()
16020 .and_then(|v| v.parse().ok())
16021 .unwrap_or(1)
16022 });
16023 }
16024 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
16035 let q5_force = q5_mode.as_deref() == Some("2");
16036 let q5_il = qtype == QT_Q5_K
16039 && m == 1
16040 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16041 if q5_il && !q5_force && out_f > 65536 {
16042 mr = 1;
16043 }
16044 if qtype == QT_Q4_0 && rp && mr != 1 {
16047 mr = 2;
16048 }
16049 if qtype == QT_Q8_0 && rp {
16053 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16054 mr = *Q80MR.get_or_init(|| {
16055 std::env::var("MEMRA_Q80_MR")
16056 .ok()
16057 .and_then(|v| v.parse().ok())
16058 .unwrap_or(1)
16059 });
16060 }
16061 let name = match (qtype, mr, rp) {
16062 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16063 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16064 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16065 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16066 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16067 (QT_Q5_K, 2, _) => {
16068 if q5_il {
16069 "qmatvec_q5_K_mmvq_mr2_il"
16070 } else {
16071 "qmatvec_q5_K_mmvq_mr2"
16072 }
16073 }
16074 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16075 (QT_Q8_0, _, true)
16080 if in_f % 1024 == 0 && {
16081 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16082 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16083 } =>
16084 {
16085 "qmatvec_q8_0_mmvq_rpca"
16086 }
16087 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16088 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16089 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16093 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16094 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16095 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16096 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16097 (QT_Q5_K, _, _) => {
16098 if q5_il {
16099 "qmatvec_q5_K_mmvq_il"
16100 } else {
16101 "qmatvec_q5_K_mmvq"
16102 }
16103 }
16104 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16105 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16106 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16107 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16108 };
16109 let f = self.func(name);
16110 let rows_per_block = ROWS_PER_BLOCK * mr;
16112 let cfg = LaunchConfig {
16113 grid_dim: (
16114 (out_f as u32 + rows_per_block - 1) / rows_per_block,
16115 m as u32,
16116 1,
16117 ),
16118 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
16121 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16122 let __s_b = self.gpu.stream();
16123 let mut b = __s_b.launch_builder(&f);
16124 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16129 if Self::pdl_on()
16132 && Self::pdl_mmvq_on()
16133 && Self::pdl_nvfp4q8_on()
16134 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16135 {
16136 use cudarc::driver::{DevicePtr, DevicePtrMut};
16137 let s = &self.gpu.stream();
16138 let (pw, _g0) = bytes.device_ptr(s);
16139 let (paq, _g1) = aq.device_ptr(s);
16140 let (pad, _g2) = ad.device_ptr(s);
16141 let (py, _g3) = y.device_ptr_mut(s);
16142 let mut ps = [
16143 &pw as *const _ as *mut std::ffi::c_void,
16144 &paq as *const _ as *mut _,
16145 &pad as *const _ as *mut _,
16146 &py as *const _ as *mut _,
16147 &inf as *const _ as *mut _,
16148 &outf as *const _ as *mut _,
16149 &mi as *const _ as *mut _,
16150 &rb as *const _ as *mut _,
16151 &scale as *const _ as *mut _,
16152 ];
16153 unsafe {
16154 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16155 }
16156 return Ok(());
16157 }
16158 b.arg(bytes)
16159 .arg(aq)
16160 .arg(ad)
16161 .arg(&mut *y)
16162 .arg(&inf)
16163 .arg(&outf)
16164 .arg(&mi)
16165 .arg(&rb)
16166 .arg(&scale);
16167 unsafe {
16168 b.launch(cfg)?;
16169 }
16170 } else if Self::pdl_on()
16171 && Self::pdl_mmvq_on()
16172 && (matches!(
16173 name,
16174 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16175 ) || (Self::pdl_nvfp4q8_on()
16176 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16177 {
16178 {
16182 use cudarc::driver::{DevicePtr, DevicePtrMut};
16183 let s = &self.gpu.stream();
16184 let (pw, _g0) = bytes.device_ptr(s);
16185 let (paq, _g1) = aq.device_ptr(s);
16186 let (pad, _g2) = ad.device_ptr(s);
16187 let (py, _g3) = y.device_ptr_mut(s);
16188 let mut ps = [
16189 &pw as *const _ as *mut std::ffi::c_void,
16190 &paq as *const _ as *mut _,
16191 &pad as *const _ as *mut _,
16192 &py as *const _ as *mut _,
16193 &inf as *const _ as *mut _,
16194 &outf as *const _ as *mut _,
16195 &mi as *const _ as *mut _,
16196 &rb as *const _ as *mut _,
16197 ];
16198 unsafe {
16199 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16200 }
16201 }
16202 if scale != 1.0 {
16203 self.scale_inplace(y, scale, m * out_f)?;
16204 }
16205 } else {
16206 b.arg(bytes)
16207 .arg(aq)
16208 .arg(ad)
16209 .arg(&mut *y)
16210 .arg(&inf)
16211 .arg(&outf)
16212 .arg(&mi)
16213 .arg(&rb);
16214 unsafe {
16215 b.launch(cfg)?;
16216 }
16217 if scale != 1.0 {
16218 self.scale_inplace(y, scale, m * out_f)?;
16219 }
16220 }
16221 Ok(())
16222 }
16223
16224 pub fn qmatvec_mmvq_raw(
16228 &self,
16229 bytes: &CudaSlice<u8>,
16230 x: &CudaSlice<f32>,
16231 m: usize,
16232 in_f: usize,
16233 out_f: usize,
16234 qtype: i32,
16235 row_bytes: usize,
16236 rp: bool,
16237 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16238 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16239 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16240 }
16241
16242 pub fn batched_supports(&self, qtype: i32) -> bool {
16246 matches!(
16247 qtype,
16248 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16249 )
16250 }
16251
16252 pub fn iq_fast_enabled() -> bool {
16260 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16261 *ON.get_or_init(|| {
16262 std::env::var("MEMRA_IQ_FAST")
16263 .map(|v| v != "0")
16264 .unwrap_or(true)
16265 })
16266 }
16267
16268 pub fn b8_enabled() -> bool {
16271 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16272 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16273 }
16274
16275 pub fn batched_mcols(m: usize) -> usize {
16277 if m == 2 {
16278 2
16279 } else if m <= 4 {
16280 4
16281 } else if m <= 8 {
16282 8
16283 } else {
16284 16
16285 }
16286 }
16287
16288 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16293 Some(match (qtype, mcols) {
16294 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16295 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16296 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16297 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16303 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16304 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16305 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16306 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16309 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16310 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16311 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16312 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16315 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16316 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16317 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16318 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16319 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16320 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16321 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16322 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16326 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16327 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16328 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16329 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16333 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16334 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16335 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16336 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16337 _ => return None,
16338 })
16339 }
16340
16341 pub fn sm_count(&self) -> i32 {
16376 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16377 *SMS.get_or_init(|| {
16378 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16379 self.gpu
16380 .ctx
16381 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16382 .unwrap_or(82)
16383 })
16384 }
16385
16386 pub fn batched_variant(
16387 &self,
16388 _m: usize,
16389 in_f: usize,
16390 out_f: usize,
16391 qtype: i32,
16392 row_bytes: usize,
16393 mcols: usize,
16394 rp: bool,
16395 ) -> &'static str {
16396 if qtype == QT_Q8_0 {
16401 return if rp { "rp" } else { "base" };
16402 }
16403 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16404 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16405 Ok("base") => "base",
16406 Ok("pf") => "pf",
16407 Ok("r2") => "r2",
16408 Ok("r2w8") => "r2w8",
16409 Ok("pfr2") => "pfr2",
16410 Ok("ca") => "ca",
16411 Ok("car2") => "car2",
16412 Ok("rp") => "rp",
16415 Ok("rpr2") => "rpr2",
16416 Ok("rpr2w8") => "rpr2w8",
16417 Ok("rpca") => "rpca",
16420 Ok("rpcar2") => "rpcar2",
16421 Ok("rpsc") => "rpsc",
16428 Ok("rpms") => "rpms",
16429 Ok("rpmsc") => "rpmsc",
16430 Ok("rpks") => "rpks",
16431 Ok("rpksc") => "rpksc",
16432 _ => "auto",
16433 });
16434 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16438 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16443 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16444 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16445 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16446 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16447 let sms = *SMS.get_or_init(|| {
16448 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16449 self.gpu
16450 .ctx
16451 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16452 .unwrap_or(82)
16453 });
16454 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16474 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16477 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16478 Ok("base") => "base",
16479 Ok("r2") => "r2",
16480 Ok("r2w8") => "r2w8",
16481 _ => "auto",
16482 });
16483 let variant: &'static str = if qtype == QT_Q4_0 {
16484 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16488 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16489 Ok("base") => "base",
16495 Ok("r2") => "r2",
16496 Ok("ms") => "ms",
16497 Ok("sm") => "sm",
16498 Ok("la") => "la",
16499 _ => "auto",
16500 });
16501 let v = if q40 != "auto" {
16502 q40
16503 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16504 "r2"
16505 } else {
16506 "base"
16507 };
16508 if rp {
16513 match v {
16514 "ms" => "r2ms_rp",
16515 "sm" => "r2sm_rp",
16516 "la" => "r2la_rp",
16517 "r2" => "r2_rp",
16518 _ => "rp",
16519 }
16520 } else if matches!(v, "ms" | "sm" | "la") {
16521 "r2"
16522 } else {
16523 v
16524 }
16525 } else if qtype != QT_NVFP4 && !kq_r2 {
16526 "base"
16527 } else if kq_r2 && rp {
16528 "rp"
16532 } else if kq_r2 {
16533 if kq_bv != "auto" {
16536 if kq_bv == "r2w8" && mcols != 4 {
16537 "r2"
16538 } else {
16539 kq_bv
16540 }
16541 } else if bv != "auto" {
16542 match bv {
16543 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16544 "r2w8" | "rpr2w8" => {
16545 if mcols != 4 {
16546 "r2"
16547 } else {
16548 "r2w8"
16549 }
16550 }
16551 _ => "base", }
16553 } else {
16554 let blocks = (out_f + 7) / 8;
16555 let waves = blocks as f64 / (7 * sms as usize) as f64;
16556 let filled = blocks >= 4 * sms as usize;
16557 let use_r2 = if qtype == QT_Q4_K {
16558 filled
16559 } else {
16560 waves >= 2.0
16561 };
16562 if use_r2 { "r2" } else { "base" }
16563 }
16564 } else if bv != "auto" {
16565 let v = if bv == "r2w8" && mcols == 2 {
16570 "r2"
16571 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16572 "pf"
16573 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16574 "r2"
16575 } else if bv == "pfr2" && mcols == 8 {
16576 "r2"
16577 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16578 "rpr2"
16579 }
16580 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16582 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16583 } else if bv == "rpcar2" && mcols == 2 {
16584 "rpca"
16585 }
16586 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16589 "rpr2"
16590 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16591 "rpr2"
16592 } else {
16593 bv
16594 };
16595 if rp {
16596 match v {
16597 "base" | "pf" | "ca" | "rp" => "rp",
16598 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16599 "r2w8" | "rpr2w8" => {
16600 if mcols == 2 {
16601 "rpr2"
16602 } else {
16603 "rpr2w8"
16604 }
16605 }
16606 other => other, }
16608 } else {
16609 v
16610 }
16611 } else if mcols == 8 {
16612 if rp {
16623 if sc_ok { "rpsc" } else { "rpr2w8" }
16624 } else {
16625 "r2w8"
16626 }
16627 } else if mcols >= 4 {
16628 let blocks = (out_f + 7) / 8;
16632 let r7 = 7 * sms as usize;
16633 let r8 = 8 * sms as usize;
16634 let waves = blocks as f64 / r7 as f64;
16635 let filled = blocks >= 4 * sms as usize;
16636 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16640 if rp { "rpr2w8" } else { "r2w8" }
16644 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16645 if rp { "rpr2" } else { "r2" }
16648 } else {
16649 if rp { "rp" } else { "pf" }
16653 }
16654 } else if in_f >= 6144 {
16655 if rp { "rpr2" } else { "r2" }
16659 } else if rp {
16660 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16665 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16666 "rpsc"
16667 } else {
16668 "rp"
16669 }
16670 } else {
16671 "base"
16672 };
16673 variant
16674 }
16675
16676 pub fn qmatvec_mmvq_batched(
16677 &self,
16678 bytes: &CudaSlice<u8>,
16679 aq: &CudaSlice<i8>,
16680 ad: &CudaSlice<f32>,
16681 m: usize,
16682 in_f: usize,
16683 out_f: usize,
16684 qtype: i32,
16685 row_bytes: usize,
16686 mcols: usize,
16687 scale: f32,
16688 rp: bool,
16689 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16690 const ROWS_PER_BLOCK: u32 = 4;
16691 let forced: Option<&'static str> = {
16696 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16697 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16698 .as_deref()
16699 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16700 };
16701 let variant = match forced {
16702 Some(v) if !rp || v.contains("rp") => v,
16703 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16704 };
16705 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16706 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16707 })?;
16708 let variant = if mcols == 16 {
16712 if rp { "rp" } else { "base" }
16713 } else {
16714 variant
16715 };
16716 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16723 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16724 if b567
16725 && qtype == QT_NVFP4
16726 && rp
16727 && mcols == 8
16728 && (5..=7).contains(&m)
16729 && matches!(variant, "rpsc" | "rpr2w8")
16730 {
16731 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16732 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16734 let cfg = LaunchConfig {
16735 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16736 block_dim: (32, ROWS_PER_BLOCK, 1),
16737 shared_mem_bytes: 0,
16738 };
16739 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16740 let __s_b = self.gpu.stream();
16741 let mut b = __s_b.launch_builder(&f);
16742 b.arg(bytes)
16743 .arg(aq)
16744 .arg(ad)
16745 .arg(&mut y)
16746 .arg(&inf)
16747 .arg(&outf)
16748 .arg(&mi)
16749 .arg(&rb);
16750 unsafe {
16751 b.launch(cfg)?;
16752 }
16753 if scale != 1.0 {
16754 self.scale_inplace(&mut y, scale, m * out_f)?;
16755 }
16756 return Ok(y);
16757 }
16758 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16759 "base" => (base_name.into(), ROWS_PER_BLOCK),
16760 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16761 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16762 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16763 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16767 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16768 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16769 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16770 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16771 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16772 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16773 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16775 debug_assert!(
16776 !rp || name.contains("_rp"),
16777 "rp weight dispatched to a GGUF-layout kernel"
16778 );
16779 let f = self.func(&name);
16780 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16781 let smem = if name.contains("_r2sm_rp") {
16783 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16784 } else {
16785 0
16786 };
16787 let cfg = LaunchConfig {
16788 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16789 block_dim: (32, ROWS_PER_BLOCK, 1),
16790 shared_mem_bytes: smem,
16791 };
16792 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16793 let __s_b = self.gpu.stream();
16794 let mut b = __s_b.launch_builder(&f);
16795 b.arg(bytes)
16796 .arg(aq)
16797 .arg(ad)
16798 .arg(&mut y)
16799 .arg(&inf)
16800 .arg(&outf)
16801 .arg(&mi)
16802 .arg(&rb);
16803 unsafe {
16804 b.launch(cfg)?;
16805 }
16806 if scale != 1.0 {
16807 self.scale_inplace(&mut y, scale, m * out_f)?;
16808 }
16809 Ok(y)
16810 }
16811
16812 pub fn qmatvec_batched_raw(
16816 &self,
16817 bytes: &CudaSlice<u8>,
16818 x: &CudaSlice<f32>,
16819 m: usize,
16820 in_f: usize,
16821 out_f: usize,
16822 qtype: i32,
16823 row_bytes: usize,
16824 mcols: usize,
16825 rp: bool,
16826 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16827 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16828 self.qmatvec_mmvq_batched(
16829 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16830 )
16831 }
16832
16833 pub fn qmatvec_nvfp4_batched_raw(
16835 &self,
16836 bytes: &CudaSlice<u8>,
16837 x: &CudaSlice<f32>,
16838 m: usize,
16839 in_f: usize,
16840 out_f: usize,
16841 row_bytes: usize,
16842 mcols: usize,
16843 rp: bool,
16844 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16845 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16846 }
16847
16848 fn try_fp4_gemm(
16852 &self,
16853 w: &crate::model::GpuTensor,
16854 x: &CudaSlice<f32>,
16855 m: usize,
16856 in_f: usize,
16857 out_f: usize,
16858 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16859 use crate::model::GpuTensor;
16860 if cfg!(memra_portable_cuda) {
16861 return Ok(None);
16862 }
16863 if std::env::var("MEMRA_FP4").is_ok() {
16866 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16867 }
16868 if std::env::var("MEMRA_FP4").is_err() {
16869 return Ok(None);
16870 }
16871 #[cfg(memra_cutlass)]
16880 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16881 if let GpuTensor::Quant {
16882 bytes,
16883 qtype,
16884 scale,
16885 row_bytes,
16886 cutlass,
16887 ..
16888 } = w
16889 {
16890 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16891 if let Some(cw) = cutlass {
16892 let y = self.cutlass_fp4_gemm(
16894 &cw.b_packed,
16895 &cw.sfb_swizzled,
16896 x,
16897 *scale,
16898 m,
16899 out_f,
16900 in_f,
16901 )?;
16902 return Ok(Some(y));
16903 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16904 let (b_packed, sfb_sw) =
16909 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16910 let y =
16911 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16912 return Ok(Some(y));
16913 }
16914 }
16915 }
16916 }
16917 if let GpuTensor::Quant {
16918 bytes,
16919 qtype,
16920 row_bytes,
16921 scale,
16922 rp,
16923 ..
16924 } = w
16925 {
16926 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16929 let y =
16930 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16931 return Ok(Some(y));
16932 }
16933 }
16934 Ok(None)
16935 }
16936
16937 pub fn rms_norm_f16out(
16941 &self,
16942 x: &CudaSlice<f32>,
16943 w: &CudaSlice<f32>,
16944 dst: &mut CudaSlice<f32>,
16945 dst16: &mut CudaSlice<u8>,
16946 ncols: usize,
16947 nrows: usize,
16948 eps: f32,
16949 ) -> Result<(), Box<dyn std::error::Error>> {
16950 let f = self.func("rms_norm_f16out_f32");
16951 let cfg = LaunchConfig {
16952 grid_dim: (nrows as u32, 1, 1),
16953 block_dim: (rms_block(), 1, 1),
16954 shared_mem_bytes: 0,
16955 };
16956 let (nc, e) = (ncols as i32, eps);
16957 let __s_b = self.gpu.stream();
16958 let mut b = __s_b.launch_builder(&f);
16959 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16960 unsafe {
16961 b.launch(cfg)?;
16962 }
16963 Ok(())
16964 }
16965
16966 #[allow(clippy::too_many_arguments)]
16969 pub fn add_rms_norm_f16out(
16970 &self,
16971 a: &CudaSlice<f32>,
16972 b: &CudaSlice<f32>,
16973 w: &CudaSlice<f32>,
16974 res: &mut CudaSlice<f32>,
16975 dst: &mut CudaSlice<f32>,
16976 dst16: &mut CudaSlice<u8>,
16977 ncols: usize,
16978 nrows: usize,
16979 eps: f32,
16980 ) -> Result<(), Box<dyn std::error::Error>> {
16981 let f = self.func("add_rms_norm_f16out_f32");
16982 let cfg = LaunchConfig {
16983 grid_dim: (nrows as u32, 1, 1),
16984 block_dim: (rms_block(), 1, 1),
16985 shared_mem_bytes: 0,
16986 };
16987 let (nc, e) = (ncols as i32, eps);
16988 let __s_lb = self.gpu.stream();
16989 let mut lb = __s_lb.launch_builder(&f);
16990 lb.arg(a)
16991 .arg(b)
16992 .arg(w)
16993 .arg(res)
16994 .arg(dst)
16995 .arg(dst16)
16996 .arg(&nc)
16997 .arg(&e);
16998 unsafe {
16999 lb.launch(cfg)?;
17000 }
17001 Ok(())
17002 }
17003
17004 pub fn matmul_group_xh(
17007 &self,
17008 ws: &[&crate::model::GpuTensor],
17009 x: &CudaSlice<f32>,
17010 xh: &CudaSlice<u8>,
17011 m: usize,
17012 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17013 let mut out = Vec::with_capacity(ws.len());
17014 let in_f = ws[0].in_features();
17015 for w in ws {
17016 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
17017 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
17018 out.push(y);
17019 continue;
17020 }
17021 }
17022 out.push(self.matmul(w, x, m)?);
17023 }
17024 Ok(out)
17025 }
17026
17027 pub fn gdn_pad_mask(
17030 &self,
17031 beta: &mut CudaSlice<f32>,
17032 g_log: &mut CudaSlice<f32>,
17033 len_d: &CudaSlice<i32>,
17034 h: usize,
17035 t: usize,
17036 ) -> Result<(), Box<dyn std::error::Error>> {
17037 let f = self.func("gdn_pad_mask_f32");
17038 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
17039 let (hi, ti) = (h as i32, t as i32);
17040 let __s_b = self.gpu.stream();
17041 let mut b = __s_b.launch_builder(&f);
17042 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17043 unsafe {
17044 b.launch(cfg)?;
17045 }
17046 Ok(())
17047 }
17048
17049 pub fn row_gather_dev(
17052 &self,
17053 src: &CudaSlice<f32>,
17054 dst: &mut CudaSlice<f32>,
17055 len_d: &CudaSlice<i32>,
17056 ncols: usize,
17057 ) -> Result<(), Box<dyn std::error::Error>> {
17058 let f = self.func("row_gather_dev_f32");
17059 let cfg = LaunchConfig::for_num_elems(ncols as u32);
17060 let nc = ncols as i32;
17061 let __s_b = self.gpu.stream();
17062 let mut b = __s_b.launch_builder(&f);
17063 b.arg(src).arg(dst).arg(len_d).arg(&nc);
17064 unsafe {
17065 b.launch(cfg)?;
17066 }
17067 Ok(())
17068 }
17069
17070 pub fn matmul_group(
17077 &self,
17078 ws: &[&crate::model::GpuTensor],
17079 x: &CudaSlice<f32>,
17080 m: usize,
17081 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17082 use crate::model::GpuTensor;
17083 let mut out = Vec::with_capacity(ws.len());
17084 let any_mirror = ws
17085 .iter()
17086 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17087 if m >= 16 && any_mirror && !self.verify_exact_on() {
17088 let in_f = ws[0].in_features();
17089 let xh = self.f16_act(x, m * in_f, in_f)?;
17090 for w in ws {
17091 if w.in_features() == in_f {
17092 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17093 out.push(y);
17094 continue;
17095 }
17096 }
17097 out.push(self.matmul(w, x, m)?);
17098 }
17099 return Ok(out);
17100 }
17101 for w in ws {
17102 out.push(self.matmul(w, x, m)?);
17103 }
17104 Ok(out)
17105 }
17106
17107 pub fn matmul_group_multi(
17114 &self,
17115 ws: &[&crate::model::GpuTensor],
17116 xs: &[&CudaSlice<f32>],
17117 ms: &[usize],
17118 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17119 assert_eq!(xs.len(), ms.len());
17120 let in_f = ws[0].in_features();
17121 let total: usize = ms.iter().sum();
17122 let mut xcat = self.uninit(total * in_f)?;
17123 let mut off = 0usize;
17124 for (x, &m) in xs.iter().zip(ms) {
17125 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17126 off += m;
17127 }
17128 let ys = self.matmul_group(ws, &xcat, total)?;
17129 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17130 for (w, y) in ws.iter().zip(ys) {
17131 let out_f = w.out_features();
17132 let mut off = 0usize;
17133 for (s, &m) in ms.iter().enumerate() {
17134 let mut ys_s = self.uninit(m * out_f)?;
17135 let src = y.slice(off * out_f..(off + m) * out_f);
17136 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17137 out[s].push(ys_s);
17138 off += m;
17139 }
17140 }
17141 Ok(out)
17142 }
17143
17144 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17154 use crate::model::GpuTensor;
17155 if !legacy_quant_gemm_allowed(
17156 cfg!(memra_portable_cuda),
17157 cfg!(memra_hopper_mma),
17158 std::env::var_os("MEMRA_NO_GEMM").is_some(),
17159 ) {
17160 return false;
17161 }
17162 match w {
17163 GpuTensor::Quant { qtype, .. } => {
17164 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17165 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17166 }
17167 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17168 }
17169 }
17170
17171 pub fn qmatvec_gemm(
17178 &self,
17179 w: &crate::model::GpuTensor,
17180 aq: &CudaSlice<i8>,
17181 ad: &CudaSlice<f32>,
17182 m: usize,
17183 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17184 use crate::model::GpuTensor;
17185 let in_f = w.in_features();
17186 let out_f = w.out_features();
17187 let (bytes, qtype, row_bytes, scale, rp) = match w {
17188 GpuTensor::Quant {
17189 bytes,
17190 qtype,
17191 row_bytes,
17192 scale,
17193 rp,
17194 ..
17195 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17196 _ => unreachable!("gemm_supports guaranteed Quant"),
17197 };
17198 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17204 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17205 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17206 if scale != 1.0 {
17207 self.scale_inplace(&mut y, scale, m * out_f)?;
17208 }
17209 return Ok(y);
17210 }
17211 }
17212 let name = match qtype {
17213 QT_Q8_0 => "qmatvec_gemm_q8_0",
17214 QT_Q4_K => "qmatvec_gemm_q4_K",
17215 QT_Q4_0 => {
17216 if rp {
17217 "qmatvec_gemm_q4_0_rp"
17218 } else {
17219 "qmatvec_gemm_q4_0"
17220 }
17221 }
17222 QT_Q5_K => "qmatvec_gemm_q5_K",
17223 QT_Q6_K => "qmatvec_gemm_q6_K",
17224 QT_NVFP4 => {
17225 if rp {
17226 "qmatvec_gemm_nvfp4_rp"
17227 } else {
17228 "qmatvec_gemm_nvfp4"
17229 }
17230 }
17231 _ => unreachable!(),
17232 };
17233 let f = self.func(name);
17234 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);
17239 let k1_tile = if is_k1 {
17241 k1_launch_override().unwrap_or((128, 128, 8))
17242 } else {
17243 (128, 128, 8)
17244 };
17245 let (bm, bn): (u32, u32) = if is_k1 {
17246 (k1_tile.0, k1_tile.1)
17247 } else {
17248 (64, 256)
17249 };
17250 let warps: u32 = if is_k1 {
17251 k1_tile.2
17252 } else {
17253 match qtype {
17254 QT_NVFP4 => 8,
17255 _ => 4,
17256 }
17257 };
17258 let cfg = LaunchConfig {
17259 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17260 block_dim: (32, warps, 1),
17261 shared_mem_bytes: 0,
17262 };
17263 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17264 let __s_b = self.gpu.stream();
17265 let mut b = __s_b.launch_builder(&f);
17266 b.arg(bytes)
17267 .arg(aq)
17268 .arg(ad)
17269 .arg(&mut y)
17270 .arg(&inf)
17271 .arg(&outf)
17272 .arg(&mi)
17273 .arg(&rb);
17274 unsafe {
17275 b.launch(cfg)?;
17276 }
17277 if scale != 1.0 {
17278 self.scale_inplace(&mut y, scale, m * out_f)?;
17279 }
17280 Ok(y)
17281 }
17282
17283 pub fn qmatvec_gemm_raw(
17288 &self,
17289 bytes: &CudaSlice<u8>,
17290 x: &CudaSlice<f32>,
17291 m: usize,
17292 in_f: usize,
17293 out_f: usize,
17294 qtype: i32,
17295 row_bytes: usize,
17296 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17297 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17298 let name = match qtype {
17299 QT_Q8_0 => "qmatvec_gemm_q8_0",
17300 QT_Q4_K => "qmatvec_gemm_q4_K",
17301 QT_Q4_0 => "qmatvec_gemm_q4_0",
17302 QT_Q5_K => "qmatvec_gemm_q5_K",
17303 QT_Q6_K => "qmatvec_gemm_q6_K",
17304 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17305 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17306 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17307 };
17308 let f = self.func(name);
17309 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);
17313 let k1_tile = if is_k1 {
17315 k1_launch_override().unwrap_or((128, 128, 8))
17316 } else {
17317 (128, 128, 8)
17318 };
17319 let (bm, bn): (u32, u32) = if is_k1 {
17320 (k1_tile.0, k1_tile.1)
17321 } else {
17322 (64, 256)
17323 };
17324 let warps: u32 = if is_k1 {
17325 k1_tile.2
17326 } else {
17327 match qtype {
17328 QT_NVFP4 | QT_NVFP4_RP => 8,
17329 _ => 4,
17330 }
17331 };
17332 let cfg = LaunchConfig {
17333 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17334 block_dim: (32, warps, 1),
17335 shared_mem_bytes: 0,
17336 };
17337 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17338 let __s_b = self.gpu.stream();
17339 let mut b = __s_b.launch_builder(&f);
17340 b.arg(bytes)
17341 .arg(&aq)
17342 .arg(&ad)
17343 .arg(&mut y)
17344 .arg(&inf)
17345 .arg(&outf)
17346 .arg(&mi)
17347 .arg(&rb);
17348 unsafe {
17349 b.launch(cfg)?;
17350 }
17351 Ok(y)
17352 }
17353
17354 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17361 &self,
17362 rp4: &CudaSlice<u8>,
17363 aq: &CudaSlice<i8>,
17364 ad: &CudaSlice<f32>,
17365 m: usize,
17366 in_f: usize,
17367 out_f: usize,
17368 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17369 assert!(
17370 out_f % 64 == 0 && in_f % 32 == 0,
17371 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17372 );
17373 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17374 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17376 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17377 block_dim: (128, 1, 1),
17378 shared_mem_bytes: 0,
17379 };
17380 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17381 let __s_b = self.gpu.stream();
17382 let mut b = __s_b.launch_builder(&f);
17383 b.arg(rp4)
17384 .arg(aq)
17385 .arg(ad)
17386 .arg(&mut y)
17387 .arg(&inf)
17388 .arg(&outf)
17389 .arg(&mi);
17390 unsafe {
17391 b.launch(cfg)?;
17392 }
17393 Ok(y)
17394 }
17395
17396 pub fn scale_inplace(
17398 &self,
17399 y: &mut CudaSlice<f32>,
17400 s: f32,
17401 n: usize,
17402 ) -> Result<(), Box<dyn std::error::Error>> {
17403 let f = self.func("scale_f32");
17404 let cfg = LaunchConfig::for_num_elems(n as u32);
17405 let (sf, ni) = (s, n as i32);
17406 let __s_b = self.gpu.stream();
17407 let mut b = __s_b.launch_builder(&f);
17408 b.arg(y).arg(&sf).arg(&ni);
17409 unsafe {
17410 b.launch(cfg)?;
17411 }
17412 Ok(())
17413 }
17414
17415 pub fn bf16_to_f32(
17420 &self,
17421 data: &cudarc::driver::CudaView<'_, u8>,
17422 n: usize,
17423 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17424 let mut out = self.alloc_uninit::<f32>(n)?;
17425 let f = self.func("bf16_to_f32");
17426 let cfg = LaunchConfig::for_num_elems(n as u32);
17427 let ni = n as i32;
17428 let __s_b = self.gpu.stream();
17429 let mut b = __s_b.launch_builder(&f);
17430 b.arg(data).arg(&mut out).arg(&ni);
17431 unsafe {
17432 b.launch(cfg)?;
17433 }
17434 Ok(out)
17435 }
17436
17437 fn linear_bf16_chunked(
17444 &self,
17445 x: &CudaSlice<f32>,
17446 data: &CudaSlice<u8>,
17447 m: usize,
17448 in_f: usize,
17449 out_f: usize,
17450 exact: bool,
17451 canonical_chunk_rows: Option<usize>,
17452 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17453 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17457 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17458 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17459 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17460 let started = timing.then(std::time::Instant::now);
17461 let result =
17462 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17463 if let Some(started) = started {
17464 use std::sync::atomic::Ordering;
17465 self.stream().synchronize()?;
17466 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17467 + started.elapsed().as_nanos() as u64;
17468 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17469 + (in_f * out_f * 2) as u64;
17470 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17471 if calls % 1024 == 0 {
17472 eprintln!(
17473 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17474 weight_gb={:.2}",
17475 ns as f64 / 1.0e6,
17476 ns as f64 / calls as f64 / 1.0e3,
17477 wb as f64 / 1.0e9,
17478 );
17479 }
17480 }
17481 result
17482 }
17483
17484 pub(crate) fn bf16_mmv_on() -> bool {
17489 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17490 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17491 }
17492
17493 fn matvec_bf16(
17496 &self,
17497 data: &CudaSlice<u8>,
17498 x: &CudaSlice<f32>,
17499 in_f: usize,
17500 out_f: usize,
17501 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17502 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17503 return Err(format!(
17504 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17505 data.len(),
17506 x.len()
17507 )
17508 .into());
17509 }
17510 let mut y = self.alloc_uninit::<f32>(out_f)?;
17511 let f = self.func("matvec_bf16_f32acc");
17512 let cfg = LaunchConfig {
17513 grid_dim: (out_f as u32, 1, 1),
17514 block_dim: (mmv_block(), 1, 1),
17515 shared_mem_bytes: 0,
17516 };
17517 let ini = in_f as i32;
17518 let __s_bld = self.gpu.stream();
17519 let mut bld = __s_bld.launch_builder(&f);
17520 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17521 unsafe {
17522 bld.launch(cfg)?;
17523 }
17524 Ok(y)
17525 }
17526
17527 #[allow(clippy::too_many_arguments)]
17531 #[allow(clippy::too_many_arguments)]
17536 #[allow(clippy::too_many_arguments)]
17541 pub fn qk_norm_rope_append_inc_dcw_rows(
17542 &self,
17543 q_raw_t: &CudaSlice<f32>,
17544 k_raw_t: &CudaSlice<f32>,
17545 v_raw_t: &CudaSlice<f32>,
17546 qw: &CudaSlice<f32>,
17547 kw: &CudaSlice<f32>,
17548 q_out_t: &mut CudaSlice<f32>,
17549 k_out_t: &mut CudaSlice<f32>,
17550 tab: &CudaSlice<u64>,
17551 pos_t: &CudaSlice<i32>,
17552 same_session: bool,
17553 t: usize,
17554 kv_dim_k: usize,
17555 kv_dim_v: usize,
17556 k_tok_bytes: usize,
17557 v_tok_bytes: usize,
17558 head_dim: usize,
17559 n_dims: usize,
17560 nh_q: usize,
17561 nh_k: usize,
17562 eps: f32,
17563 freq_base: f32,
17564 freq_scale: f32,
17565 ff: Option<&CudaSlice<f32>>,
17566 ) -> Result<(), Box<dyn std::error::Error>> {
17567 if head_dim != 128
17568 || kv_dim_v != kv_dim_k
17569 || kv_dim_k != nh_k * head_dim
17570 || t == 0
17571 || t > 32
17572 || tab.len() < t * 6
17573 || pos_t.len() < t
17574 || q_raw_t.len() < t * nh_q * head_dim
17575 || k_raw_t.len() < t * nh_k * head_dim
17576 || v_raw_t.len() < t * kv_dim_v
17577 || q_out_t.len() < t * nh_q * head_dim
17578 || k_out_t.len() < t * nh_k * head_dim
17579 {
17580 return Err(format!(
17581 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17582 nh_q={nh_q} nh_k={nh_k}"
17583 )
17584 .into());
17585 }
17586 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17587 let same_t: i32 = if same_session { t as i32 } else { 0 };
17588 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17589 let cfg = LaunchConfig {
17590 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17591 block_dim: (128, 1, 1),
17592 shared_mem_bytes: 0,
17593 };
17594 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17595 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17596 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17597 let null: u64 = 0;
17598 let __s_b = self.gpu.stream();
17599 let mut b = __s_b.launch_builder(&f);
17600 b.arg(q_raw_t)
17601 .arg(k_raw_t)
17602 .arg(v_raw_t)
17603 .arg(qw)
17604 .arg(kw)
17605 .arg(q_out_t)
17606 .arg(k_out_t)
17607 .arg(tab)
17608 .arg(pos_t)
17609 .arg(&same_t)
17610 .arg(&kvk)
17611 .arg(&kvv)
17612 .arg(&ktb)
17613 .arg(&vtb)
17614 .arg(&hd)
17615 .arg(&nd)
17616 .arg(&nq)
17617 .arg(&nk)
17618 .arg(&eps)
17619 .arg(&theta_scale)
17620 .arg(&freq_scale);
17621 match ff {
17622 Some(freqs) => {
17623 b.arg(freqs);
17624 }
17625 None => {
17626 b.arg(&null);
17627 }
17628 }
17629 unsafe {
17630 b.launch(cfg)?;
17631 }
17632 Ok(())
17633 }
17634
17635 pub fn qk_norm_rope_append_inc_dcw(
17636 &self,
17637 q_raw: &CudaSlice<f32>,
17638 k_raw: &CudaSlice<f32>,
17639 v_raw: &CudaSlice<f32>,
17640 qw: &CudaSlice<f32>,
17641 kw: &CudaSlice<f32>,
17642 q_out: &mut CudaSlice<f32>,
17643 k_out: &mut CudaSlice<f32>,
17644 pos: &CudaSlice<i32>,
17645 k_plane: &mut CudaSlice<u8>,
17646 v_plane: &mut CudaSlice<u8>,
17647 len_dev: &CudaSlice<i32>,
17650 base_dev: Option<&CudaSlice<i32>>,
17651 done_ctr: &mut CudaSlice<u32>,
17652 kv_dim_k: usize,
17653 kv_dim_v: usize,
17654 k_tok_bytes: usize,
17655 v_tok_bytes: usize,
17656 head_dim: usize,
17657 n_dims: usize,
17658 nh_q: usize,
17659 nh_k: usize,
17660 eps: f32,
17661 freq_base: f32,
17662 freq_scale: f32,
17663 ff: Option<&CudaSlice<f32>>,
17664 ) -> Result<(), Box<dyn std::error::Error>> {
17665 if head_dim != 128
17666 || kv_dim_v != kv_dim_k
17667 || kv_dim_k != nh_k * head_dim
17668 || q_raw.len() < nh_q * head_dim
17669 || k_raw.len() < nh_k * head_dim
17670 || v_raw.len() < kv_dim_v
17671 || q_out.len() < nh_q * head_dim
17672 || k_out.len() < nh_k * head_dim
17673 || pos.is_empty()
17674 || done_ctr.is_empty()
17675 {
17676 return Err(format!(
17677 "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}"
17678 )
17679 .into());
17680 }
17681 let f = self.func("qk_norm_rope_append_inc_dcw");
17682 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17683 let cfg = LaunchConfig {
17684 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17685 block_dim: (128, 1, 1),
17686 shared_mem_bytes: 0,
17687 };
17688 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17689 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17690 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17691 let null: u64 = 0;
17692 let __s_b = self.gpu.stream();
17693 let mut b = __s_b.launch_builder(&f);
17694 b.arg(q_raw)
17695 .arg(k_raw)
17696 .arg(v_raw)
17697 .arg(qw)
17698 .arg(kw)
17699 .arg(q_out)
17700 .arg(k_out)
17701 .arg(pos)
17702 .arg(&mut *k_plane)
17703 .arg(&mut *v_plane)
17704 .arg(len_dev);
17705 match base_dev {
17706 Some(base) => {
17707 b.arg(base);
17708 }
17709 None => {
17710 b.arg(&null);
17711 }
17712 }
17713 b.arg(&mut *done_ctr)
17714 .arg(&kvk)
17715 .arg(&kvv)
17716 .arg(&ktb)
17717 .arg(&vtb)
17718 .arg(&hd)
17719 .arg(&nd)
17720 .arg(&nq)
17721 .arg(&eps)
17722 .arg(&theta_scale)
17723 .arg(&freq_scale);
17724 match ff {
17725 Some(freqs) => {
17726 b.arg(freqs);
17727 }
17728 None => {
17729 b.arg(&null);
17730 }
17731 }
17732 unsafe {
17733 b.launch(cfg)?;
17734 }
17735 Ok(())
17736 }
17737
17738 pub fn qk_norm_rope_into(
17739 &self,
17740 q_raw: &CudaSlice<f32>,
17741 k_raw: &CudaSlice<f32>,
17742 qw: &CudaSlice<f32>,
17743 kw: &CudaSlice<f32>,
17744 q_out: &mut CudaSlice<f32>,
17745 k_out: &mut CudaSlice<f32>,
17746 pos: &CudaSlice<i32>,
17747 head_dim: usize,
17748 n_dims: usize,
17749 nh_q: usize,
17750 nh_k: usize,
17751 eps: f32,
17752 freq_base: f32,
17753 freq_scale: f32,
17754 ff: Option<&CudaSlice<f32>>,
17755 ) -> Result<(), Box<dyn std::error::Error>> {
17756 if head_dim > 512
17757 || q_raw.len() < nh_q * head_dim
17758 || k_raw.len() < nh_k * head_dim
17759 || q_out.len() < nh_q * head_dim
17760 || k_out.len() < nh_k * head_dim
17761 || qw.len() < head_dim
17762 || kw.len() < head_dim
17763 || pos.is_empty()
17764 {
17765 return Err(format!(
17766 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17767 )
17768 .into());
17769 }
17770 let f = self.func("qk_norm_rope_f32");
17771 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17772 let cfg = LaunchConfig {
17773 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17774 block_dim: (128, 1, 1),
17775 shared_mem_bytes: 0,
17776 };
17777 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17778 let __s_b = self.gpu.stream();
17779 let mut b = __s_b.launch_builder(&f);
17780 b.arg(q_raw)
17781 .arg(k_raw)
17782 .arg(qw)
17783 .arg(kw)
17784 .arg(q_out)
17785 .arg(k_out)
17786 .arg(pos)
17787 .arg(&hd)
17788 .arg(&nd)
17789 .arg(&nq)
17790 .arg(&eps)
17791 .arg(&theta_scale)
17792 .arg(&freq_scale);
17793 match ff {
17794 Some(ffv) => {
17795 b.arg(ffv);
17796 unsafe {
17797 b.launch(cfg)?;
17798 }
17799 }
17800 None => {
17801 let null: u64 = 0;
17802 b.arg(&null);
17803 unsafe {
17804 b.launch(cfg)?;
17805 }
17806 }
17807 }
17808 Ok(())
17809 }
17810
17811 #[allow(clippy::too_many_arguments)]
17814 pub fn matvec_f32_b4_into(
17815 &self,
17816 w: [&CudaSlice<f32>; 4],
17817 x: &CudaSlice<f32>,
17818 y: &mut CudaSlice<f32>,
17819 block_cols: usize,
17820 out_f: usize,
17821 ) -> Result<(), Box<dyn std::error::Error>> {
17822 if block_cols % 4 != 0
17823 || x.len() < 4 * block_cols
17824 || y.len() < out_f
17825 || w.iter().any(|w| w.len() != out_f * block_cols)
17826 {
17827 return Err(format!(
17828 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17829 x.len()
17830 )
17831 .into());
17832 }
17833 let f = self.func("matvec_f32_b4");
17834 let cfg = LaunchConfig {
17835 grid_dim: (out_f as u32, 1, 1),
17836 block_dim: (128, 1, 1),
17837 shared_mem_bytes: 0,
17838 };
17839 let (bc, of) = (block_cols as i32, out_f as i32);
17840 let __s_b = self.gpu.stream();
17841 let mut b = __s_b.launch_builder(&f);
17842 b.arg(w[0])
17843 .arg(w[1])
17844 .arg(w[2])
17845 .arg(w[3])
17846 .arg(x)
17847 .arg(y)
17848 .arg(&bc)
17849 .arg(&of);
17850 unsafe {
17851 b.launch(cfg)?;
17852 }
17853 Ok(())
17854 }
17855
17856 pub fn axpy_rows_seq_into(
17859 &self,
17860 x: &CudaSlice<f32>,
17861 w: &CudaSlice<f32>,
17862 y: &mut CudaSlice<f32>,
17863 width: usize,
17864 n_rows: usize,
17865 ) -> Result<(), Box<dyn std::error::Error>> {
17866 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17867 return Err(format!(
17868 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17869 x.len(),
17870 w.len(),
17871 y.len()
17872 )
17873 .into());
17874 }
17875 let f = self.func("axpy_rows_seq_f32");
17876 let cfg = LaunchConfig::for_num_elems(width as u32);
17877 let (wi, nr) = (width as i32, n_rows as i32);
17878 let __s_b = self.gpu.stream();
17879 let mut b = __s_b.launch_builder(&f);
17880 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17881 unsafe {
17882 b.launch(cfg)?;
17883 }
17884 Ok(())
17885 }
17886
17887 #[allow(clippy::too_many_arguments)]
17891 pub fn axpy_rows_seq_md_off_into(
17892 &self,
17893 x: &CudaSlice<f32>,
17894 w_route: &CudaSlice<f32>,
17895 md: &CudaSlice<f32>,
17896 sel: &CudaSlice<i32>,
17897 y: &mut CudaSlice<f32>,
17898 width: usize,
17899 n_rows: usize,
17900 row0: usize,
17901 ) -> Result<(), Box<dyn std::error::Error>> {
17902 if x.len() < (row0 + n_rows) * width
17903 || w_route.len() < row0 + n_rows
17904 || sel.len() < row0 + n_rows
17905 || y.len() < width
17906 {
17907 return Err(format!(
17908 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17909 rows={n_rows} row0={row0}",
17910 x.len(),
17911 w_route.len(),
17912 sel.len(),
17913 y.len()
17914 )
17915 .into());
17916 }
17917 let f = self.func("axpy_rows_seq_md_off_f32");
17918 let cfg = LaunchConfig::for_num_elems(width as u32);
17919 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17920 let __s_b = self.gpu.stream();
17921 let mut b = __s_b.launch_builder(&f);
17922 b.arg(x)
17923 .arg(w_route)
17924 .arg(md)
17925 .arg(sel)
17926 .arg(y)
17927 .arg(&wi)
17928 .arg(&nr)
17929 .arg(&r0);
17930 unsafe {
17931 b.launch(cfg)?;
17932 }
17933 Ok(())
17934 }
17935
17936 #[allow(clippy::too_many_arguments)]
17941 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17942 &self,
17943 gate_bank: &CudaSlice<u8>,
17944 up_bank: &CudaSlice<u8>,
17945 sel: &CudaSlice<i32>,
17946 aq: &CudaSlice<i8>,
17947 ad: &CudaSlice<f32>,
17948 yg: &mut CudaSlice<f32>,
17949 yu: &mut CudaSlice<f32>,
17950 n_sel: usize,
17951 n_sel_col: usize,
17952 in_f: usize,
17953 out_f: usize,
17954 row_bytes: usize,
17955 expert_stride: usize,
17956 act_row_stride: usize,
17957 ad_row_stride: usize,
17958 ) -> Result<(), Box<dyn std::error::Error>> {
17959 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17960 if yg.len() < n_sel * out_f
17961 || yu.len() < n_sel * out_f
17962 || sel.len() < n_sel
17963 || n_sel_col == 0
17964 || n_sel % n_sel_col != 0
17965 {
17966 return Err("NVFP4 gu tcol geometry".into());
17967 }
17968 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17969 let cfg = LaunchConfig {
17970 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17971 block_dim: (128, 1, 1),
17972 shared_mem_bytes: 0,
17973 };
17974 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17975 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17976 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17977 let __s_b = self.gpu.stream();
17978 let mut b = __s_b.launch_builder(&f);
17979 b.arg(gate_bank)
17980 .arg(up_bank)
17981 .arg(sel)
17982 .arg(aq)
17983 .arg(ad)
17984 .arg(yg)
17985 .arg(yu)
17986 .arg(&inf)
17987 .arg(&outf)
17988 .arg(&ns)
17989 .arg(&rb)
17990 .arg(&es)
17991 .arg(&ars)
17992 .arg(&adrs)
17993 .arg(&nsc);
17994 unsafe {
17995 b.launch(cfg)?;
17996 }
17997 Ok(())
17998 }
17999
18000 #[allow(clippy::too_many_arguments)]
18003 pub fn axpy_rows_seq_md_into(
18004 &self,
18005 x: &CudaSlice<f32>,
18006 w_route: &CudaSlice<f32>,
18007 md: &CudaSlice<f32>,
18008 sel: &CudaSlice<i32>,
18009 y: &mut CudaSlice<f32>,
18010 width: usize,
18011 n_rows: usize,
18012 ) -> Result<(), Box<dyn std::error::Error>> {
18013 if x.len() < n_rows * width
18014 || w_route.len() < n_rows
18015 || sel.len() < n_rows
18016 || y.len() < width
18017 {
18018 return Err(format!(
18019 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
18020 x.len(),
18021 w_route.len(),
18022 sel.len(),
18023 y.len()
18024 )
18025 .into());
18026 }
18027 let f = self.func("axpy_rows_seq_md_f32");
18028 let cfg = LaunchConfig::for_num_elems(width as u32);
18029 let (wi, nr) = (width as i32, n_rows as i32);
18030 let __s_b = self.gpu.stream();
18031 let mut b = __s_b.launch_builder(&f);
18032 b.arg(x)
18033 .arg(w_route)
18034 .arg(md)
18035 .arg(sel)
18036 .arg(y)
18037 .arg(&wi)
18038 .arg(&nr);
18039 unsafe {
18040 b.launch(cfg)?;
18041 }
18042 Ok(())
18043 }
18044
18045 #[allow(clippy::too_many_arguments)]
18047 #[allow(clippy::too_many_arguments)]
18051 pub fn matvec_bf16_qkvg_tcol_into(
18052 &self,
18053 wq: &CudaSlice<u8>,
18054 wk: &CudaSlice<u8>,
18055 wv: &CudaSlice<u8>,
18056 wg: &CudaSlice<u8>,
18057 x_t: &CudaSlice<f32>,
18058 yq: &mut CudaSlice<f32>,
18059 yk: &mut CudaSlice<f32>,
18060 yv: &mut CudaSlice<f32>,
18061 yg: &mut CudaSlice<f32>,
18062 in_f: usize,
18063 out_q: usize,
18064 out_kv: usize,
18065 out_g: usize,
18066 t: usize,
18067 ) -> Result<(), Box<dyn std::error::Error>> {
18068 if t == 0
18069 || t > 8
18070 || in_f % 8 != 0
18071 || x_t.len() < t * in_f
18072 || yq.len() < t * out_q
18073 || yk.len() < t * out_kv
18074 || yv.len() < t * out_kv
18075 || (out_g > 0 && yg.len() < t * out_g)
18076 {
18077 return Err("matvec_bf16_qkvg_tcol geometry".into());
18078 }
18079 let grid = out_q + 2 * out_kv + out_g;
18080 let cfg = LaunchConfig {
18081 grid_dim: (grid as u32, 1, 1),
18082 block_dim: (mmv_block(), 1, 1),
18083 shared_mem_bytes: 0,
18084 };
18085 let (ini, oq, okv, og, ti) = (
18086 in_f as i32,
18087 out_q as i32,
18088 out_kv as i32,
18089 out_g as i32,
18090 t as i32,
18091 );
18092 let __s_b = self.gpu.stream();
18093 let f = self.func("matvec_bf16_qkvg_tcol");
18099 let mut b = __s_b.launch_builder(&f);
18100 b.arg(wq)
18101 .arg(wk)
18102 .arg(wv)
18103 .arg(wg)
18104 .arg(x_t)
18105 .arg(yq)
18106 .arg(yk)
18107 .arg(yv)
18108 .arg(yg)
18109 .arg(&ini)
18110 .arg(&oq)
18111 .arg(&okv)
18112 .arg(&og)
18113 .arg(&ti);
18114 unsafe {
18115 b.launch(cfg)?;
18116 }
18117 Ok(())
18118 }
18119
18120 pub fn matvec_bf16_qkvg_into(
18121 &self,
18122 wq: &CudaSlice<u8>,
18123 wk: &CudaSlice<u8>,
18124 wv: &CudaSlice<u8>,
18125 wg: &CudaSlice<u8>,
18126 x: &CudaSlice<f32>,
18127 yq: &mut CudaSlice<f32>,
18128 yk: &mut CudaSlice<f32>,
18129 yv: &mut CudaSlice<f32>,
18130 yg: &mut CudaSlice<f32>,
18131 in_f: usize,
18132 out_q: usize,
18133 out_kv: usize,
18134 out_g: usize,
18135 ) -> Result<(), Box<dyn std::error::Error>> {
18136 if in_f % 8 != 0
18137 || wq.len() != out_q * in_f * 2
18138 || wk.len() != out_kv * in_f * 2
18139 || wv.len() != out_kv * in_f * 2
18140 || wg.len() < out_g * in_f * 2
18141 || x.len() < in_f
18142 || yq.len() < out_q
18143 || yk.len() < out_kv
18144 || yv.len() < out_kv
18145 || (out_g > 0 && yg.len() < out_g)
18146 {
18147 return Err(format!(
18148 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18149 )
18150 .into());
18151 }
18152 let f = self.func("matvec_bf16_qkvg");
18153 let cfg = LaunchConfig {
18154 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18155 block_dim: (mmv_block(), 1, 1),
18156 shared_mem_bytes: 0,
18157 };
18158 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18159 let __s_b = self.gpu.stream();
18160 let mut b = __s_b.launch_builder(&f);
18161 b.arg(wq)
18162 .arg(wk)
18163 .arg(wv)
18164 .arg(wg)
18165 .arg(x)
18166 .arg(yq)
18167 .arg(yk)
18168 .arg(yv)
18169 .arg(yg)
18170 .arg(&inf)
18171 .arg(&oq)
18172 .arg(&okv)
18173 .arg(&og);
18174 unsafe {
18175 b.launch(cfg)?;
18176 }
18177 Ok(())
18178 }
18179
18180 pub fn matvec_bf16_b4_into(
18182 &self,
18183 w: [&CudaSlice<u8>; 4],
18184 x: &CudaSlice<f32>,
18185 y: &mut CudaSlice<f32>,
18186 block_cols: usize,
18187 out_f: usize,
18188 ) -> Result<(), Box<dyn std::error::Error>> {
18189 if block_cols % 8 != 0
18190 || x.len() < 4 * block_cols
18191 || y.len() < out_f
18192 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18193 {
18194 return Err(format!(
18195 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18196 x.len()
18197 )
18198 .into());
18199 }
18200 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18203 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18204 let f = self.func(if x2 {
18205 "matvec_bf16_b4_x2"
18206 } else {
18207 "matvec_bf16_b4"
18208 });
18209 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18210 let cfg = LaunchConfig {
18211 grid_dim: (grid as u32, 1, 1),
18212 block_dim: (mmv_block(), 1, 1),
18213 shared_mem_bytes: 0,
18214 };
18215 let (bc, of) = (block_cols as i32, out_f as i32);
18216 let __s_b = self.gpu.stream();
18217 let mut b = __s_b.launch_builder(&f);
18218 b.arg(w[0])
18219 .arg(w[1])
18220 .arg(w[2])
18221 .arg(w[3])
18222 .arg(x)
18223 .arg(y)
18224 .arg(&bc)
18225 .arg(&of);
18226 unsafe {
18227 b.launch(cfg)?;
18228 }
18229 Ok(())
18230 }
18231
18232 pub fn matvec_bf16_b4_tcol_into(
18238 &self,
18239 w: [&CudaSlice<u8>; 4],
18240 x_t: &CudaSlice<f32>,
18241 y_t: &mut CudaSlice<f32>,
18242 block_cols: usize,
18243 out_f: usize,
18244 t: usize,
18245 ) -> Result<(), Box<dyn std::error::Error>> {
18246 if block_cols % 8 != 0
18247 || t == 0
18248 || t > 8
18249 || x_t.len() < t * 4 * block_cols
18250 || y_t.len() < t * out_f
18251 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18252 {
18253 return Err(format!(
18254 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18255 x_t.len()
18256 )
18257 .into());
18258 }
18259 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18260 return Err(
18261 "b4 tcol verify is qualified against the plain b4 kernel only \
18262 (MEMRA_B4_X2=1 is a different t=1 program)"
18263 .into(),
18264 );
18265 }
18266 let cfg = LaunchConfig {
18270 grid_dim: (out_f as u32, 1, 1),
18271 block_dim: (mmv_block(), 1, 1),
18272 shared_mem_bytes: 0,
18273 };
18274 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18275 let __s_b = self.gpu.stream();
18276 let f = self.func("matvec_bf16_b4_tcol");
18277 let mut b = __s_b.launch_builder(&f);
18278 b.arg(w[0])
18279 .arg(w[1])
18280 .arg(w[2])
18281 .arg(w[3])
18282 .arg(x_t)
18283 .arg(y_t)
18284 .arg(&bc)
18285 .arg(&of)
18286 .arg(&ti);
18287 unsafe {
18288 b.launch(cfg)?;
18289 }
18290 Ok(())
18291 }
18292
18293 pub fn q8_0_row_bytes(in_f: usize) -> usize {
18296 in_f / 32 * 34
18297 }
18298
18299 pub fn encode_q8_0_from_bf16(
18303 &self,
18304 w_bf16: &CudaSlice<u8>,
18305 out: &mut CudaSlice<u8>,
18306 in_f: usize,
18307 out_f: usize,
18308 ) -> Result<(), Box<dyn std::error::Error>> {
18309 if in_f % 32 != 0
18310 || w_bf16.len() < in_f * out_f * 2
18311 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18312 {
18313 return Err(format!(
18314 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18315 w_bf16.len(),
18316 out.len()
18317 )
18318 .into());
18319 }
18320 let f = self.func("encode_q8_0_rows_from_bf16");
18321 const PAIRS_PER_BLOCK: u32 = 4;
18324 let pairs = (out_f * (in_f / 32)) as u64;
18325 let cfg = LaunchConfig {
18326 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18327 block_dim: (32, PAIRS_PER_BLOCK, 1),
18328 shared_mem_bytes: 0,
18329 };
18330 let (ini, outi) = (in_f as i32, out_f as i32);
18331 let __s_b = self.gpu.stream();
18332 let mut b = __s_b.launch_builder(&f);
18333 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18334 unsafe {
18335 b.launch(cfg)?;
18336 }
18337 Ok(())
18338 }
18339
18340 #[allow(clippy::too_many_arguments)]
18345 pub fn qmatvec_q8_0_qkv_rp_into(
18346 &self,
18347 wq: &CudaSlice<u8>,
18348 wk: &CudaSlice<u8>,
18349 wv: &CudaSlice<u8>,
18350 aq: &CudaSlice<i8>,
18351 ad: &CudaSlice<f32>,
18352 yq: &mut CudaSlice<f32>,
18353 yk: &mut CudaSlice<f32>,
18354 yv: &mut CudaSlice<f32>,
18355 in_f: usize,
18356 out_q: usize,
18357 out_kv: usize,
18358 ) -> Result<(), Box<dyn std::error::Error>> {
18359 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
18361 let nblk = in_f / 32;
18362 if in_f % 32 != 0
18363 || aq.len() < in_f
18364 || ad.len() < nblk
18365 || yq.len() < out_q
18366 || yk.len() < out_kv
18367 || yv.len() < out_kv
18368 || wq.len() < out_q * nblk * 34
18369 || wk.len() < out_kv * nblk * 34
18370 || wv.len() < out_kv * nblk * 34
18371 {
18372 return Err(
18373 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18374 );
18375 }
18376 let f = self.func("qmatvec_q8_0_qkv_rp");
18377 let cfg = LaunchConfig {
18378 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18379 block_dim: (32, ROWS_PER_BLOCK, 1),
18380 shared_mem_bytes: 0,
18381 };
18382 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18383 let __s_b = self.gpu.stream();
18384 let mut b = __s_b.launch_builder(&f);
18385 b.arg(wq)
18386 .arg(wk)
18387 .arg(wv)
18388 .arg(aq)
18389 .arg(ad)
18390 .arg(yq)
18391 .arg(yk)
18392 .arg(yv)
18393 .arg(&ini)
18394 .arg(&oq)
18395 .arg(&okv);
18396 unsafe {
18397 b.launch(cfg)?;
18398 }
18399 Ok(())
18400 }
18401
18402 #[allow(clippy::too_many_arguments)]
18406 pub fn qmatvec_q8_0_b4_rp_into(
18407 &self,
18408 w: [&CudaSlice<u8>; 4],
18409 aq: &CudaSlice<i8>,
18410 ad: &CudaSlice<f32>,
18411 y: &mut CudaSlice<f32>,
18412 block_cols: usize,
18413 out_f: usize,
18414 ) -> Result<(), Box<dyn std::error::Error>> {
18415 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
18417 if block_cols % 32 != 0
18418 || aq.len() < 4 * block_cols
18419 || ad.len() < 4 * nblk
18420 || y.len() < out_f
18421 || w.iter().any(|p| p.len() < out_f * nblk * 34)
18422 {
18423 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18424 }
18425 let f = self.func("qmatvec_q8_0_b4_rp");
18426 let cfg = LaunchConfig {
18427 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18428 block_dim: (32, ROWS_PER_BLOCK, 1),
18429 shared_mem_bytes: 0,
18430 };
18431 let (bc, of) = (block_cols as i32, out_f as i32);
18432 let __s_b = self.gpu.stream();
18433 let mut b = __s_b.launch_builder(&f);
18434 b.arg(w[0])
18435 .arg(w[1])
18436 .arg(w[2])
18437 .arg(w[3])
18438 .arg(aq)
18439 .arg(ad)
18440 .arg(y)
18441 .arg(&bc)
18442 .arg(&of);
18443 unsafe {
18444 b.launch(cfg)?;
18445 }
18446 Ok(())
18447 }
18448
18449 fn matvec_bf16_via_q8_mirror(
18452 &self,
18453 data: &CudaSlice<u8>,
18454 x: &CudaSlice<f32>,
18455 y: &mut CudaSlice<f32>,
18456 in_f: usize,
18457 out_f: usize,
18458 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18459 use cudarc::driver::DevicePtr;
18460 let key = {
18461 let s = self.gpu.stream();
18462 let (p, _g) = data.device_ptr(&s);
18463 p as u64
18464 };
18465 {
18466 let mut mirrors = self
18467 .w8_mirrors
18468 .lock()
18469 .map_err(|_| "w8 mirror map is poisoned")?;
18470 if !mirrors.contains_key(&key) {
18471 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18472 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18473 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18474 mirrors.insert(key, planar);
18475 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18481 eprintln!(
18482 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18483 mirrors.len()
18484 );
18485 }
18486 }
18487 }
18488 let nblk = in_f / 32;
18489 {
18490 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18491 if !act.contains_key(&in_f) {
18492 let aq = self.alloc_uninit::<i8>(in_f)?;
18493 let ad = self.alloc_uninit::<f32>(nblk)?;
18494 act.insert(in_f, (aq, ad));
18495 }
18496 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18497 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18498 }
18499 let mirrors = self
18500 .w8_mirrors
18501 .lock()
18502 .map_err(|_| "w8 mirror map is poisoned")?;
18503 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18504 let mirror = mirrors.get(&key).expect("built above");
18505 let (aq, ad) = act.get(&in_f).expect("built above");
18506 self.qmatvec_mmvq_into(
18507 mirror,
18508 aq,
18509 ad,
18510 1,
18511 in_f,
18512 out_f,
18513 QT_Q8_0,
18514 Self::q8_0_row_bytes(in_f),
18515 1.0,
18516 true,
18517 y,
18518 )?;
18519 Ok(Some(()))
18520 }
18521
18522 pub fn matvec_bf16_into(
18523 &self,
18524 data: &CudaSlice<u8>,
18525 x: &CudaSlice<f32>,
18526 y: &mut CudaSlice<f32>,
18527 in_f: usize,
18528 out_f: usize,
18529 ) -> Result<(), Box<dyn std::error::Error>> {
18530 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18531 return Err(format!(
18532 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18533 data.len(),
18534 x.len(),
18535 y.len()
18536 )
18537 .into());
18538 }
18539 if step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18546 if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
18547 return Ok(());
18548 }
18549 }
18550 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18554 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
18555 && in_f <= 2048;
18556 if x4 {
18557 let f = self.func("matvec_bf16_f32acc_x4");
18558 let cfg = LaunchConfig {
18559 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
18560 block_dim: (mmv_block(), 1, 1),
18561 shared_mem_bytes: 0,
18562 };
18563 let (ini, outi) = (in_f as i32, out_f as i32);
18564 let __s_b = self.gpu.stream();
18565 let mut b = __s_b.launch_builder(&f);
18566 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
18567 unsafe {
18568 b.launch(cfg)?;
18569 }
18570 return Ok(());
18571 }
18572 let f = self.func("matvec_bf16_f32acc");
18573 let cfg = LaunchConfig {
18574 grid_dim: (out_f as u32, 1, 1),
18575 block_dim: (mmv_block(), 1, 1),
18576 shared_mem_bytes: 0,
18577 };
18578 let ini = in_f as i32;
18579 let __s_b = self.gpu.stream();
18580 let mut b = __s_b.launch_builder(&f);
18581 b.arg(data).arg(x).arg(y).arg(&ini);
18582 unsafe {
18583 b.launch(cfg)?;
18584 }
18585 Ok(())
18586 }
18587
18588 pub fn matvec_bf16_view_into(
18591 &self,
18592 data: &cudarc::driver::CudaView<'_, u8>,
18593 x: &CudaSlice<f32>,
18594 y: &mut CudaSlice<f32>,
18595 in_f: usize,
18596 out_f: usize,
18597 ) -> Result<(), Box<dyn std::error::Error>> {
18598 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
18599 return Err(format!(
18600 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
18601 data.len(),
18602 x.len(),
18603 y.len()
18604 )
18605 .into());
18606 }
18607 let f = self.func("matvec_bf16_f32acc");
18608 let cfg = LaunchConfig {
18609 grid_dim: (out_f as u32, 1, 1),
18610 block_dim: (mmv_block(), 1, 1),
18611 shared_mem_bytes: 0,
18612 };
18613 let ini = in_f as i32;
18614 let __s_b = self.gpu.stream();
18615 let mut b = __s_b.launch_builder(&f);
18616 b.arg(data).arg(x).arg(y).arg(&ini);
18617 unsafe {
18618 b.launch(cfg)?;
18619 }
18620 Ok(())
18621 }
18622
18623 pub fn matvec_bf16_raw_out(
18626 &self,
18627 w: &CudaSlice<u8>,
18628 x: &CudaSlice<f32>,
18629 y_raw: u64,
18630 in_f: usize,
18631 out_f: usize,
18632 ) -> Result<(), Box<dyn std::error::Error>> {
18633 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
18634 return Err("matvec_bf16_raw_out geometry".into());
18635 }
18636 let f = self.func("matvec_bf16_f32acc");
18637 let cfg = LaunchConfig {
18638 grid_dim: (out_f as u32, 1, 1),
18639 block_dim: (mmv_block(), 1, 1),
18640 shared_mem_bytes: 0,
18641 };
18642 let ini = in_f as i32;
18643 let __s_b = self.gpu.stream();
18644 let mut b = __s_b.launch_builder(&f);
18645 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
18646 unsafe {
18647 b.launch(cfg)?;
18648 }
18649 Ok(())
18650 }
18651
18652 pub fn add3_raw(
18656 &self,
18657 a: &CudaSlice<f32>,
18658 b: &CudaSlice<f32>,
18659 sh_raw: u64,
18660 scale_raw: u64,
18661 dst: &mut CudaSlice<f32>,
18662 n: usize,
18663 ) -> Result<(), Box<dyn std::error::Error>> {
18664 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
18665 return Err("add3_raw geometry".into());
18666 }
18667 let f = self.func("add3_f32");
18668 let cfg = LaunchConfig {
18669 grid_dim: ((n as u32).div_ceil(256), 1, 1),
18670 block_dim: (256, 1, 1),
18671 shared_mem_bytes: 0,
18672 };
18673 let ni = n as i32;
18674 let __s_b = self.gpu.stream();
18675 let mut bld = __s_b.launch_builder(&f);
18676 bld.arg(a)
18677 .arg(b)
18678 .arg(&sh_raw)
18679 .arg(&scale_raw)
18680 .arg(dst)
18681 .arg(&ni);
18682 unsafe {
18683 bld.launch(cfg)?;
18684 }
18685 Ok(())
18686 }
18687
18688 pub fn matvec_bf16_down_addscale_into(
18691 &self,
18692 w: &CudaSlice<u8>,
18693 x: &CudaSlice<f32>,
18694 scale: &CudaSlice<f32>,
18695 dst: &mut CudaSlice<f32>,
18696 in_f: usize,
18697 out_f: usize,
18698 ) -> Result<(), Box<dyn std::error::Error>> {
18699 if w.len() != in_f * out_f * 2
18700 || x.len() < in_f
18701 || in_f % 8 != 0
18702 || dst.len() < out_f
18703 || scale.is_empty()
18704 {
18705 return Err("matvec_bf16_down_addscale geometry".into());
18706 }
18707 let f = self.func("matvec_bf16_down_addscale");
18708 let cfg = LaunchConfig {
18709 grid_dim: (out_f as u32, 1, 1),
18710 block_dim: (mmv_block(), 1, 1),
18711 shared_mem_bytes: 0,
18712 };
18713 let ini = in_f as i32;
18714 let __s_b = self.gpu.stream();
18715 let mut b = __s_b.launch_builder(&f);
18716 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
18717 unsafe {
18718 b.launch(cfg)?;
18719 }
18720 Ok(())
18721 }
18722
18723 #[allow(clippy::too_many_arguments)]
18727 pub fn matvec_bf16_dual_silu_rows_into(
18728 &self,
18729 wg: &CudaSlice<u8>,
18730 wu: &CudaSlice<u8>,
18731 x: &CudaSlice<f32>,
18732 act: &mut CudaSlice<f32>,
18733 in_f: usize,
18734 out_f: usize,
18735 limit: Option<f32>,
18736 t: usize,
18737 ) -> Result<(), Box<dyn std::error::Error>> {
18738 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
18739 return Err("matvec_bf16_dual_silu_rows geometry".into());
18740 }
18741 let f = self.func("matvec_bf16_dual_silu_rows");
18742 let cfg = LaunchConfig {
18743 grid_dim: (out_f as u32, t as u32, 1),
18744 block_dim: (mmv_block(), 1, 1),
18745 shared_mem_bytes: 0,
18746 };
18747 let (ini, outi) = (in_f as i32, out_f as i32);
18748 let lim = limit.unwrap_or(0.0);
18749 let __s_b = self.gpu.stream();
18750 let mut b = __s_b.launch_builder(&f);
18751 b.arg(wg)
18752 .arg(wu)
18753 .arg(x)
18754 .arg(&mut *act)
18755 .arg(&ini)
18756 .arg(&outi)
18757 .arg(&lim);
18758 unsafe {
18759 b.launch(cfg)?;
18760 }
18761 Ok(())
18762 }
18763
18764 pub fn matvec_bf16_rows_into(
18766 &self,
18767 w: &CudaSlice<u8>,
18768 x: &CudaSlice<f32>,
18769 y: &mut CudaSlice<f32>,
18770 in_f: usize,
18771 out_f: usize,
18772 t: usize,
18773 ) -> Result<(), Box<dyn std::error::Error>> {
18774 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
18775 return Err("matvec_bf16_rows geometry".into());
18776 }
18777 if t == 1 && step_tp_w8_on() && in_f % 32 == 0 && out_f >= 64 {
18783 if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
18784 return Ok(());
18785 }
18786 }
18787 let f = self.func("matvec_bf16_f32acc_x4_rows");
18788 let cfg = LaunchConfig {
18789 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
18790 block_dim: (mmv_block(), 1, 1),
18791 shared_mem_bytes: 0,
18792 };
18793 let (ini, outi) = (in_f as i32, out_f as i32);
18794 let __s_b = self.gpu.stream();
18795 let mut b = __s_b.launch_builder(&f);
18796 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
18797 unsafe {
18798 b.launch(cfg)?;
18799 }
18800 Ok(())
18801 }
18802
18803 pub fn matvec_bf16_dual_silu_into(
18804 &self,
18805 wg: &CudaSlice<u8>,
18806 wu: &CudaSlice<u8>,
18807 x: &CudaSlice<f32>,
18808 act: &mut CudaSlice<f32>,
18809 in_f: usize,
18810 out_f: usize,
18811 limit: Option<f32>,
18812 ) -> Result<(), Box<dyn std::error::Error>> {
18813 if wg.len() != in_f * out_f * 2
18814 || wu.len() != in_f * out_f * 2
18815 || x.len() < in_f
18816 || in_f % 8 != 0
18817 || act.len() < out_f
18818 {
18819 return Err("matvec_bf16_dual_silu geometry".into());
18820 }
18821 let f = self.func("matvec_bf16_dual_silu");
18822 let cfg = LaunchConfig {
18823 grid_dim: (out_f as u32, 1, 1),
18824 block_dim: (mmv_block(), 1, 1),
18825 shared_mem_bytes: 0,
18826 };
18827 let (ini, outi) = (in_f as i32, out_f as i32);
18828 let lim = limit.unwrap_or(0.0);
18829 let __s_b = self.gpu.stream();
18830 let mut b = __s_b.launch_builder(&f);
18831 b.arg(wg)
18832 .arg(wu)
18833 .arg(x)
18834 .arg(act)
18835 .arg(&ini)
18836 .arg(&outi)
18837 .arg(&lim);
18838 unsafe {
18839 b.launch(cfg)?;
18840 }
18841 Ok(())
18842 }
18843
18844 #[allow(clippy::too_many_arguments)]
18847 pub fn matvec_bf16_dual_view_into(
18848 &self,
18849 wg: &cudarc::driver::CudaView<'_, u8>,
18850 wu: &cudarc::driver::CudaView<'_, u8>,
18851 x: &CudaSlice<f32>,
18852 yg: &mut CudaSlice<f32>,
18853 yu: &mut CudaSlice<f32>,
18854 in_f: usize,
18855 out_f: usize,
18856 ) -> Result<(), Box<dyn std::error::Error>> {
18857 if wg.len() != in_f * out_f * 2
18858 || wu.len() != in_f * out_f * 2
18859 || x.len() < in_f
18860 || in_f % 8 != 0
18861 || yg.len() < out_f
18862 || yu.len() < out_f
18863 {
18864 return Err(format!(
18865 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18866 wg.len(),
18867 wu.len(),
18868 x.len()
18869 )
18870 .into());
18871 }
18872 let f = self.func("matvec_bf16_dual");
18873 let cfg = LaunchConfig {
18874 grid_dim: ((2 * out_f) as u32, 1, 1),
18875 block_dim: (mmv_block(), 1, 1),
18876 shared_mem_bytes: 0,
18877 };
18878 let (ini, outi) = (in_f as i32, out_f as i32);
18879 let __s_b = self.gpu.stream();
18880 let mut b = __s_b.launch_builder(&f);
18881 b.arg(wg)
18882 .arg(wu)
18883 .arg(x)
18884 .arg(yg)
18885 .arg(yu)
18886 .arg(&ini)
18887 .arg(&outi);
18888 unsafe {
18889 b.launch(cfg)?;
18890 }
18891 Ok(())
18892 }
18893
18894 #[allow(clippy::too_many_arguments)]
18896 pub fn matvec_bf16_dual_into(
18897 &self,
18898 wg: &CudaSlice<u8>,
18899 wu: &CudaSlice<u8>,
18900 x: &CudaSlice<f32>,
18901 yg: &mut CudaSlice<f32>,
18902 yu: &mut CudaSlice<f32>,
18903 in_f: usize,
18904 out_f: usize,
18905 ) -> Result<(), Box<dyn std::error::Error>> {
18906 if wg.len() != in_f * out_f * 2
18907 || wu.len() != in_f * out_f * 2
18908 || x.len() < in_f
18909 || in_f % 8 != 0
18910 || yg.len() < out_f
18911 || yu.len() < out_f
18912 {
18913 return Err(format!(
18914 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
18915 wg.len(),
18916 wu.len(),
18917 x.len()
18918 )
18919 .into());
18920 }
18921 let f = self.func("matvec_bf16_dual");
18922 let cfg = LaunchConfig {
18923 grid_dim: ((2 * out_f) as u32, 1, 1),
18924 block_dim: (mmv_block(), 1, 1),
18925 shared_mem_bytes: 0,
18926 };
18927 let (ini, outi) = (in_f as i32, out_f as i32);
18928 let __s_b = self.gpu.stream();
18929 let mut b = __s_b.launch_builder(&f);
18930 b.arg(wg)
18931 .arg(wu)
18932 .arg(x)
18933 .arg(yg)
18934 .arg(yu)
18935 .arg(&ini)
18936 .arg(&outi);
18937 unsafe {
18938 b.launch(cfg)?;
18939 }
18940 Ok(())
18941 }
18942
18943 pub(crate) fn matvec_bf16_dual(
18946 &self,
18947 wg: &CudaSlice<u8>,
18948 wu: &CudaSlice<u8>,
18949 x: &CudaSlice<f32>,
18950 in_f: usize,
18951 out_f: usize,
18952 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18953 if wg.len() != in_f * out_f * 2
18954 || wu.len() != in_f * out_f * 2
18955 || x.len() < in_f
18956 || in_f % 8 != 0
18957 {
18958 return Err(format!(
18959 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18960 wg.len(),
18961 wu.len(),
18962 x.len()
18963 )
18964 .into());
18965 }
18966 let mut yg = self.alloc_uninit::<f32>(out_f)?;
18967 let mut yu = self.alloc_uninit::<f32>(out_f)?;
18968 let f = self.func("matvec_bf16_dual");
18969 let cfg = LaunchConfig {
18970 grid_dim: ((2 * out_f) as u32, 1, 1),
18971 block_dim: (mmv_block(), 1, 1),
18972 shared_mem_bytes: 0,
18973 };
18974 let (ini, outi) = (in_f as i32, out_f as i32);
18975 let __s_b = self.gpu.stream();
18976 let mut b = __s_b.launch_builder(&f);
18977 b.arg(wg)
18978 .arg(wu)
18979 .arg(x)
18980 .arg(&mut yg)
18981 .arg(&mut yu)
18982 .arg(&ini)
18983 .arg(&outi);
18984 unsafe {
18985 b.launch(cfg)?;
18986 }
18987 Ok((yg, yu))
18988 }
18989
18990 #[allow(clippy::too_many_arguments)]
18991 fn linear_bf16_chunked_inner(
18992 &self,
18993 x: &CudaSlice<f32>,
18994 data: &CudaSlice<u8>,
18995 m: usize,
18996 in_f: usize,
18997 out_f: usize,
18998 exact: bool,
18999 canonical_chunk_rows: Option<usize>,
19000 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19001 const CHUNK_BYTES: usize = 256 << 20;
19002 if m == 1
19005 && !exact
19006 && canonical_chunk_rows.is_none()
19007 && in_f % 8 == 0
19008 && Self::bf16_mmv_on()
19009 {
19010 return self.matvec_bf16(data, x, in_f, out_f);
19011 }
19012 let row_bytes = in_f
19013 .checked_mul(std::mem::size_of::<f32>())
19014 .ok_or("BF16 chunk row byte count overflow")?;
19015 if row_bytes == 0 || out_f == 0 {
19016 return Err("BF16 chunk dimensions must be nonzero".into());
19017 }
19018 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
19019 let chunk_rows = match canonical_chunk_rows {
19020 Some(rows) if rows == 0 => {
19021 return Err("canonical BF16 chunk rows must be nonzero".into());
19022 }
19023 Some(rows) if rows > max_chunk_rows => {
19024 return Err(format!(
19025 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
19026 )
19027 .into());
19028 }
19029 Some(rows) if out_f % rows != 0 => {
19030 return Err(format!(
19031 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
19032 )
19033 .into());
19034 }
19035 Some(rows) => rows,
19036 None => max_chunk_rows,
19037 };
19038 if chunk_rows >= out_f {
19039 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
19040 return if exact {
19041 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19042 } else {
19043 self.linear(x, &wf32, m, in_f, out_f)
19044 };
19045 }
19046 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19047 let mut r0 = 0usize;
19048 while r0 < out_f {
19049 let rows = chunk_rows.min(out_f - r0);
19050 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19051 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19052 let yc = if exact {
19053 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19054 } else {
19055 self.linear(x, &wf32, m, in_f, rows)?
19056 };
19057 for mi in 0..m {
19059 let src = yc.slice(mi * rows..(mi + 1) * rows);
19060 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19061 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19062 }
19063 r0 += rows;
19064 }
19065 Ok(y)
19066 }
19067
19068 pub fn linear_bf16_resident(
19072 &self,
19073 x: &CudaSlice<f32>,
19074 data: &CudaSlice<u8>,
19075 m: usize,
19076 in_f: usize,
19077 out_f: usize,
19078 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19079 if data.len() != in_f * out_f * 2 {
19080 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19081 }
19082 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19083 }
19084
19085 pub fn linear_bf16_resident_canonical_rows(
19091 &self,
19092 x: &CudaSlice<f32>,
19093 data: &CudaSlice<u8>,
19094 m: usize,
19095 in_f: usize,
19096 out_f: usize,
19097 canonical_chunk_rows: usize,
19098 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19099 if data.len() != in_f * out_f * 2 {
19100 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19101 }
19102 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19103 }
19104
19105 pub fn linear_f32_resident_canonical_rows(
19110 &self,
19111 x: &CudaSlice<f32>,
19112 data: &CudaSlice<f32>,
19113 m: usize,
19114 in_f: usize,
19115 out_f: usize,
19116 canonical_chunk_rows: usize,
19117 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19118 self.linear_f32_resident_canonical_rows_inner(
19119 x,
19120 data,
19121 m,
19122 in_f,
19123 out_f,
19124 canonical_chunk_rows,
19125 false,
19126 )
19127 }
19128
19129 pub fn linear_f32_resident_canonical_rows_strided(
19135 &self,
19136 x: &CudaSlice<f32>,
19137 data: &CudaSlice<f32>,
19138 m: usize,
19139 in_f: usize,
19140 out_f: usize,
19141 canonical_chunk_rows: usize,
19142 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19143 self.linear_f32_resident_canonical_rows_inner(
19144 x,
19145 data,
19146 m,
19147 in_f,
19148 out_f,
19149 canonical_chunk_rows,
19150 true,
19151 )
19152 }
19153
19154 fn linear_f32_resident_canonical_rows_inner(
19155 &self,
19156 x: &CudaSlice<f32>,
19157 data: &CudaSlice<f32>,
19158 m: usize,
19159 in_f: usize,
19160 out_f: usize,
19161 canonical_chunk_rows: usize,
19162 strided_output: bool,
19163 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19164 if data.len() != in_f * out_f {
19165 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19166 }
19167 if canonical_chunk_rows == 0
19168 || canonical_chunk_rows > out_f
19169 || out_f % canonical_chunk_rows != 0
19170 {
19171 return Err(format!(
19172 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19173 )
19174 .into());
19175 }
19176 if canonical_chunk_rows == out_f {
19177 return self.linear(x, data, m, in_f, out_f);
19178 }
19179
19180 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19181 let input = x.slice(0..x.len());
19182 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19183 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19184 if m == 1 {
19185 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19186 self.linear_device_into(
19187 &input,
19188 &weights,
19189 &mut destination,
19190 1,
19191 in_f,
19192 canonical_chunk_rows,
19193 )?;
19194 continue;
19195 }
19196 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19197 if strided_output {
19198 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19199 } else {
19200 for token in 0..m {
19201 let source = chunk
19202 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19203 let mut destination =
19204 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19205 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19206 }
19207 }
19208 }
19209 Ok(y)
19210 }
19211
19212 pub fn linear_f32_resident_canonical_rows_t1_into(
19217 &self,
19218 x: &CudaSlice<f32>,
19219 data: &CudaSlice<f32>,
19220 y: &mut CudaSlice<f32>,
19221 in_f: usize,
19222 out_f: usize,
19223 canonical_chunk_rows: usize,
19224 ) -> Result<(), Box<dyn std::error::Error>> {
19225 if data.len() != in_f * out_f {
19226 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19227 }
19228 if y.len() != out_f || x.len() != in_f {
19229 return Err(format!(
19230 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19231 x.len(),
19232 y.len()
19233 )
19234 .into());
19235 }
19236 if canonical_chunk_rows == 0
19237 || canonical_chunk_rows > out_f
19238 || out_f % canonical_chunk_rows != 0
19239 {
19240 return Err(format!(
19241 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19242 )
19243 .into());
19244 }
19245 let input = x.slice(0..x.len());
19246 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19247 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19248 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19249 self.linear_device_into(
19250 &input,
19251 &weights,
19252 &mut destination,
19253 1,
19254 in_f,
19255 canonical_chunk_rows,
19256 )?;
19257 }
19258 Ok(())
19259 }
19260
19261 pub fn linear_t1_into(
19264 &self,
19265 x: &cudarc::driver::CudaView<'_, f32>,
19266 w: &cudarc::driver::CudaView<'_, f32>,
19267 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19268 in_f: usize,
19269 out_f: usize,
19270 ) -> Result<(), Box<dyn std::error::Error>> {
19271 self.linear_device_into(x, w, y, 1, in_f, out_f)
19272 }
19273
19274 pub fn linear_decode_exact(
19281 &self,
19282 x: &CudaSlice<f32>,
19283 w: &CudaSlice<f32>,
19284 m_tokens: usize,
19285 in_f: usize,
19286 out_f: usize,
19287 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19288 if m_tokens == 1 {
19289 return self.linear(x, w, 1, in_f, out_f);
19290 }
19291 let xv = self.view(x, m_tokens * in_f);
19292 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19293 for t in 0..m_tokens {
19294 let row = xv.slice(t * in_f..(t + 1) * in_f);
19295 let mut xr = self.alloc_uninit::<f32>(in_f)?;
19296 self.copy_view_into(&mut xr, 0, &row, in_f)?;
19297 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19298 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19299 }
19300 Ok(y)
19301 }
19302
19303 pub fn linear(
19304 &self,
19305 x: &CudaSlice<f32>,
19306 w: &CudaSlice<f32>,
19307 m_tokens: usize,
19308 in_f: usize,
19309 out_f: usize,
19310 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19311 self.linear_device(x, w, m_tokens, in_f, out_f)
19312 }
19313
19314 fn linear_device<I>(
19315 &self,
19316 x: &I,
19317 w: &I,
19318 m_tokens: usize,
19319 in_f: usize,
19320 out_f: usize,
19321 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19322 where
19323 I: cudarc::driver::DevicePtr<f32>,
19324 {
19325 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)?;
19327 Ok(c)
19328 }
19329
19330 fn linear_device_into<I, O>(
19331 &self,
19332 x: &I,
19333 w: &I,
19334 c: &mut O,
19335 m_tokens: usize,
19336 in_f: usize,
19337 out_f: usize,
19338 ) -> Result<(), Box<dyn std::error::Error>>
19339 where
19340 I: cudarc::driver::DevicePtr<f32>,
19341 O: cudarc::driver::DevicePtrMut<f32>,
19342 {
19343 use cudarc::cublaslt::{Matmul, MatmulConfig};
19344 let cfg = MatmulConfig {
19345 transa: true,
19346 transb: false,
19347 transc: false,
19348 m: out_f as u64,
19349 n: m_tokens as u64,
19350 k: in_f as u64,
19351 alpha: 1.0,
19352 lda: in_f as i64,
19353 ldb: in_f as i64,
19354 beta: 0.0,
19355 ldc: out_f as i64,
19356 stride_a: None,
19357 stride_b: None,
19358 stride_c: None,
19359 stride_bias: None,
19360 batch_size: None,
19361 };
19362 let blas = self.gpu.blas();
19363 unsafe {
19364 blas.matmul(cfg, w, x, c, None, None)?;
19365 }
19366 Ok(())
19367 }
19368
19369 pub fn sdpa_naive(
19378 &self,
19379 q: &CudaSlice<f32>,
19380 k: &CudaSlice<f32>,
19381 v: &CudaSlice<f32>,
19382 o: &mut CudaSlice<f32>,
19383 head_dim: usize,
19384 n_head: usize,
19385 n_head_kv: usize,
19386 t: usize,
19387 t_kv: usize,
19388 scale: f32,
19389 causal: bool,
19390 ) -> Result<(), Box<dyn std::error::Error>> {
19391 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
19392 return self.sdpa_naive_gmem(
19393 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19394 );
19395 }
19396 let f = self.func("sdpa_naive_f32");
19397 let cfg = LaunchConfig {
19398 grid_dim: (n_head as u32, t as u32, 1),
19399 block_dim: (128, 1, 1),
19400 shared_mem_bytes: (t_kv * 4) as u32,
19401 };
19402 let (hd, nh, nhkv, ti, tkvi, cz) = (
19403 head_dim as i32,
19404 n_head as i32,
19405 n_head_kv as i32,
19406 t as i32,
19407 t_kv as i32,
19408 causal as i32,
19409 );
19410 let __s_b = self.gpu.stream();
19411 let mut b = __s_b.launch_builder(&f);
19412 b.arg(q)
19413 .arg(k)
19414 .arg(v)
19415 .arg(o)
19416 .arg(&hd)
19417 .arg(&nh)
19418 .arg(&nhkv)
19419 .arg(&ti)
19420 .arg(&tkvi)
19421 .arg(&scale)
19422 .arg(&cz);
19423 unsafe {
19424 b.launch(cfg)?;
19425 }
19426 Ok(())
19427 }
19428
19429 #[allow(clippy::too_many_arguments)]
19438 pub fn sdpa_naive_gmem(
19439 &self,
19440 q: &CudaSlice<f32>,
19441 k: &CudaSlice<f32>,
19442 v: &CudaSlice<f32>,
19443 o: &mut CudaSlice<f32>,
19444 head_dim: usize,
19445 n_head: usize,
19446 n_head_kv: usize,
19447 t: usize,
19448 t_kv: usize,
19449 scale: f32,
19450 causal: bool,
19451 ) -> Result<(), Box<dyn std::error::Error>> {
19452 let ws_len = n_head
19453 .checked_mul(t)
19454 .and_then(|x| x.checked_mul(t_kv))
19455 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
19456 let ws_bytes = ws_len
19457 .checked_mul(std::mem::size_of::<f32>())
19458 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
19459 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
19460 return Err(format!(
19461 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
19462 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
19463 needs a tiled/flash kernel, not the naive oracle"
19464 )
19465 .into());
19466 }
19467 let mut scores = self.uninit(ws_len)?;
19468 let f = self.func("sdpa_naive_gmem_f32");
19469 let cfg = LaunchConfig {
19470 grid_dim: (n_head as u32, t as u32, 1),
19471 block_dim: (128, 1, 1),
19472 shared_mem_bytes: 0,
19473 };
19474 let (hd, nh, nhkv, ti, tkvi, cz) = (
19475 head_dim as i32,
19476 n_head as i32,
19477 n_head_kv as i32,
19478 t as i32,
19479 t_kv as i32,
19480 causal as i32,
19481 );
19482 let __s_b = self.gpu.stream();
19483 let mut b = __s_b.launch_builder(&f);
19484 b.arg(q)
19485 .arg(k)
19486 .arg(v)
19487 .arg(o)
19488 .arg(&mut scores)
19489 .arg(&hd)
19490 .arg(&nh)
19491 .arg(&nhkv)
19492 .arg(&ti)
19493 .arg(&tkvi)
19494 .arg(&scale)
19495 .arg(&cz);
19496 unsafe {
19497 b.launch(cfg)?;
19498 }
19499 Ok(())
19500 }
19501
19502 #[allow(clippy::too_many_arguments)]
19507 pub fn sdpa_naive_island(
19508 &self,
19509 q: &CudaSlice<f32>,
19510 k: &CudaSlice<f32>,
19511 v: &CudaSlice<f32>,
19512 o: &mut CudaSlice<f32>,
19513 span_id: &CudaSlice<i32>,
19514 head_dim: usize,
19515 n_head: usize,
19516 n_head_kv: usize,
19517 t: usize,
19518 t_kv: usize,
19519 scale: f32,
19520 window: usize,
19521 ) -> Result<(), Box<dyn std::error::Error>> {
19522 let f = self.func("sdpa_naive_island_f32");
19523 let cfg = LaunchConfig {
19524 grid_dim: (n_head as u32, t as u32, 1),
19525 block_dim: (128, 1, 1),
19526 shared_mem_bytes: (t_kv * 4) as u32,
19527 };
19528 let (hd, nh, nhkv, ti, tkvi, wi) = (
19529 head_dim as i32,
19530 n_head as i32,
19531 n_head_kv as i32,
19532 t as i32,
19533 t_kv as i32,
19534 window as i32,
19535 );
19536 let __s_b = self.gpu.stream();
19537 let mut b = __s_b.launch_builder(&f);
19538 b.arg(q)
19539 .arg(k)
19540 .arg(v)
19541 .arg(o)
19542 .arg(span_id)
19543 .arg(&hd)
19544 .arg(&nh)
19545 .arg(&nhkv)
19546 .arg(&ti)
19547 .arg(&tkvi)
19548 .arg(&scale)
19549 .arg(&wi);
19550 unsafe {
19551 b.launch(cfg)?;
19552 }
19553 Ok(())
19554 }
19555
19556 #[allow(clippy::too_many_arguments)]
19558 pub fn sdpa_naive_w(
19559 &self,
19560 q: &CudaSlice<f32>,
19561 k: &CudaSlice<f32>,
19562 v: &CudaSlice<f32>,
19563 o: &mut CudaSlice<f32>,
19564 head_dim: usize,
19565 n_head: usize,
19566 n_head_kv: usize,
19567 t: usize,
19568 t_kv: usize,
19569 scale: f32,
19570 causal: bool,
19571 window: usize,
19572 ) -> Result<(), Box<dyn std::error::Error>> {
19573 let f = self.func("sdpa_naive_w_f32");
19574 let cfg = LaunchConfig {
19575 grid_dim: (n_head as u32, t as u32, 1),
19576 block_dim: (128, 1, 1),
19577 shared_mem_bytes: (t_kv * 4) as u32,
19578 };
19579 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19580 head_dim as i32,
19581 n_head as i32,
19582 n_head_kv as i32,
19583 t as i32,
19584 t_kv as i32,
19585 causal as i32,
19586 window as i32,
19587 );
19588 let __s_b = self.gpu.stream();
19589 let mut b = __s_b.launch_builder(&f);
19590 b.arg(q)
19591 .arg(k)
19592 .arg(v)
19593 .arg(o)
19594 .arg(&hd)
19595 .arg(&nh)
19596 .arg(&nhkv)
19597 .arg(&ti)
19598 .arg(&tkvi)
19599 .arg(&scale)
19600 .arg(&cz)
19601 .arg(&wi);
19602 unsafe {
19603 b.launch(cfg)?;
19604 }
19605 Ok(())
19606 }
19607
19608 #[allow(clippy::too_many_arguments)]
19618 pub fn sdpa_naive_w_lo(
19619 &self,
19620 q: &CudaSlice<f32>,
19621 k: &CudaSlice<f32>,
19622 v: &CudaSlice<f32>,
19623 o: &mut CudaSlice<f32>,
19624 head_dim: usize,
19625 n_head: usize,
19626 n_head_kv: usize,
19627 t: usize,
19628 t_kv: usize,
19629 scale: f32,
19630 causal: bool,
19631 window: usize,
19632 ) -> Result<(), Box<dyn std::error::Error>> {
19633 let kv_lo = if window > 0 {
19634 (t_kv - t + 1).saturating_sub(window)
19635 } else {
19636 0
19637 };
19638 let smem = (t_kv - kv_lo) * 4;
19639 if smem > 48 * 1024 {
19640 return Err(format!(
19641 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
19642 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
19643 a window this wide needs the multi-pass long-ctx kernel"
19644 )
19645 .into());
19646 }
19647 let f = self.func("sdpa_naive_w_lo_f32");
19648 let cfg = LaunchConfig {
19649 grid_dim: (n_head as u32, t as u32, 1),
19650 block_dim: (128, 1, 1),
19651 shared_mem_bytes: smem as u32,
19652 };
19653 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
19654 head_dim as i32,
19655 n_head as i32,
19656 n_head_kv as i32,
19657 t as i32,
19658 t_kv as i32,
19659 causal as i32,
19660 window as i32,
19661 kv_lo as i32,
19662 );
19663 let __s_b = self.gpu.stream();
19664 let mut b = __s_b.launch_builder(&f);
19665 b.arg(q)
19666 .arg(k)
19667 .arg(v)
19668 .arg(o)
19669 .arg(&hd)
19670 .arg(&nh)
19671 .arg(&nhkv)
19672 .arg(&ti)
19673 .arg(&tkvi)
19674 .arg(&scale)
19675 .arg(&cz)
19676 .arg(&wi)
19677 .arg(&lo);
19678 unsafe {
19679 b.launch(cfg)?;
19680 }
19681 Ok(())
19682 }
19683
19684 pub fn sdpa_naive_view(
19686 &self,
19687 q: &CudaSlice<f32>,
19688 k: &cudarc::driver::CudaView<f32>,
19689 v: &cudarc::driver::CudaView<f32>,
19690 o: &mut CudaSlice<f32>,
19691 head_dim: usize,
19692 n_head: usize,
19693 n_head_kv: usize,
19694 t: usize,
19695 t_kv: usize,
19696 scale: f32,
19697 causal: bool,
19698 ) -> Result<(), Box<dyn std::error::Error>> {
19699 let f = self.func("sdpa_naive_f32");
19700 let cfg = LaunchConfig {
19701 grid_dim: (n_head as u32, t as u32, 1),
19702 block_dim: (128, 1, 1),
19703 shared_mem_bytes: (t_kv * 4) as u32,
19704 };
19705 let (hd, nh, nhkv, ti, tkvi, cz) = (
19706 head_dim as i32,
19707 n_head as i32,
19708 n_head_kv as i32,
19709 t as i32,
19710 t_kv as i32,
19711 causal as i32,
19712 );
19713 let __s_b = self.gpu.stream();
19714 let mut b = __s_b.launch_builder(&f);
19715 b.arg(q)
19716 .arg(k)
19717 .arg(v)
19718 .arg(o)
19719 .arg(&hd)
19720 .arg(&nh)
19721 .arg(&nhkv)
19722 .arg(&ti)
19723 .arg(&tkvi)
19724 .arg(&scale)
19725 .arg(&cz);
19726 unsafe {
19727 b.launch(cfg)?;
19728 }
19729 Ok(())
19730 }
19731
19732 #[allow(clippy::too_many_arguments)]
19740 pub fn fa_dequant_kv_view_f32(
19741 &self,
19742 k: &cudarc::driver::CudaView<u8>,
19743 v: &cudarc::driver::CudaView<u8>,
19744 kf: &mut CudaSlice<f32>,
19745 vf: &mut CudaSlice<f32>,
19746 kv_dim_k: usize,
19747 kv_dim_v: usize,
19748 t_kv: usize,
19749 k_tok_bytes: usize,
19750 v_tok_bytes: usize,
19751 g: bool,
19752 ) -> Result<(), Box<dyn std::error::Error>> {
19753 let f = if g {
19754 self.func_g("fa_dequant_kv_ws_f32")
19755 } else {
19756 self.func("fa_dequant_kv_ws_f32")
19757 };
19758 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
19759 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19760 let cfg = LaunchConfig {
19761 grid_dim: (nblk.max(1), 1, 1),
19762 block_dim: (256, 1, 1),
19763 shared_mem_bytes: 0,
19764 };
19765 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
19766 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19767 let __s_b = self.gpu.stream();
19768 let mut b = __s_b.launch_builder(&f);
19769 b.arg(k)
19770 .arg(v)
19771 .arg(&mut *kf)
19772 .arg(&mut *vf)
19773 .arg(&kdk)
19774 .arg(&kdv)
19775 .arg(&tkvi)
19776 .arg(&ktb)
19777 .arg(&vtb);
19778 unsafe {
19779 b.launch(cfg)?;
19780 }
19781 Ok(())
19782 }
19783
19784 #[allow(clippy::too_many_arguments)]
19785 pub fn sdpa_naive_quantized_view(
19786 &self,
19787 q: &CudaSlice<f32>,
19788 k: &cudarc::driver::CudaView<u8>,
19789 v: &cudarc::driver::CudaView<u8>,
19790 o: &mut CudaSlice<f32>,
19791 head_dim: usize,
19792 n_head: usize,
19793 n_head_kv: usize,
19794 t: usize,
19795 t_kv: usize,
19796 scale: f32,
19797 causal: bool,
19798 k_tok_bytes: usize,
19799 v_tok_bytes: usize,
19800 ) -> Result<(), Box<dyn std::error::Error>> {
19801 let kv_dim = n_head_kv * head_dim;
19802 let mut kf = self.uninit(t_kv * kv_dim)?;
19803 let mut vf = self.uninit(t_kv * kv_dim)?;
19804 let f = self.func("fa_dequant_kv_ws_f32");
19805 let total = (2 * t_kv * kv_dim) as u64;
19806 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19807 let cfg = LaunchConfig {
19808 grid_dim: (nblk.max(1), 1, 1),
19809 block_dim: (256, 1, 1),
19810 shared_mem_bytes: 0,
19811 };
19812 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19813 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19814 let __s_b = self.gpu.stream();
19815 let mut b = __s_b.launch_builder(&f);
19816 b.arg(k)
19817 .arg(v)
19818 .arg(&mut kf)
19819 .arg(&mut vf)
19820 .arg(&kv_dim_i)
19821 .arg(&kv_dim_i)
19822 .arg(&t_kv_i)
19823 .arg(&k_tok_bytes_i)
19824 .arg(&v_tok_bytes_i);
19825 unsafe { b.launch(cfg)? };
19826 self.sdpa_naive(
19827 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19828 )
19829 }
19830
19831 #[allow(clippy::too_many_arguments)]
19843 pub fn sdpa_naive_w_quantized_view(
19844 &self,
19845 q: &CudaSlice<f32>,
19846 k: &cudarc::driver::CudaView<u8>,
19847 v: &cudarc::driver::CudaView<u8>,
19848 o: &mut CudaSlice<f32>,
19849 head_dim: usize,
19850 n_head: usize,
19851 n_head_kv: usize,
19852 t: usize,
19853 t_kv: usize,
19854 scale: f32,
19855 causal: bool,
19856 window: usize,
19857 k_tok_bytes: usize,
19858 v_tok_bytes: usize,
19859 ) -> Result<(), Box<dyn std::error::Error>> {
19860 let kv_dim = n_head_kv * head_dim;
19861 let mut kf = self.uninit(t_kv * kv_dim)?;
19862 let mut vf = self.uninit(t_kv * kv_dim)?;
19863 let f = self.func("fa_dequant_kv_ws_f32");
19864 let total = (2 * t_kv * kv_dim) as u64;
19865 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
19866 let cfg = LaunchConfig {
19867 grid_dim: (nblk.max(1), 1, 1),
19868 block_dim: (256, 1, 1),
19869 shared_mem_bytes: 0,
19870 };
19871 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
19872 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
19873 let __s_b = self.gpu.stream();
19874 let mut b = __s_b.launch_builder(&f);
19875 b.arg(k)
19876 .arg(v)
19877 .arg(&mut kf)
19878 .arg(&mut vf)
19879 .arg(&kv_dim_i)
19880 .arg(&kv_dim_i)
19881 .arg(&t_kv_i)
19882 .arg(&k_tok_bytes_i)
19883 .arg(&v_tok_bytes_i);
19884 unsafe { b.launch(cfg)? };
19885 self.sdpa_naive_w(
19886 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19887 )
19888 }
19889
19890 pub fn fa_prefill(
19894 &self,
19895 q: &CudaSlice<f32>,
19896 k: &CudaSlice<f32>,
19897 v: &CudaSlice<f32>,
19898 o: &mut CudaSlice<f32>,
19899 head_dim: usize,
19900 n_head: usize,
19901 n_head_kv: usize,
19902 t: usize,
19903 t_kv: usize,
19904 scale: f32,
19905 causal: bool,
19906 ) -> Result<(), Box<dyn std::error::Error>> {
19907 if portable_mma_gated() {
19908 return self.sdpa_naive(
19909 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19910 );
19911 }
19912 let fa3_on = head_dim == 256
19920 && causal
19921 && t == t_kv
19922 && match std::env::var("MEMRA_FA3").as_deref() {
19923 Ok("0") => false,
19924 Ok("1") => {
19928 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
19929 true
19930 }
19931 _ => cfg!(memra_hopper_mma),
19932 };
19933 if fa3_on {
19934 let n = t * n_head * head_dim;
19935 let nkv = t * n_head_kv * head_dim;
19936 let mut q16 = self.alloc_u8_uninit(n * 2)?;
19937 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
19938 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
19939 self.f32_to_bf16_into(q, &mut q16, n)?;
19940 self.f32_to_bf16_into(k, &mut k16, nkv)?;
19941 self.f32_to_bf16_into(v, &mut v16, nkv)?;
19942 let rc = {
19943 use cudarc::driver::{DevicePtr, DevicePtrMut};
19944 let stream = self.gpu.stream();
19945 let (qp, _g1) = q16.device_ptr(&stream);
19946 let (kp, _g2) = k16.device_ptr(&stream);
19947 let (vp, _g3) = v16.device_ptr(&stream);
19948 let (op, _g4) = o.device_ptr_mut(&stream);
19949 unsafe {
19950 memra_fa3_prefill(
19951 qp as *const core::ffi::c_void,
19952 kp as *const core::ffi::c_void,
19953 vp as *const core::ffi::c_void,
19954 op as *mut f32,
19955 t as i32,
19956 n_head as i32,
19957 n_head_kv as i32,
19958 head_dim as i32,
19959 scale,
19960 stream.cu_stream() as *mut core::ffi::c_void,
19961 )
19962 }
19963 };
19964 if rc != 0 {
19965 return Err(format!("memra_fa3_prefill rc={rc}").into());
19966 }
19967 return Ok(());
19968 }
19969 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19974 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
19975 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
19976 const BLOCK_Q: usize = 64;
19977 const BKX: usize = 32;
19978 let f = self.func("fa_prefill_bf16_p1");
19979 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
19980 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
19981 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19982 f.set_attribute(
19983 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19984 shmem as i32,
19985 )?;
19986 let cfg = LaunchConfig {
19987 grid_dim: (
19988 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19989 n_head as u32,
19990 1,
19991 ),
19992 block_dim: (32, 4, 1),
19993 shared_mem_bytes: shmem,
19994 };
19995 let (hd, nh, nhkv, ti, tkvi, cz) = (
19996 head_dim as i32,
19997 n_head as i32,
19998 n_head_kv as i32,
19999 t as i32,
20000 t_kv as i32,
20001 causal as i32,
20002 );
20003 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20004 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20005 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20006 let __s_b = self.gpu.stream();
20007 let mut b = __s_b.launch_builder(&f);
20008 b.arg(&qb)
20009 .arg(&kb)
20010 .arg(&vb)
20011 .arg(o)
20012 .arg(&hd)
20013 .arg(&nh)
20014 .arg(&nhkv)
20015 .arg(&ti)
20016 .arg(&tkvi)
20017 .arg(&scale)
20018 .arg(&cz);
20019 unsafe {
20020 b.launch(cfg)?;
20021 }
20022 return Ok(());
20023 }
20024 const BK: usize = 32;
20030 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
20033 let (block_q, warps, w2_sfx): (usize, u32, &str) =
20034 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
20035 let hd_sfx = fa_hd_suffix(head_dim)?;
20039 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20040 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20045 let (kb16, vb16) = if bf16kv {
20046 let n = t_kv * n_head_kv * head_dim;
20047 let mut kb = self.alloc_u8_uninit(n * 2)?;
20048 let mut vb = self.alloc_u8_uninit(n * 2)?;
20049 let fcv = self.func("f32_to_bf16_bulk");
20050 let ni = n as i64;
20051 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20052 let __s_b = self.gpu.stream();
20053 let mut b = __s_b.launch_builder(&fcv);
20054 b.arg(k).arg(&mut kb).arg(&ni);
20055 unsafe {
20056 b.launch(cfgc)?;
20057 }
20058 let __s_b = self.gpu.stream();
20059 let mut b = __s_b.launch_builder(&fcv);
20060 b.arg(v).arg(&mut vb).arg(&ni);
20061 unsafe {
20062 b.launch(cfgc)?;
20063 }
20064 (Some(kb), Some(vb))
20065 } else {
20066 (None, None)
20067 };
20068 let f = self.func(&if bf16kv {
20069 format!("fa_prefill_bf16kv_pp{hd_sfx}")
20070 } else {
20071 format!(
20072 "fa_prefill_f32{}{}{hd_sfx}",
20073 if floor { "" } else { "_pp" },
20074 if floor { "" } else { w2_sfx }
20075 )
20076 });
20077 let kv_stages = if bf16kv { 2 } else { 1 };
20080 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20081 + 4 * (block_q * BK + 2 * block_q)) as u32;
20082 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20083 f.set_attribute(
20084 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20085 shmem as i32,
20086 )?;
20087 let cfg = LaunchConfig {
20088 grid_dim: (
20089 (t as u32 + block_q as u32 - 1) / block_q as u32,
20090 n_head as u32,
20091 1,
20092 ),
20093 block_dim: (32, warps, 1),
20094 shared_mem_bytes: shmem,
20095 };
20096 let (hd, nh, nhkv, ti, tkvi, cz) = (
20097 head_dim as i32,
20098 n_head as i32,
20099 n_head_kv as i32,
20100 t as i32,
20101 t_kv as i32,
20102 causal as i32,
20103 );
20104 let __s_b = self.gpu.stream();
20105 let mut b = __s_b.launch_builder(&f);
20106 b.arg(q);
20107 match (&kb16, &vb16) {
20108 (Some(kb), Some(vb)) => {
20109 b.arg(kb).arg(vb);
20110 }
20111 _ => {
20112 b.arg(k).arg(v);
20113 }
20114 }
20115 b.arg(o)
20116 .arg(&hd)
20117 .arg(&nh)
20118 .arg(&nhkv)
20119 .arg(&ti)
20120 .arg(&tkvi)
20121 .arg(&scale)
20122 .arg(&cz);
20123 unsafe {
20124 b.launch(cfg)?;
20125 }
20126 Ok(())
20127 }
20128
20129 #[allow(clippy::too_many_arguments)]
20133 pub fn fa_prefill_w(
20134 &self,
20135 q: &CudaSlice<f32>,
20136 k: &CudaSlice<f32>,
20137 v: &CudaSlice<f32>,
20138 o: &mut CudaSlice<f32>,
20139 head_dim: usize,
20140 n_head: usize,
20141 n_head_kv: usize,
20142 t: usize,
20143 t_kv: usize,
20144 scale: f32,
20145 causal: bool,
20146 window: usize,
20147 ) -> Result<(), Box<dyn std::error::Error>> {
20148 if portable_mma_gated() {
20151 return self.sdpa_naive_w(
20152 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20153 );
20154 }
20155 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20159 let faw_f32 =
20160 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20161 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20162 self.fa_prefill_w_arm(
20163 q,
20164 k,
20165 v,
20166 o,
20167 head_dim,
20168 n_head,
20169 n_head_kv,
20170 t,
20171 t_kv,
20172 scale,
20173 causal,
20174 window,
20175 floor || faw_f32,
20176 floor,
20177 )
20178 }
20179
20180 #[allow(clippy::too_many_arguments)]
20183 pub fn fa_prefill_w_pre(
20184 &self,
20185 qb: &CudaSlice<u8>,
20186 kb: &CudaSlice<u8>,
20187 vb: &CudaSlice<u8>,
20188 o: &mut CudaSlice<f32>,
20189 head_dim: usize,
20190 n_head: usize,
20191 n_head_kv: usize,
20192 t: usize,
20193 t_kv: usize,
20194 scale: f32,
20195 causal: bool,
20196 window: usize,
20197 v_f16: bool,
20198 ) -> Result<(), Box<dyn std::error::Error>> {
20199 const BLOCK_Q: usize = 64;
20200 const BK: usize = 32;
20201 debug_assert_eq!(head_dim, 256);
20202 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20203 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20204 if hp {
20205 const BLOCK_QH: usize = 32;
20206 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20209 let vh: &CudaSlice<u8> = if v_f16 {
20210 vb
20211 } else {
20212 let n = t_kv * n_head_kv * head_dim;
20213 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20214 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20215 }
20216 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20217 vguard.as_ref().unwrap()
20218 };
20219 let f = self.func("fa_prefill_w_bf16_p1h2");
20220 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) 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: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20228 block_dim: (32, 4, 1),
20229 shared_mem_bytes: shmem,
20230 };
20231 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20232 head_dim as i32,
20233 n_head as i32,
20234 n_head_kv as i32,
20235 t as i32,
20236 t_kv as i32,
20237 causal as i32,
20238 window as i32,
20239 );
20240 let __s_b = self.gpu.stream();
20241 let mut b = __s_b.launch_builder(&f);
20242 b.arg(qb)
20243 .arg(kb)
20244 .arg(vh)
20245 .arg(o)
20246 .arg(&hd)
20247 .arg(&nh)
20248 .arg(&nhkv)
20249 .arg(&ti)
20250 .arg(&tkvi)
20251 .arg(&scale)
20252 .arg(&cz)
20253 .arg(&wi);
20254 unsafe {
20255 b.launch(cfg)?;
20256 }
20257 return Ok(());
20258 }
20259 let f = self.func("fa_prefill_w_bf16_p1");
20260 let shmem =
20261 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20262 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20263 f.set_attribute(
20264 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20265 shmem as i32,
20266 )?;
20267 let cfg = LaunchConfig {
20268 grid_dim: (
20269 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20270 n_head as u32,
20271 1,
20272 ),
20273 block_dim: (32, 4, 1),
20274 shared_mem_bytes: shmem,
20275 };
20276 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20277 head_dim as i32,
20278 n_head as i32,
20279 n_head_kv as i32,
20280 t as i32,
20281 t_kv as i32,
20282 causal as i32,
20283 window as i32,
20284 );
20285 let __s_b = self.gpu.stream();
20286 let mut b = __s_b.launch_builder(&f);
20287 b.arg(qb)
20288 .arg(kb)
20289 .arg(vb)
20290 .arg(o)
20291 .arg(&hd)
20292 .arg(&nh)
20293 .arg(&nhkv)
20294 .arg(&ti)
20295 .arg(&tkvi)
20296 .arg(&scale)
20297 .arg(&cz)
20298 .arg(&wi);
20299 unsafe {
20300 b.launch(cfg)?;
20301 }
20302 Ok(())
20303 }
20304
20305 #[allow(clippy::too_many_arguments)]
20307 pub fn fa_prefill_w_arm(
20308 &self,
20309 q: &CudaSlice<f32>,
20310 k: &CudaSlice<f32>,
20311 v: &CudaSlice<f32>,
20312 o: &mut CudaSlice<f32>,
20313 head_dim: usize,
20314 n_head: usize,
20315 n_head_kv: usize,
20316 t: usize,
20317 t_kv: usize,
20318 scale: f32,
20319 causal: bool,
20320 window: usize,
20321 f32_stage: bool,
20322 floor: bool,
20323 ) -> Result<(), Box<dyn std::error::Error>> {
20324 const BLOCK_Q: usize = 64;
20325 const BK: usize = 32;
20326 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20327 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20331 let p1 = !floor
20332 && !f32_stage
20333 && *P1_ON.get_or_init(|| {
20334 std::env::var("MEMRA_FAW_P1")
20335 .map(|v| v != "0")
20336 .unwrap_or(true)
20337 });
20338 let hp =
20339 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20340 if hp {
20341 const BLOCK_QH: usize = 32;
20342 let f = self.func("fa_prefill_w_bf16_p1h2");
20343 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20344 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20345 f.set_attribute(
20346 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20347 shmem as i32,
20348 )?;
20349 let cfg = LaunchConfig {
20350 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20351 block_dim: (32, 4, 1),
20352 shared_mem_bytes: shmem,
20353 };
20354 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20355 head_dim as i32,
20356 n_head as i32,
20357 n_head_kv as i32,
20358 t as i32,
20359 t_kv as i32,
20360 causal as i32,
20361 window as i32,
20362 );
20363 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20364 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20365 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
20366 let __s_b = self.gpu.stream();
20367 let mut b = __s_b.launch_builder(&f);
20368 b.arg(&qb)
20369 .arg(&kb)
20370 .arg(&vh)
20371 .arg(o)
20372 .arg(&hd)
20373 .arg(&nh)
20374 .arg(&nhkv)
20375 .arg(&ti)
20376 .arg(&tkvi)
20377 .arg(&scale)
20378 .arg(&cz)
20379 .arg(&wi);
20380 unsafe {
20381 b.launch(cfg)?;
20382 }
20383 return Ok(());
20384 }
20385 if p1 {
20386 let f = self.func("fa_prefill_w_bf16_p1");
20387 let shmem =
20388 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20389 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20390 f.set_attribute(
20391 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20392 shmem as i32,
20393 )?;
20394 let cfg = LaunchConfig {
20395 grid_dim: (
20396 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20397 n_head as u32,
20398 1,
20399 ),
20400 block_dim: (32, 4, 1),
20401 shared_mem_bytes: shmem,
20402 };
20403 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20404 head_dim as i32,
20405 n_head as i32,
20406 n_head_kv as i32,
20407 t as i32,
20408 t_kv as i32,
20409 causal as i32,
20410 window as i32,
20411 );
20412 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20413 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20414 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20415 let __s_b = self.gpu.stream();
20416 let mut b = __s_b.launch_builder(&f);
20417 b.arg(&qb)
20418 .arg(&kb)
20419 .arg(&vb)
20420 .arg(o)
20421 .arg(&hd)
20422 .arg(&nh)
20423 .arg(&nhkv)
20424 .arg(&ti)
20425 .arg(&tkvi)
20426 .arg(&scale)
20427 .arg(&cz)
20428 .arg(&wi);
20429 unsafe {
20430 b.launch(cfg)?;
20431 }
20432 return Ok(());
20433 }
20434 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20437 let g4 = !floor
20438 && !f32_stage
20439 && n_head_kv == 1
20440 && n_head % 4 == 0
20441 && *G4_ON.get_or_init(|| {
20442 std::env::var("MEMRA_FAW_G4")
20443 .map(|v| v != "0")
20444 .unwrap_or(true)
20445 });
20446 if g4 {
20447 const SP_M: usize = 16;
20448 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20451 let o2 = *O2_ON.get_or_init(|| {
20452 std::env::var("MEMRA_FAW_O2")
20453 .map(|v| v != "0")
20454 .unwrap_or(true)
20455 });
20456 let f = self.func(if o2 {
20457 "fa_prefill_w_bf16_g4o2"
20458 } else {
20459 "fa_prefill_w_bf16_g4"
20460 });
20461 let shmem = if o2 {
20462 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
20463 } else {
20464 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
20465 as u32
20466 };
20467 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20468 f.set_attribute(
20469 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20470 shmem as i32,
20471 )?;
20472 let cfg = LaunchConfig {
20473 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
20474 block_dim: (32, 4, 1),
20475 shared_mem_bytes: shmem,
20476 };
20477 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20478 head_dim as i32,
20479 n_head as i32,
20480 n_head_kv as i32,
20481 t as i32,
20482 t_kv as i32,
20483 causal as i32,
20484 window as i32,
20485 );
20486 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20487 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20488 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20489 let __s_b = self.gpu.stream();
20490 let mut b = __s_b.launch_builder(&f);
20491 b.arg(&qb)
20492 .arg(&kb)
20493 .arg(&vb)
20494 .arg(o)
20495 .arg(&hd)
20496 .arg(&nh)
20497 .arg(&nhkv)
20498 .arg(&ti)
20499 .arg(&tkvi)
20500 .arg(&scale)
20501 .arg(&cz)
20502 .arg(&wi);
20503 unsafe {
20504 b.launch(cfg)?;
20505 }
20506 return Ok(());
20507 }
20508 let f = self.func(if floor {
20509 "fa_prefill_w_f32"
20510 } else if f32_stage {
20511 "fa_prefill_w_f32_pp"
20512 } else {
20513 "fa_prefill_w_bf16_pp"
20514 });
20515 let shmem =
20516 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20517 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20518 f.set_attribute(
20519 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20520 shmem as i32,
20521 )?;
20522 let cfg = LaunchConfig {
20523 grid_dim: (
20524 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20525 n_head as u32,
20526 1,
20527 ),
20528 block_dim: (32, 4, 1),
20529 shared_mem_bytes: shmem,
20530 };
20531 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20532 head_dim as i32,
20533 n_head as i32,
20534 n_head_kv as i32,
20535 t as i32,
20536 t_kv as i32,
20537 causal as i32,
20538 window as i32,
20539 );
20540 if f32_stage {
20541 let __s_b = self.gpu.stream();
20542 let mut b = __s_b.launch_builder(&f);
20543 b.arg(q)
20544 .arg(k)
20545 .arg(v)
20546 .arg(o)
20547 .arg(&hd)
20548 .arg(&nh)
20549 .arg(&nhkv)
20550 .arg(&ti)
20551 .arg(&tkvi)
20552 .arg(&scale)
20553 .arg(&cz)
20554 .arg(&wi);
20555 unsafe {
20556 b.launch(cfg)?;
20557 }
20558 } else {
20559 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20560 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20561 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20562 let __s_b = self.gpu.stream();
20563 let mut b = __s_b.launch_builder(&f);
20564 b.arg(&qb)
20565 .arg(&kb)
20566 .arg(&vb)
20567 .arg(o)
20568 .arg(&hd)
20569 .arg(&nh)
20570 .arg(&nhkv)
20571 .arg(&ti)
20572 .arg(&tkvi)
20573 .arg(&scale)
20574 .arg(&cz)
20575 .arg(&wi);
20576 unsafe {
20577 b.launch(cfg)?;
20578 }
20579 }
20580 Ok(())
20581 }
20582
20583 #[allow(clippy::too_many_arguments)]
20587 pub fn fa_prefill_hd512(
20588 &self,
20589 q: &CudaSlice<f32>,
20590 k: &CudaSlice<f32>,
20591 v: &CudaSlice<f32>,
20592 o: &mut CudaSlice<f32>,
20593 head_dim: usize,
20594 n_head: usize,
20595 n_head_kv: usize,
20596 t: usize,
20597 t_kv: usize,
20598 scale: f32,
20599 causal: bool,
20600 ) -> Result<(), Box<dyn std::error::Error>> {
20601 if portable_mma_gated() {
20603 return self.sdpa_naive(
20604 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20605 );
20606 }
20607 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20613 let f32_stage =
20614 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
20615 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20619 let sp = !f32_stage
20620 && *SP_ON.get_or_init(|| {
20621 std::env::var("MEMRA_FA512_SP")
20622 .map(|v| v != "0")
20623 .unwrap_or(true)
20624 });
20625 self.fa_prefill_hd512_arm(
20626 q,
20627 k,
20628 v,
20629 o,
20630 head_dim,
20631 n_head,
20632 n_head_kv,
20633 t,
20634 t_kv,
20635 scale,
20636 causal,
20637 f32_stage,
20638 sp,
20639 sp && fa_f16pv_on(),
20640 )
20641 }
20642
20643 #[allow(clippy::too_many_arguments)]
20645 pub fn fa_prefill_hd512_pre(
20646 &self,
20647 qb: &CudaSlice<u8>,
20648 kb: &CudaSlice<u8>,
20649 vb: &CudaSlice<u8>,
20650 o: &mut CudaSlice<f32>,
20651 head_dim: usize,
20652 n_head: usize,
20653 n_head_kv: usize,
20654 t: usize,
20655 t_kv: usize,
20656 scale: f32,
20657 causal: bool,
20658 v_f16: bool,
20659 ) -> Result<(), Box<dyn std::error::Error>> {
20660 debug_assert_eq!(head_dim, 512);
20661 const SP_M: usize = 16;
20662 const BKS: usize = 32;
20663 let f16pv = fa_f16pv_on();
20667 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20668 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20669 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
20670 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20671 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
20672 let n = t_kv * n_head_kv * head_dim;
20674 let need = n * 2;
20675 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
20676 *vguard = Some(self.alloc_uninit::<u8>(need)?);
20677 }
20678 let dst = vguard.as_mut().unwrap();
20679 self.bf16_to_f16_into(vb, n, dst)?;
20680 vguard.as_ref().unwrap()
20681 } else {
20682 vb
20683 };
20684 let f = self.func(if hp {
20685 "fa_prefill_bf16_hd512_sp16h2"
20686 } else {
20687 match (f16pv, nw) {
20688 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20689 (true, _) => "fa_prefill_bf16_hd512_sp16",
20690 _ => "fa_prefill_bf16_hd512_sp",
20691 }
20692 });
20693 let (nwarp, npart) = if hp {
20694 (4usize, 4usize)
20695 } else if nw > 2 {
20696 (nw, nw)
20697 } else {
20698 (2, 1)
20699 };
20700 let shmem = if hp {
20702 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
20703 as u32
20704 } else {
20705 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20706 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20707 };
20708 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20709 f.set_attribute(
20710 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20711 shmem as i32,
20712 )?;
20713 let grid_y = if hp {
20714 (n_head / 2) as u32
20715 } else {
20716 n_head as u32
20717 };
20718 let cfg = LaunchConfig {
20719 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20720 block_dim: (32, nwarp as u32, 1),
20721 shared_mem_bytes: shmem,
20722 };
20723 let (hd, nh, nhkv, ti, tkvi, cz) = (
20724 head_dim as i32,
20725 n_head as i32,
20726 n_head_kv as i32,
20727 t as i32,
20728 t_kv as i32,
20729 causal as i32,
20730 );
20731 let __s_b = self.gpu.stream();
20732 let mut b = __s_b.launch_builder(&f);
20733 b.arg(qb)
20734 .arg(kb)
20735 .arg(vref)
20736 .arg(o)
20737 .arg(&hd)
20738 .arg(&nh)
20739 .arg(&nhkv)
20740 .arg(&ti)
20741 .arg(&tkvi)
20742 .arg(&scale)
20743 .arg(&cz);
20744 unsafe {
20745 b.launch(cfg)?;
20746 }
20747 Ok(())
20748 }
20749
20750 #[allow(clippy::too_many_arguments)]
20753 pub fn fa_prefill_hd512_arm(
20754 &self,
20755 q: &CudaSlice<f32>,
20756 k: &CudaSlice<f32>,
20757 v: &CudaSlice<f32>,
20758 o: &mut CudaSlice<f32>,
20759 head_dim: usize,
20760 n_head: usize,
20761 n_head_kv: usize,
20762 t: usize,
20763 t_kv: usize,
20764 scale: f32,
20765 causal: bool,
20766 f32_stage: bool,
20767 sp: bool,
20768 f16pv: bool,
20769 ) -> Result<(), Box<dyn std::error::Error>> {
20770 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
20771 if sp && !f32_stage {
20772 const SP_M: usize = 16;
20776 const BKS: usize = 32;
20777 let nw = if f16pv { fa512_wide_warps() } else { 2 };
20778 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20779 let f = self.func(if hp {
20780 "fa_prefill_bf16_hd512_sp16h2"
20781 } else {
20782 match (f16pv, nw) {
20783 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
20784 (true, _) => "fa_prefill_bf16_hd512_sp16",
20785 _ => "fa_prefill_bf16_hd512_sp",
20786 }
20787 });
20788 let (nwarp, npart) = if hp {
20789 (4usize, 4usize)
20790 } else if nw > 2 {
20791 (nw, nw)
20792 } else {
20793 (2, 1)
20794 };
20795 let shmem = if hp {
20796 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
20797 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
20798 } else {
20799 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
20800 + 4 * (npart * SP_M * BKS + SP_M)) as u32
20801 };
20802 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20803 f.set_attribute(
20804 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20805 shmem as i32,
20806 )?;
20807 let grid_y = if hp {
20808 (n_head / 2) as u32
20809 } else {
20810 n_head as u32
20811 };
20812 let cfg = LaunchConfig {
20813 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
20814 block_dim: (32, nwarp as u32, 1),
20815 shared_mem_bytes: shmem,
20816 };
20817 let (hd, nh, nhkv, ti, tkvi, cz) = (
20818 head_dim as i32,
20819 n_head as i32,
20820 n_head_kv as i32,
20821 t as i32,
20822 t_kv as i32,
20823 causal as i32,
20824 );
20825 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20826 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20827 let vb = if f16pv {
20828 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
20829 } else {
20830 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
20831 };
20832 let __s_b = self.gpu.stream();
20833 let mut b = __s_b.launch_builder(&f);
20834 b.arg(&qb)
20835 .arg(&kb)
20836 .arg(&vb)
20837 .arg(o)
20838 .arg(&hd)
20839 .arg(&nh)
20840 .arg(&nhkv)
20841 .arg(&ti)
20842 .arg(&tkvi)
20843 .arg(&scale)
20844 .arg(&cz);
20845 unsafe {
20846 b.launch(cfg)?;
20847 }
20848 return Ok(());
20849 }
20850 const BLOCK_Q: usize = 32;
20851 const BK: usize = 32;
20852 const HALF: usize = 256;
20853 let f = self.func(if f32_stage {
20854 "fa_prefill_f32_hd512"
20855 } else {
20856 "fa_prefill_bf16_hd512"
20857 });
20858 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
20860 + 4 * BLOCK_Q) as u32;
20861 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20862 f.set_attribute(
20863 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20864 shmem as i32,
20865 )?;
20866 let cfg = LaunchConfig {
20867 grid_dim: (
20868 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20869 n_head as u32,
20870 2,
20871 ),
20872 block_dim: (32, 2, 1),
20873 shared_mem_bytes: shmem,
20874 };
20875 let (hd, nh, nhkv, ti, tkvi, cz) = (
20876 head_dim as i32,
20877 n_head as i32,
20878 n_head_kv as i32,
20879 t as i32,
20880 t_kv as i32,
20881 causal as i32,
20882 );
20883 if f32_stage {
20884 let __s_b = self.gpu.stream();
20885 let mut b = __s_b.launch_builder(&f);
20886 b.arg(q)
20887 .arg(k)
20888 .arg(v)
20889 .arg(o)
20890 .arg(&hd)
20891 .arg(&nh)
20892 .arg(&nhkv)
20893 .arg(&ti)
20894 .arg(&tkvi)
20895 .arg(&scale)
20896 .arg(&cz);
20897 unsafe {
20898 b.launch(cfg)?;
20899 }
20900 } else {
20901 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20902 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20903 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20904 let __s_b = self.gpu.stream();
20905 let mut b = __s_b.launch_builder(&f);
20906 b.arg(&qb)
20907 .arg(&kb)
20908 .arg(&vb)
20909 .arg(o)
20910 .arg(&hd)
20911 .arg(&nh)
20912 .arg(&nhkv)
20913 .arg(&ti)
20914 .arg(&tkvi)
20915 .arg(&scale)
20916 .arg(&cz);
20917 unsafe {
20918 b.launch(cfg)?;
20919 }
20920 }
20921 Ok(())
20922 }
20923
20924 #[allow(clippy::too_many_arguments)]
20928 pub fn rope_neox2_bf16e(
20929 &self,
20930 q: &mut CudaSlice<f32>,
20931 k: &mut CudaSlice<f32>,
20932 qb: &mut CudaSlice<u8>,
20933 kb: &mut CudaSlice<u8>,
20934 pos: &CudaSlice<i32>,
20935 head_dim: usize,
20936 n_dims: usize,
20937 nh_q: usize,
20938 nh_k: usize,
20939 n_tokens: usize,
20940 base: f32,
20941 freq_scale: f32,
20942 ff: Option<&CudaSlice<f32>>,
20943 ) -> Result<(), Box<dyn std::error::Error>> {
20944 let f = self.func("rope_neox2_bf16e_f32");
20945 let rows = ((nh_q + nh_k) * n_tokens) as u32;
20946 let cfg = LaunchConfig {
20947 grid_dim: (rows, 1, 1),
20948 block_dim: ((head_dim / 2) as u32, 1, 1),
20949 shared_mem_bytes: 0,
20950 };
20951 let theta_scale = base.powf(-2.0 / n_dims as f32);
20952 let (hd, nd, nhq, nhk, nt) = (
20953 head_dim as i32,
20954 n_dims as i32,
20955 nh_q as i32,
20956 nh_k as i32,
20957 n_tokens as i32,
20958 );
20959 let __s_b = self.gpu.stream();
20960 let mut b = __s_b.launch_builder(&f);
20961 match ff {
20962 Some(t) => {
20963 b.arg(&mut *q)
20964 .arg(&mut *k)
20965 .arg(&mut *qb)
20966 .arg(&mut *kb)
20967 .arg(pos)
20968 .arg(&hd)
20969 .arg(&nd)
20970 .arg(&nhq)
20971 .arg(&nhk)
20972 .arg(&nt)
20973 .arg(&theta_scale)
20974 .arg(&freq_scale)
20975 .arg(t);
20976 unsafe {
20977 b.launch(cfg)?;
20978 }
20979 }
20980 None => {
20981 let null: u64 = 0;
20982 b.arg(&mut *q)
20983 .arg(&mut *k)
20984 .arg(&mut *qb)
20985 .arg(&mut *kb)
20986 .arg(pos)
20987 .arg(&hd)
20988 .arg(&nd)
20989 .arg(&nhq)
20990 .arg(&nhk)
20991 .arg(&nt)
20992 .arg(&theta_scale)
20993 .arg(&freq_scale)
20994 .arg(&null);
20995 unsafe {
20996 b.launch(cfg)?;
20997 }
20998 }
20999 }
21000 Ok(())
21001 }
21002
21003 pub fn f32_to_bf16(
21006 &self,
21007 x: &CudaSlice<f32>,
21008 n: usize,
21009 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21010 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
21011 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21012 let f = self.func("f32_to_bf16_flat");
21013 let n_i = n as i64;
21014 let cfg = LaunchConfig {
21015 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21016 block_dim: (256, 1, 1),
21017 shared_mem_bytes: 0,
21018 };
21019 let __s_b = self.gpu.stream();
21020 let mut b = __s_b.launch_builder(&f);
21021 b.arg(x).arg(&mut y).arg(&n_i);
21022 unsafe {
21023 b.launch(cfg)?;
21024 }
21025 Ok(y)
21026 }
21027
21028 pub fn f32_to_f16(
21029 &self,
21030 x: &CudaSlice<f32>,
21031 n: usize,
21032 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21033 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
21034 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21035 let f = self.func("f32_to_f16_flat");
21036 let n_i = n as i64;
21037 let cfg = LaunchConfig {
21038 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21039 block_dim: (256, 1, 1),
21040 shared_mem_bytes: 0,
21041 };
21042 let __s_b = self.gpu.stream();
21043 let mut b = __s_b.launch_builder(&f);
21044 b.arg(x).arg(&mut y).arg(&n_i);
21045 unsafe {
21046 b.launch(cfg)?;
21047 }
21048 Ok(y)
21049 }
21050
21051 pub fn bf16_to_f16(
21053 &self,
21054 xb: &CudaSlice<u8>,
21055 n: usize,
21056 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21057 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21058 self.bf16_to_f16_into(xb, n, &mut y)?;
21059 Ok(y)
21060 }
21061
21062 pub fn bf16_to_f16_into(
21064 &self,
21065 xb: &CudaSlice<u8>,
21066 n: usize,
21067 y: &mut CudaSlice<u8>,
21068 ) -> Result<(), Box<dyn std::error::Error>> {
21069 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21070 assert!(y.len() >= n * 2);
21071 let f = self.func("bf16_to_f16_flat");
21072 let n2 = (n / 2) as i64;
21073 let cfg = LaunchConfig {
21074 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21075 block_dim: (256, 1, 1),
21076 shared_mem_bytes: 0,
21077 };
21078 let __s_b = self.gpu.stream();
21079 let mut b = __s_b.launch_builder(&f);
21080 b.arg(xb).arg(y).arg(&n2);
21081 unsafe {
21082 b.launch(cfg)?;
21083 }
21084 Ok(())
21085 }
21086
21087 #[allow(clippy::too_many_arguments)]
21092 pub fn fa_prefill_vl8(
21093 &self,
21094 seqs: &[FaSeqVl],
21095 head_dim: usize,
21096 n_head: usize,
21097 n_head_kv: usize,
21098 scale: f32,
21099 ) -> Result<(), Box<dyn std::error::Error>> {
21100 const BK: usize = 32;
21101 let b = seqs.len();
21102 assert!(b >= 1 && b <= 8);
21103 let mut packed = [FaSeqVl::default(); 8];
21104 packed[..b].copy_from_slice(seqs);
21105 let v = FaVl8(packed);
21106 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21107 let ept = (n_head_kv * head_dim) as i32;
21108 {
21109 let f = self.func("fa_mirror_vl");
21110 let max_n = (max_t as i64) * ept as i64;
21111 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21112 for which in 0..2i32 {
21113 let cfg = LaunchConfig {
21114 grid_dim: (blocks, 1, b as u32),
21115 block_dim: (256, 1, 1),
21116 shared_mem_bytes: 0,
21117 };
21118 let __s_lb = self.gpu.stream();
21119 let mut lb = __s_lb.launch_builder(&f);
21120 lb.arg(&v).arg(&ept).arg(&which);
21121 unsafe {
21122 lb.launch(cfg)?;
21123 }
21124 }
21125 }
21126 let hd_sfx = fa_hd_suffix(head_dim)?;
21127 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21128 let block_q = 64usize;
21129 let kv_stages = 2usize;
21130 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21131 + 4 * (block_q * BK + 2 * block_q)) as u32;
21132 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21133 f.set_attribute(
21134 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21135 shmem as i32,
21136 )?;
21137 let cfg = LaunchConfig {
21138 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21139 block_dim: (32, 4, 1),
21140 shared_mem_bytes: shmem,
21141 };
21142 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21143 let __s_lb = self.gpu.stream();
21144 let mut lb = __s_lb.launch_builder(&f);
21145 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21146 unsafe {
21147 lb.launch(cfg)?;
21148 }
21149 Ok(())
21150 }
21151
21152 #[allow(clippy::too_many_arguments)]
21156 pub fn attn_pre_vl8(
21157 &self,
21158 seqs: &[AttnPreVl],
21159 wq: &CudaSlice<f32>,
21160 wk: &CudaSlice<f32>,
21161 head_dim: usize,
21162 rope_dims: usize,
21163 n_head: usize,
21164 n_head_kv: usize,
21165 eps: f32,
21166 freq_base: f32,
21167 freq_scale: f32,
21168 kv_dim_k: usize,
21169 kv_dim_v: usize,
21170 k_tok_bytes: usize,
21171 v_tok_bytes: usize,
21172 ) -> Result<(), Box<dyn std::error::Error>> {
21173 let b = seqs.len();
21174 assert!(b >= 1 && b <= 8);
21175 let mut packed = [AttnPreVl::default(); 8];
21176 packed[..b].copy_from_slice(seqs);
21177 let v = AttnPreVl8(packed);
21178 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21179 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21180 {
21181 let f = self.func("q_gate_split_vl");
21182 let n = max_t * (n_head * head_dim) as u32;
21183 let cfg = LaunchConfig {
21184 grid_dim: (n.div_ceil(256), 1, b as u32),
21185 block_dim: (256, 1, 1),
21186 shared_mem_bytes: 0,
21187 };
21188 let __s_lb = self.gpu.stream();
21189 let mut lb = __s_lb.launch_builder(&f);
21190 lb.arg(&v).arg(&hd).arg(&nh);
21191 unsafe {
21192 lb.launch(cfg)?;
21193 }
21194 }
21195 {
21196 let f = self.func("attn_rms_vl");
21197 let cfg = LaunchConfig {
21198 grid_dim: (max_t * n_head as u32, 2, b as u32),
21199 block_dim: (rms_block(), 1, 1),
21200 shared_mem_bytes: 0,
21201 };
21202 let __s_lb = self.gpu.stream();
21203 let mut lb = __s_lb.launch_builder(&f);
21204 lb.arg(&v)
21205 .arg(wq)
21206 .arg(wk)
21207 .arg(&hd)
21208 .arg(&nh)
21209 .arg(&nhkv)
21210 .arg(&eps);
21211 unsafe {
21212 lb.launch(cfg)?;
21213 }
21214 }
21215 {
21216 let f = self.func("attn_rope_vl");
21217 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21218 let nd = rope_dims as i32;
21219 let cfg = LaunchConfig {
21220 grid_dim: (max_t * n_head as u32, 2, b as u32),
21221 block_dim: ((head_dim / 2) as u32, 1, 1),
21222 shared_mem_bytes: 0,
21223 };
21224 let __s_lb = self.gpu.stream();
21225 let mut lb = __s_lb.launch_builder(&f);
21226 lb.arg(&v)
21227 .arg(&hd)
21228 .arg(&nd)
21229 .arg(&nh)
21230 .arg(&nhkv)
21231 .arg(&theta_scale)
21232 .arg(&freq_scale);
21233 unsafe {
21234 lb.launch(cfg)?;
21235 }
21236 }
21237 {
21238 let f = self.func("append_kv_vl");
21239 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21240 let cfg = LaunchConfig {
21241 grid_dim: (nblk, max_t, b as u32),
21242 block_dim: (32, 1, 1),
21243 shared_mem_bytes: 0,
21244 };
21245 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21246 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21247 let __s_lb = self.gpu.stream();
21248 let mut lb = __s_lb.launch_builder(&f);
21249 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21250 unsafe {
21251 lb.launch(cfg)?;
21252 }
21253 }
21254 Ok(())
21255 }
21256
21257 pub fn fa_prefill_view(
21262 &self,
21263 q: &CudaSlice<f32>,
21264 k: &cudarc::driver::CudaView<u8>,
21265 v: &cudarc::driver::CudaView<u8>,
21266 o: &mut CudaSlice<f32>,
21267 head_dim: usize,
21268 n_head: usize,
21269 n_head_kv: usize,
21270 t: usize,
21271 t_kv: usize,
21272 scale: f32,
21273 causal: bool,
21274 k_tok_bytes: usize,
21275 v_tok_bytes: usize,
21276 g: bool,
21277 ) -> Result<(), Box<dyn std::error::Error>> {
21278 if portable_mma_gated() {
21279 return self.sdpa_naive_quantized_view(
21280 q,
21281 k,
21282 v,
21283 o,
21284 head_dim,
21285 n_head,
21286 n_head_kv,
21287 t,
21288 t_kv,
21289 scale,
21290 causal,
21291 k_tok_bytes,
21292 v_tok_bytes,
21293 );
21294 }
21295 const BLOCK_Q: usize = 64;
21296 const BK: usize = 32;
21297 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21300 let f = if g {
21301 self.func_g(&name)
21302 } else {
21303 self.func(&name)
21304 };
21305 let shmem =
21306 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21307 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21308 f.set_attribute(
21309 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21310 shmem as i32,
21311 )?;
21312 let cfg = LaunchConfig {
21313 grid_dim: (
21314 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21315 n_head as u32,
21316 1,
21317 ),
21318 block_dim: (32, 4, 1),
21319 shared_mem_bytes: shmem,
21320 };
21321 let (hd, nh, nhkv, ti, tkvi, cz) = (
21322 head_dim as i32,
21323 n_head as i32,
21324 n_head_kv as i32,
21325 t as i32,
21326 t_kv as i32,
21327 causal as i32,
21328 );
21329 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21330 let __s_b = self.gpu.stream();
21331 let mut b = __s_b.launch_builder(&f);
21332 b.arg(q)
21333 .arg(k)
21334 .arg(v)
21335 .arg(o)
21336 .arg(&hd)
21337 .arg(&nh)
21338 .arg(&nhkv)
21339 .arg(&ti)
21340 .arg(&tkvi)
21341 .arg(&scale)
21342 .arg(&cz)
21343 .arg(&ktb)
21344 .arg(&vtb);
21345 unsafe {
21346 b.launch(cfg)?;
21347 }
21348 Ok(())
21349 }
21350
21351 #[allow(clippy::too_many_arguments)]
21361 pub fn fa_prefill_view_ws(
21362 &self,
21363 q: &CudaSlice<f32>,
21364 k: &cudarc::driver::CudaView<u8>,
21365 v: &cudarc::driver::CudaView<u8>,
21366 o: &mut CudaSlice<f32>,
21367 head_dim: usize,
21368 n_head: usize,
21369 n_head_kv: usize,
21370 t: usize,
21371 t_kv: usize,
21372 scale: f32,
21373 causal: bool,
21374 k_tok_bytes: usize,
21375 v_tok_bytes: usize,
21376 g: bool,
21377 ) -> Result<(), Box<dyn std::error::Error>> {
21378 if portable_mma_gated() {
21379 return self.sdpa_naive_quantized_view(
21380 q,
21381 k,
21382 v,
21383 o,
21384 head_dim,
21385 n_head,
21386 n_head_kv,
21387 t,
21388 t_kv,
21389 scale,
21390 causal,
21391 k_tok_bytes,
21392 v_tok_bytes,
21393 );
21394 }
21395 const BLOCK_Q: usize = 64;
21396 const BK: usize = 32;
21397 let kv_dim_k = n_head_kv * head_dim;
21398 let kv_dim_v = n_head_kv * head_dim;
21399 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21401 let mut guard = self.prime_deqw_ws.lock().unwrap();
21403 let need_grow = match guard.as_ref() {
21404 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21405 None => true,
21406 };
21407 if need_grow {
21408 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21409 let (ck, cv) = guard
21410 .as_ref()
21411 .map(|(a, b)| (a.len(), b.len()))
21412 .unwrap_or((0, 0));
21413 *guard = Some((
21414 self.alloc_u8(grow(ck, k_ws_bytes))?,
21415 self.alloc_u8(grow(cv, v_ws_bytes))?,
21416 ));
21417 }
21418 let (kw, vw) = guard.as_mut().unwrap();
21419 {
21421 let f = if g {
21423 self.func_g("fa_dequant_kv_ws_bf16")
21424 } else {
21425 self.func("fa_dequant_kv_ws_bf16")
21426 };
21427 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21428 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21429 let cfg = LaunchConfig {
21430 grid_dim: (nblk.max(1), 1, 1),
21431 block_dim: (256, 1, 1),
21432 shared_mem_bytes: 0,
21433 };
21434 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21435 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21436 let __s_b = self.gpu.stream();
21437 let mut b = __s_b.launch_builder(&f);
21438 b.arg(k)
21439 .arg(v)
21440 .arg(&mut *kw)
21441 .arg(&mut *vw)
21442 .arg(&kdk)
21443 .arg(&kdv)
21444 .arg(&tkvi)
21445 .arg(&ktb)
21446 .arg(&vtb);
21447 unsafe {
21448 b.launch(cfg)?;
21449 }
21450 }
21451 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21459 .map(|v| v != "0")
21460 .unwrap_or(true);
21461 {
21462 let hd_sfx = fa_hd_suffix(head_dim)?;
21463 let f = self.func(&format!(
21464 "fa_prefill_qw{}{hd_sfx}",
21465 if db { "_db" } else { "" }
21466 ));
21467 let shmem = if db {
21468 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21470 } else {
21471 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21472 };
21473 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21474 f.set_attribute(
21475 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21476 shmem as i32,
21477 )?;
21478 let cfg = LaunchConfig {
21479 grid_dim: (
21480 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21481 n_head as u32,
21482 1,
21483 ),
21484 block_dim: (32, 4, 1),
21485 shared_mem_bytes: shmem,
21486 };
21487 let (hd, nh, nhkv, ti, tkvi, cz) = (
21488 head_dim as i32,
21489 n_head as i32,
21490 n_head_kv as i32,
21491 t as i32,
21492 t_kv as i32,
21493 causal as i32,
21494 );
21495 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21496 let __s_b = self.gpu.stream();
21497 let mut b = __s_b.launch_builder(&f);
21498 b.arg(q)
21499 .arg(&*kw)
21500 .arg(&*vw)
21501 .arg(o)
21502 .arg(&hd)
21503 .arg(&nh)
21504 .arg(&nhkv)
21505 .arg(&ti)
21506 .arg(&tkvi)
21507 .arg(&scale)
21508 .arg(&cz)
21509 .arg(&kdk)
21510 .arg(&kdv);
21511 unsafe {
21512 b.launch(cfg)?;
21513 }
21514 }
21515 Ok(())
21516 }
21517
21518 #[allow(clippy::too_many_arguments)]
21534 pub fn fa_prefill_view_ws_w_hd128(
21535 &self,
21536 q: &CudaSlice<f32>,
21537 k: &cudarc::driver::CudaView<u8>,
21538 v: &cudarc::driver::CudaView<u8>,
21539 o: &mut CudaSlice<f32>,
21540 head_dim: usize,
21541 n_head: usize,
21542 n_head_kv: usize,
21543 t: usize,
21544 t_kv: usize,
21545 scale: f32,
21546 causal: bool,
21547 window: usize,
21548 k_tok_bytes: usize,
21549 v_tok_bytes: usize,
21550 ) -> Result<(), Box<dyn std::error::Error>> {
21551 assert_eq!(
21552 head_dim, 128,
21553 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
21554 );
21555 if portable_mma_gated() {
21556 return self.sdpa_naive_w_quantized_view(
21557 q,
21558 k,
21559 v,
21560 o,
21561 head_dim,
21562 n_head,
21563 n_head_kv,
21564 t,
21565 t_kv,
21566 scale,
21567 causal,
21568 window,
21569 k_tok_bytes,
21570 v_tok_bytes,
21571 );
21572 }
21573 const BLOCK_Q: usize = 64;
21574 const BK: usize = 32;
21575 let kv_dim_k = n_head_kv * head_dim;
21576 let kv_dim_v = n_head_kv * head_dim;
21577 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
21579 let mut guard = self.prime_deqw_ws.lock().unwrap();
21580 let need_grow = match guard.as_ref() {
21581 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
21582 None => true,
21583 };
21584 if need_grow {
21585 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
21586 let (ck, cv) = guard
21587 .as_ref()
21588 .map(|(a, b)| (a.len(), b.len()))
21589 .unwrap_or((0, 0));
21590 *guard = Some((
21591 self.alloc_u8(grow(ck, k_ws_bytes))?,
21592 self.alloc_u8(grow(cv, v_ws_bytes))?,
21593 ));
21594 }
21595 let (kw, vw) = guard.as_mut().unwrap();
21596 {
21599 let f = self.func("fa_dequant_kv_ws_bf16");
21600 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
21601 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
21602 let cfg = LaunchConfig {
21603 grid_dim: (nblk.max(1), 1, 1),
21604 block_dim: (256, 1, 1),
21605 shared_mem_bytes: 0,
21606 };
21607 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
21608 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21609 let __s_b = self.gpu.stream();
21610 let mut b = __s_b.launch_builder(&f);
21611 b.arg(k)
21612 .arg(v)
21613 .arg(&mut *kw)
21614 .arg(&mut *vw)
21615 .arg(&kdk)
21616 .arg(&kdv)
21617 .arg(&tkvi)
21618 .arg(&ktb)
21619 .arg(&vtb);
21620 unsafe {
21621 b.launch(cfg)?;
21622 }
21623 }
21624 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
21626 .map(|v| v != "0")
21627 .unwrap_or(true);
21628 {
21629 let f = self.func(if db {
21630 "fa_prefill_qw_db_w_hd128"
21631 } else {
21632 "fa_prefill_qw_w_hd128"
21633 });
21634 let shmem = if db {
21635 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
21636 } else {
21637 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
21638 };
21639 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21640 f.set_attribute(
21641 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21642 shmem as i32,
21643 )?;
21644 let cfg = LaunchConfig {
21645 grid_dim: (
21646 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21647 n_head as u32,
21648 1,
21649 ),
21650 block_dim: (32, 4, 1),
21651 shared_mem_bytes: shmem,
21652 };
21653 let (hd, nh, nhkv, ti, tkvi, cz) = (
21654 head_dim as i32,
21655 n_head as i32,
21656 n_head_kv as i32,
21657 t as i32,
21658 t_kv as i32,
21659 causal as i32,
21660 );
21661 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
21662 let __s_b = self.gpu.stream();
21663 let mut b = __s_b.launch_builder(&f);
21664 b.arg(q)
21665 .arg(&*kw)
21666 .arg(&*vw)
21667 .arg(o)
21668 .arg(&hd)
21669 .arg(&nh)
21670 .arg(&nhkv)
21671 .arg(&ti)
21672 .arg(&tkvi)
21673 .arg(&scale)
21674 .arg(&cz)
21675 .arg(&kdk)
21676 .arg(&kdv)
21677 .arg(&wnd);
21678 unsafe {
21679 b.launch(cfg)?;
21680 }
21681 }
21682 Ok(())
21683 }
21684
21685 pub fn fa_decode(
21689 &self,
21690 q: &CudaSlice<f32>,
21691 k: &cudarc::driver::CudaView<u8>,
21692 v: &cudarc::driver::CudaView<u8>,
21693 o: &mut CudaSlice<f32>,
21694 head_dim: usize,
21695 n_head: usize,
21696 n_head_kv: usize,
21697 t_kv: usize,
21698 scale: f32,
21699 k_tok_bytes: usize,
21700 v_tok_bytes: usize,
21701 ) -> Result<(), Box<dyn std::error::Error>> {
21702 self.fa_decode_kvmod(
21703 q,
21704 k,
21705 v,
21706 o,
21707 head_dim,
21708 n_head,
21709 n_head_kv,
21710 t_kv,
21711 scale,
21712 k_tok_bytes,
21713 v_tok_bytes,
21714 false,
21715 )
21716 }
21717
21718 #[allow(clippy::too_many_arguments)]
21722 #[allow(clippy::too_many_arguments)]
21726 #[allow(clippy::too_many_arguments)]
21727 fn fa_decode_scalar_unified(
21728 &self,
21729 q: &cudarc::driver::CudaView<f32>,
21730 k: &cudarc::driver::CudaView<u8>,
21731 v: &cudarc::driver::CudaView<u8>,
21732 o: &mut cudarc::driver::CudaViewMut<f32>,
21733 head_dim: usize,
21734 n_head: usize,
21735 n_head_kv: usize,
21736 t_kv_host: usize,
21737 t_kv_dev: Option<&CudaSlice<i32>>,
21738 scale: f32,
21739 n_splits: usize,
21740 split_keys: usize,
21741 k_tok_bytes: usize,
21742 v_tok_bytes: usize,
21743 g: bool,
21744 part_o: &mut CudaSlice<f32>,
21745 part_m: &mut CudaSlice<f32>,
21746 part_l: &mut CudaSlice<f32>,
21747 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21748 ) -> Result<(), Box<dyn std::error::Error>> {
21749 let f = if g {
21750 self.func_g("fa_decode_f32")
21751 } else {
21752 self.fa_func("fa_decode_f32", head_dim)
21753 };
21754 let cfg = LaunchConfig {
21755 grid_dim: (n_head as u32, n_splits as u32, 1),
21756 block_dim: (head_dim as u32, 1, 1),
21757 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
21758 };
21759 let (hd, nh, nhkv, nsp) = (
21760 head_dim as i32,
21761 n_head as i32,
21762 n_head_kv as i32,
21763 n_splits as i32,
21764 );
21765 let (ktb, vtb, tkvi, ski) = (
21766 k_tok_bytes as i64,
21767 v_tok_bytes as i64,
21768 t_kv_host as i32,
21769 split_keys as i32,
21770 );
21771 let __s_b = self.gpu.stream();
21772 let mut b = __s_b.launch_builder(&f);
21773 match t_kv_dev {
21774 Some(d) => {
21775 b.arg(q)
21776 .arg(k)
21777 .arg(v)
21778 .arg(&mut *part_o)
21779 .arg(&mut *part_m)
21780 .arg(&mut *part_l)
21781 .arg(&hd)
21782 .arg(&nh)
21783 .arg(&nhkv)
21784 .arg(&tkvi)
21785 .arg(d)
21786 .arg(&scale)
21787 .arg(&nsp)
21788 .arg(&ski)
21789 .arg(&ktb)
21790 .arg(&vtb);
21791 unsafe {
21792 b.launch(cfg)?;
21793 }
21794 }
21795 None => {
21796 let null: u64 = 0;
21797 b.arg(q)
21798 .arg(k)
21799 .arg(v)
21800 .arg(&mut *part_o)
21801 .arg(&mut *part_m)
21802 .arg(&mut *part_l)
21803 .arg(&hd)
21804 .arg(&nh)
21805 .arg(&nhkv)
21806 .arg(&tkvi)
21807 .arg(&null)
21808 .arg(&scale)
21809 .arg(&nsp)
21810 .arg(&ski)
21811 .arg(&ktb)
21812 .arg(&vtb);
21813 unsafe {
21814 b.launch(cfg)?;
21815 }
21816 }
21817 }
21818 let cfg2 = LaunchConfig {
21819 grid_dim: (n_head as u32, 1, 1),
21820 block_dim: (head_dim as u32, 1, 1),
21821 shared_mem_bytes: 0,
21822 };
21823 if let Some((oq, od)) = q8_out {
21824 let fc = if g {
21826 self.func_g("fa_decode_combine_q8_1")
21827 } else {
21828 self.fa_func("fa_decode_combine_q8_1", head_dim)
21829 };
21830 let __s_b2 = self.gpu.stream();
21831 let mut b2 = __s_b2.launch_builder(&fc);
21832 b2.arg(&*part_o)
21833 .arg(&*part_m)
21834 .arg(&*part_l)
21835 .arg(oq)
21836 .arg(od)
21837 .arg(&hd)
21838 .arg(&nh)
21839 .arg(&nsp);
21840 unsafe {
21841 b2.launch(cfg2)?;
21842 }
21843 return Ok(());
21844 }
21845 let fc = if g {
21846 self.func_g("fa_decode_combine_f32")
21847 } else {
21848 self.fa_func("fa_decode_combine_f32", head_dim)
21849 };
21850 let __s_b2 = self.gpu.stream();
21851 let mut b2 = __s_b2.launch_builder(&fc);
21852 b2.arg(&*part_o)
21853 .arg(&*part_m)
21854 .arg(&*part_l)
21855 .arg(o)
21856 .arg(&hd)
21857 .arg(&nh)
21858 .arg(&nsp);
21859 unsafe {
21860 b2.launch(cfg2)?;
21861 }
21862 Ok(())
21863 }
21864
21865 pub fn fa_decode_kvmod(
21866 &self,
21867 q: &CudaSlice<f32>,
21868 k: &cudarc::driver::CudaView<u8>,
21869 v: &cudarc::driver::CudaView<u8>,
21870 o: &mut CudaSlice<f32>,
21871 head_dim: usize,
21872 n_head: usize,
21873 n_head_kv: usize,
21874 t_kv: usize,
21875 scale: f32,
21876 k_tok_bytes: usize,
21877 v_tok_bytes: usize,
21878 g: bool,
21879 ) -> Result<(), Box<dyn std::error::Error>> {
21880 let q_view = q.as_view();
21881 let mut o_view = o.as_view_mut();
21882 self.fa_decode_kvmod_view(
21883 &q_view,
21884 k,
21885 v,
21886 &mut o_view,
21887 head_dim,
21888 n_head,
21889 n_head_kv,
21890 t_kv,
21891 scale,
21892 k_tok_bytes,
21893 v_tok_bytes,
21894 g,
21895 )
21896 }
21897
21898 #[allow(clippy::too_many_arguments)]
21903 pub fn fa_decode_kvmod_view(
21904 &self,
21905 q: &cudarc::driver::CudaView<f32>,
21906 k: &cudarc::driver::CudaView<u8>,
21907 v: &cudarc::driver::CudaView<u8>,
21908 o: &mut cudarc::driver::CudaViewMut<f32>,
21909 head_dim: usize,
21910 n_head: usize,
21911 n_head_kv: usize,
21912 t_kv: usize,
21913 scale: f32,
21914 k_tok_bytes: usize,
21915 v_tok_bytes: usize,
21916 g: bool,
21917 ) -> Result<(), Box<dyn std::error::Error>> {
21918 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
21939 if g && head_dim == 256 && !fa_v4_at(t_kv) {
21943 fa_vec = false;
21944 }
21945 let sp = fa_split_keys(t_kv, n_head_kv);
21946 let n_splits = if fa_vec {
21947 ((t_kv + sp - 1) / sp).max(1)
21948 } else {
21949 ((t_kv + 255) / 256).max(1)
21950 };
21951 let o_len = n_head * n_splits * head_dim;
21952 let ml_len = n_head * n_splits;
21953 let mut part_guard = self.fa_part_pool.lock().unwrap();
21954 if part_guard
21955 .as_ref()
21956 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21957 .unwrap_or(true)
21958 {
21959 let old = part_guard.take();
21970 let (co, cm) = old
21971 .as_ref()
21972 .map(|pp| (pp.0.len(), pp.1.len()))
21973 .unwrap_or((0, 0));
21974 if let Some(old) = old {
21975 self.fa_part_retired.lock().unwrap().push(old);
21976 }
21977 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21978 eprintln!(
21979 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21980 co, o_len, cm, ml_len
21981 );
21982 }
21983 *part_guard = Some((
21984 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21985 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21986 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21987 ));
21988 }
21989 let pg = part_guard.as_mut().unwrap();
21990 self.gpu
21991 .stream()
21992 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21993 self.gpu
21994 .stream()
21995 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21996 self.gpu
21997 .stream()
21998 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21999 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22000 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22001 let (hd, nh, nhkv, tkvi, nsp) = (
22002 head_dim as i32,
22003 n_head as i32,
22004 n_head_kv as i32,
22005 t_kv as i32,
22006 n_splits as i32,
22007 );
22008 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22009 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22013 let fa512_min = fa512_min_tkv();
22018 let deep = fa_vec
22021 && head_dim == 256
22022 && fa_v4_at(t_kv)
22023 && !g
22024 && fa_deep_at(t_kv)
22025 && !matches!(fa_v4_mode(), "noB3" | "stage");
22026 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
22027 let gqa = (n_head / n_head_kv).max(1) as u32;
22030 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
22031 (
22032 fv,
22033 LaunchConfig {
22034 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22035 block_dim: (32, gqa, 1),
22036 shared_mem_bytes: 0,
22037 },
22038 )
22039 } else if fa_vec && head_dim <= 256 {
22040 let gqa = (n_head / n_head_kv).max(1) as u32;
22041 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22052 let smem_tkv = *SMEM_TKV.get_or_init(|| {
22053 std::env::var("MEMRA_FA_SMEM_TKV")
22054 .ok()
22055 .and_then(|v| v.parse().ok())
22056 .unwrap_or_else(|| {
22057 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22058 })
22059 });
22060 if fa_v4_at(t_kv) && head_dim == 256 {
22061 let v4name = match fa_v4_mode() {
22065 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
22068 _ => "fa_decode_vec_q_v4",
22069 };
22070 let fv = if g {
22071 self.func_g(v4name)
22072 } else {
22073 self.func(v4name)
22074 };
22075 let shmem = (if deep { 12160 } else { 11520 }
22078 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22079 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22080 fv.set_attribute(
22081 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22082 shmem as i32,
22083 )?;
22084 (
22085 fv,
22086 LaunchConfig {
22087 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22088 block_dim: (32, gqa, 1),
22089 shared_mem_bytes: shmem,
22090 },
22091 )
22092 } else if fa_v3_active(head_dim) {
22093 let fv = if g {
22096 self.func_g("fa_decode_vec_q_v3")
22097 } else {
22098 self.func("fa_decode_vec_q_v3")
22099 };
22100 let shmem = (32 * head_dim * 2) as u32; (
22102 fv,
22103 LaunchConfig {
22104 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22105 block_dim: (32, gqa, 1),
22106 shared_mem_bytes: shmem,
22107 },
22108 )
22109 } else if fa_v2_on() {
22110 let fv = if g {
22114 self.func_g("fa_decode_vec_q_v2")
22115 } else {
22116 self.func("fa_decode_vec_q_v2")
22117 };
22118 let shmem = (2 * 32 * head_dim * 2) as u32; (
22120 fv,
22121 LaunchConfig {
22122 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22123 block_dim: (32, gqa, 1),
22124 shared_mem_bytes: shmem,
22125 },
22126 )
22127 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22128 {
22129 let fv = if g {
22133 self.func_g("fa_decode_vec_q_smem")
22134 } else {
22135 self.func("fa_decode_vec_q_smem")
22136 };
22137 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
22139 fv.set_attribute(
22140 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22141 shmem as i32,
22142 )?;
22143 (
22144 fv,
22145 LaunchConfig {
22146 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22147 block_dim: (32, gqa, 1),
22148 shared_mem_bytes: shmem,
22149 },
22150 )
22151 } else {
22152 let fv = if g {
22155 self.func_g("fa_decode_vec_q")
22156 } else {
22157 self.func("fa_decode_vec_q")
22158 };
22159 (
22160 fv,
22161 LaunchConfig {
22162 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22163 block_dim: (32, gqa, 1),
22164 shared_mem_bytes: 0,
22165 },
22166 )
22167 }
22168 } else {
22169 return self.fa_decode_scalar_unified(
22172 q,
22173 k,
22174 v,
22175 o,
22176 head_dim,
22177 n_head,
22178 n_head_kv,
22179 t_kv,
22180 None,
22181 scale,
22182 n_splits,
22183 if fa_vec { sp } else { 256 },
22184 k_tok_bytes,
22185 v_tok_bytes,
22186 g,
22187 part_o,
22188 part_m,
22189 part_l,
22190 None,
22191 );
22192 };
22193 let __s_b = self.gpu.stream();
22194 let mut b = __s_b.launch_builder(&f);
22195 b.arg(q)
22196 .arg(k)
22197 .arg(v)
22198 .arg(&mut *part_o)
22199 .arg(&mut *part_m)
22200 .arg(&mut *part_l)
22201 .arg(&hd)
22202 .arg(&nh)
22203 .arg(&nhkv)
22204 .arg(&tkvi)
22205 .arg(&scale)
22206 .arg(&nsp)
22207 .arg(&ktb)
22208 .arg(&vtb);
22209 unsafe {
22210 b.launch(cfg)?;
22211 }
22212 let (fc, cfg2) = (
22215 if g {
22216 self.func_g("fa_decode_combine_f32")
22217 } else {
22218 self.fa_func("fa_decode_combine_f32", head_dim)
22219 },
22220 LaunchConfig {
22221 grid_dim: (n_head as u32, 1, 1),
22222 block_dim: (head_dim as u32, 1, 1),
22223 shared_mem_bytes: 0,
22224 },
22225 );
22226 let __s_b2 = self.gpu.stream();
22227 let mut b2 = __s_b2.launch_builder(&fc);
22228 b2.arg(&*part_o)
22229 .arg(&*part_m)
22230 .arg(&*part_l)
22231 .arg(o)
22232 .arg(&hd)
22233 .arg(&nh)
22234 .arg(&nsp);
22235 unsafe {
22236 b2.launch(cfg2)?;
22237 }
22238 Ok(())
22239 }
22240
22241 #[allow(clippy::too_many_arguments)]
22252 pub fn fa_decode_batch_seqs_v4(
22253 &self,
22254 q: &CudaSlice<f32>,
22255 kv_ptrs: &cudarc::driver::CudaView<u64>,
22256 pos_seq: &CudaSlice<i32>,
22257 o: &mut CudaSlice<f32>,
22258 head_dim: usize,
22259 n_head: usize,
22260 n_head_kv: usize,
22261 b_n: usize,
22262 t_kv_max: usize,
22263 scale: f32,
22264 split_keys: usize,
22265 k_tok_bytes: usize,
22266 v_tok_bytes: usize,
22267 ) -> Result<(), Box<dyn std::error::Error>> {
22268 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22269 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22270 let o_len = b_n * n_head * n_splits_max * head_dim;
22271 let ml_len = b_n * n_head * n_splits_max;
22272 let mut part_guard = self.fa_part_pool.lock().unwrap();
22273 if part_guard
22274 .as_ref()
22275 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22276 .unwrap_or(true)
22277 {
22278 let old = part_guard.take();
22289 let (co, cm) = old
22290 .as_ref()
22291 .map(|pp| (pp.0.len(), pp.1.len()))
22292 .unwrap_or((0, 0));
22293 if let Some(old) = old {
22294 self.fa_part_retired.lock().unwrap().push(old);
22295 }
22296 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22297 eprintln!(
22298 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22299 co, o_len, cm, ml_len
22300 );
22301 }
22302 *part_guard = Some((
22303 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22304 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22305 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22306 ));
22307 }
22308 let pg = part_guard.as_mut().unwrap();
22309 self.gpu
22310 .stream()
22311 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22312 self.gpu
22313 .stream()
22314 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22315 self.gpu
22316 .stream()
22317 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22318 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22319 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22320 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22321 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22322 let gqa = (n_head / n_head_kv).max(1) as u32;
22323 let f = self.func("fa_decode_vec_q_seqs_v4");
22324 let shmem = (11520 + 32 * head_dim * 2) as u32;
22326 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22327 f.set_attribute(
22328 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22329 shmem as i32,
22330 )?;
22331 let cfg = LaunchConfig {
22332 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22333 block_dim: (32, gqa, 1),
22334 shared_mem_bytes: shmem,
22335 };
22336 {
22337 let __s_b = self.gpu.stream();
22338 let mut b = __s_b.launch_builder(&f);
22339 b.arg(q)
22340 .arg(kv_ptrs)
22341 .arg(pos_seq)
22342 .arg(&mut *part_o)
22343 .arg(&mut *part_m)
22344 .arg(&mut *part_l)
22345 .arg(&hd)
22346 .arg(&nh)
22347 .arg(&nhkv)
22348 .arg(&scale)
22349 .arg(&nspm)
22350 .arg(&spk)
22351 .arg(&ktb)
22352 .arg(&vtb);
22353 unsafe {
22354 b.launch(cfg)?;
22355 }
22356 }
22357 let fc = self.func("fa_decode_combine_seqs");
22358 let cfg2 = LaunchConfig {
22359 grid_dim: (n_head as u32, b_n as u32, 1),
22360 block_dim: (head_dim as u32, 1, 1),
22361 shared_mem_bytes: 0,
22362 };
22363 let __s_b2 = self.gpu.stream();
22364 let mut b2 = __s_b2.launch_builder(&fc);
22365 b2.arg(&*part_o)
22366 .arg(&*part_m)
22367 .arg(&*part_l)
22368 .arg(o)
22369 .arg(&hd)
22370 .arg(&nh)
22371 .arg(pos_seq)
22372 .arg(&nspm)
22373 .arg(&spk);
22374 unsafe {
22375 b2.launch(cfg2)?;
22376 }
22377 Ok(())
22378 }
22379
22380 #[allow(clippy::too_many_arguments)]
22387 pub fn append_kv_quantized_seqs(
22388 &self,
22389 k_rows: &CudaSlice<f32>,
22390 v_rows: &CudaSlice<f32>,
22391 kv_ptrs: &cudarc::driver::CudaView<u64>,
22392 pos_seq: &CudaSlice<i32>,
22393 b_n: usize,
22394 kv_dim_k: usize,
22395 kv_dim_v: usize,
22396 k_tok_bytes: usize,
22397 v_tok_bytes: usize,
22398 ) -> Result<(), Box<dyn std::error::Error>> {
22399 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
22400 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22401 let cfg = LaunchConfig {
22402 grid_dim: (nblk, b_n as u32, 1),
22403 block_dim: (32, 1, 1),
22404 shared_mem_bytes: 0,
22405 };
22406 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22407 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22408 let __s_b = self.gpu.stream();
22409 let mut b = __s_b.launch_builder(&f);
22410 b.arg(k_rows)
22411 .arg(v_rows)
22412 .arg(kv_ptrs)
22413 .arg(pos_seq)
22414 .arg(&kdk)
22415 .arg(&kdv)
22416 .arg(&ktb)
22417 .arg(&vtb);
22418 unsafe {
22419 b.launch(cfg)?;
22420 }
22421 Ok(())
22422 }
22423
22424 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
22430 std::env::var("MEMRA_NO_FA_VEC").is_err()
22431 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
22432 && base_len + 1 >= fa_vec_min_tkv()
22433 && head_dim <= 256
22434 && head_dim % 32 == 0
22435 }
22436
22437 #[allow(clippy::too_many_arguments)]
22446 pub fn fa_decode_rows(
22447 &self,
22448 q: &CudaSlice<f32>,
22449 k: &cudarc::driver::CudaView<u8>,
22450 v: &cudarc::driver::CudaView<u8>,
22451 o: &mut CudaSlice<f32>,
22452 head_dim: usize,
22453 n_head: usize,
22454 n_head_kv: usize,
22455 base_len: usize,
22456 t: usize,
22457 scale: f32,
22458 k_tok_bytes: usize,
22459 v_tok_bytes: usize,
22460 base_dev: Option<(&CudaSlice<i32>, i32)>,
22464 kv_shared: bool,
22467 g: bool,
22471 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22474 ) -> Result<(), Box<dyn std::error::Error>> {
22475 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
22476 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
22483 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22484 let v = *SP512.get_or_init(|| {
22487 std::env::var("MEMRA_FA_SP512")
22488 .ok()
22489 .and_then(|x| x.parse().ok())
22490 .unwrap_or(0)
22491 });
22492 sp = if v >= 8 {
22493 v
22494 } else {
22495 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22496 };
22497 }
22498 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22499 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22500 let gqa = (n_head / n_head_kv).max(1) as u32;
22501 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
22512 groups.push((0, t, sp));
22513 } else {
22514 let mut r0 = 0usize;
22515 while r0 < t {
22516 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
22517 let mut r1 = r0 + 1;
22518 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
22519 r1 += 1;
22520 }
22521 groups.push((r0, r1 - r0, sp_g));
22522 r0 = r1;
22523 }
22524 }
22525 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22529 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
22530 std::env::var("MEMRA_FA_SMEM_TKV")
22531 .ok()
22532 .and_then(|v| v.parse().ok())
22533 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22534 });
22535 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
22536 let v3 = fa_v3_active(head_dim);
22537 let smem_rows =
22538 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
22539 let _ = kv_shared;
22544 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
22547 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22561 let tb512 = head_dim == 512
22563 && sp <= 32
22564 && n_head / n_head_kv.max(1) <= 16
22565 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
22566 let fname = if tb512 {
22567 "fa_decode_vec_q_rows_v4_512_tb"
22568 } else if i2 {
22569 "fa_decode_vec_q_rows_dpl16_i2"
22570 } else if head_dim == 512 {
22571 "fa_decode_vec_q_rows_dpl16"
22572 }
22573 else if v4 {
22575 "fa_decode_vec_q_rows_v4"
22576 } else if v3 {
22577 "fa_decode_vec_q_rows_v3"
22578 } else if fa_v2_on() {
22579 "fa_decode_vec_q_rows_v2"
22580 } else if smem_rows {
22581 "fa_decode_vec_q_rows_smem"
22582 } else {
22583 "fa_decode_vec_q_rows"
22584 };
22585 let f = if head_dim == 512 {
22586 self.fa_func(fname, head_dim)
22587 } else if g {
22588 self.func_g(if smem_rows {
22596 "fa_decode_vec_q_rows"
22597 } else {
22598 fname
22599 })
22600 } else {
22601 self.func(fname)
22602 };
22603 let shmem = if tb512 {
22604 let gk = Self::gkv_on();
22606 let sh =
22607 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
22608 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22609 f.set_attribute(
22610 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22611 sh as i32,
22612 )?;
22613 sh
22614 } else if v4 || v3 || smem_rows || fa_v2_on() {
22615 let sh = (if v4 {
22617 11520 + 32 * head_dim * if g { 1 } else { 2 }
22618 } else if v3 {
22619 32 * head_dim * 2
22620 } else {
22621 2 * 32 * head_dim * 2
22622 }) as u32;
22623 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22624 f.set_attribute(
22625 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22626 sh as i32,
22627 )?;
22628 sh
22629 } else {
22630 0
22631 };
22632 for &(r0, t_g, sp_g) in &groups {
22636 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
22637 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
22638 let base_i = (base_len + r0) as i32;
22639 let o_len = t_g * n_head * n_splits_g * head_dim;
22640 let ml_len = t_g * n_head * n_splits_g;
22641 let mut part_guard = self.fa_part_pool.lock().unwrap();
22642 if part_guard
22643 .as_ref()
22644 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22645 .unwrap_or(true)
22646 {
22647 let old = part_guard.take();
22658 let (co, cm) = old
22659 .as_ref()
22660 .map(|pp| (pp.0.len(), pp.1.len()))
22661 .unwrap_or((0, 0));
22662 if let Some(old) = old {
22663 self.fa_part_retired.lock().unwrap().push(old);
22664 }
22665 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22666 eprintln!(
22667 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22668 co, o_len, cm, ml_len
22669 );
22670 }
22671 *part_guard = Some((
22672 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22673 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22674 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22675 ));
22676 }
22677 let pg = part_guard.as_mut().unwrap();
22678 self.gpu
22679 .stream()
22680 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22681 self.gpu
22682 .stream()
22683 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22684 self.gpu
22685 .stream()
22686 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22687 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22688 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22689 let qv = self.view(q, t * n_head * head_dim);
22690 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22691 let cfg = LaunchConfig {
22692 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
22693 block_dim: (32, gqa, 1),
22694 shared_mem_bytes: shmem,
22695 };
22696 {
22697 let __s_b = self.gpu.stream();
22698 let mut b = __s_b.launch_builder(&f);
22699 if tb512 {
22700 let (bd, plus) =
22702 base_dev.expect("hd512 rows twin requires a device base counter");
22703 let plus_g = plus + r0 as i32;
22704 let nr = t_g as i32;
22705 if Self::pdl_on() && Self::pdl_wb_on() {
22706 use cudarc::driver::{DevicePtr, DevicePtrMut};
22708 let s = &self.gpu.stream();
22709 let (pq, _b0) = q_g.device_ptr(s);
22710 let (pk, _b1) = k.device_ptr(s);
22711 let (pv, _b2) = v.device_ptr(s);
22712 let (po, _b3) = part_o.device_ptr_mut(s);
22713 let (pm, _b4) = part_m.device_ptr_mut(s);
22714 let (pl, _b5) = part_l.device_ptr_mut(s);
22715 let (pb, _b6) = bd.device_ptr(s);
22716 let mut ps = [
22717 &pq as *const _ as *mut std::ffi::c_void,
22718 &pk as *const _ as *mut _,
22719 &pv as *const _ as *mut _,
22720 &po as *const _ as *mut _,
22721 &pm as *const _ as *mut _,
22722 &pl as *const _ as *mut _,
22723 &hd as *const _ as *mut _,
22724 &nh as *const _ as *mut _,
22725 &nhkv as *const _ as *mut _,
22726 &pb as *const _ as *mut _,
22727 &plus_g as *const _ as *mut _,
22728 &scale as *const _ as *mut _,
22729 &nspm as *const _ as *mut _,
22730 &spk as *const _ as *mut _,
22731 &ktb as *const _ as *mut _,
22732 &vtb as *const _ as *mut _,
22733 &nr as *const _ as *mut _,
22734 ];
22735 unsafe {
22736 self.launch_pdl_flash(
22737 Self::gkv_on(),
22738 "fa_decode_vec_q_rows_v4_512_tb",
22739 (n_head_kv as u32, n_splits_g as u32, 1),
22740 (32, gqa, 1),
22741 shmem,
22742 &mut ps,
22743 )?;
22744 }
22745 } else {
22746 let cfg_tb = LaunchConfig {
22747 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
22748 block_dim: (32, gqa, 1),
22749 shared_mem_bytes: shmem,
22750 };
22751 b.arg(&q_g)
22752 .arg(k)
22753 .arg(v)
22754 .arg(&mut *part_o)
22755 .arg(&mut *part_m)
22756 .arg(&mut *part_l)
22757 .arg(&hd)
22758 .arg(&nh)
22759 .arg(&nhkv)
22760 .arg(bd)
22761 .arg(&plus_g)
22762 .arg(&scale)
22763 .arg(&nspm)
22764 .arg(&spk)
22765 .arg(&ktb)
22766 .arg(&vtb)
22767 .arg(&nr);
22768 unsafe {
22769 b.launch(cfg_tb)?;
22770 }
22771 }
22772 } else if head_dim == 512 {
22773 let (bd, plus) =
22774 base_dev.expect("hd512 rows twin requires a device base counter");
22775 let plus_g = plus + r0 as i32;
22776 b.arg(&q_g)
22777 .arg(k)
22778 .arg(v)
22779 .arg(&mut *part_o)
22780 .arg(&mut *part_m)
22781 .arg(&mut *part_l)
22782 .arg(&hd)
22783 .arg(&nh)
22784 .arg(&nhkv)
22785 .arg(bd)
22786 .arg(&plus_g)
22787 .arg(&scale)
22788 .arg(&nspm)
22789 .arg(&spk)
22790 .arg(&ktb)
22791 .arg(&vtb);
22792 unsafe {
22793 b.launch(cfg)?;
22794 }
22795 } else {
22796 b.arg(&q_g)
22797 .arg(k)
22798 .arg(v)
22799 .arg(&mut *part_o)
22800 .arg(&mut *part_m)
22801 .arg(&mut *part_l)
22802 .arg(&hd)
22803 .arg(&nh)
22804 .arg(&nhkv)
22805 .arg(&base_i)
22806 .arg(&scale)
22807 .arg(&nspm)
22808 .arg(&spk)
22809 .arg(&ktb)
22810 .arg(&vtb);
22811 unsafe {
22812 b.launch(cfg)?;
22813 }
22814 }
22815 }
22816 let cfg2 = LaunchConfig {
22817 grid_dim: (n_head as u32, t_g as u32, 1),
22818 block_dim: (head_dim as u32, 1, 1),
22819 shared_mem_bytes: 0,
22820 };
22821 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
22822 if head_dim == 512 {
22823 let (bd, plus) = base_dev.unwrap();
22826 let plus_g = plus + r0 as i32;
22827 if let Some((oq, od)) = q8_out.as_mut() {
22828 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
22830 if Self::pdl_on() && Self::pdl_wb_on() {
22831 use cudarc::driver::{DevicePtr, DevicePtrMut};
22833 let s = &self.gpu.stream();
22834 let (po, _g0) = part_o.device_ptr(s);
22835 let (pm, _g1) = part_m.device_ptr(s);
22836 let (pl, _g2) = part_l.device_ptr(s);
22837 let (pq, _g3) = oq.device_ptr_mut(s);
22838 let (pd, _g4) = od.device_ptr_mut(s);
22839 let (pb, _g5) = bd.device_ptr(s);
22840 let mut ps = [
22841 &po as *const _ as *mut std::ffi::c_void,
22842 &pm as *const _ as *mut _,
22843 &pl as *const _ as *mut _,
22844 &pq as *const _ as *mut _,
22845 &pd as *const _ as *mut _,
22846 &hd as *const _ as *mut _,
22847 &nh as *const _ as *mut _,
22848 &pb as *const _ as *mut _,
22849 &plus_g as *const _ as *mut _,
22850 &nspm as *const _ as *mut _,
22851 &spk as *const _ as *mut _,
22852 ];
22853 unsafe {
22854 self.launch_pdl_flash(
22855 Self::gkv_on(),
22856 "fa_decode_combine_rows_dc_q8_1",
22857 cfg2.grid_dim,
22858 cfg2.block_dim,
22859 0,
22860 &mut ps,
22861 )?;
22862 }
22863 continue;
22864 }
22865 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
22866 let __s_b2 = self.gpu.stream();
22867 let mut b2 = __s_b2.launch_builder(&fc);
22868 b2.arg(&*part_o)
22869 .arg(&*part_m)
22870 .arg(&*part_l)
22871 .arg(&mut **oq)
22872 .arg(&mut **od)
22873 .arg(&hd)
22874 .arg(&nh)
22875 .arg(bd)
22876 .arg(&plus_g)
22877 .arg(&nspm)
22878 .arg(&spk);
22879 unsafe {
22880 b2.launch(cfg2)?;
22881 }
22882 continue;
22883 }
22884 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
22885 let __s_b2 = self.gpu.stream();
22886 let mut b2 = __s_b2.launch_builder(&fc);
22887 b2.arg(&*part_o)
22888 .arg(&*part_m)
22889 .arg(&*part_l)
22890 .arg(&mut o_g)
22891 .arg(&hd)
22892 .arg(&nh)
22893 .arg(bd)
22894 .arg(&plus_g)
22895 .arg(&nspm)
22896 .arg(&spk);
22897 unsafe {
22898 b2.launch(cfg2)?;
22899 }
22900 } else {
22901 assert!(
22904 q8_out.is_none(),
22905 "rows q8 emit requires the hd512 dc combine"
22906 );
22907 let fc = self.func("fa_decode_combine_rows");
22908 let __s_b2 = self.gpu.stream();
22909 let mut b2 = __s_b2.launch_builder(&fc);
22910 b2.arg(&*part_o)
22911 .arg(&*part_m)
22912 .arg(&*part_l)
22913 .arg(&mut o_g)
22914 .arg(&hd)
22915 .arg(&nh)
22916 .arg(&base_i)
22917 .arg(&nspm)
22918 .arg(&spk);
22919 unsafe {
22920 b2.launch(cfg2)?;
22921 }
22922 }
22923 }
22924 Ok(())
22925 }
22926
22927 #[allow(clippy::too_many_arguments)]
22931 pub fn fa_decode_rows_w(
22932 &self,
22933 q: &CudaSlice<f32>,
22934 k: &cudarc::driver::CudaView<u8>,
22935 v: &cudarc::driver::CudaView<u8>,
22936 o: &mut CudaSlice<f32>,
22937 head_dim: usize,
22938 n_head: usize,
22939 n_head_kv: usize,
22940 base_dev: &CudaSlice<i32>,
22941 base_plus: i32,
22942 t: usize,
22943 scale: f32,
22944 window: usize,
22945 k_tok_bytes: usize,
22946 v_tok_bytes: usize,
22947 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22948 ) -> Result<(), Box<dyn std::error::Error>> {
22949 debug_assert!(head_dim == 256);
22954 let sp = {
22962 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22963 let v = *SPW.get_or_init(|| {
22964 std::env::var("MEMRA_FA_SPW")
22965 .ok()
22966 .and_then(|x| x.parse().ok())
22967 .unwrap_or(0)
22968 });
22969 if v >= 8 {
22970 v
22971 } else {
22972 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22973 }
22974 };
22975 let n_splits_max = (window + sp - 1) / sp;
22976 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22977 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
22978 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22979 let gqa = (n_head / n_head_kv).max(1) as u32;
22980 let o_len = t * n_head * n_splits_max * head_dim;
22981 let ml_len = t * n_head * n_splits_max;
22982 let mut part_guard = self.fa_part_pool.lock().unwrap();
22983 if part_guard
22984 .as_ref()
22985 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22986 .unwrap_or(true)
22987 {
22988 let old = part_guard.take();
22999 let (co, cm) = old
23000 .as_ref()
23001 .map(|pp| (pp.0.len(), pp.1.len()))
23002 .unwrap_or((0, 0));
23003 if let Some(old) = old {
23004 self.fa_part_retired.lock().unwrap().push(old);
23005 }
23006 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23007 eprintln!(
23008 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23009 co, o_len, cm, ml_len
23010 );
23011 }
23012 *part_guard = Some((
23013 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23014 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23015 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23016 ));
23017 }
23018 let pg = part_guard.as_mut().unwrap();
23019 self.gpu
23020 .stream()
23021 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23022 self.gpu
23023 .stream()
23024 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23025 self.gpu
23026 .stream()
23027 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23028 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23029 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23035 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
23036 std::env::var("MEMRA_FA_SMEM_TKV")
23037 .ok()
23038 .and_then(|v| v.parse().ok())
23039 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23040 });
23041 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23047 let wg = Self::wkv_on();
23052 let sp2 =
23055 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23056 if sp2 {
23057 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23058 if Self::pdl_on() && Self::pdl_wb_on() {
23059 use cudarc::driver::{DevicePtr, DevicePtrMut};
23061 let s = &self.gpu.stream();
23062 let (pq, _b0) = q.device_ptr(s);
23063 let (pk, _b1) = k.device_ptr(s);
23064 let (pv, _b2) = v.device_ptr(s);
23065 let (po, _b3) = part_o.device_ptr_mut(s);
23066 let (pm, _b4) = part_m.device_ptr_mut(s);
23067 let (pl, _b5) = part_l.device_ptr_mut(s);
23068 let (pb, _b6) = base_dev.device_ptr(s);
23069 let mut ps = [
23070 &pq as *const _ as *mut std::ffi::c_void,
23071 &pk as *const _ as *mut _,
23072 &pv as *const _ as *mut _,
23073 &po as *const _ as *mut _,
23074 &pm as *const _ as *mut _,
23075 &pl as *const _ as *mut _,
23076 &hd as *const _ as *mut _,
23077 &nh as *const _ as *mut _,
23078 &nhkv as *const _ as *mut _,
23079 &pb as *const _ as *mut _,
23080 &base_plus as *const _ as *mut _,
23081 &scale as *const _ as *mut _,
23082 &nspm as *const _ as *mut _,
23083 &spk as *const _ as *mut _,
23084 &ktb as *const _ as *mut _,
23085 &vtb as *const _ as *mut _,
23086 &wini as *const _ as *mut _,
23087 ];
23088 unsafe {
23089 self.launch_pdl_flash(
23090 wg,
23091 "fa_decode_vec_q_rows_v4_w_sp",
23092 (n_head_kv as u32, n_splits_max as u32, t as u32),
23093 (32, gqa + 1, 1),
23094 sh,
23095 &mut ps,
23096 )?;
23097 }
23098 } else {
23099 let f = if wg {
23100 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23101 } else {
23102 self.func("fa_decode_vec_q_rows_v4_w_sp")
23103 };
23104 f.set_attribute(
23105 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23106 sh as i32,
23107 )?;
23108 let cfg = LaunchConfig {
23109 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23110 block_dim: (32, gqa + 1, 1),
23111 shared_mem_bytes: sh,
23112 };
23113 let __s_b = self.gpu.stream();
23114 let mut b = __s_b.launch_builder(&f);
23115 b.arg(q)
23116 .arg(k)
23117 .arg(v)
23118 .arg(&mut *part_o)
23119 .arg(&mut *part_m)
23120 .arg(&mut *part_l)
23121 .arg(&hd)
23122 .arg(&nh)
23123 .arg(&nhkv)
23124 .arg(base_dev)
23125 .arg(&base_plus)
23126 .arg(&scale)
23127 .arg(&nspm)
23128 .arg(&spk)
23129 .arg(&ktb)
23130 .arg(&vtb)
23131 .arg(&wini);
23132 unsafe {
23133 b.launch(cfg)?;
23134 }
23135 }
23136 } else {
23137 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23138 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23140 use cudarc::driver::{DevicePtr, DevicePtrMut};
23141 let s = &self.gpu.stream();
23142 let (pq, _b0) = q.device_ptr(s);
23143 let (pk, _b1) = k.device_ptr(s);
23144 let (pv, _b2) = v.device_ptr(s);
23145 let (po, _b3) = part_o.device_ptr_mut(s);
23146 let (pm, _b4) = part_m.device_ptr_mut(s);
23147 let (pl, _b5) = part_l.device_ptr_mut(s);
23148 let (pb, _b6) = base_dev.device_ptr(s);
23149 let mut ps = [
23150 &pq as *const _ as *mut std::ffi::c_void,
23151 &pk as *const _ as *mut _,
23152 &pv as *const _ as *mut _,
23153 &po as *const _ as *mut _,
23154 &pm as *const _ as *mut _,
23155 &pl as *const _ as *mut _,
23156 &hd as *const _ as *mut _,
23157 &nh as *const _ as *mut _,
23158 &nhkv as *const _ as *mut _,
23159 &pb as *const _ as *mut _,
23160 &base_plus as *const _ as *mut _,
23161 &scale as *const _ as *mut _,
23162 &nspm as *const _ as *mut _,
23163 &spk as *const _ as *mut _,
23164 &ktb as *const _ as *mut _,
23165 &vtb as *const _ as *mut _,
23166 &wini as *const _ as *mut _,
23167 ];
23168 unsafe {
23169 self.launch_pdl_flash(
23170 wg,
23171 "fa_decode_vec_q_rows_v4_w",
23172 (n_head_kv as u32, n_splits_max as u32, t as u32),
23173 (32, gqa, 1),
23174 sh,
23175 &mut ps,
23176 )?;
23177 }
23178 } else {
23179 let pick = |name: &str| {
23180 if wg {
23181 self.func_g(name)
23182 } else {
23183 self.func(name)
23184 }
23185 };
23186 let (f, sh) = if fa_v4_at(window) {
23187 let f = pick("fa_decode_vec_q_rows_v4_w");
23188 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23189 } else if smem_tkv > 0 && window >= smem_tkv {
23190 (
23193 pick("fa_decode_vec_q_rows_smem_w"),
23194 (2 * 32 * head_dim * 2) as u32,
23195 )
23196 } else {
23197 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23198 };
23199 f.set_attribute(
23200 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23201 sh as i32,
23202 )?;
23203 let cfg = LaunchConfig {
23204 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23205 block_dim: (32, gqa, 1),
23206 shared_mem_bytes: sh,
23207 };
23208 let __s_b = self.gpu.stream();
23209 let mut b = __s_b.launch_builder(&f);
23210 b.arg(q)
23211 .arg(k)
23212 .arg(v)
23213 .arg(&mut *part_o)
23214 .arg(&mut *part_m)
23215 .arg(&mut *part_l)
23216 .arg(&hd)
23217 .arg(&nh)
23218 .arg(&nhkv)
23219 .arg(base_dev)
23220 .arg(&base_plus)
23221 .arg(&scale)
23222 .arg(&nspm)
23223 .arg(&spk)
23224 .arg(&ktb)
23225 .arg(&vtb)
23226 .arg(&wini);
23227 unsafe {
23228 b.launch(cfg)?;
23229 }
23230 }
23231 }
23232 let cfg2 = LaunchConfig {
23233 grid_dim: (n_head as u32, t as u32, 1),
23234 block_dim: (head_dim as u32, 1, 1),
23235 shared_mem_bytes: 0,
23236 };
23237 if let Some((oq, od)) = q8_out {
23238 if Self::pdl_on() && Self::pdl_wb_on() {
23241 use cudarc::driver::{DevicePtr, DevicePtrMut};
23243 let s = &self.gpu.stream();
23244 let (po, _g0) = part_o.device_ptr(s);
23245 let (pm, _g1) = part_m.device_ptr(s);
23246 let (pl, _g2) = part_l.device_ptr(s);
23247 let (pq, _g3) = oq.device_ptr_mut(s);
23248 let (pd, _g4) = od.device_ptr_mut(s);
23249 let mut ps = [
23250 &po as *const _ as *mut std::ffi::c_void,
23251 &pm as *const _ as *mut _,
23252 &pl as *const _ as *mut _,
23253 &pq as *const _ as *mut _,
23254 &pd as *const _ as *mut _,
23255 &hd as *const _ as *mut _,
23256 &nh as *const _ as *mut _,
23257 &nspm as *const _ as *mut _,
23258 &spk as *const _ as *mut _,
23259 &wini as *const _ as *mut _,
23260 ];
23261 unsafe {
23262 self.launch_pdl_flash(
23263 wg,
23264 "fa_decode_combine_rows_w_q8_1",
23265 cfg2.grid_dim,
23266 cfg2.block_dim,
23267 0,
23268 &mut ps,
23269 )?;
23270 }
23271 return Ok(());
23272 }
23273 let fc = if wg {
23274 self.func_g("fa_decode_combine_rows_w_q8_1")
23275 } else {
23276 self.func("fa_decode_combine_rows_w_q8_1")
23277 };
23278 let __s_b2 = self.gpu.stream();
23279 let mut b2 = __s_b2.launch_builder(&fc);
23280 b2.arg(&*part_o)
23281 .arg(&*part_m)
23282 .arg(&*part_l)
23283 .arg(oq)
23284 .arg(od)
23285 .arg(&hd)
23286 .arg(&nh)
23287 .arg(&nspm)
23288 .arg(&spk)
23289 .arg(&wini);
23290 unsafe {
23291 b2.launch(cfg2)?;
23292 }
23293 return Ok(());
23294 }
23295 let fc = if wg {
23296 self.func_g("fa_decode_combine_rows_w")
23297 } else {
23298 self.func("fa_decode_combine_rows_w")
23299 };
23300 let __s_b2 = self.gpu.stream();
23301 let mut b2 = __s_b2.launch_builder(&fc);
23302 b2.arg(&*part_o)
23303 .arg(&*part_m)
23304 .arg(&*part_l)
23305 .arg(o)
23306 .arg(&hd)
23307 .arg(&nh)
23308 .arg(&nspm)
23309 .arg(&spk)
23310 .arg(&wini);
23311 unsafe {
23312 b2.launch(cfg2)?;
23313 }
23314 Ok(())
23315 }
23316
23317 #[allow(clippy::too_many_arguments)]
23323 pub fn fa_decode_rows_dc(
23324 &self,
23325 q: &CudaSlice<f32>,
23326 k: &cudarc::driver::CudaView<u8>,
23327 v: &cudarc::driver::CudaView<u8>,
23328 o: &mut CudaSlice<f32>,
23329 head_dim: usize,
23330 n_head: usize,
23331 n_head_kv: usize,
23332 base_dev: &CudaSlice<i32>,
23333 t_kv_upper: usize,
23334 t: usize,
23335 scale: f32,
23336 k_tok_bytes: usize,
23337 v_tok_bytes: usize,
23338 base_plus: i32,
23339 g: bool,
23340 ) -> Result<(), Box<dyn std::error::Error>> {
23341 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23342 assert!(
23343 v4 || fa_v3_active(head_dim),
23344 "stream fa rows requires the v3 or v4 lane"
23345 );
23346 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23347 if v4 {
23348 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23349 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23350 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23351 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23352 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23353 let gqa = (n_head / n_head_kv).max(1) as u32;
23354 let o_len = t * n_head * n_splits_max * head_dim;
23355 let ml_len = t * n_head * n_splits_max;
23356 let mut part_guard = self.fa_part_pool.lock().unwrap();
23357 if part_guard
23358 .as_ref()
23359 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23360 .unwrap_or(true)
23361 {
23362 let old = part_guard.take();
23373 let (co, cm) = old
23374 .as_ref()
23375 .map(|pp| (pp.0.len(), pp.1.len()))
23376 .unwrap_or((0, 0));
23377 if let Some(old) = old {
23378 self.fa_part_retired.lock().unwrap().push(old);
23379 }
23380 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23381 eprintln!(
23382 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23383 co, o_len, cm, ml_len
23384 );
23385 }
23386 *part_guard = Some((
23387 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23388 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23389 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23390 ));
23391 }
23392 let pg = part_guard.as_mut().unwrap();
23393 self.gpu
23394 .stream()
23395 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23396 self.gpu
23397 .stream()
23398 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23399 self.gpu
23400 .stream()
23401 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23402 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23403 let f = if g {
23404 self.func_g("fa_decode_vec_q_rows_v4_dc")
23405 } else {
23406 self.func("fa_decode_vec_q_rows_v4_dc")
23407 };
23408 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23409 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23410 f.set_attribute(
23411 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23412 sh as i32,
23413 )?;
23414 let cfg = LaunchConfig {
23415 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23416 block_dim: (32, gqa, 1),
23417 shared_mem_bytes: sh,
23418 };
23419 let __s_b = self.gpu.stream();
23420 let mut b = __s_b.launch_builder(&f);
23421 b.arg(q)
23422 .arg(k)
23423 .arg(v)
23424 .arg(&mut *part_o)
23425 .arg(&mut *part_m)
23426 .arg(&mut *part_l)
23427 .arg(&hd)
23428 .arg(&nh)
23429 .arg(&nhkv)
23430 .arg(base_dev)
23431 .arg(&base_plus)
23432 .arg(&scale)
23433 .arg(&nspm)
23434 .arg(&spk)
23435 .arg(&ktb)
23436 .arg(&vtb);
23437 unsafe {
23438 b.launch(cfg)?;
23439 }
23440 let fc = self.func("fa_decode_combine_rows_dc");
23441 let cfg2 = LaunchConfig {
23442 grid_dim: (n_head as u32, t as u32, 1),
23443 block_dim: (head_dim as u32, 1, 1),
23444 shared_mem_bytes: 0,
23445 };
23446 let __s_b2 = self.gpu.stream();
23447 let mut b2 = __s_b2.launch_builder(&fc);
23448 b2.arg(&*part_o)
23449 .arg(&*part_m)
23450 .arg(&*part_l)
23451 .arg(o)
23452 .arg(&hd)
23453 .arg(&nh)
23454 .arg(base_dev)
23455 .arg(&base_plus)
23456 .arg(&nspm)
23457 .arg(&spk);
23458 unsafe {
23459 b2.launch(cfg2)?;
23460 }
23461 return Ok(());
23462 }
23463 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23464 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23465 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23466 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23467 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23468 let gqa = (n_head / n_head_kv).max(1) as u32;
23469 let o_len = t * n_head * n_splits_max * head_dim;
23470 let ml_len = t * n_head * n_splits_max;
23471 let mut part_guard = self.fa_part_pool.lock().unwrap();
23472 if part_guard
23473 .as_ref()
23474 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23475 .unwrap_or(true)
23476 {
23477 let old = part_guard.take();
23488 let (co, cm) = old
23489 .as_ref()
23490 .map(|pp| (pp.0.len(), pp.1.len()))
23491 .unwrap_or((0, 0));
23492 if let Some(old) = old {
23493 self.fa_part_retired.lock().unwrap().push(old);
23494 }
23495 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23496 eprintln!(
23497 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23498 co, o_len, cm, ml_len
23499 );
23500 }
23501 *part_guard = Some((
23502 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23503 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23504 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23505 ));
23506 }
23507 let pg = part_guard.as_mut().unwrap();
23508 self.gpu
23509 .stream()
23510 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23511 self.gpu
23512 .stream()
23513 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23514 self.gpu
23515 .stream()
23516 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23517 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23518 let f = self.func("fa_decode_vec_q_rows_v3_dc");
23519 let sh = (32 * head_dim * 2) as u32;
23520 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23521 f.set_attribute(
23522 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23523 sh as i32,
23524 )?;
23525 let cfg = LaunchConfig {
23526 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23527 block_dim: (32, gqa, 1),
23528 shared_mem_bytes: sh,
23529 };
23530 let __s_b = self.gpu.stream();
23531 let mut b = __s_b.launch_builder(&f);
23532 b.arg(q)
23533 .arg(k)
23534 .arg(v)
23535 .arg(&mut *part_o)
23536 .arg(&mut *part_m)
23537 .arg(&mut *part_l)
23538 .arg(&hd)
23539 .arg(&nh)
23540 .arg(&nhkv)
23541 .arg(base_dev)
23542 .arg(&scale)
23543 .arg(&nspm)
23544 .arg(&spk)
23545 .arg(&ktb)
23546 .arg(&vtb);
23547 unsafe {
23548 b.launch(cfg)?;
23549 }
23550 let fc = self.func("fa_decode_combine_rows_dc");
23551 let cfg2 = LaunchConfig {
23552 grid_dim: (n_head as u32, t as u32, 1),
23553 block_dim: (head_dim as u32, 1, 1),
23554 shared_mem_bytes: 0,
23555 };
23556 let plus0 = 0i32;
23557 let __s_b2 = self.gpu.stream();
23558 let mut b2 = __s_b2.launch_builder(&fc);
23559 b2.arg(&*part_o)
23560 .arg(&*part_m)
23561 .arg(&*part_l)
23562 .arg(o)
23563 .arg(&hd)
23564 .arg(&nh)
23565 .arg(base_dev)
23566 .arg(&plus0)
23567 .arg(&nspm)
23568 .arg(&spk);
23569 unsafe {
23570 b2.launch(cfg2)?;
23571 }
23572 Ok(())
23573 }
23574
23575 pub fn fa_decode_dc(
23586 &self,
23587 q: &CudaSlice<f32>,
23588 k: &cudarc::driver::CudaView<u8>,
23589 v: &cudarc::driver::CudaView<u8>,
23590 o: &mut CudaSlice<f32>,
23591 head_dim: usize,
23592 n_head: usize,
23593 n_head_kv: usize,
23594 t_kv_dev: &CudaSlice<i32>,
23595 bucket_max: usize,
23596 scale: f32,
23597 k_tok_bytes: usize,
23598 v_tok_bytes: usize,
23599 g: bool,
23600 ) -> Result<(), Box<dyn std::error::Error>> {
23601 self.fa_decode_dc_q8(
23602 q,
23603 k,
23604 v,
23605 o,
23606 head_dim,
23607 n_head,
23608 n_head_kv,
23609 t_kv_dev,
23610 bucket_max,
23611 scale,
23612 k_tok_bytes,
23613 v_tok_bytes,
23614 g,
23615 None,
23616 )
23617 }
23618
23619 #[allow(clippy::too_many_arguments)]
23622 pub fn fa_decode_dc_q8(
23623 &self,
23624 q: &CudaSlice<f32>,
23625 k: &cudarc::driver::CudaView<u8>,
23626 v: &cudarc::driver::CudaView<u8>,
23627 o: &mut CudaSlice<f32>,
23628 head_dim: usize,
23629 n_head: usize,
23630 n_head_kv: usize,
23631 t_kv_dev: &CudaSlice<i32>,
23632 bucket_max: usize,
23633 scale: f32,
23634 k_tok_bytes: usize,
23635 v_tok_bytes: usize,
23636 g: bool,
23637 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23638 ) -> Result<(), Box<dyn std::error::Error>> {
23639 let mut fa_vec =
23647 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23648 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
23649 fa_vec = false;
23650 } let sp = fa_split_keys(bucket_max, n_head_kv);
23652 let n_splits = if fa_vec {
23653 ((bucket_max + sp - 1) / sp).max(1)
23654 } else {
23655 ((bucket_max + 255) / 256).max(1)
23656 };
23657 let o_len = n_head * n_splits * head_dim;
23658 let ml_len = n_head * n_splits;
23659 let mut part_guard = self.fa_part_pool.lock().unwrap();
23660 if part_guard
23661 .as_ref()
23662 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23663 .unwrap_or(true)
23664 {
23665 let old = part_guard.take();
23676 let (co, cm) = old
23677 .as_ref()
23678 .map(|pp| (pp.0.len(), pp.1.len()))
23679 .unwrap_or((0, 0));
23680 if let Some(old) = old {
23681 self.fa_part_retired.lock().unwrap().push(old);
23682 }
23683 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23684 eprintln!(
23685 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23686 co, o_len, cm, ml_len
23687 );
23688 }
23689 *part_guard = Some((
23690 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23691 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23692 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23693 ));
23694 }
23695 let pg = part_guard.as_mut().unwrap();
23696 self.gpu
23697 .stream()
23698 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23699 self.gpu
23700 .stream()
23701 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23702 self.gpu
23703 .stream()
23704 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23705 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23706 let (hd, nh, nhkv, nsp) = (
23707 head_dim as i32,
23708 n_head as i32,
23709 n_head_kv as i32,
23710 n_splits as i32,
23711 );
23712 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23713 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
23714 let deep = fa_vec
23717 && head_dim == 256
23718 && fa_v4_at(bucket_max)
23719 && !g
23720 && fa_deep_at(bucket_max)
23721 && !matches!(fa_v4_mode(), "noB3" | "stage");
23722 let (f, cfg) = if fa_vec
23723 && head_dim == 512
23724 && bucket_max >= {
23725 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23726 *FA512_MIN_DC.get_or_init(|| {
23727 std::env::var("MEMRA_FA512_MIN")
23728 .ok()
23729 .and_then(|v| v.parse().ok())
23730 .unwrap_or(512)
23731 })
23732 } {
23733 let gqa = (n_head / n_head_kv).max(1) as u32;
23735 (
23736 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
23737 LaunchConfig {
23738 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23739 block_dim: (32, gqa, 1),
23740 shared_mem_bytes: 0,
23741 },
23742 )
23743 } else if fa_vec && head_dim == 512 {
23744 let q_view = q.as_view();
23747 let mut o_view = o.as_view_mut();
23748 return self.fa_decode_scalar_unified(
23749 &q_view,
23750 k,
23751 v,
23752 &mut o_view,
23753 head_dim,
23754 n_head,
23755 n_head_kv,
23756 0,
23757 Some(t_kv_dev),
23758 scale,
23759 n_splits,
23760 sp,
23761 k_tok_bytes,
23762 v_tok_bytes,
23763 g,
23764 &mut *part_o,
23765 &mut *part_m,
23766 &mut *part_l,
23767 q8_out,
23768 );
23769 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
23770 let gqa = (n_head / n_head_kv).max(1) as u32;
23773 let fv = if g {
23774 self.func_g("fa_decode_vec_q_v4_dc")
23775 } else if deep {
23776 self.func("fa_decode_vec_q_v4_deep_dc")
23777 } else {
23778 self.func("fa_decode_vec_q_v4_dc")
23779 };
23780 let shmem =
23781 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
23782 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23783 fv.set_attribute(
23784 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23785 shmem as i32,
23786 )?;
23787 (
23788 fv,
23789 LaunchConfig {
23790 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23791 block_dim: (32, gqa, 1),
23792 shared_mem_bytes: shmem,
23793 },
23794 )
23795 } else if fa_vec && fa_v3_active(head_dim) {
23796 let gqa = (n_head / n_head_kv).max(1) as u32;
23799 let fv = if g {
23800 self.func_g("fa_decode_vec_q_v3_dc")
23801 } else {
23802 self.func("fa_decode_vec_q_v3_dc")
23803 };
23804 let shmem = (32 * head_dim * 2) as u32; (
23806 fv,
23807 LaunchConfig {
23808 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23809 block_dim: (32, gqa, 1),
23810 shared_mem_bytes: shmem,
23811 },
23812 )
23813 } else if fa_vec && fa_v2_on() {
23814 let gqa = (n_head / n_head_kv).max(1) as u32;
23818 let fv = if g {
23819 self.func_g("fa_decode_vec_q_v2_dc")
23820 } else {
23821 self.func("fa_decode_vec_q_v2_dc")
23822 };
23823 let shmem = (2 * 32 * head_dim * 2) as u32; (
23825 fv,
23826 LaunchConfig {
23827 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23828 block_dim: (32, gqa, 1),
23829 shared_mem_bytes: shmem,
23830 },
23831 )
23832 } else if fa_vec {
23833 let gqa = (n_head / n_head_kv).max(1) as u32;
23834 let fv = if g {
23836 self.func_g("fa_decode_vec_q_dc")
23837 } else {
23838 self.func("fa_decode_vec_q_dc")
23839 };
23840 (
23841 fv,
23842 LaunchConfig {
23843 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23844 block_dim: (32, gqa, 1),
23845 shared_mem_bytes: 0,
23846 },
23847 )
23848 } else {
23849 let q_view = q.as_view();
23850 let mut o_view = o.as_view_mut();
23851 return self.fa_decode_scalar_unified(
23852 &q_view,
23853 k,
23854 v,
23855 &mut o_view,
23856 head_dim,
23857 n_head,
23858 n_head_kv,
23859 0,
23860 Some(t_kv_dev),
23861 scale,
23862 n_splits,
23863 if fa_vec { sp } else { 256 },
23864 k_tok_bytes,
23865 v_tok_bytes,
23866 g,
23867 &mut *part_o,
23868 &mut *part_m,
23869 &mut *part_l,
23870 q8_out,
23871 );
23872 };
23873 let ski = sp as i32; let __s_b = self.gpu.stream();
23875 let mut b = __s_b.launch_builder(&f);
23876 b.arg(q)
23877 .arg(k)
23878 .arg(v)
23879 .arg(&mut *part_o)
23880 .arg(&mut *part_m)
23881 .arg(&mut *part_l)
23882 .arg(&hd)
23883 .arg(&nh)
23884 .arg(&nhkv)
23885 .arg(t_kv_dev)
23886 .arg(&scale)
23887 .arg(&nsp)
23888 .arg(&ski)
23889 .arg(&ktb)
23890 .arg(&vtb);
23891 unsafe {
23892 b.launch(cfg)?;
23893 }
23894 let cfg2 = LaunchConfig {
23895 grid_dim: (n_head as u32, 1, 1),
23896 block_dim: (head_dim as u32, 1, 1),
23897 shared_mem_bytes: 0,
23898 };
23899 if let Some((oq, od)) = q8_out {
23900 let fc = if g {
23901 self.func_g("fa_decode_combine_q8_1")
23902 } else {
23903 self.fa_func("fa_decode_combine_q8_1", head_dim)
23904 };
23905 let __s_b2 = self.gpu.stream();
23906 let mut b2 = __s_b2.launch_builder(&fc);
23907 b2.arg(&*part_o)
23908 .arg(&*part_m)
23909 .arg(&*part_l)
23910 .arg(oq)
23911 .arg(od)
23912 .arg(&hd)
23913 .arg(&nh)
23914 .arg(&nsp);
23915 unsafe {
23916 b2.launch(cfg2)?;
23917 }
23918 return Ok(());
23919 }
23920 let fc = if g {
23921 self.func_g("fa_decode_combine_f32")
23922 } else {
23923 self.fa_func("fa_decode_combine_f32", head_dim)
23924 };
23925 let __s_b2 = self.gpu.stream();
23926 let mut b2 = __s_b2.launch_builder(&fc);
23927 b2.arg(&*part_o)
23928 .arg(&*part_m)
23929 .arg(&*part_l)
23930 .arg(o)
23931 .arg(&hd)
23932 .arg(&nh)
23933 .arg(&nsp);
23934 unsafe {
23935 b2.launch(cfg2)?;
23936 }
23937 Ok(())
23938 }
23939
23940 #[allow(clippy::too_many_arguments)]
23944 pub fn append_kv_quantized_dcw(
23945 &self,
23946 k_row: &CudaSlice<f32>,
23947 v_row: &CudaSlice<f32>,
23948 kc: &mut CudaSlice<u8>,
23949 vc: &mut CudaSlice<u8>,
23950 len_dev: &CudaSlice<i32>,
23951 base_dev: Option<&CudaSlice<i32>>,
23952 kv_dim_k: usize,
23953 kv_dim_v: usize,
23954 k_tok_bytes: usize,
23955 v_tok_bytes: usize,
23956 ) -> Result<(), Box<dyn std::error::Error>> {
23957 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
23958 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23959 let cfg = LaunchConfig {
23960 grid_dim: (nblk, 1, 1),
23961 block_dim: (32, 1, 1),
23962 shared_mem_bytes: 0,
23963 };
23964 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23965 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23966 let null: u64 = 0;
23967 let __s_b = self.gpu.stream();
23968 let mut b = __s_b.launch_builder(&f);
23969 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
23970 match base_dev {
23971 Some(base) => {
23972 b.arg(base);
23973 }
23974 None => {
23975 b.arg(&null);
23976 }
23977 }
23978 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
23979 unsafe {
23980 b.launch(cfg)?;
23981 }
23982 Ok(())
23983 }
23984
23985 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
23987 let f = self.func("inc_i32");
23988 let cfg = LaunchConfig {
23989 grid_dim: (1, 1, 1),
23990 block_dim: (1, 1, 1),
23991 shared_mem_bytes: 0,
23992 };
23993 let __s_b = self.gpu.stream();
23994 let mut b = __s_b.launch_builder(&f);
23995 b.arg(counter);
23996 unsafe {
23997 b.launch(cfg)?;
23998 }
23999 Ok(())
24000 }
24001
24002 #[allow(clippy::too_many_arguments)]
24011 fn fa_part_pool_grow(
24015 &self,
24016 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
24017 o_len: usize,
24018 ml_len: usize,
24019 ) -> Result<(), Box<dyn std::error::Error>> {
24020 if part_guard
24021 .as_ref()
24022 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24023 .unwrap_or(true)
24024 {
24025 let old = part_guard.take();
24026 let (co, cm) = old
24027 .as_ref()
24028 .map(|pp| (pp.0.len(), pp.1.len()))
24029 .unwrap_or((0, 0));
24030 if let Some(old) = old {
24031 self.fa_part_retired.lock().unwrap().push(old);
24032 }
24033 *part_guard = Some((
24034 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
24035 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
24036 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
24037 ));
24038 }
24039 Ok(())
24040 }
24041
24042 pub fn fa_dcw_pool_ensure(
24045 &self,
24046 head_dim: usize,
24047 n_head: usize,
24048 n_head_kv: usize,
24049 bucket_max: usize,
24050 ) -> Result<(), Box<dyn std::error::Error>> {
24051 let sp = fa_split_keys(bucket_max, n_head_kv);
24052 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24053 let o_len = n_head * n_splits * head_dim;
24054 let ml_len = n_head * n_splits;
24055 let mut part_guard = self.fa_part_pool.lock().unwrap();
24056 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24057 }
24058
24059 #[allow(clippy::too_many_arguments)]
24067 pub fn fa_decode_dcw2(
24068 &self,
24069 q2: &CudaSlice<f32>,
24070 k_ring: &cudarc::driver::CudaView<u8>,
24071 v_ring: &cudarc::driver::CudaView<u8>,
24072 o2: &mut CudaSlice<f32>,
24073 head_dim: usize,
24074 n_head: usize,
24075 n_head_kv: usize,
24076 len_dev: &CudaSlice<i32>,
24077 base_dev: Option<&CudaSlice<i32>>,
24078 window: usize,
24079 bucket_max: usize,
24080 scale: f32,
24081 k_tok_bytes: usize,
24082 v_tok_bytes: usize,
24083 gate2: &CudaSlice<f32>,
24084 ) -> Result<(), Box<dyn std::error::Error>> {
24085 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24086 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24087 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24088 }
24089 let sp = fa_split_keys(bucket_max, n_head_kv);
24090 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24091 let o_len = 2 * n_head * n_splits * head_dim;
24093 let ml_len = 2 * n_head * n_splits;
24094 let mut part_guard = self.fa_part_pool.lock().unwrap();
24095 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24096 let pg = part_guard.as_mut().unwrap();
24097 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24098 let (hd, nh, nhkv, nsp) = (
24099 head_dim as i32,
24100 n_head as i32,
24101 n_head_kv as i32,
24102 n_splits as i32,
24103 );
24104 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24105 let (ski, win) = (sp as i32, window as i32);
24106 let gqa = (n_head / n_head_kv).max(1) as u32;
24107 let smem = (32 * head_dim * 2) as u32;
24108 let f = self.func("fa_decode_vec_q_v3_dcw2");
24109 let cfg = LaunchConfig {
24110 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24111 block_dim: (32, gqa, 1),
24112 shared_mem_bytes: smem,
24113 };
24114 let null: u64 = 0;
24115 {
24116 let __s_b = self.gpu.stream();
24117 let mut b = __s_b.launch_builder(&f);
24118 b.arg(q2)
24119 .arg(k_ring)
24120 .arg(v_ring)
24121 .arg(&mut *part_o)
24122 .arg(&mut *part_m)
24123 .arg(&mut *part_l)
24124 .arg(&hd)
24125 .arg(&nh)
24126 .arg(&nhkv)
24127 .arg(len_dev);
24128 match base_dev {
24129 Some(base) => {
24130 b.arg(base);
24131 }
24132 None => {
24133 b.arg(&null);
24134 }
24135 }
24136 b.arg(&win)
24137 .arg(&scale)
24138 .arg(&nsp)
24139 .arg(&ski)
24140 .arg(&ktb)
24141 .arg(&vtb);
24142 unsafe {
24143 b.launch(cfg)?;
24144 }
24145 }
24146 let fc = {
24150 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24151 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24152 self.func("fa_decode_combine_gate_f32_s")
24153 } else {
24154 self.func("fa_decode_combine_gate_f32")
24155 }
24156 };
24157 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24158 let nh2 = (2 * n_head) as i32;
24159 let cfg2 = LaunchConfig {
24160 grid_dim: ((2 * n_head) as u32, 1, 1),
24161 block_dim: (head_dim as u32, 1, 1),
24162 shared_mem_bytes: if combine_shared {
24163 (2 * n_splits * 4) as u32
24164 } else {
24165 0
24166 },
24167 };
24168 let __s_b2 = self.gpu.stream();
24169 let mut b2 = __s_b2.launch_builder(&fc);
24170 b2.arg(&*part_o)
24171 .arg(&*part_m)
24172 .arg(&*part_l)
24173 .arg(gate2)
24174 .arg(o2)
24175 .arg(&hd)
24176 .arg(&nh2)
24177 .arg(&nsp);
24178 unsafe {
24179 b2.launch(cfg2)?;
24180 }
24181 Ok(())
24182 }
24183
24184 #[allow(clippy::too_many_arguments)]
24193 pub fn fa_decode_dcw_rows(
24194 &self,
24195 q_rows: &CudaSlice<f32>,
24196 tab: &CudaSlice<u64>,
24197 o_rows: &mut CudaSlice<f32>,
24198 t: usize,
24199 head_dim: usize,
24200 n_head: usize,
24201 n_head_kv: usize,
24202 window: usize,
24203 max_ns: usize,
24204 scale: f32,
24205 k_tok_bytes: usize,
24206 v_tok_bytes: usize,
24207 gate_rows: &CudaSlice<f32>,
24208 ) -> Result<(), Box<dyn std::error::Error>> {
24209 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24210 || head_dim > 256
24211 || head_dim % 32 != 0
24212 || !fa_v3_on()
24213 {
24214 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24215 }
24216 if fa_sm_count() < 128
24217 || std::env::var("MEMRA_FA_SPLIT").is_ok()
24218 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24219 || std::env::var("MEMRA_FA_SP16").is_ok()
24220 {
24221 return Err(
24222 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24223 (or a <128-SM rig) keep the per-row path"
24224 .into(),
24225 );
24226 }
24227 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24228 return Err("fa_decode_dcw_rows geometry".into());
24229 }
24230 let o_len = t * n_head * max_ns * head_dim;
24231 let ml_len = t * n_head * max_ns;
24232 let mut part_guard = self.fa_part_pool.lock().unwrap();
24233 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24234 let pg = part_guard.as_mut().unwrap();
24235 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24236 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24237 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24238 let (win, mns) = (window as i32, max_ns as i32);
24239 let gqa = (n_head / n_head_kv).max(1) as u32;
24240 let smem = (32 * head_dim * 2) as u32;
24241 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24242 let cfg = LaunchConfig {
24243 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24244 block_dim: (32, gqa, 1),
24245 shared_mem_bytes: smem,
24246 };
24247 {
24248 let __s_b = self.gpu.stream();
24249 let mut b = __s_b.launch_builder(&f);
24250 b.arg(q_rows)
24251 .arg(tab)
24252 .arg(&mut *part_o)
24253 .arg(&mut *part_m)
24254 .arg(&mut *part_l)
24255 .arg(&hd)
24256 .arg(&nh)
24257 .arg(&nhkv)
24258 .arg(&win)
24259 .arg(&scale)
24260 .arg(&mns)
24261 .arg(&ktb)
24262 .arg(&vtb);
24263 unsafe {
24264 b.launch(cfg)?;
24265 }
24266 }
24267 let fc = {
24271 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24272 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24273 self.func("fa_decode_combine_gate_f32_s")
24274 } else {
24275 self.func("fa_decode_combine_gate_f32")
24276 }
24277 };
24278 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24279 let nht = (t * n_head) as i32;
24280 let cfg2 = LaunchConfig {
24281 grid_dim: ((t * n_head) as u32, 1, 1),
24282 block_dim: (head_dim as u32, 1, 1),
24283 shared_mem_bytes: if combine_shared {
24284 (2 * max_ns * 4) as u32
24285 } else {
24286 0
24287 },
24288 };
24289 let __s_b2 = self.gpu.stream();
24290 let mut b2 = __s_b2.launch_builder(&fc);
24291 b2.arg(&*part_o)
24292 .arg(&*part_m)
24293 .arg(&*part_l)
24294 .arg(gate_rows)
24295 .arg(o_rows)
24296 .arg(&hd)
24297 .arg(&nht)
24298 .arg(&mns);
24299 unsafe {
24300 b2.launch(cfg2)?;
24301 }
24302 Ok(())
24303 }
24304
24305 pub fn fa_decode_dcw(
24306 &self,
24307 q: &CudaSlice<f32>,
24308 k_ring: &cudarc::driver::CudaView<u8>,
24309 v_ring: &cudarc::driver::CudaView<u8>,
24310 o: &mut CudaSlice<f32>,
24311 head_dim: usize,
24312 n_head: usize,
24313 n_head_kv: usize,
24314 len_dev: &CudaSlice<i32>,
24315 base_dev: Option<&CudaSlice<i32>>,
24316 window: usize,
24317 bucket_max: usize,
24318 scale: f32,
24319 k_tok_bytes: usize,
24320 v_tok_bytes: usize,
24321 fused_gate: Option<&CudaSlice<f32>>,
24325 ) -> Result<(), Box<dyn std::error::Error>> {
24326 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24327 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24328 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"
24329 .into());
24330 }
24331 let sp = fa_split_keys(bucket_max, n_head_kv);
24332 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24333 let o_len = n_head * n_splits * head_dim;
24334 let ml_len = n_head * n_splits;
24335 let mut part_guard = self.fa_part_pool.lock().unwrap();
24336 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24337 let pg = part_guard.as_mut().unwrap();
24338 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24343 let memset_on = *MEMSET_ON
24348 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
24349 || crate::tp::token_graph_building();
24350 if memset_on {
24351 self.gpu
24352 .stream()
24353 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24354 self.gpu
24355 .stream()
24356 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24357 self.gpu
24358 .stream()
24359 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24360 }
24361 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24362 let (hd, nh, nhkv, nsp) = (
24363 head_dim as i32,
24364 n_head as i32,
24365 n_head_kv as i32,
24366 n_splits as i32,
24367 );
24368 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24369 let (ski, win) = (sp as i32, window as i32);
24370 let gqa = (n_head / n_head_kv).max(1) as u32;
24371 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24375 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
24376 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
24377 Ok("2") => 2,
24378 Ok("1") => 1,
24379 _ => 0,
24380 });
24381 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24386 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
24387 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
24388 std::sync::Mutex::new(None);
24389 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24393 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
24394 && (n_head / n_head_kv) % 2 == 0
24395 && (n_head / n_head_kv) >= 2;
24396 let f = if fprof {
24397 self.func("fa_decode_vec_q_v3_dcw_prof")
24398 } else if hs2 {
24399 self.func("fa_decode_vec_q_v3_dcw_hs2")
24400 } else if hoist == 2 {
24401 self.func("fa_decode_vec_q_v3_dcw_hc")
24403 } else if hoist == 1 {
24404 self.func("fa_decode_vec_q_v3_dcw_h")
24406 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
24407 self.func("fa_decode_vec_q_v3_dcw_u8")
24408 } else {
24409 self.func("fa_decode_vec_q_v3_dcw")
24410 };
24411 let cfg = LaunchConfig {
24412 grid_dim: if hs2 {
24413 ((2 * n_head_kv) as u32, n_splits as u32, 1)
24414 } else {
24415 (n_head_kv as u32, n_splits as u32, 1)
24416 },
24417 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
24418 shared_mem_bytes: smem,
24419 };
24420 let null: u64 = 0;
24421 let __s_b = self.gpu.stream();
24422 let mut b = __s_b.launch_builder(&f);
24423 b.arg(q)
24424 .arg(k_ring)
24425 .arg(v_ring)
24426 .arg(&mut *part_o)
24427 .arg(&mut *part_m)
24428 .arg(&mut *part_l)
24429 .arg(&hd)
24430 .arg(&nh)
24431 .arg(&nhkv)
24432 .arg(len_dev);
24433 match base_dev {
24434 Some(base) => {
24435 b.arg(base);
24436 }
24437 None => {
24438 b.arg(&null);
24439 }
24440 }
24441 b.arg(&win)
24442 .arg(&scale)
24443 .arg(&nsp)
24444 .arg(&ski)
24445 .arg(&ktb)
24446 .arg(&vtb);
24447 if fprof {
24448 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
24449 if guard
24450 .as_ref()
24451 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
24452 {
24453 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
24454 }
24455 let (_, buf) = guard.as_mut().expect("armed above");
24456 b.arg(&*buf);
24457 unsafe {
24458 b.launch(cfg)?;
24459 }
24460 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
24461 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
24462 if n % 430 == 0 {
24463 self.stream().synchronize()?;
24464 let h = self.dtoh_u64(buf)?;
24465 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
24466 let tot: u64 = h[..6].iter().sum();
24467 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
24468 for (i, name) in phases.iter().enumerate() {
24469 let pct = if tot > 0 {
24470 h[i] as f64 / tot as f64 * 100.0
24471 } else {
24472 0.0
24473 };
24474 line.push_str(&format!(" {name}={pct:.1}%"));
24475 }
24476 if h[6] > 0 {
24477 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
24478 }
24479 eprintln!("{line}");
24480 }
24481 } else {
24482 unsafe {
24483 b.launch(cfg)?;
24484 }
24485 }
24486 let mut combine_shared = false;
24487 let fc = if fused_gate.is_some() {
24488 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24491 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24492 combine_shared = true;
24493 self.func("fa_decode_combine_gate_f32_s")
24494 } else {
24495 self.func("fa_decode_combine_gate_f32")
24496 }
24497 } else {
24498 self.fa_func("fa_decode_combine_f32", head_dim)
24499 };
24500 let cfg2 = LaunchConfig {
24501 grid_dim: (n_head as u32, 1, 1),
24502 block_dim: (head_dim as u32, 1, 1),
24503 shared_mem_bytes: if combine_shared {
24504 (2 * n_splits * 4) as u32
24505 } else {
24506 0
24507 },
24508 };
24509 let __s_b2 = self.gpu.stream();
24510 let mut b2 = __s_b2.launch_builder(&fc);
24511 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
24512 if let Some(gate_row) = fused_gate {
24513 b2.arg(gate_row);
24514 }
24515 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
24516 unsafe {
24517 b2.launch(cfg2)?;
24518 }
24519 Ok(())
24520 }
24521
24522 pub fn fa_geom_eager(
24528 &self,
24529 t_kv: usize,
24530 head_dim: usize,
24531 n_head_kv: usize,
24532 g: bool,
24533 ) -> (bool, usize) {
24534 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
24538 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
24544 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
24545 if g && head_dim == 256 && !fa_v4_at(t_kv) {
24551 fa_vec = false;
24552 }
24553 let sp = fa_split_keys(t_kv, n_head_kv);
24554 let n_splits = if fa_vec {
24555 ((t_kv + sp - 1) / sp).max(1)
24556 } else {
24557 ((t_kv + 255) / 256).max(1)
24558 };
24559 (fa_vec, n_splits)
24560 }
24561
24562 pub fn fa_bucket_key(
24568 &self,
24569 t_kv: usize,
24570 head_dim: usize,
24571 n_head_kv: usize,
24572 g: bool,
24573 ) -> (bool, usize) {
24574 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
24575 }
24576
24577 pub fn capture_graph_retained<F>(
24589 &self,
24590 step: F,
24591 ) -> Result<
24592 (
24593 cudarc::driver::CudaGraph,
24594 Vec<Box<dyn std::any::Any + Send>>,
24595 ),
24596 Box<dyn std::error::Error>,
24597 >
24598 where
24599 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24600 {
24601 use cudarc::driver::sys::CUgraphInstantiate_flags;
24602 self.capture_graph_retained_flags(
24603 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24604 step,
24605 )
24606 }
24607
24608 pub fn capture_graph_retained_flags<F>(
24613 &self,
24614 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
24615 mut step: F,
24616 ) -> Result<
24617 (
24618 cudarc::driver::CudaGraph,
24619 Vec<Box<dyn std::any::Any + Send>>,
24620 ),
24621 Box<dyn std::error::Error>,
24622 >
24623 where
24624 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24625 {
24626 use cudarc::driver::sys::CUstreamCaptureMode;
24627 self.capture_keep.lock().unwrap().clear();
24635 let was_tracking = self.gpu.ctx.is_event_tracking();
24636 if was_tracking {
24637 unsafe {
24638 self.gpu.ctx.disable_event_tracking();
24639 }
24640 }
24641 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24642 self.capture_keep_on
24643 .store(true, std::sync::atomic::Ordering::Relaxed);
24644 let w = (|| {
24645 step(self)?;
24646 step(self)
24647 })();
24648 self.capture_keep_on
24649 .store(false, std::sync::atomic::Ordering::Relaxed);
24650 w?;
24651 self.gpu.stream().synchronize()?;
24652 self.gpu
24653 .stream()
24654 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24655 let r = step(self);
24656 let g = self.gpu.stream().end_capture(flags);
24657 r?;
24658 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24659 graph.upload()?;
24660 Ok(graph)
24661 };
24662 let result = run();
24663 self.capture_keep_on
24664 .store(false, std::sync::atomic::Ordering::Relaxed);
24665 if was_tracking {
24666 unsafe {
24667 self.gpu.ctx.enable_event_tracking();
24668 }
24669 }
24670 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
24671 Ok((result?, keeper))
24672 }
24673
24674 pub fn capture_graph_retained_nowarm<F>(
24680 &self,
24681 mut step: F,
24682 ) -> Result<
24683 (
24684 cudarc::driver::CudaGraph,
24685 Vec<Box<dyn std::any::Any + Send>>,
24686 ),
24687 Box<dyn std::error::Error>,
24688 >
24689 where
24690 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24691 {
24692 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24693 let was_tracking = self.gpu.ctx.is_event_tracking();
24694 if was_tracking {
24695 unsafe {
24696 self.gpu.ctx.disable_event_tracking();
24697 }
24698 }
24699 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24700 self.gpu.stream().synchronize()?;
24701 self.gpu
24702 .stream()
24703 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24704 let r = step(self);
24705 let g = self.gpu.stream().end_capture(
24706 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24707 );
24708 r?;
24709 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24710 graph.upload()?;
24711 Ok(graph)
24712 };
24713 let result = run();
24714 if was_tracking {
24715 unsafe {
24716 self.gpu.ctx.enable_event_tracking();
24717 }
24718 }
24719 Ok((result?, Vec::new()))
24720 }
24721
24722 pub fn capture_graph<F>(
24723 &self,
24724 mut step: F,
24725 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
24726 where
24727 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
24728 {
24729 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
24730 let was_tracking = self.gpu.ctx.is_event_tracking();
24738 if was_tracking {
24739 unsafe {
24740 self.gpu.ctx.disable_event_tracking();
24741 }
24742 }
24743 let iflag = {
24750 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
24751 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
24752 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
24755 Ok("priority") => {
24756 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
24757 }
24758 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
24759 })
24760 };
24761 let ct = {
24768 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24769 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
24770 };
24771 let warmups = {
24794 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24795 *W.get_or_init(|| {
24796 std::env::var("MEMRA_GRAPH_WARMUPS")
24797 .ok()
24798 .and_then(|v| v.parse().ok())
24799 .filter(|n| *n >= 1)
24800 .unwrap_or(1)
24801 })
24802 };
24803 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
24804 let t_w = std::time::Instant::now();
24805 for _ in 0..warmups {
24807 step(self)?;
24808 }
24809 self.gpu.stream().synchronize()?;
24810 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
24811 let t_c = std::time::Instant::now();
24813 self.gpu
24814 .stream()
24815 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
24816 let r = step(self);
24819 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
24820 let t_i = std::time::Instant::now();
24821 let g = self.gpu.stream().end_capture(iflag);
24822 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
24823 r?;
24824 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
24825 let t_u = std::time::Instant::now();
24826 graph.upload()?;
24827 if ct {
24828 println!(
24829 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
24830 instantiate {ms_inst:.2} ms upload {:.2} ms",
24831 t_u.elapsed().as_secs_f64() * 1e3
24832 );
24833 }
24834 Ok(graph)
24835 };
24836 let result = run();
24837 if was_tracking {
24838 unsafe {
24839 self.gpu.ctx.enable_event_tracking();
24840 }
24841 }
24842 result
24843 }
24844
24845 pub fn gdn_scan_s128_view(
24847 &self,
24848 q: &CudaSlice<f32>,
24849 k: &CudaSlice<f32>,
24850 v: &CudaSlice<f32>,
24851 g: &CudaSlice<f32>,
24852 beta: &CudaSlice<f32>,
24853 state_in: &cudarc::driver::CudaView<f32>,
24854 state_out: &mut cudarc::driver::CudaViewMut<f32>,
24855 o: &mut CudaSlice<f32>,
24856 n_head: usize,
24857 t: usize,
24858 scale: f32,
24859 ) -> Result<(), Box<dyn std::error::Error>> {
24860 let f = self.func("gdn_scan_s128");
24861 const S_V: u32 = 128;
24862 const WARP: u32 = 32;
24863 const COLS: u32 = 4;
24864 let cfg = LaunchConfig {
24865 grid_dim: (n_head as u32, 1, S_V / COLS),
24866 block_dim: (WARP, COLS, 1),
24867 shared_mem_bytes: 0,
24868 };
24869 let (h, ti) = (n_head as i32, t as i32);
24870 let __s_b = self.gpu.stream();
24871 let mut b = __s_b.launch_builder(&f);
24872 b.arg(q)
24873 .arg(k)
24874 .arg(v)
24875 .arg(g)
24876 .arg(beta)
24877 .arg(state_in)
24878 .arg(state_out)
24879 .arg(o)
24880 .arg(&h)
24881 .arg(&ti)
24882 .arg(&scale);
24883 unsafe {
24884 b.launch(cfg)?;
24885 }
24886 Ok(())
24887 }
24888
24889 pub fn ssm_conv1d_view(
24891 &self,
24892 x: &cudarc::driver::CudaView<f32>,
24893 w: &CudaSlice<f32>,
24894 y: &mut CudaSlice<f32>,
24895 conv_dim: usize,
24896 t: usize,
24897 d_conv: usize,
24898 silu: bool,
24899 ) -> Result<(), Box<dyn std::error::Error>> {
24900 let f = self.func("ssm_conv1d_silu_f32");
24901 let cfg = LaunchConfig {
24903 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24904 block_dim: (256, 1, 1),
24905 shared_mem_bytes: 0,
24906 };
24907 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24908 let __s_b = self.gpu.stream();
24909 let mut b = __s_b.launch_builder(&f);
24910 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24911 unsafe {
24912 b.launch(cfg)?;
24913 }
24914 Ok(())
24915 }
24916
24917 pub fn ssm_conv1d_tm(
24924 &self,
24925 qkv_tm: &CudaSlice<f32>,
24926 w: &CudaSlice<f32>,
24927 y: &mut CudaSlice<f32>,
24928 conv_dim: usize,
24929 t: usize,
24930 d_conv: usize,
24931 ) -> Result<(), Box<dyn std::error::Error>> {
24932 let f = self.func("ssm_conv1d_tm_f32");
24933 let cfg = LaunchConfig {
24934 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24935 block_dim: (256, 1, 1),
24936 shared_mem_bytes: 0,
24937 };
24938 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24939 let __s_b = self.gpu.stream();
24940 let mut b = __s_b.launch_builder(&f);
24941 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
24942 unsafe {
24943 b.launch(cfg)?;
24944 }
24945 Ok(())
24946 }
24947
24948 pub fn ssm_conv1d_tm_state(
24956 &self,
24957 qkv_tm: &CudaSlice<f32>,
24958 conv_state: &mut CudaSlice<f32>,
24959 w: &CudaSlice<f32>,
24960 y: &mut CudaSlice<f32>,
24961 conv_dim: usize,
24962 t: usize,
24963 d_conv: usize,
24964 ) -> Result<(), Box<dyn std::error::Error>> {
24965 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
24966 }
24967
24968 #[allow(clippy::too_many_arguments)]
24971 pub fn ssm_conv1d_tm_state_pad(
24972 &self,
24973 qkv_tm: &CudaSlice<f32>,
24974 conv_state: &mut CudaSlice<f32>,
24975 w: &CudaSlice<f32>,
24976 y: &mut CudaSlice<f32>,
24977 conv_dim: usize,
24978 t: usize,
24979 d_conv: usize,
24980 pad_len: Option<&CudaSlice<i32>>,
24981 ) -> Result<(), Box<dyn std::error::Error>> {
24982 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
24983 let ring_old = if t < d_conv - 1 {
24987 Some(self.clone_dtod(conv_state)?)
24988 } else {
24989 None
24990 };
24991 {
24992 let f = self.func("ssm_conv1d_tm_state_f32");
24993 let cfg = LaunchConfig {
24994 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24995 block_dim: (256, 1, 1),
24996 shared_mem_bytes: 0,
24997 };
24998 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24999 let __s_b = self.gpu.stream();
25000 let mut b = __s_b.launch_builder(&f);
25001 b.arg(qkv_tm)
25002 .arg(&*conv_state)
25003 .arg(w)
25004 .arg(y)
25005 .arg(&cd)
25006 .arg(&ti)
25007 .arg(&dc);
25008 unsafe {
25009 b.launch(cfg)?;
25010 }
25011 }
25012 match (ring_old, pad_len) {
25013 (None, Some(len_d)) => {
25014 let f = self.func("ssm_conv_ring_update_dev_f32");
25015 let n = conv_dim * (d_conv - 1);
25016 let cfg = LaunchConfig::for_num_elems(n as u32);
25017 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25018 let __s_b = self.gpu.stream();
25019 let mut b = __s_b.launch_builder(&f);
25020 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25021 unsafe {
25022 b.launch(cfg)?;
25023 }
25024 }
25025 (None, None) => {
25026 let f = self.func("ssm_conv_ring_update_f32");
25027 let n = conv_dim * (d_conv - 1);
25028 let cfg = LaunchConfig::for_num_elems(n as u32);
25029 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25030 let __s_b = self.gpu.stream();
25031 let mut b = __s_b.launch_builder(&f);
25032 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25033 unsafe {
25034 b.launch(cfg)?;
25035 }
25036 }
25037 (Some(old), _) => {
25038 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
25039 }
25040 }
25041 Ok(())
25042 }
25043
25044 pub fn ssm_conv1d_tm_state_pad_v(
25046 &self,
25047 qkv_tm: &cudarc::driver::CudaView<f32>,
25048 conv_state: &mut CudaSlice<f32>,
25049 w: &CudaSlice<f32>,
25050 y: &mut CudaSlice<f32>,
25051 conv_dim: usize,
25052 t: usize,
25053 d_conv: usize,
25054 pad_len: Option<&CudaSlice<i32>>,
25055 ) -> Result<(), Box<dyn std::error::Error>> {
25056 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25057 let ring_old = if t < d_conv - 1 {
25061 Some(self.clone_dtod(conv_state)?)
25062 } else {
25063 None
25064 };
25065 {
25066 let f = self.func("ssm_conv1d_tm_state_f32");
25067 let cfg = LaunchConfig {
25068 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25069 block_dim: (256, 1, 1),
25070 shared_mem_bytes: 0,
25071 };
25072 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25073 let __s_b = self.gpu.stream();
25074 let mut b = __s_b.launch_builder(&f);
25075 b.arg(qkv_tm)
25076 .arg(&*conv_state)
25077 .arg(w)
25078 .arg(y)
25079 .arg(&cd)
25080 .arg(&ti)
25081 .arg(&dc);
25082 unsafe {
25083 b.launch(cfg)?;
25084 }
25085 }
25086 match (ring_old, pad_len) {
25087 (None, Some(len_d)) => {
25088 let f = self.func("ssm_conv_ring_update_dev_f32");
25089 let n = conv_dim * (d_conv - 1);
25090 let cfg = LaunchConfig::for_num_elems(n as u32);
25091 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25092 let __s_b = self.gpu.stream();
25093 let mut b = __s_b.launch_builder(&f);
25094 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25095 unsafe {
25096 b.launch(cfg)?;
25097 }
25098 }
25099 (None, None) => {
25100 let f = self.func("ssm_conv_ring_update_f32");
25101 let n = conv_dim * (d_conv - 1);
25102 let cfg = LaunchConfig::for_num_elems(n as u32);
25103 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25104 let __s_b = self.gpu.stream();
25105 let mut b = __s_b.launch_builder(&f);
25106 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25107 unsafe {
25108 b.launch(cfg)?;
25109 }
25110 }
25111 (Some(_), _) => unreachable!(
25112 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25113 ),
25114 }
25115 Ok(())
25116 }
25117
25118 pub fn ssm_conv_ring_rebuild(
25123 &self,
25124 qkv_tm: &CudaSlice<f32>,
25125 ring_old: &CudaSlice<f32>,
25126 conv_state: &mut CudaSlice<f32>,
25127 conv_dim: usize,
25128 tc: usize,
25129 d_conv: usize,
25130 ) -> Result<(), Box<dyn std::error::Error>> {
25131 let f = self.func("ssm_conv_ring_rebuild_f32");
25132 let n = conv_dim * (d_conv - 1);
25133 let cfg = LaunchConfig::for_num_elems(n as u32);
25134 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25135 let __s_b = self.gpu.stream();
25136 let mut b = __s_b.launch_builder(&f);
25137 b.arg(qkv_tm)
25138 .arg(ring_old)
25139 .arg(conv_state)
25140 .arg(&cd)
25141 .arg(&ti)
25142 .arg(&dc);
25143 unsafe {
25144 b.launch(cfg)?;
25145 }
25146 Ok(())
25147 }
25148
25149 #[allow(clippy::too_many_arguments)]
25154 pub fn gdn_prep_decode(
25155 &self,
25156 conv_out: &CudaSlice<f32>,
25157 beta_raw: &CudaSlice<f32>,
25158 alpha: &CudaSlice<f32>,
25159 dt_bias: &CudaSlice<f32>,
25160 a: &CudaSlice<f32>,
25161 q_l2: &mut CudaSlice<f32>,
25162 k_l2: &mut CudaSlice<f32>,
25163 v_g: &mut CudaSlice<f32>,
25164 beta: &mut CudaSlice<f32>,
25165 g_log: &mut CudaSlice<f32>,
25166 d_state: usize,
25167 num_v: usize,
25168 num_k: usize,
25169 key_dim: usize,
25170 eps: f32,
25171 ) -> Result<(), Box<dyn std::error::Error>> {
25172 let f = self.func("gdn_prep_decode_f32");
25173 let cfg = LaunchConfig {
25174 grid_dim: (num_v as u32, 1, 1),
25175 block_dim: (32, 4, 1),
25176 shared_mem_bytes: 0,
25177 };
25178 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25179 let __s_b = self.gpu.stream();
25180 let mut b = __s_b.launch_builder(&f);
25181 b.arg(conv_out)
25182 .arg(beta_raw)
25183 .arg(alpha)
25184 .arg(dt_bias)
25185 .arg(a)
25186 .arg(q_l2)
25187 .arg(k_l2)
25188 .arg(v_g)
25189 .arg(beta)
25190 .arg(g_log)
25191 .arg(&ds)
25192 .arg(&nv)
25193 .arg(&nk)
25194 .arg(&kd)
25195 .arg(&eps);
25196 unsafe {
25197 b.launch(cfg)?;
25198 }
25199 Ok(())
25200 }
25201
25202 #[allow(clippy::too_many_arguments)]
25206 pub fn ssm_conv1d_gdn(
25207 &self,
25208 qkv_tm: &CudaSlice<f32>,
25209 w: &CudaSlice<f32>,
25210 q_g: &mut CudaSlice<f32>,
25211 k_g: &mut CudaSlice<f32>,
25212 v_g: &mut CudaSlice<f32>,
25213 conv_dim: usize,
25214 t: usize,
25215 d_conv: usize,
25216 d_state: usize,
25217 num_v: usize,
25218 num_k: usize,
25219 key_dim: usize,
25220 ) -> Result<(), Box<dyn std::error::Error>> {
25221 let f = self.func("ssm_conv1d_gdn_f32");
25222 let cfg = LaunchConfig {
25223 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25224 block_dim: (256, 1, 1),
25225 shared_mem_bytes: 0,
25226 };
25227 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25228 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25229 let __s_b = self.gpu.stream();
25230 let mut b = __s_b.launch_builder(&f);
25231 b.arg(qkv_tm)
25232 .arg(w)
25233 .arg(q_g)
25234 .arg(k_g)
25235 .arg(v_g)
25236 .arg(&cd)
25237 .arg(&ti)
25238 .arg(&dc)
25239 .arg(&ds)
25240 .arg(&nv)
25241 .arg(&nk)
25242 .arg(&kd);
25243 unsafe {
25244 b.launch(cfg)?;
25245 }
25246 Ok(())
25247 }
25248
25249 pub fn ssm_conv1d(
25250 &self,
25251 x: &CudaSlice<f32>,
25252 w: &CudaSlice<f32>,
25253 y: &mut CudaSlice<f32>,
25254 conv_dim: usize,
25255 t: usize,
25256 d_conv: usize,
25257 silu: bool,
25258 ) -> Result<(), Box<dyn std::error::Error>> {
25259 let f = self.func("ssm_conv1d_silu_f32");
25260 let cfg = LaunchConfig {
25261 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25262 block_dim: (256, 1, 1),
25263 shared_mem_bytes: 0,
25264 };
25265 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25266 let __s_b = self.gpu.stream();
25267 let mut b = __s_b.launch_builder(&f);
25268 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25269 unsafe {
25270 b.launch(cfg)?;
25271 }
25272 Ok(())
25273 }
25274
25275 pub fn gdn_scan_s128(
25278 &self,
25279 q: &CudaSlice<f32>,
25280 k: &CudaSlice<f32>,
25281 v: &CudaSlice<f32>,
25282 g: &CudaSlice<f32>,
25283 beta: &CudaSlice<f32>,
25284 state_in: &CudaSlice<f32>,
25285 state_out: &mut CudaSlice<f32>,
25286 o: &mut CudaSlice<f32>,
25287 n_head: usize,
25288 t: usize,
25289 scale: f32,
25290 ) -> Result<(), Box<dyn std::error::Error>> {
25291 let f = self.func("gdn_scan_s128");
25292 const S_V: u32 = 128;
25293 const WARP: u32 = 32;
25294 const COLS_PER_BLOCK: u32 = 4;
25295 let cfg = LaunchConfig {
25296 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25297 block_dim: (WARP, COLS_PER_BLOCK, 1),
25298 shared_mem_bytes: 0,
25299 };
25300 let (h, ti) = (n_head as i32, t as i32);
25301 let __s_b = self.gpu.stream();
25302 let mut b = __s_b.launch_builder(&f);
25303 b.arg(q)
25304 .arg(k)
25305 .arg(v)
25306 .arg(g)
25307 .arg(beta)
25308 .arg(state_in)
25309 .arg(state_out)
25310 .arg(o)
25311 .arg(&h)
25312 .arg(&ti)
25313 .arg(&scale);
25314 unsafe {
25315 b.launch(cfg)?;
25316 }
25317 Ok(())
25318 }
25319
25320 #[allow(clippy::too_many_arguments)]
25325 pub fn ssm_conv1d_fused_decode_b(
25326 &self,
25327 qkv_cols: &CudaSlice<f32>,
25328 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25329 w: &CudaSlice<f32>,
25330 conv_outs: &mut CudaSlice<f32>,
25331 conv_dim: usize,
25332 d_conv: usize,
25333 b_n: usize,
25334 ) -> Result<(), Box<dyn std::error::Error>> {
25335 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25336 let cfg = LaunchConfig {
25337 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25338 block_dim: (256, 1, 1),
25339 shared_mem_bytes: 0,
25340 };
25341 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25342 let __s_b = self.gpu.stream();
25343 let mut b = __s_b.launch_builder(&f);
25344 b.arg(qkv_cols)
25345 .arg(conv_state_ptrs)
25346 .arg(w)
25347 .arg(conv_outs)
25348 .arg(&cd)
25349 .arg(&dc);
25350 unsafe {
25351 b.launch(cfg)?;
25352 }
25353 Ok(())
25354 }
25355
25356 #[allow(clippy::too_many_arguments)]
25357 pub fn gdn_prep_decode_b(
25358 &self,
25359 conv_outs: &CudaSlice<f32>,
25360 beta_raws: &CudaSlice<f32>,
25361 alphas: &CudaSlice<f32>,
25362 dt_bias: &CudaSlice<f32>,
25363 a: &CudaSlice<f32>,
25364 q_l2: &mut CudaSlice<f32>,
25365 k_l2: &mut CudaSlice<f32>,
25366 v_g: &mut CudaSlice<f32>,
25367 beta: &mut CudaSlice<f32>,
25368 g_log: &mut CudaSlice<f32>,
25369 d_state: usize,
25370 num_v: usize,
25371 num_k: usize,
25372 key_dim: usize,
25373 eps: f32,
25374 conv_dim: usize,
25375 b_n: usize,
25376 ) -> Result<(), Box<dyn std::error::Error>> {
25377 let f = self.func("gdn_prep_decode_b_f32");
25378 let cfg = LaunchConfig {
25379 grid_dim: (num_v as u32, 1, b_n as u32),
25380 block_dim: (32, 4, 1),
25381 shared_mem_bytes: 0,
25382 };
25383 let (ds, nv, nk, kd, cd) = (
25384 d_state as i32,
25385 num_v as i32,
25386 num_k as i32,
25387 key_dim as i32,
25388 conv_dim as i32,
25389 );
25390 let __s_b = self.gpu.stream();
25391 let mut b = __s_b.launch_builder(&f);
25392 b.arg(conv_outs)
25393 .arg(beta_raws)
25394 .arg(alphas)
25395 .arg(dt_bias)
25396 .arg(a)
25397 .arg(q_l2)
25398 .arg(k_l2)
25399 .arg(v_g)
25400 .arg(beta)
25401 .arg(g_log)
25402 .arg(&ds)
25403 .arg(&nv)
25404 .arg(&nk)
25405 .arg(&kd)
25406 .arg(&eps)
25407 .arg(&cd);
25408 unsafe {
25409 b.launch(cfg)?;
25410 }
25411 Ok(())
25412 }
25413
25414 #[allow(clippy::too_many_arguments)]
25415 pub fn gdn_scan_s128_batched(
25416 &self,
25417 q: &CudaSlice<f32>,
25418 k: &CudaSlice<f32>,
25419 v: &CudaSlice<f32>,
25420 g: &CudaSlice<f32>,
25421 beta: &CudaSlice<f32>,
25422 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25423 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25424 o: &mut CudaSlice<f32>,
25425 n_head: usize,
25426 b_n: usize,
25427 scale: f32,
25428 ) -> Result<(), Box<dyn std::error::Error>> {
25429 let f = self.func("gdn_scan_s128_b");
25430 const S_V: u32 = 128;
25431 const WARP: u32 = 32;
25432 const COLS_PER_BLOCK: u32 = 4;
25433 let cfg = LaunchConfig {
25434 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25435 block_dim: (WARP, COLS_PER_BLOCK, 1),
25436 shared_mem_bytes: 0,
25437 };
25438 let h = n_head as i32;
25439 let __s_b = self.gpu.stream();
25440 let mut b = __s_b.launch_builder(&f);
25441 b.arg(q)
25442 .arg(k)
25443 .arg(v)
25444 .arg(g)
25445 .arg(beta)
25446 .arg(state_in_ptrs)
25447 .arg(state_out_ptrs)
25448 .arg(o)
25449 .arg(&h)
25450 .arg(&scale);
25451 unsafe {
25452 b.launch(cfg)?;
25453 }
25454 Ok(())
25455 }
25456
25457 #[allow(clippy::too_many_arguments)]
25463 pub fn ssm_conv1d_fused_decode_b_view(
25464 &self,
25465 qkv_cols: &cudarc::driver::CudaView<f32>,
25466 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
25467 w: &CudaSlice<f32>,
25468 conv_outs: &mut CudaSlice<f32>,
25469 conv_dim: usize,
25470 d_conv: usize,
25471 b_n: usize,
25472 ) -> Result<(), Box<dyn std::error::Error>> {
25473 let f = self.func("ssm_conv1d_fused_decode_b_f32");
25474 let cfg = LaunchConfig {
25475 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
25476 block_dim: (256, 1, 1),
25477 shared_mem_bytes: 0,
25478 };
25479 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25480 let __s_b = self.gpu.stream();
25481 let mut b = __s_b.launch_builder(&f);
25482 b.arg(qkv_cols)
25483 .arg(conv_state_ptrs)
25484 .arg(w)
25485 .arg(conv_outs)
25486 .arg(&cd)
25487 .arg(&dc);
25488 unsafe {
25489 b.launch(cfg)?;
25490 }
25491 Ok(())
25492 }
25493
25494 #[allow(clippy::too_many_arguments)]
25495 pub fn gdn_prep_decode_b_view(
25496 &self,
25497 conv_outs: &CudaSlice<f32>,
25498 beta_raws: &cudarc::driver::CudaView<f32>,
25499 alphas: &cudarc::driver::CudaView<f32>,
25500 dt_bias: &CudaSlice<f32>,
25501 a: &CudaSlice<f32>,
25502 q_l2: &mut CudaSlice<f32>,
25503 k_l2: &mut CudaSlice<f32>,
25504 v_g: &mut CudaSlice<f32>,
25505 beta: &mut CudaSlice<f32>,
25506 g_log: &mut CudaSlice<f32>,
25507 d_state: usize,
25508 num_v: usize,
25509 num_k: usize,
25510 key_dim: usize,
25511 eps: f32,
25512 conv_dim: usize,
25513 b_n: usize,
25514 ) -> Result<(), Box<dyn std::error::Error>> {
25515 let f = self.func("gdn_prep_decode_b_f32");
25516 let cfg = LaunchConfig {
25517 grid_dim: (num_v as u32, 1, b_n as u32),
25518 block_dim: (32, 4, 1),
25519 shared_mem_bytes: 0,
25520 };
25521 let (ds, nv, nk, kd, cd) = (
25522 d_state as i32,
25523 num_v as i32,
25524 num_k as i32,
25525 key_dim as i32,
25526 conv_dim as i32,
25527 );
25528 let __s_b = self.gpu.stream();
25529 let mut b = __s_b.launch_builder(&f);
25530 b.arg(conv_outs)
25531 .arg(beta_raws)
25532 .arg(alphas)
25533 .arg(dt_bias)
25534 .arg(a)
25535 .arg(q_l2)
25536 .arg(k_l2)
25537 .arg(v_g)
25538 .arg(beta)
25539 .arg(g_log)
25540 .arg(&ds)
25541 .arg(&nv)
25542 .arg(&nk)
25543 .arg(&kd)
25544 .arg(&eps)
25545 .arg(&cd);
25546 unsafe {
25547 b.launch(cfg)?;
25548 }
25549 Ok(())
25550 }
25551
25552 #[allow(clippy::too_many_arguments)]
25553 pub fn gdn_scan_s128_batched_view(
25554 &self,
25555 q: &CudaSlice<f32>,
25556 k: &CudaSlice<f32>,
25557 v: &CudaSlice<f32>,
25558 g: &CudaSlice<f32>,
25559 beta: &CudaSlice<f32>,
25560 state_in_ptrs: &cudarc::driver::CudaView<u64>,
25561 state_out_ptrs: &cudarc::driver::CudaView<u64>,
25562 o: &mut cudarc::driver::CudaViewMut<f32>,
25563 n_head: usize,
25564 b_n: usize,
25565 scale: f32,
25566 ) -> Result<(), Box<dyn std::error::Error>> {
25567 let f = self.func("gdn_scan_s128_b");
25568 const S_V: u32 = 128;
25569 const WARP: u32 = 32;
25570 const COLS_PER_BLOCK: u32 = 4;
25571 let cfg = LaunchConfig {
25572 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
25573 block_dim: (WARP, COLS_PER_BLOCK, 1),
25574 shared_mem_bytes: 0,
25575 };
25576 let h = n_head as i32;
25577 let __s_b = self.gpu.stream();
25578 let mut b = __s_b.launch_builder(&f);
25579 b.arg(q)
25580 .arg(k)
25581 .arg(v)
25582 .arg(g)
25583 .arg(beta)
25584 .arg(state_in_ptrs)
25585 .arg(state_out_ptrs)
25586 .arg(o)
25587 .arg(&h)
25588 .arg(&scale);
25589 unsafe {
25590 b.launch(cfg)?;
25591 }
25592 Ok(())
25593 }
25594
25595 pub fn gdn_chunked_enabled() -> bool {
25604 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25605 *E.get_or_init(|| {
25606 std::env::var("MEMRA_GDN_CHUNKED")
25607 .map(|v| v != "0")
25608 .unwrap_or(true)
25609 })
25610 }
25611
25612 pub fn gdn_chunk_size() -> usize {
25617 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25618 *C.get_or_init(|| {
25619 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
25620 .ok()
25621 .and_then(|v| v.parse().ok())
25622 .unwrap_or(32);
25623 c.clamp(32, 128) / 32 * 32
25624 })
25625 }
25626
25627 #[allow(clippy::too_many_arguments)]
25632 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
25635 #[allow(clippy::too_many_arguments)]
25636 pub fn gdn_chunk_k123(
25637 &self,
25638 q: &CudaSlice<f32>,
25639 k: &CudaSlice<f32>,
25640 v: &CudaSlice<f32>,
25641 g: &CudaSlice<f32>,
25642 beta: &CudaSlice<f32>,
25643 wb16: Option<&mut CudaSlice<u8>>,
25644 n_head: usize,
25645 t: usize,
25646 c: usize,
25647 hk: usize,
25648 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
25649 ) -> Result<
25650 (
25651 CudaSlice<f32>,
25652 CudaSlice<f32>,
25653 CudaSlice<f32>,
25654 CudaSlice<f32>,
25655 ),
25656 Box<dyn std::error::Error>,
25657 > {
25658 const D: usize = 128;
25659 let h = n_head;
25660 let nc = (t + c - 1) / c;
25661 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25662 let mut gcum = self.uninit(t * h)?;
25663 let mut a = self.uninit(nc * h * c * c)?;
25664 let mut p = self.uninit(nc * h * c * c)?;
25665 let mut u = self.uninit(nc * h * c * D)?;
25666 let mut w = self.uninit(nc * h * c * D)?;
25667 {
25668 let f = self.func("gdn_chunk_cumgate_f32");
25670 let cfg = LaunchConfig {
25671 grid_dim: (nc as u32, h as u32, 1),
25672 block_dim: (32, 1, 1),
25673 shared_mem_bytes: 0,
25674 };
25675 let __s_b = self.gpu.stream();
25676 let mut b = __s_b.launch_builder(&f);
25677 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
25678 unsafe {
25679 b.launch(cfg)?;
25680 }
25681 }
25682 if let Some((qb, kb, pb)) = k2w {
25683 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
25686 let f = self.func("gdn_k2_wgmma");
25687 let cfg = LaunchConfig {
25688 grid_dim: (nc as u32, h as u32, 1),
25689 block_dim: (128, 1, 1),
25690 shared_mem_bytes: 0,
25691 };
25692 let hki = hk as i32;
25693 let __s_b = self.gpu.stream();
25694 let mut b = __s_b.launch_builder(&f);
25695 b.arg(qb)
25696 .arg(kb)
25697 .arg(&gcum)
25698 .arg(beta)
25699 .arg(&mut a)
25700 .arg(&mut *pb)
25701 .arg(&hi)
25702 .arg(&ti)
25703 .arg(&ci)
25704 .arg(&hki);
25705 unsafe {
25706 b.launch(cfg)?;
25707 }
25708 } else if c <= 64 && !portable_mma_gated() {
25709 let f = self.func("gdn_chunk_attn_f32");
25711 f.set_attribute(
25712 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
25713 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
25714 )?;
25715 let jt = ((c + 31) / 32) as u32;
25716 let cfg = LaunchConfig {
25717 grid_dim: (nc as u32, h as u32, jt),
25718 block_dim: (256, 1, 1),
25719 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
25720 };
25721 let hki = hk as i32;
25722 let __s_b = self.gpu.stream();
25723 let mut b = __s_b.launch_builder(&f);
25724 b.arg(q)
25725 .arg(k)
25726 .arg(&gcum)
25727 .arg(beta)
25728 .arg(&mut a)
25729 .arg(&mut p)
25730 .arg(&hi)
25731 .arg(&ti)
25732 .arg(&ci)
25733 .arg(&hki);
25734 unsafe {
25735 b.launch(cfg)?;
25736 }
25737 } else {
25738 assert!(
25740 hk == h,
25741 "generic K2 is broadcast-only (de-broadcast rides C==32)"
25742 );
25743 let f = self.func("gdn_chunk_attn_g_f32");
25744 let cfg = LaunchConfig {
25745 grid_dim: (nc as u32, h as u32, 1),
25746 block_dim: (32, 8, 1),
25747 shared_mem_bytes: 0,
25748 };
25749 let __s_b = self.gpu.stream();
25750 let mut b = __s_b.launch_builder(&f);
25751 b.arg(q)
25752 .arg(k)
25753 .arg(&gcum)
25754 .arg(beta)
25755 .arg(&mut a)
25756 .arg(&mut p)
25757 .arg(&hi)
25758 .arg(&ti)
25759 .arg(&ci);
25760 unsafe {
25761 b.launch(cfg)?;
25762 }
25763 }
25764 {
25765 let cfg = LaunchConfig {
25767 grid_dim: (nc as u32, h as u32, 1),
25768 block_dim: (256, 1, 1),
25769 shared_mem_bytes: 0,
25770 };
25771 match c {
25772 32 | 64 => {
25773 let f = self.func(if c == 32 {
25774 "gdn_chunk_solve32_f32"
25775 } else {
25776 "gdn_chunk_solve64_f32"
25777 });
25778 let wb: u64 = match wb16 {
25780 Some(d) => self.addr_u8(d),
25781 None => 0,
25782 };
25783 let hki = hk as i32;
25784 let __s_b = self.gpu.stream();
25785 let mut b = __s_b.launch_builder(&f);
25786 b.arg(v)
25787 .arg(k)
25788 .arg(&a)
25789 .arg(&gcum)
25790 .arg(&mut u)
25791 .arg(&mut w)
25792 .arg(&wb)
25793 .arg(&hi)
25794 .arg(&ti)
25795 .arg(&hki);
25796 unsafe {
25797 b.launch(cfg)?;
25798 }
25799 }
25800 _ => {
25801 assert!(hk == h, "generic K3 is broadcast-only");
25802 let f = self.func("gdn_chunk_solve_f32");
25803 let __s_b = self.gpu.stream();
25804 let mut b = __s_b.launch_builder(&f);
25805 b.arg(v)
25806 .arg(k)
25807 .arg(&a)
25808 .arg(&gcum)
25809 .arg(&mut u)
25810 .arg(&mut w)
25811 .arg(&hi)
25812 .arg(&ti)
25813 .arg(&ci);
25814 unsafe {
25815 b.launch(cfg)?;
25816 }
25817 }
25818 }
25819 }
25820 Ok((gcum, p, u, w))
25821 }
25822
25823 pub fn gdn_db_on() -> bool {
25827 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
25828 }
25829
25830 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
25839 !portable_mma_gated()
25840 && c == 32
25841 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25842 Ok("1") => true,
25843 Ok("0") => false,
25844 _ => gdn_mma_default_on(),
25845 }
25846 }
25847
25848 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
25855 cfg!(memra_hopper_mma)
25856 && self.gdn_mma_enabled(c)
25857 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
25858 }
25859
25860 #[allow(clippy::too_many_arguments)]
25865 pub fn ssm_conv1d_gdn_state_pad(
25866 &self,
25867 qkv_tm: &cudarc::driver::CudaView<f32>,
25868 conv_state: &mut CudaSlice<f32>,
25869 w: &CudaSlice<f32>,
25870 q_g: &mut CudaSlice<f32>,
25871 k_g: &mut CudaSlice<f32>,
25872 v_g: &mut CudaSlice<f32>,
25873 conv_dim: usize,
25874 t: usize,
25875 d_conv: usize,
25876 d_state: usize,
25877 num_v: usize,
25878 num_k: usize,
25879 key_dim: usize,
25880 hk: usize,
25881 pad_len: Option<&CudaSlice<i32>>,
25882 ) -> Result<(), Box<dyn std::error::Error>> {
25883 assert!(
25884 t >= d_conv - 1,
25885 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
25886 );
25887 {
25888 let f = self.func("ssm_conv1d_gdn_state_f32");
25889 let cfg = LaunchConfig {
25890 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25891 block_dim: (256, 1, 1),
25892 shared_mem_bytes: 0,
25893 };
25894 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25895 let (ds, nv, nk, kd, hki) = (
25896 d_state as i32,
25897 num_v as i32,
25898 num_k as i32,
25899 key_dim as i32,
25900 hk as i32,
25901 );
25902 let __s_b = self.gpu.stream();
25903 let mut b = __s_b.launch_builder(&f);
25904 b.arg(qkv_tm)
25905 .arg(&*conv_state)
25906 .arg(w)
25907 .arg(q_g)
25908 .arg(k_g)
25909 .arg(v_g)
25910 .arg(&cd)
25911 .arg(&ti)
25912 .arg(&dc)
25913 .arg(&ds)
25914 .arg(&nv)
25915 .arg(&nk)
25916 .arg(&kd)
25917 .arg(&hki);
25918 unsafe {
25919 b.launch(cfg)?;
25920 }
25921 }
25922 match pad_len {
25923 Some(len_d) => {
25924 let f = self.func("ssm_conv_ring_update_dev_f32");
25925 let n = conv_dim * (d_conv - 1);
25926 let cfg = LaunchConfig::for_num_elems(n as u32);
25927 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25928 let __s_b = self.gpu.stream();
25929 let mut b = __s_b.launch_builder(&f);
25930 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25931 unsafe {
25932 b.launch(cfg)?;
25933 }
25934 }
25935 None => {
25936 let f = self.func("ssm_conv_ring_update_f32");
25937 let n = conv_dim * (d_conv - 1);
25938 let cfg = LaunchConfig::for_num_elems(n as u32);
25939 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25940 let __s_b = self.gpu.stream();
25941 let mut b = __s_b.launch_builder(&f);
25942 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25943 unsafe {
25944 b.launch(cfg)?;
25945 }
25946 }
25947 }
25948 Ok(())
25949 }
25950
25951 pub fn gdn_chunk_alloc(
25955 &self,
25956 n_head: usize,
25957 t: usize,
25958 c: usize,
25959 hk: usize,
25960 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
25961 const D: usize = 128;
25962 assert!(
25963 c == 32,
25964 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
25965 );
25966 let h = n_head;
25967 let nc = (t + c - 1) / c;
25968 Ok(GdnChunkBufs {
25969 gcum: self.uninit(t * h)?,
25970 a: self.uninit(nc * h * c * c)?,
25971 p: self.uninit(nc * h * c * c)?,
25972 u: self.uninit(nc * h * c * D)?,
25973 w: self.uninit(nc * h * c * D)?,
25974 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25975 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25976 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
25977 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
25978 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
25979 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
25980 o: self.uninit(D * h * t)?,
25981 t,
25982 nc,
25983 })
25984 }
25985
25986 pub fn f32_to_bf16_v(
25988 &self,
25989 x: &cudarc::driver::CudaView<f32>,
25990 dst: &mut CudaSlice<u8>,
25991 n: usize,
25992 ) -> Result<(), Box<dyn std::error::Error>> {
25993 let f = self.func("f32_to_bf16_bulk");
25994 let ni = n as i64;
25995 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
25996 let __s_b = self.gpu.stream();
25997 let mut b = __s_b.launch_builder(&f);
25998 b.arg(x).arg(dst).arg(&ni);
25999 unsafe {
26000 b.launch(cfg)?;
26001 }
26002 Ok(())
26003 }
26004
26005 pub fn f32_to_bf16_into(
26007 &self,
26008 x: &CudaSlice<f32>,
26009 dst: &mut CudaSlice<u8>,
26010 n: usize,
26011 ) -> Result<(), Box<dyn std::error::Error>> {
26012 let f = self.func("f32_to_bf16_bulk");
26013 let ni = n as i64;
26014 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26015 let __s_b = self.gpu.stream();
26016 let mut b = __s_b.launch_builder(&f);
26017 b.arg(x).arg(dst).arg(&ni);
26018 unsafe {
26019 b.launch(cfg)?;
26020 }
26021 Ok(())
26022 }
26023
26024 pub fn gdn_chunk_k123_vl8(
26027 &self,
26028 seqs: &[GdnSeqVl],
26029 n_head: usize,
26030 hk: usize,
26031 wq: Option<&GdnWVl8>,
26032 ) -> Result<(), Box<dyn std::error::Error>> {
26033 let b = seqs.len();
26034 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
26035 let mut packed = [GdnSeqVl::default(); 8];
26036 packed[..b].copy_from_slice(seqs);
26037 let v = GdnVl8(packed);
26038 let (hi, ci) = (n_head as i32, 32i32);
26039 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26040 {
26041 let f = self.func("gdn_chunk_cumgate_vl");
26042 let cfg = LaunchConfig {
26043 grid_dim: (max_nc, n_head as u32, b as u32),
26044 block_dim: (32, 1, 1),
26045 shared_mem_bytes: 0,
26046 };
26047 let __s_lb = self.gpu.stream();
26048 let mut lb = __s_lb.launch_builder(&f);
26049 lb.arg(&v).arg(&hi).arg(&ci);
26050 unsafe {
26051 lb.launch(cfg)?;
26052 }
26053 }
26054 let hki = hk as i32;
26055 if let Some(w) = wq {
26056 let f = self.func("gdn_k2_wgmma_vl");
26058 let cfg = LaunchConfig {
26059 grid_dim: (max_nc, n_head as u32, b as u32),
26060 block_dim: (128, 1, 1),
26061 shared_mem_bytes: 0,
26062 };
26063 let __s_lb = self.gpu.stream();
26064 let mut lb = __s_lb.launch_builder(&f);
26065 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26066 unsafe {
26067 lb.launch(cfg)?;
26068 }
26069 } else {
26070 let f = self.func("gdn_chunk_attn_vl");
26071 f.set_attribute(
26072 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26073 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26074 )?;
26075 let cfg = LaunchConfig {
26076 grid_dim: (max_nc, n_head as u32, b as u32),
26077 block_dim: (256, 1, 1),
26078 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26079 };
26080 let __s_lb = self.gpu.stream();
26081 let mut lb = __s_lb.launch_builder(&f);
26082 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26083 unsafe {
26084 lb.launch(cfg)?;
26085 }
26086 }
26087 {
26088 let f = self.func("gdn_chunk_solve32_vl");
26089 let cfg = LaunchConfig {
26090 grid_dim: (max_nc, n_head as u32, b as u32),
26091 block_dim: (256, 1, 1),
26092 shared_mem_bytes: 0,
26093 };
26094 let __s_lb = self.gpu.stream();
26095 let mut lb = __s_lb.launch_builder(&f);
26096 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26097 unsafe {
26098 lb.launch(cfg)?;
26099 }
26100 }
26101 Ok(())
26102 }
26103
26104 #[allow(clippy::too_many_arguments)]
26108 pub fn gdn_prep_vl8(
26109 &self,
26110 seqs: &[GdnPrepVl],
26111 conv_w: &CudaSlice<f32>,
26112 dt_bias: &CudaSlice<f32>,
26113 a: &CudaSlice<f32>,
26114 conv_dim: usize,
26115 d_conv: usize,
26116 d_state: usize,
26117 num_v: usize,
26118 num_k: usize,
26119 key_dim: usize,
26120 hk: usize,
26121 eps: f32,
26122 ) -> Result<(), Box<dyn std::error::Error>> {
26123 let b = seqs.len();
26124 assert!(b >= 1 && b <= 8);
26125 let mut packed = [GdnPrepVl::default(); 8];
26126 packed[..b].copy_from_slice(seqs);
26127 let v = GdnPrepVl8(packed);
26128 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26129 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26130 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26131 assert!(
26132 conv_fuse || hk == num_v,
26133 "de-broadcast requires the fused conv"
26134 );
26135 if conv_fuse {
26136 let f = self.func("ssm_conv1d_gdn_state_vl");
26137 let cfg = LaunchConfig {
26138 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26139 block_dim: (256, 1, 1),
26140 shared_mem_bytes: 0,
26141 };
26142 let (dsi, nvi, nki, kdi, hki) = (
26143 d_state as i32,
26144 num_v as i32,
26145 num_k as i32,
26146 key_dim as i32,
26147 hk as i32,
26148 );
26149 let __s_lb = self.gpu.stream();
26150 let mut lb = __s_lb.launch_builder(&f);
26151 lb.arg(&v)
26152 .arg(conv_w)
26153 .arg(&cdi)
26154 .arg(&dci)
26155 .arg(&dsi)
26156 .arg(&nvi)
26157 .arg(&nki)
26158 .arg(&kdi)
26159 .arg(&hki);
26160 unsafe {
26161 lb.launch(cfg)?;
26162 }
26163 } else {
26164 let f = self.func("ssm_conv1d_tm_state_vl");
26165 let cfg = LaunchConfig {
26166 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26167 block_dim: (256, 1, 1),
26168 shared_mem_bytes: 0,
26169 };
26170 let __s_lb = self.gpu.stream();
26171 let mut lb = __s_lb.launch_builder(&f);
26172 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26173 unsafe {
26174 lb.launch(cfg)?;
26175 }
26176 }
26177 {
26178 let f = self.func("ssm_conv_ring_update_vl");
26179 let n = (conv_dim * (d_conv - 1)) as u32;
26180 let cfg = LaunchConfig {
26181 grid_dim: (n.div_ceil(256), 1, b as u32),
26182 block_dim: (256, 1, 1),
26183 shared_mem_bytes: 0,
26184 };
26185 let __s_lb = self.gpu.stream();
26186 let mut lb = __s_lb.launch_builder(&f);
26187 lb.arg(&v).arg(&cdi).arg(&dci);
26188 unsafe {
26189 lb.launch(cfg)?;
26190 }
26191 }
26192 if !conv_fuse {
26193 let f = self.func("qkv_to_gdn_repack_vl");
26194 let n = max_t * (num_v * d_state) as u32;
26195 let cfg = LaunchConfig {
26196 grid_dim: (n.div_ceil(256), 1, b as u32),
26197 block_dim: (256, 1, 1),
26198 shared_mem_bytes: 0,
26199 };
26200 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26201 let __s_lb = self.gpu.stream();
26202 let mut lb = __s_lb.launch_builder(&f);
26203 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26204 unsafe {
26205 lb.launch(cfg)?;
26206 }
26207 }
26208 if Self::l2_v2_on(d_state) {
26209 let f = self.func("gdn_l2_v2_vl");
26210 let cfg = LaunchConfig {
26211 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26212 block_dim: (256, 1, 1),
26213 shared_mem_bytes: 0,
26214 };
26215 let (dsi, nvi) = (d_state as i32, hk as i32);
26216 let __s_lb = self.gpu.stream();
26217 let mut lb = __s_lb.launch_builder(&f);
26218 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26219 unsafe {
26220 lb.launch(cfg)?;
26221 }
26222 } else {
26223 let f = self.func("gdn_l2_vl");
26224 let cfg = LaunchConfig {
26225 grid_dim: (max_t * hk as u32, 2, b as u32),
26226 block_dim: (256, 1, 1),
26227 shared_mem_bytes: 0,
26228 };
26229 let (dsi, nvi) = (d_state as i32, hk as i32);
26230 let __s_lb = self.gpu.stream();
26231 let mut lb = __s_lb.launch_builder(&f);
26232 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26233 unsafe {
26234 lb.launch(cfg)?;
26235 }
26236 }
26237 {
26238 let f = self.func("gdn_gate_prep_vl");
26239 let n = max_t * num_v as u32;
26240 let cfg = LaunchConfig {
26241 grid_dim: (n.div_ceil(256), 1, b as u32),
26242 block_dim: (256, 1, 1),
26243 shared_mem_bytes: 0,
26244 };
26245 let nvi = num_v as i32;
26246 let __s_lb = self.gpu.stream();
26247 let mut lb = __s_lb.launch_builder(&f);
26248 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26249 unsafe {
26250 lb.launch(cfg)?;
26251 }
26252 }
26253 Ok(())
26254 }
26255
26256 pub fn gdn_mirror_vl8(
26258 &self,
26259 seqs: &[GdnSeqVl],
26260 n_head: usize,
26261 which: i32,
26262 hk: usize,
26263 ) -> Result<(), Box<dyn std::error::Error>> {
26264 let b = seqs.len();
26265 assert!(b >= 1 && b <= 8);
26266 let mut packed = [GdnSeqVl::default(); 8];
26267 packed[..b].copy_from_slice(seqs);
26268 let v = GdnVl8(packed);
26269 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26270 let max_n = seqs
26271 .iter()
26272 .map(|s| {
26273 if which == 0 {
26274 s.t as i64 * ept as i64
26275 } else {
26276 s.nc as i64 * ept as i64 * 32
26277 }
26278 })
26279 .max()
26280 .unwrap();
26281 let f = self.func("gdn_mirror_vl");
26282 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26283 let cfg = LaunchConfig {
26284 grid_dim: (blocks, 1, b as u32),
26285 block_dim: (256, 1, 1),
26286 shared_mem_bytes: 0,
26287 };
26288 let __s_lb = self.gpu.stream();
26289 let mut lb = __s_lb.launch_builder(&f);
26290 lb.arg(&v).arg(&ept).arg(&which);
26291 unsafe {
26292 lb.launch(cfg)?;
26293 }
26294 Ok(())
26295 }
26296
26297 pub fn gdn_tail_vl8(
26299 &self,
26300 seqs: &[GdnPrepVl],
26301 norm_w: &CudaSlice<f32>,
26302 d_state: usize,
26303 num_v: usize,
26304 eps: f32,
26305 ) -> Result<(), Box<dyn std::error::Error>> {
26306 let b = seqs.len();
26307 assert!(b >= 1 && b <= 8);
26308 let mut packed = [GdnPrepVl::default(); 8];
26309 packed[..b].copy_from_slice(seqs);
26310 let v = GdnPrepVl8(packed);
26311 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26312 let f = self.func("gated_rmsnorm_f16out_vl");
26313 let cfg = LaunchConfig {
26315 grid_dim: (max_t * num_v as u32, 1, b as u32),
26316 block_dim: (128, 1, 1),
26317 shared_mem_bytes: 0,
26318 };
26319 let (dsi, nvi) = (d_state as i32, num_v as i32);
26320 let __s_lb = self.gpu.stream();
26321 let mut lb = __s_lb.launch_builder(&f);
26322 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26323 unsafe {
26324 lb.launch(cfg)?;
26325 }
26326 Ok(())
26327 }
26328
26329 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
26332 use cudarc::driver::DevicePtr;
26333 let s = self.gpu.stream();
26334 let (p, _g) = x.device_ptr(&s);
26335 p as u64
26336 }
26337 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
26338 use cudarc::driver::DevicePtrMut;
26339 let s = self.gpu.stream();
26340 let (p, _g) = x.device_ptr_mut(&s);
26341 p as u64
26342 }
26343 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
26344 use cudarc::driver::DevicePtr;
26345 let s = self.gpu.stream();
26346 let (p, _g) = x.device_ptr(&s);
26347 p as u64
26348 }
26349 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
26350 use cudarc::driver::DevicePtr;
26351 let s = self.gpu.stream();
26352 let (p, _g) = x.device_ptr(&s);
26353 p as u64
26354 }
26355
26356 pub fn gdn_chunk_vl8(
26360 &self,
26361 seqs: &[GdnSeqVl],
26362 n_head: usize,
26363 scale: f32,
26364 hk: usize,
26365 wq: Option<&GdnWVl8>,
26366 ) -> Result<(), Box<dyn std::error::Error>> {
26367 const NSPLIT: u32 = 4;
26368 let b = seqs.len();
26369 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
26370 let mut packed = [GdnSeqVl::default(); 8];
26371 packed[..b].copy_from_slice(seqs);
26372 let v = GdnVl8(packed);
26373 let (hi, ci) = (n_head as i32, 32i32);
26374 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26375 let hki = hk as i32;
26376 if let Some(w) = wq {
26377 let f = self.func("gdn_k45_wgmma_vl");
26379 let cfg = LaunchConfig {
26380 grid_dim: (n_head as u32, NSPLIT, b as u32),
26381 block_dim: (256, 1, 1),
26382 shared_mem_bytes: 0,
26383 };
26384 let __s_lb = self.gpu.stream();
26385 let mut lb = __s_lb.launch_builder(&f);
26386 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
26387 unsafe {
26388 lb.launch(cfg)?;
26389 }
26390 let _ = max_nc;
26391 return Ok(());
26392 }
26393 {
26394 let f = self.func("gdn_chunk_state_mma_vl");
26395 let cfg = LaunchConfig {
26396 grid_dim: (n_head as u32, NSPLIT, b as u32),
26397 block_dim: (256, 1, 1),
26398 shared_mem_bytes: 0,
26399 };
26400 let __s_lb = self.gpu.stream();
26401 let mut lb = __s_lb.launch_builder(&f);
26402 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26403 unsafe {
26404 lb.launch(cfg)?;
26405 }
26406 }
26407 {
26408 let f = self.func("gdn_chunk_output_mma_vl");
26409 let cfg = LaunchConfig {
26410 grid_dim: (max_nc, n_head as u32, b as u32),
26411 block_dim: (256, 1, 1),
26412 shared_mem_bytes: 0,
26413 };
26414 let __s_lb = self.gpu.stream();
26415 let mut lb = __s_lb.launch_builder(&f);
26416 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
26417 unsafe {
26418 lb.launch(cfg)?;
26419 }
26420 }
26421 Ok(())
26422 }
26423 pub fn gdn_scan_chunked(
26424 &self,
26425 q: &CudaSlice<f32>,
26426 k: &CudaSlice<f32>,
26427 v: &CudaSlice<f32>,
26428 g: &CudaSlice<f32>,
26429 beta: &CudaSlice<f32>,
26430 kb16_pre: Option<&CudaSlice<u8>>,
26431 qb16_pre: Option<&CudaSlice<u8>>,
26432 state_in: &CudaSlice<f32>,
26433 state_out: &mut CudaSlice<f32>,
26434 o: &mut CudaSlice<f32>,
26435 n_head: usize,
26436 t: usize,
26437 scale: f32,
26438 c: usize,
26439 hk: usize,
26440 ) -> Result<(), Box<dyn std::error::Error>> {
26441 const D: usize = 128;
26442 const NSPLIT: u32 = 4;
26443 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
26444 let h = n_head;
26445 let nc = (t + c - 1) / c;
26446 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26447 let gdn_mma_pre = !portable_mma_gated()
26452 && c == 32
26453 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26454 Ok("1") => true,
26455 Ok("0") => false,
26456 _ => gdn_mma_default_on(),
26457 };
26458 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
26459 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
26460 } else {
26461 None
26462 };
26463 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
26468 && gdn_mma_pre
26469 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
26470 let nk = t * hk * D;
26471 let mut kb16_local: Option<CudaSlice<u8>> = None;
26472 if gdn_mma_pre && kb16_pre.is_none() {
26473 let mut kb = self.alloc_u8_uninit(nk * 2)?;
26474 let f = self.func("f32_to_bf16_bulk");
26475 let n2 = nk as i64;
26476 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26477 let __s_b = self.gpu.stream();
26478 let mut b = __s_b.launch_builder(&f);
26479 b.arg(k).arg(&mut kb).arg(&n2);
26480 unsafe {
26481 b.launch(cfg2)?;
26482 }
26483 kb16_local = Some(kb);
26484 }
26485 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
26486 if let Some(kb) = kb16_pre {
26487 assert!(kb.len() >= nk * 2, "kb16_pre too small");
26488 }
26489 let mut qb16: Option<CudaSlice<u8>> = None;
26490 let mut pb16: Option<CudaSlice<u8>> = None;
26491 if gdn_wgmma_pre {
26492 if qb16_pre.is_none() {
26495 let mut qb = self.alloc_u8_uninit(nk * 2)?;
26496 let f = self.func("f32_to_bf16_bulk");
26497 let n2 = nk as i64;
26498 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
26499 let __s_b = self.gpu.stream();
26500 let mut b = __s_b.launch_builder(&f);
26501 b.arg(q).arg(&mut qb).arg(&n2);
26502 unsafe {
26503 b.launch(cfg2)?;
26504 }
26505 qb16 = Some(qb);
26506 } else if let Some(qb) = qb16_pre {
26507 assert!(qb.len() >= nk * 2, "qb16_pre too small");
26508 }
26509 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
26510 }
26511 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
26512 let k2w = if gdn_wgmma_pre {
26513 Some((
26514 *qb16_ref0.as_ref().unwrap(),
26515 *kb16_ref0.as_ref().unwrap(),
26516 pb16.as_mut().unwrap(),
26517 ))
26518 } else {
26519 None
26520 };
26521 let (gcum, p, u, w) =
26522 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
26523 let _ = &w;
26524 let mut y = self.uninit(nc * h * c * D)?;
26525 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
26539 && c == 32
26540 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26541 Ok("1") => true,
26542 Ok("0") => false,
26543 _ => gdn_mma_default_on(),
26544 };
26545 if gdn_mma {
26546 let wb16 = wb16_pre
26547 .take()
26548 .expect("mma path pre-allocates wb16 (K3 store fold)");
26549 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
26550 if gdn_wgmma_pre {
26562 let qb16 = qb16_ref0.unwrap();
26564 let pb16 = pb16.as_ref().unwrap();
26565 {
26566 let f = self.func("gdn_k45_wgmma");
26567 let cfg = LaunchConfig {
26568 grid_dim: (h as u32, 4, 1),
26569 block_dim: (256, 1, 1),
26570 shared_mem_bytes: 0,
26571 };
26572 let hki = hk as i32;
26573 let __s_b = self.gpu.stream();
26574 let mut b = __s_b.launch_builder(&f);
26575 b.arg(kb16_ref)
26576 .arg(&gcum)
26577 .arg(beta)
26578 .arg(&u)
26579 .arg(&wb16)
26580 .arg(qb16)
26581 .arg(pb16)
26582 .arg(o)
26583 .arg(&scale)
26584 .arg(state_in)
26585 .arg(&mut *state_out)
26586 .arg(&hi)
26587 .arg(&ti)
26588 .arg(&ci)
26589 .arg(&hki);
26590 unsafe {
26591 b.launch(cfg)?;
26592 }
26593 }
26594 return Ok(());
26595 }
26596 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
26600 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
26601 {
26602 let f = self.func("gdn_chunk_state_mma");
26603 let cfg = LaunchConfig {
26604 grid_dim: (h as u32, NSPLIT, 1),
26605 block_dim: (256, 1, 1),
26606 shared_mem_bytes: 0,
26607 };
26608 let hki = hk as i32;
26609 let __s_b = self.gpu.stream();
26610 let mut b = __s_b.launch_builder(&f);
26611 b.arg(kb16_ref)
26612 .arg(&gcum)
26613 .arg(beta)
26614 .arg(&u)
26615 .arg(&wb16)
26616 .arg(&mut y16)
26617 .arg(&mut ssnap16)
26618 .arg(state_in)
26619 .arg(&mut *state_out)
26620 .arg(&hi)
26621 .arg(&ti)
26622 .arg(&ci)
26623 .arg(&hki);
26624 unsafe {
26625 b.launch(cfg)?;
26626 }
26627 }
26628 {
26629 let f = self.func("gdn_chunk_output_mma");
26631 let jt = ((c + 31) / 32) as u32;
26632 let cfg = LaunchConfig {
26633 grid_dim: (nc as u32, h as u32, jt),
26634 block_dim: (256, 1, 1),
26635 shared_mem_bytes: 0,
26636 };
26637 let hki = hk as i32;
26638 let __s_b = self.gpu.stream();
26639 let mut b = __s_b.launch_builder(&f);
26640 b.arg(q)
26641 .arg(&gcum)
26642 .arg(&p)
26643 .arg(&y16)
26644 .arg(&ssnap16)
26645 .arg(o)
26646 .arg(&hi)
26647 .arg(&ti)
26648 .arg(&ci)
26649 .arg(&scale)
26650 .arg(&hki);
26651 unsafe {
26652 b.launch(cfg)?;
26653 }
26654 }
26655 return Ok(());
26656 }
26657 {
26658 let f = self.func("gdn_chunk_state_f32");
26660 let cfg = LaunchConfig {
26661 grid_dim: (h as u32, NSPLIT, 1),
26662 block_dim: (256, 1, 1),
26663 shared_mem_bytes: 0,
26664 };
26665 let __s_b = self.gpu.stream();
26666 let mut b = __s_b.launch_builder(&f);
26667 b.arg(k)
26668 .arg(&gcum)
26669 .arg(beta)
26670 .arg(&u)
26671 .arg(&w)
26672 .arg(&mut y)
26673 .arg(&mut ssnap)
26674 .arg(state_in)
26675 .arg(&mut *state_out)
26676 .arg(&hi)
26677 .arg(&ti)
26678 .arg(&ci);
26679 unsafe {
26680 b.launch(cfg)?;
26681 }
26682 }
26683 {
26684 let f = self.func("gdn_chunk_output_f32");
26686 let jt = ((c + 31) / 32) as u32;
26687 let cfg = LaunchConfig {
26688 grid_dim: (nc as u32, h as u32, jt),
26689 block_dim: (256, 1, 1),
26690 shared_mem_bytes: 0,
26691 };
26692 let __s_b = self.gpu.stream();
26693 let mut b = __s_b.launch_builder(&f);
26694 b.arg(q)
26695 .arg(&gcum)
26696 .arg(&p)
26697 .arg(&y)
26698 .arg(&ssnap)
26699 .arg(o)
26700 .arg(&hi)
26701 .arg(&ti)
26702 .arg(&ci)
26703 .arg(&scale);
26704 unsafe {
26705 b.launch(cfg)?;
26706 }
26707 }
26708 Ok(())
26709 }
26710
26711 #[allow(clippy::too_many_arguments)]
26720 #[allow(clippy::too_many_arguments)]
26721 pub fn gdn_scan_prefill(
26722 &self,
26723 q: &CudaSlice<f32>,
26724 k: &CudaSlice<f32>,
26725 v: &CudaSlice<f32>,
26726 g: &CudaSlice<f32>,
26727 beta: &CudaSlice<f32>,
26728 kb16_pre: Option<&CudaSlice<u8>>,
26729 qb16_pre: Option<&CudaSlice<u8>>,
26730 state_in: &CudaSlice<f32>,
26731 state_out: &mut CudaSlice<f32>,
26732 o: &mut CudaSlice<f32>,
26733 n_head: usize,
26734 t: usize,
26735 scale: f32,
26736 hk: usize,
26737 ) -> Result<(), Box<dyn std::error::Error>> {
26738 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
26739 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
26740 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
26741 }
26742 if Self::gdn_chunked_enabled() && t >= 16 {
26743 self.gdn_scan_chunked(
26744 q,
26745 k,
26746 v,
26747 g,
26748 beta,
26749 kb16_pre,
26750 qb16_pre,
26751 state_in,
26752 state_out,
26753 o,
26754 n_head,
26755 t,
26756 scale,
26757 Self::gdn_chunk_size(),
26758 hk,
26759 )
26760 } else {
26761 assert!(
26762 hk == n_head,
26763 "s128 scan is broadcast-only (prep guarantees by predicate)"
26764 );
26765 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
26766 }
26767 }
26768
26769 #[allow(clippy::too_many_arguments)]
26771 fn gdn_scan_diff(
26772 &self,
26773 q: &CudaSlice<f32>,
26774 k: &CudaSlice<f32>,
26775 v: &CudaSlice<f32>,
26776 g: &CudaSlice<f32>,
26777 beta: &CudaSlice<f32>,
26778 state_in: &CudaSlice<f32>,
26779 state_out: &mut CudaSlice<f32>,
26780 o: &mut CudaSlice<f32>,
26781 n_head: usize,
26782 t: usize,
26783 scale: f32,
26784 ) -> Result<(), Box<dyn std::error::Error>> {
26785 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
26786 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
26787 let mut o_c = self.uninit(o.len())?;
26788 let mut st_c = self.uninit(state_out.len())?;
26789 self.gdn_scan_chunked(
26790 q,
26791 k,
26792 v,
26793 g,
26794 beta,
26795 None,
26796 None,
26797 state_in,
26798 &mut st_c,
26799 &mut o_c,
26800 n_head,
26801 t,
26802 scale,
26803 Self::gdn_chunk_size(),
26804 n_head,
26805 )?;
26806 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
26807 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
26808 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
26809 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
26810 let mut max_abs = 0f32;
26811 let mut max_rel = 0f32;
26812 let mut sum_rel = 0f64;
26813 for (x, y) in a.iter().zip(b) {
26814 let ad = (x - y).abs();
26815 let rel = ad / x.abs().max(y.abs()).max(1e-3);
26816 if ad > max_abs {
26817 max_abs = ad;
26818 }
26819 if rel > max_rel {
26820 max_rel = rel;
26821 }
26822 sum_rel += rel as f64;
26823 }
26824 (max_abs, max_rel, sum_rel / a.len() as f64)
26825 };
26826 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
26827 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
26828 println!(
26829 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
26830 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
26831 Self::gdn_chunk_size()
26832 );
26833 Ok(())
26834 }
26835
26836 pub fn gdn_glog(
26838 &self,
26839 alpha: &CudaSlice<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_v(
26861 &self,
26862 x: &cudarc::driver::CudaView<f32>,
26863 y: &mut CudaSlice<f32>,
26864 n: usize,
26865 ) -> Result<(), Box<dyn std::error::Error>> {
26866 let f = self.func("sigmoid_f32");
26867 let cfg = LaunchConfig::for_num_elems(n as u32);
26868 let ni = n as i32;
26869 let __s_b = self.gpu.stream();
26870 let mut b = __s_b.launch_builder(&f);
26871 b.arg(x).arg(y).arg(&ni);
26872 unsafe {
26873 b.launch(cfg)?;
26874 }
26875 Ok(())
26876 }
26877
26878 pub fn gdn_glog_v(
26879 &self,
26880 alpha: &cudarc::driver::CudaView<f32>,
26881 dt_bias: &CudaSlice<f32>,
26882 a: &CudaSlice<f32>,
26883 g_log: &mut CudaSlice<f32>,
26884 n_head: usize,
26885 t: usize,
26886 ) -> Result<(), Box<dyn std::error::Error>> {
26887 let f = self.func("gdn_glog_f32");
26888 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
26889 let (h, ti) = (n_head as i32, t as i32);
26890 let __s_b = self.gpu.stream();
26891 let mut b = __s_b.launch_builder(&f);
26892 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
26893 unsafe {
26894 b.launch(cfg)?;
26895 }
26896 Ok(())
26897 }
26898
26899 pub fn sigmoid(
26900 &self,
26901 x: &CudaSlice<f32>,
26902 y: &mut CudaSlice<f32>,
26903 n: usize,
26904 ) -> Result<(), Box<dyn std::error::Error>> {
26905 let f = self.func("sigmoid_f32");
26906 let cfg = LaunchConfig::for_num_elems(n as u32);
26907 let ni = n as i32;
26908 let __s_b = self.gpu.stream();
26909 let mut b = __s_b.launch_builder(&f);
26910 b.arg(x).arg(y).arg(&ni);
26911 unsafe {
26912 b.launch(cfg)?;
26913 }
26914 Ok(())
26915 }
26916
26917 pub fn sig_mul_f16out(
26920 &self,
26921 a: &CudaSlice<f32>,
26922 g: &CudaSlice<f32>,
26923 dst: &mut CudaSlice<f32>,
26924 dst16: &mut CudaSlice<u8>,
26925 n: usize,
26926 ) -> Result<(), Box<dyn std::error::Error>> {
26927 let f = self.func("sig_mul_f16out_f32");
26928 let cfg = LaunchConfig::for_num_elems(n as u32);
26929 let ni = n as i32;
26930 let __s_b = self.gpu.stream();
26931 let mut b = __s_b.launch_builder(&f);
26932 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
26933 unsafe {
26934 b.launch(cfg)?;
26935 }
26936 Ok(())
26937 }
26938
26939 #[allow(clippy::too_many_arguments)]
26948 pub fn attn_head_gate(
26949 &self,
26950 a: &CudaSlice<f32>,
26951 g: &CudaSlice<f32>,
26952 dst: &mut CudaSlice<f32>,
26953 dst16: Option<&mut CudaSlice<u8>>,
26954 head_dim: usize,
26955 n_head: usize,
26956 t: usize,
26957 ) -> Result<(), Box<dyn std::error::Error>> {
26958 let f = self.func("attn_head_gate_f32");
26959 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26960 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26961 let d16: u64 = match dst16 {
26963 Some(d) => self.addr_u8(d),
26964 None => 0,
26965 };
26966 let __s_b = self.gpu.stream();
26967 let mut b = __s_b.launch_builder(&f);
26968 b.arg(a)
26969 .arg(g)
26970 .arg(dst)
26971 .arg(&d16)
26972 .arg(&hd)
26973 .arg(&nh)
26974 .arg(&ti);
26975 unsafe {
26976 b.launch(cfg)?;
26977 }
26978 Ok(())
26979 }
26980
26981 #[allow(clippy::too_many_arguments)]
26990 pub fn swiglu_clamped_mul_scaled(
26991 &self,
26992 gate: &CudaSlice<f32>,
26993 up: &CudaSlice<f32>,
26994 gs: f32,
26995 us: f32,
26996 limit: f32,
26997 dst: &mut CudaSlice<f32>,
26998 n: usize,
26999 ) -> Result<(), Box<dyn std::error::Error>> {
27000 debug_assert!(
27001 limit > 1e-6,
27002 "swiglu_clamped needs a live limit; use silu_mul_scaled"
27003 );
27004 let f = self.func("swiglu_clamped_mul_scaled_f32");
27005 let cfg = LaunchConfig::for_num_elems(n as u32);
27006 let ni = n as i32;
27007 let __s_b = self.gpu.stream();
27008 let mut b = __s_b.launch_builder(&f);
27009 b.arg(gate)
27010 .arg(up)
27011 .arg(&gs)
27012 .arg(&us)
27013 .arg(&limit)
27014 .arg(dst)
27015 .arg(&ni);
27016 unsafe {
27017 b.launch(cfg)?;
27018 }
27019 Ok(())
27020 }
27021
27022 pub fn gated_rmsnorm(
27024 &self,
27025 o: &CudaSlice<f32>,
27026 w: &CudaSlice<f32>,
27027 z: &CudaSlice<f32>,
27028 dst: &mut CudaSlice<f32>,
27029 ncols: usize,
27030 nrows: usize,
27031 eps: f32,
27032 ) -> Result<(), Box<dyn std::error::Error>> {
27033 let f = self.func("gated_rmsnorm_f32");
27034 let cfg = LaunchConfig {
27035 grid_dim: (nrows as u32, 1, 1),
27036 block_dim: (128, 1, 1),
27037 shared_mem_bytes: 0,
27038 };
27039 let (nc, e) = (ncols as i32, eps);
27040 let __s_b = self.gpu.stream();
27041 let mut b = __s_b.launch_builder(&f);
27042 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27043 unsafe {
27044 b.launch(cfg)?;
27045 }
27046 Ok(())
27047 }
27048
27049 pub fn gated_rmsnorm_f16out(
27052 &self,
27053 o: &CudaSlice<f32>,
27054 w: &CudaSlice<f32>,
27055 z: &CudaSlice<f32>,
27056 dst: &mut CudaSlice<f32>,
27057 dst16: &mut CudaSlice<u8>,
27058 ncols: usize,
27059 nrows: usize,
27060 eps: f32,
27061 ) -> Result<(), Box<dyn std::error::Error>> {
27062 let f = self.func("gated_rmsnorm_f16out_f32");
27063 let cfg = LaunchConfig {
27065 grid_dim: (nrows as u32, 1, 1),
27066 block_dim: (128, 1, 1),
27067 shared_mem_bytes: 0,
27068 };
27069 let (nc, e) = (ncols as i32, eps);
27070 let __s_b = self.gpu.stream();
27071 let mut b = __s_b.launch_builder(&f);
27072 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27073 unsafe {
27074 b.launch(cfg)?;
27075 }
27076 Ok(())
27077 }
27078
27079 #[allow(clippy::too_many_arguments)]
27083 pub fn add_rms_norm_zq8(
27084 &self,
27085 a: &CudaSlice<f32>,
27086 b_in: &CudaSlice<f32>,
27087 w: &CudaSlice<f32>,
27088 res: &mut CudaSlice<f32>,
27089 z: &mut CudaSlice<f32>,
27090 ncols: usize,
27091 nrows: usize,
27092 eps: f32,
27093 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27094 assert!(ncols % 32 == 0);
27095 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27096 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27097 let f = self.func("add_rms_norm_zq8");
27098 let cfg = LaunchConfig {
27099 grid_dim: (nrows as u32, 1, 1),
27100 block_dim: (1024, 1, 1),
27101 shared_mem_bytes: 0,
27102 };
27103 let (nc, ep) = (ncols as i32, eps);
27104 let __s_b = self.gpu.stream();
27105 let mut b = __s_b.launch_builder(&f);
27106 b.arg(a)
27107 .arg(b_in)
27108 .arg(w)
27109 .arg(res)
27110 .arg(z)
27111 .arg(&mut q)
27112 .arg(&mut d)
27113 .arg(&nc)
27114 .arg(&ep);
27115 unsafe {
27116 b.launch(cfg)?;
27117 }
27118 Ok((q, d))
27119 }
27120
27121 pub fn gated_rmsnorm_zv(
27126 &self,
27127 o: &CudaSlice<f32>,
27128 w: &CudaSlice<f32>,
27129 z: &cudarc::driver::CudaView<f32>,
27130 dst: &mut CudaSlice<f32>,
27131 ncols: usize,
27132 nrows: usize,
27133 eps: f32,
27134 ) -> Result<(), Box<dyn std::error::Error>> {
27135 let f = self.func("gated_rmsnorm_f32");
27136 let cfg = LaunchConfig {
27137 grid_dim: (nrows as u32, 1, 1),
27138 block_dim: (128, 1, 1),
27139 shared_mem_bytes: 0,
27140 };
27141 let (nc, e) = (ncols as i32, eps);
27142 let __s_b = self.gpu.stream();
27143 let mut b = __s_b.launch_builder(&f);
27144 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27145 unsafe {
27146 b.launch(cfg)?;
27147 }
27148 Ok(())
27149 }
27150
27151 pub fn gated_rmsnorm_f16out_zv(
27152 &self,
27153 o: &CudaSlice<f32>,
27154 w: &CudaSlice<f32>,
27155 z: &cudarc::driver::CudaView<f32>,
27156 dst: &mut CudaSlice<f32>,
27157 dst16: &mut CudaSlice<u8>,
27158 ncols: usize,
27159 nrows: usize,
27160 eps: f32,
27161 ) -> Result<(), Box<dyn std::error::Error>> {
27162 let f = self.func("gated_rmsnorm_f16out_f32");
27163 let cfg = LaunchConfig {
27165 grid_dim: (nrows as u32, 1, 1),
27166 block_dim: (128, 1, 1),
27167 shared_mem_bytes: 0,
27168 };
27169 let (nc, e) = (ncols as i32, eps);
27170 let __s_b = self.gpu.stream();
27171 let mut b = __s_b.launch_builder(&f);
27172 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27173 unsafe {
27174 b.launch(cfg)?;
27175 }
27176 Ok(())
27177 }
27178
27179 pub fn gated_rmsnorm_q8_1(
27180 &self,
27181 o: &CudaSlice<f32>,
27182 w: &CudaSlice<f32>,
27183 z: &CudaSlice<f32>,
27184 ncols: usize,
27185 nrows: usize,
27186 eps: f32,
27187 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27188 assert!(ncols % 32 == 0);
27189 let f = self.func("gated_rmsnorm_q8_1");
27190 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27191 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27192 let cfg = LaunchConfig {
27193 grid_dim: (nrows as u32, 1, 1),
27194 block_dim: (128, 1, 1),
27195 shared_mem_bytes: 0,
27196 };
27197 let (nc, ep) = (ncols as i32, eps);
27198 let __s_b = self.gpu.stream();
27199 let mut b = __s_b.launch_builder(&f);
27200 b.arg(o)
27201 .arg(w)
27202 .arg(z)
27203 .arg(&mut out_q)
27204 .arg(&mut out_d)
27205 .arg(&nc)
27206 .arg(&ep);
27207 unsafe {
27208 b.launch(cfg)?;
27209 }
27210 Ok((out_q, out_d))
27211 }
27212
27213 pub fn transpose(
27215 &self,
27216 inp: &CudaSlice<f32>,
27217 rows: usize,
27218 cols: usize,
27219 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27220 let f = self.func("transpose_f32");
27221 let mut out = self.zeros(rows * cols)?;
27222 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27223 let (r, c) = (rows as i32, cols as i32);
27224 let __s_b = self.gpu.stream();
27225 let mut b = __s_b.launch_builder(&f);
27226 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27227 unsafe {
27228 b.launch(cfg)?;
27229 }
27230 Ok(out)
27231 }
27232
27233 pub fn repeat_heads(
27235 &self,
27236 inp: &CudaSlice<f32>,
27237 out: &mut CudaSlice<f32>,
27238 head_dim: usize,
27239 n_in: usize,
27240 n_out: usize,
27241 t: usize,
27242 ) -> Result<(), Box<dyn std::error::Error>> {
27243 let f = self.func("repeat_heads_f32");
27244 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27245 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27246 let __s_b = self.gpu.stream();
27247 let mut b = __s_b.launch_builder(&f);
27248 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27249 unsafe {
27250 b.launch(cfg)?;
27251 }
27252 Ok(())
27253 }
27254
27255 pub fn q_gate_split(
27262 &self,
27263 qf: &CudaSlice<f32>,
27264 q_out: &mut CudaSlice<f32>,
27265 gate_out: &mut CudaSlice<f32>,
27266 head_dim: usize,
27267 n_head: usize,
27268 t: usize,
27269 ) -> Result<(), Box<dyn std::error::Error>> {
27270 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27271 let out_need = head_dim * n_head * t;
27272 if q_out.len() < out_need || gate_out.len() < out_need {
27273 return Err(format!(
27274 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27275 q_out.len(),
27276 gate_out.len()
27277 )
27278 .into());
27279 }
27280 let f = self.func("q_gate_split_f32");
27281 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27282 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27283 let __s_b = self.gpu.stream();
27284 let mut b = __s_b.launch_builder(&f);
27285 b.arg(qf)
27286 .arg(q_out)
27287 .arg(gate_out)
27288 .arg(&hd)
27289 .arg(&nh)
27290 .arg(&ti);
27291 unsafe {
27292 b.launch(cfg)?;
27293 }
27294 Ok(())
27295 }
27296
27297 pub fn qkv_to_gdn_repack(
27301 &self,
27302 conv_out: &CudaSlice<f32>,
27303 q_g: &mut CudaSlice<f32>,
27304 k_g: &mut CudaSlice<f32>,
27305 v_g: &mut CudaSlice<f32>,
27306 d_state: usize,
27307 num_v: usize,
27308 num_k: usize,
27309 key_dim: usize,
27310 t: usize,
27311 ) -> Result<(), Box<dyn std::error::Error>> {
27312 let f = self.func("qkv_to_gdn_repack_f32");
27313 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27314 let (ds, nv, nk, kd, ti) = (
27315 d_state as i32,
27316 num_v as i32,
27317 num_k as i32,
27318 key_dim as i32,
27319 t as i32,
27320 );
27321 let __s_b = self.gpu.stream();
27322 let mut b = __s_b.launch_builder(&f);
27323 b.arg(conv_out)
27324 .arg(q_g)
27325 .arg(k_g)
27326 .arg(v_g)
27327 .arg(&ds)
27328 .arg(&nv)
27329 .arg(&nk)
27330 .arg(&kd)
27331 .arg(&ti);
27332 unsafe {
27333 b.launch(cfg)?;
27334 }
27335 Ok(())
27336 }
27337
27338 pub fn conv_left_pad(
27341 &self,
27342 src: &CudaSlice<f32>,
27343 dst: &mut CudaSlice<f32>,
27344 conv_dim: usize,
27345 t: usize,
27346 pad: usize,
27347 ) -> Result<(), Box<dyn std::error::Error>> {
27348 let f = self.func("conv_left_pad_f32");
27349 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
27350 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
27351 let __s_b = self.gpu.stream();
27352 let mut b = __s_b.launch_builder(&f);
27353 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
27354 unsafe {
27355 b.launch(cfg)?;
27356 }
27357 Ok(())
27358 }
27359
27360 pub fn conv_assemble_and_roll(
27364 &self,
27365 qkv_col: &CudaSlice<f32>,
27366 conv_state: &mut CudaSlice<f32>,
27367 conv_in: &mut CudaSlice<f32>,
27368 conv_dim: usize,
27369 pad: usize,
27370 ) -> Result<(), Box<dyn std::error::Error>> {
27371 let f = self.func("conv_assemble_and_roll_f32");
27372 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27373 let (cd, p) = (conv_dim as i32, pad as i32);
27374 let __s_b = self.gpu.stream();
27375 let mut b = __s_b.launch_builder(&f);
27376 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
27377 unsafe {
27378 b.launch(cfg)?;
27379 }
27380 Ok(())
27381 }
27382
27383 pub fn ssm_conv1d_fused_decode(
27389 &self,
27390 qkv_col: &CudaSlice<f32>,
27391 conv_state: &mut CudaSlice<f32>,
27392 w: &CudaSlice<f32>,
27393 conv_out: &mut CudaSlice<f32>,
27394 conv_dim: usize,
27395 d_conv: usize,
27396 ) -> Result<(), Box<dyn std::error::Error>> {
27397 let f = self.func("ssm_conv1d_fused_decode_f32");
27398 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
27399 let (cd, dc) = (conv_dim as i32, d_conv as i32);
27400 let __s_b = self.gpu.stream();
27401 let mut b = __s_b.launch_builder(&f);
27402 b.arg(qkv_col)
27403 .arg(conv_state)
27404 .arg(w)
27405 .arg(conv_out)
27406 .arg(&cd)
27407 .arg(&dc);
27408 unsafe {
27409 b.launch(cfg)?;
27410 }
27411 Ok(())
27412 }
27413
27414 pub fn slice_range(
27417 &self,
27418 src: &CudaSlice<f32>,
27419 start: usize,
27420 len: usize,
27421 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27422 let host = self.gpu.stream().clone_dtoh(src)?;
27423 self.gpu.stream().synchronize()?;
27424 Ok(self.htod(&host[start..start + len])?)
27425 }
27426}
27427
27428#[cfg(test)]
27429mod target_dispatch_tests {
27430 use super::legacy_quant_gemm_allowed;
27431
27432 #[test]
27433 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
27434 assert!(legacy_quant_gemm_allowed(false, false, false));
27436 assert!(!legacy_quant_gemm_allowed(false, false, true));
27437 assert!(!legacy_quant_gemm_allowed(true, false, false));
27439 assert!(!legacy_quant_gemm_allowed(true, false, true));
27440 assert!(legacy_quant_gemm_allowed(true, true, false));
27442 assert!(!legacy_quant_gemm_allowed(true, true, true));
27443 }
27444
27445 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
27446 #[test]
27447 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
27448 assert!(!legacy_quant_gemm_allowed(
27449 cfg!(memra_portable_cuda),
27450 cfg!(memra_hopper_mma),
27451 false
27452 ));
27453 }
27454
27455 #[cfg(memra_hopper_mma)]
27456 #[test]
27457 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
27458 assert!(legacy_quant_gemm_allowed(
27459 cfg!(memra_portable_cuda),
27460 cfg!(memra_hopper_mma),
27461 false
27462 ));
27463 assert!(super::portable_mma_gated() == false);
27464 }
27465}
27466
27467impl memra_kv::KvDev for Engine {
27470 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27471 Engine::zeros(self, n)
27472 }
27473 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27474 Engine::uninit(self, n)
27475 }
27476 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
27477 Engine::alloc_u8(self, n)
27478 }
27479 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
27480 Engine::htod_i32(self, v)
27481 }
27482 fn clone_dtod(
27483 &self,
27484 src: &CudaSlice<f32>,
27485 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27486 Engine::clone_dtod(self, src)
27487 }
27488 fn copy_into(
27489 &self,
27490 dst: &mut CudaSlice<f32>,
27491 off: usize,
27492 src: &CudaSlice<f32>,
27493 len: usize,
27494 ) -> Result<(), Box<dyn std::error::Error>> {
27495 Engine::copy_into(self, dst, off, src, len)
27496 }
27497 fn set_i32_one(
27498 &self,
27499 d: &mut CudaSlice<i32>,
27500 v: i32,
27501 ) -> Result<(), Box<dyn std::error::Error>> {
27502 Engine::set_i32_one(self, d, v)
27503 }
27504}
27505
27506#[cfg(test)]
27507mod fused_gate_bounds_tests {
27508 use super::*;
27509
27510 #[test]
27523 #[ignore = "requires a CUDA GPU"]
27524 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
27525 let e = Engine::new(0).unwrap();
27526 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
27527 let fused = 2 * head_dim * n_head * t;
27528 let out_n = head_dim * n_head * t;
27529
27530 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
27532 let mut q = e.uninit(out_n).unwrap();
27533 let mut gate = e.uninit(out_n).unwrap();
27534 let err = e
27535 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
27536 .expect_err("half-width wq must be refused, not read past")
27537 .to_string();
27538 assert!(err.contains("NO fused gate"), "{err}");
27539 assert!(err.contains(&format!("{fused}")), "{err}");
27540
27541 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
27544 let wide = e.htod(&host).unwrap();
27545 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
27546 .expect("full-width wq splits");
27547 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
27548 for tok in 0..t {
27549 for hh in 0..n_head {
27550 for d in 0..head_dim {
27551 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
27552 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
27553 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
27554 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
27555 }
27556 }
27557 }
27558
27559 let mut small = e.uninit(out_n - 1).unwrap();
27561 assert!(
27562 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
27563 .is_err()
27564 );
27565 }
27566}
27567
27568#[cfg(test)]
27572mod fused_rope_width_tests {
27573 use super::Engine;
27574
27575 #[test]
27578 fn full_width_is_accepted() {
27579 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
27580 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
27581 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
27582 }
27583
27584 #[test]
27597 fn gemma4_official_artifact_widths_pass() {
27598 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
27599 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
27600 }
27601
27602 #[test]
27605 fn partial_rotary_is_refused_with_the_geometry_named() {
27606 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
27608 .expect_err("partial rotary must refuse");
27609 let msg = err.to_string();
27610 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
27611 assert!(msg.contains("n_rot 64"), "{msg}");
27612 assert!(msg.contains("head_dim 256"), "{msg}");
27613 assert!(
27614 msg.contains("64..256"),
27615 "names the band it would corrupt: {msg}"
27616 );
27617 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
27619 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
27621 }
27622}