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
13#[cfg(debug_assertions)]
14pub(crate) fn debug_assert_tensor_stream_device<T>(
15 tensor: &CudaSlice<T>,
16 stream: &CudaStream,
17 site: &str,
18) {
19 let tensor_dev = tensor.ordinal();
20 let stream_dev = stream.context().ordinal();
21 assert_eq!(
22 tensor_dev, stream_dev,
23 "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
24 );
25}
26
27fn ensure_tensor_stream_device<T>(
28 tensor: &impl DeviceSlice<T>,
29 stream: &CudaStream,
30 site: &str,
31) -> Result<(), Box<dyn std::error::Error>> {
32 let tensor_dev = tensor.stream().context().ordinal();
33 let stream_dev = stream.context().ordinal();
34 if tensor_dev != stream_dev {
35 return Err(format!(
36 "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
37 stream on dev{stream_dev}"
38 )
39 .into());
40 }
41 Ok(())
42}
43
44pub use memra_gguf;
45pub use memra_runtime;
46
47pub mod forward;
48pub mod hybrid;
49pub mod hybrid_forward;
50pub mod model;
51pub mod sigrouter_contract;
52pub mod vision;
53pub mod vision_gemma;
54pub mod vision_pre;
55pub mod cache {
58 pub use memra_kv::*;
59}
60pub mod decode;
61pub mod decode_batch;
62pub mod dflash;
63pub mod eagle;
64pub mod gemma_spec;
65pub mod graph_update;
66pub mod mla;
70pub mod moesd;
71pub mod parallel;
72pub mod plan_backend;
73pub mod pp;
74pub mod round_stream;
75pub mod spec;
76pub mod tp;
77pub use memra_sampling as sampler;
78
79pub fn moe_f16g_mode() -> u8 {
123 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
124 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
125 Ok("0") => 0,
126 Ok("2") => 2,
127 Ok("3") => 3,
128 Ok(_) => 1,
129 Err(_) => 2,
132 })
133}
134pub fn moe_f16g_sk_params() -> (i32, i32) {
148 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
149 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
150 Ok("0") => (-1, 0),
151 Ok("32") => (0, i32::MAX),
152 Ok("128") => (0, 1),
153 _ => {
154 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
155 .ok()
156 .and_then(|v| v.parse().ok())
157 .unwrap_or(64);
158 (0, cross)
159 }
160 })
161}
162pub fn moe_f16g_direct_on(qtype: i32) -> bool {
173 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
174 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
175 Ok("0") => 0,
176 Ok("kq") => 1,
177 _ => 2,
178 });
179 match m {
180 0 => false,
181 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
182 _ => true,
183 }
184}
185pub fn moe_f16g_tail_on() -> bool {
194 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
195 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
196}
197
198pub fn moe_f16g_gemma_on() -> bool {
205 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
207}
208
209pub fn moe_fuse_actq_on() -> bool {
213 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
214 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
215}
216
217pub fn router_prefill_exact_on() -> bool {
227 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
228 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
229}
230
231pub fn router_kernel_on() -> bool {
232 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
233 *ON.get_or_init(|| {
234 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
235 if !on {
236 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
237 }
238 on
239 })
240}
241
242pub const ROUTER_BATCH_MIN_T: usize = 8;
257pub fn router_batch_on() -> bool {
258 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
259 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
260}
261mod cpu_experts;
262#[cfg(memra_cutlass)]
263pub mod cutlass_ffi;
264pub mod dsv4_ffi;
265pub mod dsv4_gpu;
266pub mod f16_ffi;
267pub mod fp8_ffi;
268pub mod mmq_ffi;
269pub mod moe_cache;
270pub mod prime_graph;
271pub mod spill;
272mod spill_pread;
273
274const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
281const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
282const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
283const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
284const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
285const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
286const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
288
289fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
295 assert!(
296 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
297 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
298 );
299 match std::env::var("MEMRA_GEMM_FATBIN") {
300 Ok(path) => std::borrow::Cow::Owned(
301 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
302 ),
303 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
304 }
305}
306
307pub(crate) const fn portable_mma_gated() -> bool {
314 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
315}
316
317#[track_caller]
334pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
335 assert!(
336 !portable_mma_gated(),
337 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
338 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
339 );
340}
341
342pub(crate) const fn gdn_mma_default_on() -> bool {
351 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
352}
353
354const fn konst_eq(a: &str, b: &str) -> bool {
356 let (a, b) = (a.as_bytes(), b.as_bytes());
357 if a.len() != b.len() {
358 return false;
359 }
360 let mut i = 0;
361 while i < a.len() {
362 if a[i] != b[i] {
363 return false;
364 }
365 i += 1;
366 }
367 true
368}
369
370const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
375 (!portable_cuda || hopper_mma) && !no_gemm
376}
377
378const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
386const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
387const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
388const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
389const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
390
391pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
394
395fn flash_fatbin_bytes() -> &'static [u8] {
397 match kv_cache_formats() {
398 ("q8_0", "q5_1") => FLASH_FATBIN,
399 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
400 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
401 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
402 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
403 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
404 other => unreachable!("kv_cache_formats returned {other:?}"),
405 }
406}
407
408fn k1_launch_override() -> Option<(u32, u32, u32)> {
415 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
416 *K1.get_or_init(|| {
417 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
418 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
419 match p.as_slice() {
420 [bm, bn, w] => Some((*bm, *bn, *w)),
421 _ => None,
422 }
423 })
424}
425
426pub(crate) fn wgmma_gemm_enabled() -> bool {
433 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
434 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
435}
436
437pub const FA_VEC_MIN_TKV: usize = 96;
452pub fn fa_vec_min_tkv() -> usize {
456 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
457 *V.get_or_init(|| {
458 std::env::var("MEMRA_FA_VEC_MIN")
459 .ok()
460 .and_then(|v| v.parse().ok())
461 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
462 })
463}
464
465pub fn fa_f16pv_on() -> bool {
476 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
477 *ON.get_or_init(|| {
478 std::env::var("MEMRA_FA_F16PV")
479 .map(|v| v != "0")
480 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
481 })
482}
483
484pub fn fa512_hp_on() -> bool {
488 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
489 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
490}
491
492pub fn faw_hp_on() -> bool {
496 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
497 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
498}
499
500pub fn fa512_wide_warps() -> usize {
504 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
505 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
506 Ok("1") => 4,
507 _ => 2,
508 })
509}
510
511pub fn fa512_min_tkv() -> usize {
514 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
515 *FA512_MIN.get_or_init(|| {
516 std::env::var("MEMRA_FA512_MIN")
517 .ok()
518 .and_then(|v| v.parse().ok())
519 .unwrap_or(512)
520 })
521}
522pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
526 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
527pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
531pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
537 std::sync::atomic::AtomicBool::new(false);
538pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
545 std::sync::atomic::AtomicBool::new(true);
546pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
547 std::sync::atomic::AtomicUsize::new(16);
548pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
553pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
555pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
560pub use memra_kv::KV_FP8_FORCE;
563pub(crate) fn mmv_block() -> u32 {
568 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
569 *V.get_or_init(|| {
570 std::env::var("MEMRA_MMV_BLOCK")
571 .ok()
572 .and_then(|v| v.parse().ok())
573 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
574 .unwrap_or(128)
575 })
576}
577
578pub(crate) fn sig_expf_dev_on() -> bool {
583 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
584 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
585}
586
587pub(crate) fn topk_fast_on() -> bool {
588 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
589 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
590}
591
592pub(crate) fn rms_block() -> u32 {
593 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
594 *V.get_or_init(|| {
595 std::env::var("MEMRA_RMS_BLOCK")
596 .ok()
597 .and_then(|v| v.parse().ok())
598 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
599 })
600}
601
602pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
603 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
604 if let Some(forced) = *S.get_or_init(|| {
605 std::env::var("MEMRA_FA_SPLIT")
606 .ok()
607 .and_then(|v| v.parse().ok())
608 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
609 }) {
610 return forced;
611 }
612 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
630 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
631 {
632 return if t_kv <= 8192 {
633 16
634 } else if t_kv <= 16384 {
635 64
636 } else {
637 128
638 };
639 }
640 let big_rig = fa_sm_count() >= 128;
641 if big_rig {
642 let _ = n_head_kv;
643 if t_kv <= 2048 {
644 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
653 if let Some(sp) = *SHORT.get_or_init(|| {
654 std::env::var("MEMRA_FA_SP_SHORT")
655 .ok()
656 .and_then(|v| v.parse().ok())
657 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
658 }) {
659 return sp;
660 }
661 16
662 } else if t_kv <= 16384 {
663 64
664 } else {
665 128
666 }
667 } else if n_head_kv <= 4 {
668 if t_kv <= 512 {
689 8
690 } else if t_kv <= 16384 {
691 64
692 } else {
693 128
694 }
695 } else {
696 if t_kv <= 8192 {
697 32
698 } else if t_kv <= 16384 {
699 64
700 } else {
701 128
702 }
703 }
704}
705
706fn fa_sm_count() -> i32 {
709 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
710 *N.get_or_init(|| {
711 cudarc::driver::result::init().ok();
712 cudarc::driver::result::device::get(0)
713 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
714 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
715 .unwrap_or(82)
716 })
717}
718
719fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
723 match head_dim {
724 256 => Ok(""),
725 128 => Ok("_hd128"),
726 d => Err(format!(
727 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
728 callers must gate to sdpa_naive"
729 )
730 .into()),
731 }
732}
733
734pub const QT_Q8_0: i32 = 0;
736pub const QT_Q4_K: i32 = 1;
737pub const QT_Q6_K: i32 = 2;
738pub const QT_Q5_K: i32 = 3;
739pub const QT_Q3_K: i32 = 4;
740pub const QT_IQ4_XS: i32 = 5;
741pub const QT_IQ3_S: i32 = 6;
742pub const QT_NVFP4: i32 = 7;
743pub const QT_F8_E4M3: i32 = 10;
749pub const QT_NVFP4_RP: i32 = 9;
752pub const QT_F32: i32 = 8;
754pub const QT_BF16: i32 = 11;
755pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
760pub const QT_F8_E4M3_BLK: i32 = 14;
776
777pub struct Engine {
779 pub gpu: memra_runtime::Gpu,
780 module: Arc<CudaModule>,
781 hybrid: Arc<CudaModule>,
782 qmatvec: Arc<CudaModule>,
783 flash: Arc<CudaModule>,
784 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
788 gemm: Arc<CudaModule>,
789 router: Arc<CudaModule>,
790 sample: Arc<CudaModule>,
792 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
796 moe_cache_layout: Mutex<Option<Vec<usize>>>,
800 capture_keep_on: std::sync::atomic::AtomicBool,
806 verify_exact: std::sync::atomic::AtomicBool,
811 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
812 pub copy_stream: Arc<CudaStream>,
814 #[cfg(memra_cutlass)]
821 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
822 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
826 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
829 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
833 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
837 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
839 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
840 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
845 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
850 router_stage: Mutex<Option<PinnedStage>>,
854}
855
856fn fa_v2_on() -> bool {
866 std::env::var("MEMRA_FA_V2")
872 .map(|v| v != "0")
873 .unwrap_or(true)
874}
875
876fn fa_v3_on() -> bool {
884 std::env::var("MEMRA_FA_V3")
888 .map(|v| v != "0")
889 .unwrap_or(true)
890}
891
892fn fa_v4_mode() -> &'static str {
897 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
898 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
899}
900fn fa_v4_on() -> bool {
901 fa_v4_mode() != "0"
902} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
911 std::sync::atomic::AtomicUsize::new(1024);
912pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
913 std::sync::atomic::AtomicUsize::new(usize::MAX);
914pub fn fa_v4_at_pub(t_kv: usize) -> bool {
915 fa_v4_at(t_kv)
916}
917fn fa_v4_at(t_kv: usize) -> bool {
918 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
919 let mx = *M.get_or_init(|| {
920 std::env::var("MEMRA_FA_V4_MAX")
921 .ok()
922 .and_then(|v| v.parse().ok())
923 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
924 });
925 fa_v4_on() && t_kv < mx
926}
927pub const FA_DEEP_MIN_DEFAULT: usize = 0;
941fn fa_deep_at(t_kv: usize) -> bool {
942 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
943 return false;
944 }
945 let min = std::env::var("MEMRA_FA_DEEP_MIN")
946 .ok()
947 .and_then(|v| v.parse().ok())
948 .unwrap_or(FA_DEEP_MIN_DEFAULT);
949 t_kv >= min
950}
951pub fn fa_deep_at_pub(t_kv: usize) -> bool {
953 fa_deep_at(t_kv)
954}
955
956fn fa_v3_active(head_dim: usize) -> bool {
957 fa_v3_on()
960 && head_dim % 128 == 0
961 && kv_cache_formats() == ("q8_0", "q5_1")
962 && !Engine::kv_fp8_on()
963}
964
965pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
973 std::env::var("MEMRA_NO_FA_VEC").is_err()
974 && t_kv >= fa_vec_min_tkv()
975 && head_dim == 256
976 && fa_v4_at(t_kv)
977 && !matches!(fa_v4_mode(), "noB3" | "stage")
978 && !Engine::kv_fp8_on()
979}
980pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
982 fa_split_keys(t_kv, n_head_kv)
983}
984
985struct PinnedStage {
990 ptr: *mut u8,
991 cap: usize,
992}
993unsafe impl Send for PinnedStage {}
994impl PinnedStage {
995 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
996 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
997 Ok(PinnedStage { ptr, cap })
998 }
999}
1000impl Drop for PinnedStage {
1001 fn drop(&mut self) {
1002 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1003 }
1004}
1005
1006pub const ARGMAX_NB: usize = 256;
1009
1010pub(crate) use memra_fa3_vl as fa3_vl_raw;
1012
1013unsafe extern "C" {
1014 fn memra_fa3_prefill(
1016 q16: *const core::ffi::c_void,
1017 k16: *const core::ffi::c_void,
1018 v16: *const core::ffi::c_void,
1019 o: *mut f32,
1020 t: i32,
1021 h: i32,
1022 hkv: i32,
1023 d: i32,
1024 scale: f32,
1025 stream: *mut core::ffi::c_void,
1026 ) -> i32;
1027 pub(crate) fn memra_fa3_vl(
1029 q16s: *const *const core::ffi::c_void,
1030 k16s: *const *const core::ffi::c_void,
1031 v16s: *const *const core::ffi::c_void,
1032 os: *const *mut f32,
1033 ts: *const i32,
1034 b: i32,
1035 h: i32,
1036 hkv: i32,
1037 d: i32,
1038 scale: f32,
1039 stream: *mut core::ffi::c_void,
1040 ) -> i32;
1041}
1042
1043#[repr(C)]
1048#[derive(Clone, Copy)]
1049pub struct WPtr8(pub [u64; 8]);
1050unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1051
1052#[repr(C)]
1057#[derive(Clone, Copy, Default)]
1058pub struct GdnSeqVl {
1059 pub kb16: u64,
1060 pub gcum: u64,
1061 pub beta: u64,
1062 pub u: u64,
1063 pub wb16: u64,
1064 pub y: u64,
1065 pub ssnap: u64,
1066 pub state_in: u64,
1067 pub state_out: u64,
1068 pub q: u64,
1069 pub p: u64,
1070 pub o: u64,
1071 pub k: u64,
1072 pub v: u64,
1073 pub g: u64,
1074 pub a: u64,
1075 pub w: u64,
1076 pub t: i32,
1077 pub nc: i32,
1078}
1079unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1080#[repr(C)]
1081#[derive(Clone, Copy)]
1082pub struct GdnVl8(pub [GdnSeqVl; 8]);
1083unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1084
1085#[repr(C)]
1088#[derive(Clone, Copy, Default)]
1089pub struct GdnWVl {
1090 pub qb16: u64,
1091 pub pb16: u64,
1092}
1093unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1094#[repr(C)]
1095#[derive(Clone, Copy)]
1096pub struct GdnWVl8(pub [GdnWVl; 8]);
1097unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1098
1099#[repr(C)]
1101#[derive(Clone, Copy, Default)]
1102pub struct GdnPrepVl {
1103 pub qkv: u64,
1104 pub conv_state: u64,
1105 pub conv_out: u64,
1106 pub q_g: u64,
1107 pub k_g: u64,
1108 pub v_g: u64,
1109 pub q_l2: u64,
1110 pub k_l2: u64,
1111 pub beta_raw: u64,
1112 pub alpha: u64,
1113 pub beta: u64,
1114 pub g_log: u64,
1115 pub o: u64,
1116 pub z: u64,
1117 pub gn: u64,
1118 pub gn16: u64,
1119 pub kb16: u64,
1120 pub qb16: u64,
1121 pub t: i32,
1122 pub pad: i32,
1123}
1124unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1125#[repr(C)]
1126#[derive(Clone, Copy)]
1127pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1128unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1129
1130#[repr(C)]
1132#[derive(Clone, Copy, Default)]
1133pub struct FaSeqVl {
1134 pub q: u64,
1135 pub k16: u64,
1136 pub v16: u64,
1137 pub o: u64,
1138 pub kf: u64,
1139 pub vf: u64,
1140 pub t: i32,
1141 pub pad: i32,
1142}
1143unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1144#[repr(C)]
1145#[derive(Clone, Copy)]
1146pub struct FaVl8(pub [FaSeqVl; 8]);
1147unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1148
1149#[repr(C)]
1151#[derive(Clone, Copy, Default)]
1152pub struct AttnPreVl {
1153 pub qf: u64,
1154 pub kf: u64,
1155 pub vf: u64,
1156 pub q: u64,
1157 pub gate: u64,
1158 pub qn: u64,
1159 pub kn: u64,
1160 pub kc: u64,
1161 pub vc: u64,
1162 pub t: i32,
1163 pub pad: i32,
1164}
1165unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1166#[repr(C)]
1167#[derive(Clone, Copy)]
1168pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1169unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1170
1171pub struct GdnChunkBufs {
1174 pub gcum: CudaSlice<f32>,
1175 pub a: CudaSlice<f32>,
1176 pub p: CudaSlice<f32>,
1177 pub u: CudaSlice<f32>,
1178 pub w: CudaSlice<f32>,
1179 pub kb16: CudaSlice<u8>,
1180 pub wb16: CudaSlice<u8>,
1181 pub y16: CudaSlice<u8>,
1182 pub ssnap16: CudaSlice<u8>,
1183 pub qb16: CudaSlice<u8>,
1184 pub pb16: CudaSlice<u8>,
1185 pub o: CudaSlice<f32>,
1186 pub t: usize,
1187 pub nc: usize,
1188}
1189
1190#[repr(C)]
1192#[derive(Clone, Copy)]
1193pub struct F32x8(pub [f32; 8]);
1194unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1195
1196pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1200
1201impl Engine {
1202 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1203 let gpu = memra_runtime::Gpu::new(ordinal)?;
1204 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1208 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1209 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1210 .and_then(|d| unsafe {
1211 Ok((
1212 cudarc::driver::result::device::get_attribute(
1213 d,
1214 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1215 )?,
1216 cudarc::driver::result::device::get_attribute(
1217 d,
1218 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1219 )?,
1220 ))
1221 })
1222 .unwrap_or((0, 0));
1223 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1224 let ok = matches!(
1225 (built, maj, min),
1226 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1227 );
1228 if !ok {
1229 return Err(format!(
1230 "memra was built for sm_{built} but device {ordinal} reports compute \
1231 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1232 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1233 )
1234 .into());
1235 }
1236 }
1237 unsafe {
1242 use cudarc::driver::sys;
1243 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1244 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1245 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1246 let mut thresh: u64 = u64::MAX;
1247 let _ = sys::cuMemPoolSetAttribute(
1248 pool,
1249 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1250 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1251 );
1252 }
1253 }
1254 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1255 let hybrid = gpu
1256 .ctx
1257 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1258 let qmatvec = gpu
1259 .ctx
1260 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1261 let flash = gpu
1262 .ctx
1263 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1264 let gemm = gpu
1265 .ctx
1266 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1267 let router = gpu
1268 .ctx
1269 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1270 let sample = gpu
1271 .ctx
1272 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1273 let copy_stream = gpu.ctx.new_stream()?;
1274 if std::env::var("MEMRA_EVT")
1290 .map(|v| v == "1")
1291 .unwrap_or(false)
1292 {
1293 } else {
1295 unsafe {
1296 gpu.ctx.disable_event_tracking();
1297 }
1298 }
1299 Ok(Self {
1300 gpu,
1301 module,
1302 hybrid,
1303 qmatvec,
1304 flash,
1305 flash_g: std::sync::OnceLock::new(),
1306 gemm,
1307 router,
1308 sample,
1309 moe_cache: Mutex::new(None),
1310 moe_cache_layout: Mutex::new(None),
1311 copy_stream,
1312 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1313 verify_exact: std::sync::atomic::AtomicBool::new(false),
1314 capture_keep: Mutex::new(Vec::new()),
1315 argmax_partials: Mutex::new(None),
1316 prime_deqw_ws: Mutex::new(None),
1317 router_stage: Mutex::new(None),
1318 fp8_scratch: Mutex::new(None),
1319 fa_vf16_scratch: Mutex::new(None),
1320 fa_part_pool: Mutex::new(None),
1321 fa_part_retired: Mutex::new(Vec::new()),
1322 fn_cache: Mutex::new(Default::default()),
1323 f16_scratch: Mutex::new(None),
1324 #[cfg(memra_cutlass)]
1325 cutlass_scratch: Mutex::new(None),
1326 })
1327 }
1328
1329 pub fn ctx(&self) -> &Arc<CudaContext> {
1330 &self.gpu.ctx
1331 }
1332
1333 pub fn pool_cached_bytes(&self) -> usize {
1351 let (reserved, used) = self.pool_reserved_used();
1352 reserved.saturating_sub(used)
1353 }
1354
1355 pub fn pool_reserved_used(&self) -> (usize, usize) {
1362 use cudarc::driver::sys;
1363 unsafe {
1364 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1365 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1366 != sys::CUresult::CUDA_SUCCESS
1367 {
1368 return (0, 0);
1369 }
1370 let (mut reserved, mut used) = (0u64, 0u64);
1371 if sys::cuMemPoolGetAttribute(
1372 pool,
1373 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1374 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1375 ) != sys::CUresult::CUDA_SUCCESS
1376 {
1377 return (0, 0);
1378 }
1379 if sys::cuMemPoolGetAttribute(
1380 pool,
1381 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1382 &mut used as *mut u64 as *mut core::ffi::c_void,
1383 ) != sys::CUresult::CUDA_SUCCESS
1384 {
1385 return (0, 0);
1386 }
1387 (reserved as usize, used as usize)
1388 }
1389 }
1390
1391 pub fn stream(&self) -> Arc<CudaStream> {
1394 self.gpu.stream()
1395 }
1396 pub fn gkv_on() -> bool {
1399 memra_kv::gkv_on()
1400 }
1401
1402 pub fn wkv_on() -> bool {
1414 memra_kv::wkv_on()
1415 }
1416
1417 pub fn kv_fp8_on() -> bool {
1423 memra_kv::kv_fp8_on()
1424 }
1425
1426 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1429 if head_dim == 512 && Self::gkv_on() {
1430 self.func_g(name)
1431 } else {
1432 self.func(name)
1433 }
1434 }
1435
1436 fn func_g(&self, name: &str) -> CudaFunction {
1440 let m = self.flash_g.get_or_init(|| {
1441 self.gpu
1442 .ctx
1443 .load_module(cudarc::nvrtc::Ptx::from_binary(
1444 FLASH_FATBIN_KF8VF8.to_vec(),
1445 ))
1446 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1447 });
1448 let key = format!("g:{name}");
1449 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1450 return f.clone();
1451 }
1452 let f = match m.load_function(name) {
1453 Ok(f) => f,
1454 Err(_) => self.func(name),
1455 };
1456 self.fn_cache.lock().unwrap().insert(key, f.clone());
1457 f
1458 }
1459
1460 fn func(&self, name: &str) -> CudaFunction {
1461 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1464 return f.clone();
1465 }
1466 let f = self
1467 .module
1468 .load_function(name)
1469 .or_else(|_| self.hybrid.load_function(name))
1470 .or_else(|_| self.qmatvec.load_function(name))
1471 .or_else(|_| self.flash.load_function(name))
1472 .or_else(|_| self.gemm.load_function(name))
1473 .or_else(|_| self.router.load_function(name))
1474 .or_else(|_| self.sample.load_function(name))
1475 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1476 self.fn_cache
1477 .lock()
1478 .unwrap()
1479 .insert(name.to_string(), f.clone());
1480 f
1481 }
1482
1483 pub fn scatter_trim_logits(
1486 &self,
1487 src: &CudaSlice<f32>,
1488 d2t: &CudaSlice<u32>,
1489 dst: &mut CudaSlice<f32>,
1490 d_vocab: usize,
1491 n_vocab: usize,
1492 ) -> Result<(), Box<dyn std::error::Error>> {
1493 let f1 = self.func("scatter_trim_logits_f32");
1494 let f2 = self.func("scatter_trim_logits_pass2_f32");
1495 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1496 let cfg1 = LaunchConfig {
1497 grid_dim: (256, 1, 1),
1498 block_dim: (256, 1, 1),
1499 shared_mem_bytes: 0,
1500 };
1501 let __s_b1 = self.gpu.stream();
1502 let mut b1 = __s_b1.launch_builder(&f1);
1503 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1504 unsafe {
1505 b1.launch(cfg1)?;
1506 }
1507 let cfg2 = LaunchConfig {
1508 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1509 block_dim: (256, 1, 1),
1510 shared_mem_bytes: 0,
1511 };
1512 let __s_b2 = self.gpu.stream();
1513 let mut b2 = __s_b2.launch_builder(&f2);
1514 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1515 unsafe {
1516 b2.launch(cfg2)?;
1517 }
1518 Ok(())
1519 }
1520
1521 #[allow(clippy::too_many_arguments)]
1527 pub fn filter_stats(
1528 &self,
1529 x: &CudaSlice<f32>,
1530 row_stride: usize,
1531 rows: &CudaSlice<i32>,
1532 out_th: &mut CudaSlice<f32>,
1533 out_z: &mut CudaSlice<f32>,
1534 out_max: &mut CudaSlice<f32>,
1535 n: usize,
1536 nrow: usize,
1537 temp: f32,
1538 top_k: i32,
1539 top_p: f32,
1540 min_p: f32,
1541 ) -> Result<(), Box<dyn std::error::Error>> {
1542 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1557 let coop_on =
1558 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1559 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1560 if coop_on && 16 * nrow <= self.sm_count() as usize {
1561 let f = self.func("filter_stats_coop_f32");
1562 let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1563 let cfg = LaunchConfig {
1564 grid_dim: (16, nrow as u32, 1),
1565 block_dim: (512, 1, 1),
1566 shared_mem_bytes: 0,
1567 };
1568 let __s_b = self.gpu.stream();
1569 let mut b = __s_b.launch_builder(&f);
1570 b.arg(x)
1571 .arg(&rs)
1572 .arg(rows)
1573 .arg(&mut *out_th)
1574 .arg(&mut *out_z)
1575 .arg(&mut *out_max)
1576 .arg(&mut ws)
1577 .arg(&ni)
1578 .arg(&nr)
1579 .arg(&temp)
1580 .arg(&top_k)
1581 .arg(&top_p)
1582 .arg(&min_p);
1583 unsafe {
1584 b.launch_cooperative(cfg)?;
1585 }
1586 return Ok(());
1587 }
1588 let f = self.func("filter_stats_f32");
1589 let cfg = LaunchConfig {
1590 grid_dim: (nrow as u32, 1, 1),
1591 block_dim: (1024, 1, 1),
1592 shared_mem_bytes: 0,
1593 };
1594 let __s_b = self.gpu.stream();
1595 let mut b = __s_b.launch_builder(&f);
1596 b.arg(x)
1597 .arg(&rs)
1598 .arg(rows)
1599 .arg(&mut *out_th)
1600 .arg(&mut *out_z)
1601 .arg(&mut *out_max)
1602 .arg(&ni)
1603 .arg(&nr)
1604 .arg(&temp)
1605 .arg(&top_k)
1606 .arg(&top_p)
1607 .arg(&min_p);
1608 unsafe {
1609 b.launch(cfg)?;
1610 }
1611 Ok(())
1612 }
1613
1614 #[allow(clippy::too_many_arguments)]
1616 pub fn softmax_gather_filtered(
1617 &self,
1618 x: &CudaSlice<f32>,
1619 row_stride: usize,
1620 ids: &CudaSlice<u32>,
1621 rows: &CudaSlice<i32>,
1622 th: &CudaSlice<f32>,
1623 z: &CudaSlice<f32>,
1624 out: &mut CudaSlice<f32>,
1625 n: usize,
1626 npair: usize,
1627 temp: f32,
1628 ) -> Result<(), Box<dyn std::error::Error>> {
1629 let f = self.func("softmax_gather_filtered_f32");
1630 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1631 let cfg = LaunchConfig {
1632 grid_dim: (npair as u32, 1, 1),
1633 block_dim: (256, 1, 1),
1634 shared_mem_bytes: 0,
1635 };
1636 let __s_b = self.gpu.stream();
1637 let mut b = __s_b.launch_builder(&f);
1638 b.arg(x)
1639 .arg(&rs)
1640 .arg(ids)
1641 .arg(rows)
1642 .arg(th)
1643 .arg(z)
1644 .arg(&mut *out)
1645 .arg(&ni)
1646 .arg(&np)
1647 .arg(&temp);
1648 unsafe {
1649 b.launch(cfg)?;
1650 }
1651 Ok(())
1652 }
1653
1654 #[allow(clippy::too_many_arguments)]
1656 pub fn residual_sample_filtered(
1657 &self,
1658 p: &CudaSlice<f32>,
1659 q: Option<&CudaSlice<f32>>,
1660 n: usize,
1661 temp: f32,
1662 seed: u64,
1663 stream_pos: u32,
1664 p_stats: (f32, f32, f32),
1665 q_stats: (f32, f32, f32),
1666 out_tok: &mut CudaSlice<u32>,
1667 ) -> Result<(), Box<dyn std::error::Error>> {
1668 let f = self.func("residual_sample_filtered_f32");
1669 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1670 let has_q: i32 = q.is_some() as i32;
1671 let qbuf = q.unwrap_or(p);
1672 let (pm, pth, pz) = p_stats;
1673 let (qm, qth, qz) = q_stats;
1674 let cfg = LaunchConfig {
1675 grid_dim: (1, 1, 1),
1676 block_dim: (1024, 1, 1),
1677 shared_mem_bytes: 0,
1678 };
1679 let __s_b = self.gpu.stream();
1680 let mut b = __s_b.launch_builder(&f);
1681 b.arg(p)
1682 .arg(qbuf)
1683 .arg(&has_q)
1684 .arg(&ni)
1685 .arg(&temp)
1686 .arg(&slo)
1687 .arg(&shi)
1688 .arg(&stream_pos)
1689 .arg(&pm)
1690 .arg(&pth)
1691 .arg(&pz)
1692 .arg(&qm)
1693 .arg(&qth)
1694 .arg(&qz)
1695 .arg(&mut *out_tok);
1696 unsafe {
1697 b.launch(cfg)?;
1698 }
1699 Ok(())
1700 }
1701
1702 #[allow(clippy::too_many_arguments)]
1708 pub fn residual_sample_sparse_q(
1709 &self,
1710 p: &CudaSlice<f32>,
1711 cand_ids: &CudaSlice<u32>,
1712 q_probs: &CudaSlice<f32>,
1713 n_cand: usize,
1714 n: usize,
1715 temp: f32,
1716 seed: u64,
1717 stream_pos: u32,
1718 p_stats: (f32, f32, f32),
1719 out_tok: &mut CudaSlice<u32>,
1720 ) -> Result<(), Box<dyn std::error::Error>> {
1721 assert!(
1722 n_cand >= 1 && n_cand <= 32,
1723 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1724 );
1725 let f = self.func("residual_sample_sparse_q_f32");
1726 let (ni, nc) = (n as i32, n_cand as i32);
1727 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1728 let (pm, pth, pz) = p_stats;
1729 let cfg = LaunchConfig {
1730 grid_dim: (1, 1, 1),
1731 block_dim: (1024, 1, 1),
1732 shared_mem_bytes: 0,
1733 };
1734 let __s_b = self.gpu.stream();
1735 let mut b = __s_b.launch_builder(&f);
1736 b.arg(p)
1737 .arg(cand_ids)
1738 .arg(q_probs)
1739 .arg(&nc)
1740 .arg(&ni)
1741 .arg(&temp)
1742 .arg(&slo)
1743 .arg(&shi)
1744 .arg(&stream_pos)
1745 .arg(&pm)
1746 .arg(&pth)
1747 .arg(&pz)
1748 .arg(&mut *out_tok);
1749 unsafe {
1750 b.launch(cfg)?;
1751 }
1752 Ok(())
1753 }
1754
1755 #[allow(clippy::too_many_arguments)]
1757 pub fn gumbel_perturb_filtered(
1758 &self,
1759 x: &CudaSlice<f32>,
1760 y: &mut CudaSlice<f32>,
1761 n: usize,
1762 seed: u64,
1763 stream_pos: u32,
1764 temp: f32,
1765 row_max: f32,
1766 th: f32,
1767 ) -> Result<(), Box<dyn std::error::Error>> {
1768 let f = self.func("gumbel_perturb_filtered_f32");
1769 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1770 let cfg = LaunchConfig {
1771 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1772 block_dim: (256, 1, 1),
1773 shared_mem_bytes: 0,
1774 };
1775 let __s_b = self.gpu.stream();
1776 let mut b = __s_b.launch_builder(&f);
1777 b.arg(x)
1778 .arg(&mut *y)
1779 .arg(&ni)
1780 .arg(&slo)
1781 .arg(&shi)
1782 .arg(&stream_pos)
1783 .arg(&temp)
1784 .arg(&row_max)
1785 .arg(&th);
1786 unsafe {
1787 b.launch(cfg)?;
1788 }
1789 Ok(())
1790 }
1791
1792 #[allow(clippy::too_many_arguments)]
1796 pub fn penalize_logits(
1797 &self,
1798 x: &mut CudaSlice<f32>,
1799 hist: &CudaSlice<u32>,
1800 n_hist: usize,
1801 rep: f32,
1802 freq: f32,
1803 present: f32,
1804 n: usize,
1805 ) -> Result<(), Box<dyn std::error::Error>> {
1806 if n_hist == 0 {
1807 return Ok(());
1808 }
1809 let f = self.func("penalize_logits_f32");
1810 let (nh, ni) = (n_hist as i32, n as i32);
1811 let cfg = LaunchConfig {
1812 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1813 block_dim: (128, 1, 1),
1814 shared_mem_bytes: 0,
1815 };
1816 let __s_b = self.gpu.stream();
1817 let mut b = __s_b.launch_builder(&f);
1818 b.arg(&mut *x)
1819 .arg(hist)
1820 .arg(&nh)
1821 .arg(&rep)
1822 .arg(&freq)
1823 .arg(&present)
1824 .arg(&ni);
1825 unsafe {
1826 b.launch(cfg)?;
1827 }
1828 Ok(())
1829 }
1830
1831 #[allow(clippy::too_many_arguments)]
1833 pub fn penalize_logits_rows(
1834 &self,
1835 x: &mut CudaSlice<f32>,
1836 hist: &CudaSlice<u32>,
1837 n_hist: usize,
1838 rep: f32,
1839 freq: f32,
1840 present: f32,
1841 n: usize,
1842 nrow: usize,
1843 ) -> Result<(), Box<dyn std::error::Error>> {
1844 if n_hist == 0 || nrow == 0 {
1845 return Ok(());
1846 }
1847 let f = self.func("penalize_logits_rows_f32");
1848 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1849 let cfg = LaunchConfig {
1850 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1851 block_dim: (128, 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(&mut *x)
1857 .arg(hist)
1858 .arg(&nh)
1859 .arg(&rep)
1860 .arg(&freq)
1861 .arg(&present)
1862 .arg(&ni)
1863 .arg(&nr);
1864 unsafe {
1865 b.launch(cfg)?;
1866 }
1867 Ok(())
1868 }
1869
1870 #[allow(clippy::too_many_arguments)]
1878 pub fn penalize_logits_rows_inc(
1879 &self,
1880 x: &mut CudaSlice<f32>,
1881 hist: &CudaSlice<u32>,
1882 n_hist0: usize,
1883 rep: f32,
1884 freq: f32,
1885 present: f32,
1886 n: usize,
1887 nrow: usize,
1888 win: usize,
1889 ) -> Result<(), Box<dyn std::error::Error>> {
1890 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1891 return Ok(());
1892 }
1893 debug_assert!(
1894 hist.len() >= n_hist0 + nrow - 1,
1895 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1896 );
1897 let f = self.func("penalize_logits_rows_inc_f32");
1898 let max_len = win.min(n_hist0 + nrow - 1).max(1);
1899 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1900 let cfg = LaunchConfig {
1901 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1902 block_dim: (128, 1, 1),
1903 shared_mem_bytes: 0,
1904 };
1905 let __s_b = self.gpu.stream();
1906 let mut b = __s_b.launch_builder(&f);
1907 b.arg(&mut *x)
1908 .arg(hist)
1909 .arg(&nh)
1910 .arg(&rep)
1911 .arg(&freq)
1912 .arg(&present)
1913 .arg(&ni)
1914 .arg(&nr)
1915 .arg(&wi);
1916 unsafe {
1917 b.launch(cfg)?;
1918 }
1919 Ok(())
1920 }
1921
1922 pub fn wpf_level() -> u32 {
1930 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1931 *ON.get_or_init(|| {
1932 std::env::var("MEMRA_WPF")
1933 .ok()
1934 .and_then(|v| v.parse().ok())
1935 .unwrap_or(1)
1936 })
1937 }
1938
1939 pub fn set_verify_exact(&self, on: bool) {
1951 self.verify_exact
1952 .store(on, std::sync::atomic::Ordering::Relaxed);
1953 }
1954 pub(crate) fn verify_exact_on(&self) -> bool {
1955 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1956 }
1957
1958 pub fn qkv_append_on() -> bool {
1961 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1962 *ON.get_or_init(|| {
1963 std::env::var("MEMRA_QKV_APPEND")
1964 .map(|v| v != "0")
1965 .unwrap_or(true)
1966 })
1967 }
1968
1969 pub fn pdl_wb_on() -> bool {
1972 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1973 *ON.get_or_init(|| {
1974 std::env::var("MEMRA_PDL_WB")
1975 .map(|v| v != "0")
1976 .unwrap_or(true)
1977 })
1978 }
1979
1980 pub fn norm_ilp_on() -> bool {
1988 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1989 *ON.get_or_init(|| {
1990 std::env::var("MEMRA_NORM_ILP")
1991 .map(|v| v != "0")
1992 .unwrap_or(true)
1993 })
1994 }
1995
1996 pub fn tk_ffn_dual_on() -> bool {
2004 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2005 *ON.get_or_init(|| {
2006 std::env::var("MEMRA_TK_FFN_DUAL")
2007 .map(|v| v != "0")
2008 .unwrap_or(true)
2009 })
2010 }
2011
2012 pub fn pdl_mmvq_on() -> bool {
2016 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2017 *ON.get_or_init(|| {
2018 std::env::var("MEMRA_PDL_MMVQ")
2019 .map(|v| v != "0")
2020 .unwrap_or(true)
2021 })
2022 }
2023
2024 pub fn pdl_on() -> bool {
2025 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2026 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2027 }
2028
2029 pub fn pdl_nvfp4q8_on() -> bool {
2035 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2036 *ON.get_or_init(|| {
2037 std::env::var("MEMRA_PDL_NVFP4")
2038 .map(|v| v != "0")
2039 .unwrap_or(true)
2040 })
2041 }
2042
2043 fn q40_mr1_on() -> bool {
2049 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2050 match *Q40MR.get_or_init(|| {
2051 std::env::var("MEMRA_Q40_MR")
2052 .ok()
2053 .and_then(|v| v.parse().ok())
2054 }) {
2055 Some(v) => v == 1,
2056 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2057 }
2058 }
2059
2060 fn pdl_func_flash(
2065 &self,
2066 g: bool,
2067 name: &'static str,
2068 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2069 use cudarc::driver::sys as cu;
2070 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2077 std::sync::Mutex::new(None);
2078 static FNS: std::sync::Mutex<
2079 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2080 > = std::sync::Mutex::new(None);
2081 let ctx_key = self.ctx().cu_ctx() as usize;
2082 if let Some(&f) = FNS
2083 .lock()
2084 .unwrap()
2085 .get_or_insert_with(Default::default)
2086 .get(&(ctx_key, g, name))
2087 {
2088 return Ok(f as cu::CUfunction);
2089 }
2090 let module = {
2091 let mut mods = MODS.lock().unwrap();
2092 let map = mods.get_or_insert_with(Default::default);
2093 match map.get(&(ctx_key, g)) {
2094 Some(&m) => m,
2095 None => {
2096 let m = self.pdl_load_module_in_ctx(if g {
2097 FLASH_FATBIN_KF8VF8
2098 } else {
2099 FLASH_FATBIN
2100 })?;
2101 map.insert((ctx_key, g), m);
2102 m
2103 }
2104 }
2105 };
2106 let cname = std::ffi::CString::new(name)?;
2107 let mut f: cu::CUfunction = std::ptr::null_mut();
2108 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2109 if r != cu::CUresult::CUDA_SUCCESS {
2110 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2111 }
2112 FNS.lock()
2113 .unwrap()
2114 .get_or_insert_with(Default::default)
2115 .insert((ctx_key, g, name), f as usize);
2116 Ok(f)
2117 }
2118
2119 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2124 use cudarc::driver::sys as cu;
2125 let mut prev: cu::CUcontext = std::ptr::null_mut();
2126 unsafe {
2127 cu::cuCtxGetCurrent(&mut prev).result()?;
2128 }
2129 self.ctx().bind_to_thread()?;
2130 let mut m: cu::CUmodule = std::ptr::null_mut();
2131 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2132 let restore = if prev.is_null() {
2133 cu::CUresult::CUDA_SUCCESS
2134 } else {
2135 unsafe { cu::cuCtxSetCurrent(prev) }
2136 };
2137 if r != cu::CUresult::CUDA_SUCCESS {
2138 return Err(format!("pdl module load: {r:?}").into());
2139 }
2140 if restore != cu::CUresult::CUDA_SUCCESS {
2141 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2142 }
2143 Ok(m as usize)
2144 }
2145
2146 pub fn raw_kernel_function(
2149 &self,
2150 name: &'static str,
2151 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2152 self.pdl_func(name)
2153 }
2154
2155 fn pdl_func(
2156 &self,
2157 name: &'static str,
2158 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2159 use cudarc::driver::sys as cu;
2160 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2163 std::sync::Mutex::new(None);
2164 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2167 std::sync::Mutex::new(None);
2168 static FNS: std::sync::Mutex<
2169 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2170 > = std::sync::Mutex::new(None);
2171 let ctx_key = self.ctx().cu_ctx() as usize;
2172 if let Some(&f) = FNS
2173 .lock()
2174 .unwrap()
2175 .get_or_insert_with(Default::default)
2176 .get(&(ctx_key, name))
2177 {
2178 return Ok(f as cu::CUfunction);
2179 }
2180 let module = {
2181 let mut mods = MODULES.lock().unwrap();
2182 let map = mods.get_or_insert_with(Default::default);
2183 match map.get(&ctx_key) {
2184 Some(&m) => m,
2185 None => {
2186 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2187 map.insert(ctx_key, m);
2188 m
2189 }
2190 }
2191 };
2192 let cname = std::ffi::CString::new(name)?;
2193 let mut f: cu::CUfunction = std::ptr::null_mut();
2194 let mut r =
2195 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2196 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2197 let qmodule = {
2198 let mut mods = QMODULES.lock().unwrap();
2199 let map = mods.get_or_insert_with(Default::default);
2200 match map.get(&ctx_key) {
2201 Some(&m) => m,
2202 None => {
2203 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2204 map.insert(ctx_key, m);
2205 m
2206 }
2207 }
2208 };
2209 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2210 }
2211 if r != cu::CUresult::CUDA_SUCCESS {
2212 return Err(format!("pdl_func {name}: {r:?}").into());
2213 }
2214 FNS.lock()
2215 .unwrap()
2216 .get_or_insert_with(Default::default)
2217 .insert((ctx_key, name), f as usize);
2218 Ok(f)
2219 }
2220
2221 unsafe fn launch_pdl_flash(
2233 &self,
2234 g: bool,
2235 name: &'static str,
2236 grid: (u32, u32, u32),
2237 block: (u32, u32, u32),
2238 smem: u32,
2239 params: &mut [*mut std::ffi::c_void],
2240 ) -> Result<(), Box<dyn std::error::Error>> {
2241 use cudarc::driver::sys as cu;
2242 let f = self.pdl_func_flash(g, name)?;
2243 if smem > 0 {
2244 let r =
2246 unsafe {
2247 cu::cuFuncSetAttribute(f,
2248 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2249 smem as i32)
2250 };
2251 if r != cu::CUresult::CUDA_SUCCESS {
2252 return Err(format!("pdl smem attr {name}: {r:?}").into());
2253 }
2254 }
2255 let mut attr = cu::CUlaunchAttribute {
2256 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2257 pad: [0; 4],
2258 value: cu::CUlaunchAttributeValue {
2259 programmaticStreamSerializationAllowed: 1,
2260 },
2261 };
2262 let cfg = cu::CUlaunchConfig {
2263 gridDimX: grid.0,
2264 gridDimY: grid.1,
2265 gridDimZ: grid.2,
2266 blockDimX: block.0,
2267 blockDimY: block.1,
2268 blockDimZ: block.2,
2269 sharedMemBytes: smem,
2270 hStream: self.gpu.stream().cu_stream(),
2271 attrs: &mut attr,
2272 numAttrs: 1,
2273 };
2274 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2275 if r != cu::CUresult::CUDA_SUCCESS {
2276 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2277 }
2278 Ok(())
2279 }
2280
2281 unsafe fn launch_pdl(
2282 &self,
2283 name: &'static str,
2284 grid: (u32, u32, u32),
2285 block: (u32, u32, u32),
2286 params: &mut [*mut std::ffi::c_void],
2287 ) -> Result<(), Box<dyn std::error::Error>> {
2288 use cudarc::driver::sys as cu;
2289 let f = self.pdl_func(name)?;
2290 let mut attr = cu::CUlaunchAttribute {
2291 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2292 pad: [0; 4],
2293 value: cu::CUlaunchAttributeValue {
2294 programmaticStreamSerializationAllowed: 1,
2295 },
2296 };
2297 let cfg = cu::CUlaunchConfig {
2298 gridDimX: grid.0,
2299 gridDimY: grid.1,
2300 gridDimZ: grid.2,
2301 blockDimX: block.0,
2302 blockDimY: block.1,
2303 blockDimZ: block.2,
2304 sharedMemBytes: 0,
2305 hStream: self.gpu.stream().cu_stream(),
2306 attrs: &mut attr,
2307 numAttrs: 1,
2308 };
2309 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2310 if r != cu::CUresult::CUDA_SUCCESS {
2311 return Err(format!("launch_pdl {name}: {r:?}").into());
2312 }
2313 Ok(())
2314 }
2315
2316 pub fn prefetch_weight_l2(
2319 &self,
2320 w: &crate::model::GpuTensor,
2321 ) -> Result<(), Box<dyn std::error::Error>> {
2322 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2323 let p = rp4.as_ref().unwrap_or(bytes);
2324 self.prefetch_l2(p, p.len())?;
2325 }
2326 Ok(())
2327 }
2328
2329 pub fn gather_row_bf16(
2332 &self,
2333 table: &CudaSlice<u8>,
2334 tok: &CudaSlice<u32>,
2335 idx: usize,
2336 dst: &mut CudaSlice<f32>,
2337 ncols: usize,
2338 ) -> Result<(), Box<dyn std::error::Error>> {
2339 let f = self.func("gather_row_bf16_f32");
2340 let cfg = LaunchConfig {
2341 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2342 block_dim: (256, 1, 1),
2343 shared_mem_bytes: 0,
2344 };
2345 let (nc, ix) = (ncols as i32, idx as i32);
2346 let __s_b = self.gpu.stream();
2347 let mut b = __s_b.launch_builder(&f);
2348 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2349 unsafe {
2350 b.launch(cfg)?;
2351 }
2352 Ok(())
2353 }
2354
2355 #[allow(clippy::too_many_arguments)]
2361 pub fn dflash2_dynconv(
2362 &self,
2363 x: &CudaSlice<f32>,
2364 dyn_: &CudaSlice<f32>,
2365 base: &CudaSlice<f32>,
2366 out: &mut CudaSlice<f32>,
2367 rows: usize,
2368 hidden: usize,
2369 group_size: usize,
2370 ksize: usize,
2371 half: usize,
2372 ) -> Result<(), Box<dyn std::error::Error>> {
2373 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2374 let f = self.func("dflash2_dynconv_f32");
2375 let n = rows * hidden;
2376 let cfg = LaunchConfig {
2377 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2378 block_dim: (256, 1, 1),
2379 shared_mem_bytes: 0,
2380 };
2381 let (ri, hi, gi, ki, hf) = (
2382 rows as i32,
2383 hidden as i32,
2384 group_size as i32,
2385 ksize as i32,
2386 half as i32,
2387 );
2388 let __s_b = self.gpu.stream();
2389 let mut b = __s_b.launch_builder(&f);
2390 b.arg(x)
2391 .arg(dyn_)
2392 .arg(base)
2393 .arg(out)
2394 .arg(&ri)
2395 .arg(&hi)
2396 .arg(&gi)
2397 .arg(&ki)
2398 .arg(&hf);
2399 unsafe {
2400 b.launch(cfg)?;
2401 }
2402 Ok(())
2403 }
2404
2405 pub fn topk_rows(
2409 &self,
2410 logits: &CudaSlice<f32>,
2411 n_rows: usize,
2412 n_cols: usize,
2413 k: usize,
2414 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2415 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2416 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2417 let f = self.func("topk_rows_f32");
2418 let nth = 256usize;
2419 let mut vals = self.uninit(n_rows * k)?;
2420 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2421 let cfg = LaunchConfig {
2422 grid_dim: (n_rows as u32, 1, 1),
2423 block_dim: (nth as u32, 1, 1),
2424 shared_mem_bytes: (nth * k * 8) as u32,
2425 };
2426 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2427 let __s_b = self.gpu.stream();
2428 let mut b = __s_b.launch_builder(&f);
2429 b.arg(logits)
2430 .arg(&nr)
2431 .arg(&nc)
2432 .arg(&ki)
2433 .arg(&mut vals)
2434 .arg(&mut idxs);
2435 unsafe {
2436 b.launch(cfg)?;
2437 }
2438 Ok((vals, idxs))
2439 }
2440
2441 pub fn add_row_inplace(
2443 &self,
2444 logits: &mut CudaSlice<f32>,
2445 bias: &CudaSlice<f32>,
2446 n: usize,
2447 row_off: usize,
2448 ) -> Result<(), Box<dyn std::error::Error>> {
2449 let f = self.func("add_row_inplace_f32");
2450 let cfg = LaunchConfig {
2451 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2452 block_dim: (256, 1, 1),
2453 shared_mem_bytes: 0,
2454 };
2455 let (ni, off) = (n as i32, row_off as i64);
2456 let __s_b = self.gpu.stream();
2457 let mut b = __s_b.launch_builder(&f);
2458 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2459 unsafe {
2460 b.launch(cfg)?;
2461 }
2462 Ok(())
2463 }
2464
2465 pub fn prefetch_l2(
2467 &self,
2468 p: &CudaSlice<u8>,
2469 n: usize,
2470 ) -> Result<(), Box<dyn std::error::Error>> {
2471 let f = self.func("prefetch_l2_bytes");
2472 let lines = n.div_ceil(128);
2473 let ni = n as i64;
2474 let cfg = LaunchConfig {
2475 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2476 block_dim: (256, 1, 1),
2477 shared_mem_bytes: 0,
2478 };
2479 let __s_b = self.gpu.stream();
2480 let mut b = __s_b.launch_builder(&f);
2481 b.arg(p).arg(&ni);
2482 unsafe {
2483 b.launch(cfg)?;
2484 }
2485 Ok(())
2486 }
2487
2488 pub fn router_gemv(
2491 &self,
2492 w: &CudaSlice<f32>,
2493 x: &CudaSlice<f32>,
2494 n_embd: usize,
2495 n_experts: usize,
2496 t: usize,
2497 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2498 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2504 Ok("0") => false,
2505 Ok(_) => true,
2506 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2507 };
2508 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2518 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2519 }
2520
2521 pub fn router_gemv_form(
2524 &self,
2525 w: &CudaSlice<f32>,
2526 x: &CudaSlice<f32>,
2527 n_embd: usize,
2528 n_experts: usize,
2529 t: usize,
2530 w8: bool,
2531 batch: bool,
2532 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2533 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2534 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2535 let f = if batch {
2536 self.func("router_gemv_f32_w8_batch")
2537 } else if w8 {
2538 self.func("router_gemv_f32_w8")
2539 } else {
2540 self.func("router_gemv_f32")
2541 };
2542 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2543 let cfg = if batch {
2544 LaunchConfig {
2545 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2546 block_dim: (32, 8, 1),
2547 shared_mem_bytes: 0,
2548 }
2549 } else {
2550 LaunchConfig {
2551 grid_dim: (n_experts as u32, t as u32, 1),
2552 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2553 shared_mem_bytes: 0,
2554 }
2555 };
2556 let __s_b = self.gpu.stream();
2557 let mut b = __s_b.launch_builder(&f);
2558 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2559 unsafe {
2560 b.launch(cfg)?;
2561 }
2562 Ok(y)
2563 }
2564
2565 pub fn router_gemv_into(
2568 &self,
2569 w: &CudaSlice<f32>,
2570 x: &CudaSlice<f32>,
2571 y: &mut CudaSlice<f32>,
2572 n_embd: usize,
2573 n_experts: usize,
2574 t: usize,
2575 ) -> Result<(), Box<dyn std::error::Error>> {
2576 if y.len() < t * n_experts {
2577 return Err("router_gemv_into output too small".into());
2578 }
2579 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2580 Ok("0") => false,
2581 Ok(_) => true,
2582 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2583 };
2584 let f = if w8 {
2585 self.func("router_gemv_f32_w8")
2586 } else {
2587 self.func("router_gemv_f32")
2588 };
2589 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2590 let cfg = LaunchConfig {
2591 grid_dim: (n_experts as u32, t as u32, 1),
2592 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2593 shared_mem_bytes: 0,
2594 };
2595 let __s_b = self.gpu.stream();
2596 let mut b = __s_b.launch_builder(&f);
2597 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2598 unsafe {
2599 b.launch(cfg)?;
2600 }
2601 Ok(())
2602 }
2603
2604 pub fn rows_permute(
2606 &self,
2607 src: &CudaSlice<f32>,
2608 idx: &CudaSlice<i32>,
2609 nrows: usize,
2610 ncols: usize,
2611 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2612 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2613 let f = self.func("rows_permute_f32");
2614 let (nc, nr) = (ncols as i32, nrows as i32);
2615 let cfg = LaunchConfig {
2616 grid_dim: (nrows as u32, 1, 1),
2617 block_dim: (256, 1, 1),
2618 shared_mem_bytes: 0,
2619 };
2620 let __s_b = self.gpu.stream();
2621 let mut b = __s_b.launch_builder(&f);
2622 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2623 unsafe {
2624 b.launch(cfg)?;
2625 }
2626 Ok(dst)
2627 }
2628
2629 pub fn sigmoid_dot_rows(
2634 &self,
2635 x: &CudaSlice<f32>,
2636 w: &CudaSlice<f32>,
2637 n_embd: usize,
2638 t: usize,
2639 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2640 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2643 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2644 let gs = self.linear(x, w, t, n_embd, 1)?;
2645 let mut g = self.uninit(t)?;
2646 self.sigmoid(&gs, &mut g, t)?;
2647 return Ok(g);
2648 }
2649 let mut g = self.alloc_uninit::<f32>(t)?;
2655 let f = self.func("sigmoid_dot_rows_f32");
2656 let (ne, ti) = (n_embd as i32, t as i32);
2657 let cfg = LaunchConfig {
2658 grid_dim: (t as u32, 1, 1),
2659 block_dim: (32, 8, 1),
2660 shared_mem_bytes: 0,
2661 };
2662 let __s_b = self.gpu.stream();
2663 let mut b = __s_b.launch_builder(&f);
2664 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2665 unsafe {
2666 b.launch(cfg)?;
2667 }
2668 Ok(g)
2669 }
2670
2671 pub fn sigmoid_dot_rows_into(
2673 &self,
2674 x: &CudaSlice<f32>,
2675 w: &CudaSlice<f32>,
2676 g: &mut CudaSlice<f32>,
2677 n_embd: usize,
2678 t: usize,
2679 ) -> Result<(), Box<dyn std::error::Error>> {
2680 if g.len() < t {
2681 return Err("sigmoid_dot_rows_into output too small".into());
2682 }
2683 let f = self.func("sigmoid_dot_rows_f32");
2684 let (ne, ti) = (n_embd as i32, t as i32);
2685 let cfg = LaunchConfig {
2686 grid_dim: (t as u32, 1, 1),
2687 block_dim: (32, 8, 1),
2688 shared_mem_bytes: 0,
2689 };
2690 let __s_b = self.gpu.stream();
2691 let mut b = __s_b.launch_builder(&f);
2692 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2693 unsafe {
2694 b.launch(cfg)?;
2695 }
2696 Ok(())
2697 }
2698
2699 pub fn spec_rollback_stream(
2701 &self,
2702 len_ptrs: &CudaSlice<u64>,
2703 pos_start: &CudaSlice<i32>,
2704 acc: &CudaSlice<u32>,
2705 base: usize,
2706 n_rows: usize,
2707 ) -> Result<(), Box<dyn std::error::Error>> {
2708 let f = self.func("spec_rollback_stream");
2709 let (b, nr) = (base as i32, n_rows as i32);
2710 let cfg = LaunchConfig {
2711 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2712 block_dim: (64, 1, 1),
2713 shared_mem_bytes: 0,
2714 };
2715 let __s_bl = self.gpu.stream();
2716 let mut bl = __s_bl.launch_builder(&f);
2717 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2718 unsafe {
2719 bl.launch(cfg)?;
2720 }
2721 Ok(())
2722 }
2723
2724 pub fn plain_tok_ring(
2726 &self,
2727 vam: &CudaSlice<u32>,
2728 pos_start: &CudaSlice<i32>,
2729 base: usize,
2730 ring: &mut CudaSlice<u32>,
2731 ) -> Result<(), Box<dyn std::error::Error>> {
2732 let f = self.func("plain_tok_ring");
2733 let (b, cap) = (base as i32, ring.len() as i32);
2734 let cfg = LaunchConfig {
2735 grid_dim: (1, 1, 1),
2736 block_dim: (32, 1, 1),
2737 shared_mem_bytes: 0,
2738 };
2739 let __s_bl = self.gpu.stream();
2740 let mut bl = __s_bl.launch_builder(&f);
2741 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2742 unsafe {
2743 bl.launch(cfg)?;
2744 }
2745 Ok(())
2746 }
2747
2748 pub fn spec_ring_commit(
2750 &self,
2751 vtok: &CudaSlice<u32>,
2752 acc: &CudaSlice<u32>,
2753 brk: &CudaSlice<u32>,
2754 ring: &mut CudaSlice<u32>,
2755 pend: &mut CudaSlice<u32>,
2756 ) -> Result<(), Box<dyn std::error::Error>> {
2757 let f = self.func("spec_ring_commit");
2758 let cfg = LaunchConfig {
2759 grid_dim: (1, 1, 1),
2760 block_dim: (32, 1, 1),
2761 shared_mem_bytes: 0,
2762 };
2763 let __s_b = self.gpu.stream();
2764 let mut b = __s_b.launch_builder(&f);
2765 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2766 unsafe {
2767 b.launch(cfg)?;
2768 }
2769 Ok(())
2770 }
2771 pub fn i32_copy_add(
2772 &self,
2773 src: &CudaSlice<i32>,
2774 dst: &mut CudaSlice<i32>,
2775 delta: i32,
2776 ) -> Result<(), Box<dyn std::error::Error>> {
2777 let f = self.func("i32_copy_add");
2778 let cfg = LaunchConfig {
2779 grid_dim: (1, 1, 1),
2780 block_dim: (32, 1, 1),
2781 shared_mem_bytes: 0,
2782 };
2783 let __s_b = self.gpu.stream();
2784 let mut b = __s_b.launch_builder(&f);
2785 b.arg(src).arg(dst).arg(&delta);
2786 unsafe {
2787 b.launch(cfg)?;
2788 }
2789 Ok(())
2790 }
2791 pub fn u32_copy(
2792 &self,
2793 src: &CudaSlice<u32>,
2794 dst: &mut CudaSlice<u32>,
2795 ) -> Result<(), Box<dyn std::error::Error>> {
2796 let f = self.func("u32_copy");
2797 let cfg = LaunchConfig {
2798 grid_dim: (1, 1, 1),
2799 block_dim: (32, 1, 1),
2800 shared_mem_bytes: 0,
2801 };
2802 let __s_b = self.gpu.stream();
2803 let mut b = __s_b.launch_builder(&f);
2804 b.arg(src).arg(dst);
2805 unsafe {
2806 b.launch(cfg)?;
2807 }
2808 Ok(())
2809 }
2810
2811 pub fn spec_adapt_k(
2815 &self,
2816 acc: &CudaSlice<u32>,
2817 brk: &mut CudaSlice<u32>,
2818 floor: usize,
2819 cap: usize,
2820 ) -> Result<(), Box<dyn std::error::Error>> {
2821 let f = self.func("spec_adapt_k");
2822 let (fl, cp) = (floor as i32, cap as i32);
2823 let cfg = LaunchConfig {
2824 grid_dim: (1, 1, 1),
2825 block_dim: (32, 1, 1),
2826 shared_mem_bytes: 0,
2827 };
2828 let __s_b = self.gpu.stream();
2829 let mut b = __s_b.launch_builder(&f);
2830 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2831 unsafe {
2832 b.launch(cfg)?;
2833 }
2834 Ok(())
2835 }
2836
2837 pub fn spec_accept_greedy_dc(
2839 &self,
2840 preds: &CudaSlice<u32>,
2841 vtok: &CudaSlice<u32>,
2842 last_pred: &CudaSlice<u32>,
2843 brk: &CudaSlice<u32>,
2844 out: &mut CudaSlice<u32>,
2845 ) -> Result<(), Box<dyn std::error::Error>> {
2846 let f = self.func("spec_accept_greedy_dc");
2847 let cfg = LaunchConfig {
2848 grid_dim: (1, 1, 1),
2849 block_dim: (32, 1, 1),
2850 shared_mem_bytes: 0,
2851 };
2852 let __s_b = self.gpu.stream();
2853 let mut b = __s_b.launch_builder(&f);
2854 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2855 unsafe {
2856 b.launch(cfg)?;
2857 }
2858 Ok(())
2859 }
2860
2861 pub fn pos_iota(
2863 &self,
2864 pos0: &CudaSlice<i32>,
2865 out: &mut CudaSlice<i32>,
2866 t: usize,
2867 ) -> Result<(), Box<dyn std::error::Error>> {
2868 let f = self.func("pos_iota_i32");
2869 let ti = t as i32;
2870 let cfg = LaunchConfig {
2871 grid_dim: (1, 1, 1),
2872 block_dim: (t.max(1) as u32, 1, 1),
2873 shared_mem_bytes: 0,
2874 };
2875 let __s_b = self.gpu.stream();
2876 let mut b = __s_b.launch_builder(&f);
2877 b.arg(pos0).arg(out).arg(&ti);
2878 unsafe {
2879 b.launch(cfg)?;
2880 }
2881 Ok(())
2882 }
2883 #[allow(clippy::too_many_arguments)]
2884 pub fn append_kv_quantized_rows_dc(
2885 &self,
2886 k_rows: &CudaSlice<f32>,
2887 v_rows: &CudaSlice<f32>,
2888 kc: &mut CudaSlice<u8>,
2889 vc: &mut CudaSlice<u8>,
2890 t0_dev: &CudaSlice<i32>,
2891 t: usize,
2892 kv_dim_k: usize,
2893 kv_dim_v: usize,
2894 k_tok_bytes: usize,
2895 v_tok_bytes: usize,
2896 g: bool,
2897 ) -> Result<(), Box<dyn std::error::Error>> {
2898 let f = if g {
2899 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2900 } else {
2901 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2902 };
2903 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2904 let cfg = LaunchConfig {
2905 grid_dim: (nblk, t as u32, 1),
2906 block_dim: (32, 1, 1),
2907 shared_mem_bytes: 0,
2908 };
2909 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2910 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2911 let __s_b = self.gpu.stream();
2912 let mut b = __s_b.launch_builder(&f);
2913 b.arg(k_rows)
2914 .arg(v_rows)
2915 .arg(kc)
2916 .arg(vc)
2917 .arg(t0_dev)
2918 .arg(&kdk)
2919 .arg(&kdv)
2920 .arg(&ktb)
2921 .arg(&vtb);
2922 unsafe {
2923 b.launch(cfg)?;
2924 }
2925 Ok(())
2926 }
2927
2928 #[allow(clippy::too_many_arguments)]
2931 pub fn append_kv_quantized_row_dc_inc(
2932 &self,
2933 k_row: &CudaSlice<f32>,
2934 v_row: &CudaSlice<f32>,
2935 kc: &mut CudaSlice<u8>,
2936 vc: &mut CudaSlice<u8>,
2937 t0_dev: &mut CudaSlice<i32>,
2938 kv_dim_k: usize,
2939 kv_dim_v: usize,
2940 k_tok_bytes: usize,
2941 v_tok_bytes: usize,
2942 g: bool,
2943 ) -> Result<(), Box<dyn std::error::Error>> {
2944 let f = if g {
2945 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2946 } else {
2947 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2948 };
2949 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2950 let cfg = LaunchConfig {
2951 grid_dim: (1, 1, 1),
2952 block_dim: (nthreads, 1, 1),
2953 shared_mem_bytes: 0,
2954 };
2955 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2956 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2957 let __s_b = self.gpu.stream();
2958 let mut b = __s_b.launch_builder(&f);
2959 b.arg(k_row)
2960 .arg(v_row)
2961 .arg(kc)
2962 .arg(vc)
2963 .arg(t0_dev)
2964 .arg(&kdk)
2965 .arg(&kdv)
2966 .arg(&ktb)
2967 .arg(&vtb);
2968 unsafe {
2969 b.launch(cfg)?;
2970 }
2971 Ok(())
2972 }
2973
2974 pub fn pack_tok_p(
2976 &self,
2977 tok: &CudaSlice<u32>,
2978 p: &CudaSlice<f32>,
2979 out: &mut CudaSlice<u32>,
2980 slot: usize,
2981 ) -> Result<(), Box<dyn std::error::Error>> {
2982 let f = self.func("pack_tok_p");
2983 let sl = slot as i32;
2984 let cfg = LaunchConfig {
2985 grid_dim: (1, 1, 1),
2986 block_dim: (32, 1, 1),
2987 shared_mem_bytes: 0,
2988 };
2989 let __s_b = self.gpu.stream();
2990 let mut b = __s_b.launch_builder(&f);
2991 b.arg(tok).arg(p).arg(out).arg(&sl);
2992 unsafe {
2993 b.launch(cfg)?;
2994 }
2995 Ok(())
2996 }
2997 pub fn tok_map_u32(
2998 &self,
2999 tok: &mut CudaSlice<u32>,
3000 map: &CudaSlice<u32>,
3001 ) -> Result<(), Box<dyn std::error::Error>> {
3002 let f = self.func("tok_map_u32");
3003 let cfg = LaunchConfig {
3004 grid_dim: (1, 1, 1),
3005 block_dim: (32, 1, 1),
3006 shared_mem_bytes: 0,
3007 };
3008 let __s_b = self.gpu.stream();
3009 let mut b = __s_b.launch_builder(&f);
3010 b.arg(tok).arg(map);
3011 unsafe {
3012 b.launch(cfg)?;
3013 }
3014 Ok(())
3015 }
3016
3017 #[allow(clippy::too_many_arguments)]
3019 pub fn spec_assemble_verify(
3020 &self,
3021 tokp: &CudaSlice<u32>,
3022 pend: &CudaSlice<u32>,
3023 d2t: Option<&CudaSlice<u32>>,
3024 vtok: &mut CudaSlice<u32>,
3025 brk: &mut CudaSlice<u32>,
3026 p_min: f32,
3027 k: usize,
3028 pmin0: bool,
3029 ) -> Result<(), Box<dyn std::error::Error>> {
3030 let f = self.func("spec_assemble_verify");
3031 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3032 let cfg = LaunchConfig {
3033 grid_dim: (1, 1, 1),
3034 block_dim: (32, 1, 1),
3035 shared_mem_bytes: 0,
3036 };
3037 let __s_b = self.gpu.stream();
3038 let mut b = __s_b.launch_builder(&f);
3039 match d2t {
3040 Some(m) => {
3041 b.arg(tokp)
3042 .arg(pend)
3043 .arg(m)
3044 .arg(vtok)
3045 .arg(brk)
3046 .arg(&p_min)
3047 .arg(&ki)
3048 .arg(&pm);
3049 unsafe {
3050 b.launch(cfg)?;
3051 }
3052 }
3053 None => {
3054 let null: u64 = 0;
3055 b.arg(tokp)
3056 .arg(pend)
3057 .arg(&null)
3058 .arg(vtok)
3059 .arg(brk)
3060 .arg(&p_min)
3061 .arg(&ki)
3062 .arg(&pm);
3063 unsafe {
3064 b.launch(cfg)?;
3065 }
3066 }
3067 }
3068 Ok(())
3069 }
3070
3071 #[allow(clippy::too_many_arguments)]
3073 pub fn ssm_conv_ring_rebuild_dc(
3074 &self,
3075 qkv_tm: &CudaSlice<f32>,
3076 ring_old: &CudaSlice<f32>,
3077 conv_state: &mut CudaSlice<f32>,
3078 conv_dim: usize,
3079 acc: &CudaSlice<u32>,
3080 base: usize,
3081 t_v: usize,
3082 d_conv: usize,
3083 ) -> Result<(), Box<dyn std::error::Error>> {
3084 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3085 let n = conv_dim * (d_conv - 1);
3086 let cfg = LaunchConfig::for_num_elems(n as u32);
3087 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3088 let __s_b = self.gpu.stream();
3089 let mut b = __s_b.launch_builder(&f);
3090 b.arg(qkv_tm)
3091 .arg(ring_old)
3092 .arg(conv_state)
3093 .arg(&cd)
3094 .arg(acc)
3095 .arg(&b0)
3096 .arg(&tv)
3097 .arg(&dc);
3098 unsafe {
3099 b.launch(cfg)?;
3100 }
3101 Ok(())
3102 }
3103 #[allow(clippy::too_many_arguments)]
3104 pub fn gdn_scan_s128_dc(
3105 &self,
3106 q: &CudaSlice<f32>,
3107 k: &CudaSlice<f32>,
3108 v: &CudaSlice<f32>,
3109 g: &CudaSlice<f32>,
3110 beta: &CudaSlice<f32>,
3111 state_in: &CudaSlice<f32>,
3112 state_out: &mut CudaSlice<f32>,
3113 o: &mut CudaSlice<f32>,
3114 n_head: usize,
3115 acc: &CudaSlice<u32>,
3116 base: usize,
3117 t_v: usize,
3118 scale: f32,
3119 ) -> Result<(), Box<dyn std::error::Error>> {
3120 let f = self.func("gdn_scan_s128_dc");
3121 const S_V: u32 = 128;
3122 const WARP: u32 = 32;
3123 const COLS_PER_BLOCK: u32 = 4;
3124 let cfg = LaunchConfig {
3125 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3126 block_dim: (WARP, COLS_PER_BLOCK, 1),
3127 shared_mem_bytes: 0,
3128 };
3129 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3130 let __s_b = self.gpu.stream();
3131 let mut b = __s_b.launch_builder(&f);
3132 b.arg(q)
3133 .arg(k)
3134 .arg(v)
3135 .arg(g)
3136 .arg(beta)
3137 .arg(state_in)
3138 .arg(state_out)
3139 .arg(o)
3140 .arg(&h)
3141 .arg(acc)
3142 .arg(&b0)
3143 .arg(&tv)
3144 .arg(&scale);
3145 unsafe {
3146 b.launch(cfg)?;
3147 }
3148 Ok(())
3149 }
3150
3151 pub fn spec_rollback_kv(
3153 &self,
3154 len_ptrs: &CudaSlice<u64>,
3155 saved: &CudaSlice<i32>,
3156 acc: &CudaSlice<u32>,
3157 base: usize,
3158 n_layer: usize,
3159 ) -> Result<(), Box<dyn std::error::Error>> {
3160 let f = self.func("spec_rollback_kv");
3161 let (b, nl) = (base as i32, n_layer as i32);
3162 let cfg = LaunchConfig {
3163 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3164 block_dim: (64, 1, 1),
3165 shared_mem_bytes: 0,
3166 };
3167 let __s_bl = self.gpu.stream();
3168 let mut bl = __s_bl.launch_builder(&f);
3169 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3170 unsafe {
3171 bl.launch(cfg)?;
3172 }
3173 Ok(())
3174 }
3175
3176 pub fn spec_fork_valid(
3178 &self,
3179 acc: &CudaSlice<u32>,
3180 optimistic_pending: u32,
3181 valid: &mut CudaSlice<u32>,
3182 ) -> Result<(), Box<dyn std::error::Error>> {
3183 let f = self.func("spec_fork_valid");
3184 let cfg = LaunchConfig {
3185 grid_dim: (1, 1, 1),
3186 block_dim: (1, 1, 1),
3187 shared_mem_bytes: 0,
3188 };
3189 let __s_bl = self.gpu.stream();
3190 let mut bl = __s_bl.launch_builder(&f);
3191 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3192 unsafe {
3193 bl.launch(cfg)?;
3194 }
3195 Ok(())
3196 }
3197
3198 pub fn spec_fork_reconcile_kv(
3200 &self,
3201 len_ptrs: &CudaSlice<u64>,
3202 saved: &CudaSlice<i32>,
3203 acc: &CudaSlice<u32>,
3204 valid: &CudaSlice<u32>,
3205 base: usize,
3206 n_layer: usize,
3207 ) -> Result<(), Box<dyn std::error::Error>> {
3208 let f = self.func("spec_fork_reconcile_kv");
3209 let (b, nl) = (base as i32, n_layer as i32);
3210 let cfg = LaunchConfig {
3211 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3212 block_dim: (64, 1, 1),
3213 shared_mem_bytes: 0,
3214 };
3215 let __s_bl = self.gpu.stream();
3216 let mut bl = __s_bl.launch_builder(&f);
3217 bl.arg(len_ptrs)
3218 .arg(saved)
3219 .arg(acc)
3220 .arg(valid)
3221 .arg(&b)
3222 .arg(&nl);
3223 unsafe {
3224 bl.launch(cfg)?;
3225 }
3226 Ok(())
3227 }
3228
3229 pub fn spec_fork_restore_f32(
3231 &self,
3232 snapshot: &CudaSlice<f32>,
3233 state: &mut CudaSlice<f32>,
3234 valid: &CudaSlice<u32>,
3235 ) -> Result<(), Box<dyn std::error::Error>> {
3236 assert_eq!(
3237 snapshot.len(),
3238 state.len(),
3239 "fork recurrent snapshot shape mismatch"
3240 );
3241 let f = self.func("spec_fork_restore_f32");
3242 let n = state.len() as i32;
3243 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3244 let cfg = LaunchConfig {
3245 grid_dim: (blocks, 1, 1),
3246 block_dim: (256, 1, 1),
3247 shared_mem_bytes: 0,
3248 };
3249 let __s_bl = self.gpu.stream();
3250 let mut bl = __s_bl.launch_builder(&f);
3251 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3252 unsafe {
3253 bl.launch(cfg)?;
3254 }
3255 Ok(())
3256 }
3257
3258 pub fn spec_seed_gather(
3261 &self,
3262 vx: &CudaSlice<f32>,
3263 fill_prev: &CudaSlice<f32>,
3264 acc: &CudaSlice<u32>,
3265 h_seed: &mut CudaSlice<f32>,
3266 base: usize,
3267 n_embd: usize,
3268 ) -> Result<(), Box<dyn std::error::Error>> {
3269 let f = self.func("spec_seed_gather");
3270 let (b, ne) = (base as i32, n_embd as i32);
3271 let cfg = LaunchConfig {
3272 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3273 block_dim: (256, 1, 1),
3274 shared_mem_bytes: 0,
3275 };
3276 let __s_bl = self.gpu.stream();
3277 let mut bl = __s_bl.launch_builder(&f);
3278 bl.arg(vx)
3279 .arg(fill_prev)
3280 .arg(acc)
3281 .arg(h_seed)
3282 .arg(&b)
3283 .arg(&ne);
3284 unsafe {
3285 bl.launch(cfg)?;
3286 }
3287 Ok(())
3288 }
3289
3290 pub fn spec_accept_greedy(
3292 &self,
3293 preds: &CudaSlice<u32>,
3294 draft: &CudaSlice<u32>,
3295 last_pred: u32,
3296 base: usize,
3297 k_round: usize,
3298 out: &mut CudaSlice<u32>,
3299 ) -> Result<(), Box<dyn std::error::Error>> {
3300 let f = self.func("spec_accept_greedy");
3301 let (b, k) = (base as i32, k_round as i32);
3302 let cfg = LaunchConfig {
3303 grid_dim: (1, 1, 1),
3304 block_dim: (32, 1, 1),
3305 shared_mem_bytes: 0,
3306 };
3307 let __s_bl = self.gpu.stream();
3308 let mut bl = __s_bl.launch_builder(&f);
3309 bl.arg(preds)
3310 .arg(draft)
3311 .arg(&last_pred)
3312 .arg(&b)
3313 .arg(&k)
3314 .arg(out);
3315 unsafe {
3316 bl.launch(cfg)?;
3317 }
3318 Ok(())
3319 }
3320
3321 pub fn gumbel_perturb(
3328 &self,
3329 x: &CudaSlice<f32>,
3330 y: &mut CudaSlice<f32>,
3331 n: usize,
3332 seed: u64,
3333 stream_pos: u32,
3334 temp: f32,
3335 ) -> Result<(), Box<dyn std::error::Error>> {
3336 let f = self.func("gumbel_perturb_f32");
3337 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3338 let cfg = LaunchConfig {
3339 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3340 block_dim: (256, 1, 1),
3341 shared_mem_bytes: 0,
3342 };
3343 let __s_b = self.gpu.stream();
3344 let mut b = __s_b.launch_builder(&f);
3345 b.arg(x)
3346 .arg(&mut *y)
3347 .arg(&ni)
3348 .arg(&slo)
3349 .arg(&shi)
3350 .arg(&stream_pos)
3351 .arg(&temp);
3352 unsafe {
3353 b.launch(cfg)?;
3354 }
3355 Ok(())
3356 }
3357
3358 pub fn mask_logits_col(
3366 &self,
3367 logits: &mut CudaSlice<f32>,
3368 mask: &CudaSlice<u32>,
3369 col: usize,
3370 n: usize,
3371 mask_words: usize,
3372 ) -> Result<(), Box<dyn std::error::Error>> {
3373 let f = self.func("mask_logits_f32");
3374 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3375 let cfg = LaunchConfig {
3376 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3377 block_dim: (256, 1, 1),
3378 shared_mem_bytes: 0,
3379 };
3380 let __s_b = self.gpu.stream();
3381 let mut b = __s_b.launch_builder(&f);
3382 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3383 unsafe {
3384 b.launch(cfg)?;
3385 }
3386 Ok(())
3387 }
3388
3389 pub fn gumbel_perturb_col(
3396 &self,
3397 x: &CudaSlice<f32>,
3398 col: usize,
3399 y: &mut CudaSlice<f32>,
3400 n: usize,
3401 seed: u64,
3402 stream_pos: u32,
3403 temp: f32,
3404 ) -> Result<(), Box<dyn std::error::Error>> {
3405 let f = self.func("gumbel_perturb_f32");
3406 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3407 let col_view = x.slice(col * n..(col + 1) * n);
3408 let cfg = LaunchConfig {
3409 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3410 block_dim: (256, 1, 1),
3411 shared_mem_bytes: 0,
3412 };
3413 let __s_b = self.gpu.stream();
3414 let mut b = __s_b.launch_builder(&f);
3415 b.arg(&col_view)
3416 .arg(&mut *y)
3417 .arg(&ni)
3418 .arg(&slo)
3419 .arg(&shi)
3420 .arg(&stream_pos)
3421 .arg(&temp);
3422 unsafe {
3423 b.launch(cfg)?;
3424 }
3425 Ok(())
3426 }
3427
3428 #[allow(clippy::too_many_arguments)]
3434 pub fn gumbel_perturb_filtered_col(
3435 &self,
3436 x: &CudaSlice<f32>,
3437 col: usize,
3438 y: &mut CudaSlice<f32>,
3439 n: usize,
3440 seed: u64,
3441 stream_pos: u32,
3442 temp: f32,
3443 stat_max: &CudaSlice<f32>,
3444 stat_th: &CudaSlice<f32>,
3445 stat_idx: usize,
3446 ) -> Result<(), Box<dyn std::error::Error>> {
3447 let f = self.func("gumbel_perturb_filtered_col_f32");
3448 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3449 let (ci, si) = (col as i32, stat_idx as i32);
3450 let cfg = LaunchConfig {
3451 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3452 block_dim: (256, 1, 1),
3453 shared_mem_bytes: 0,
3454 };
3455 let __s_b = self.gpu.stream();
3456 let mut b = __s_b.launch_builder(&f);
3457 b.arg(x)
3458 .arg(&ci)
3459 .arg(&mut *y)
3460 .arg(&ni)
3461 .arg(&slo)
3462 .arg(&shi)
3463 .arg(&stream_pos)
3464 .arg(&temp)
3465 .arg(stat_max)
3466 .arg(stat_th)
3467 .arg(&si);
3468 unsafe {
3469 b.launch(cfg)?;
3470 }
3471 Ok(())
3472 }
3473
3474 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3479 let f = self.func("memra_sctr_inc");
3480 let cfg = LaunchConfig {
3481 grid_dim: (1, 1, 1),
3482 block_dim: (1, 1, 1),
3483 shared_mem_bytes: 0,
3484 };
3485 let __s_b = self.gpu.stream();
3486 let mut b = __s_b.launch_builder(&f);
3487 b.arg(&mut *ctr);
3488 unsafe {
3489 b.launch(cfg)?;
3490 }
3491 Ok(())
3492 }
3493
3494 pub fn gumbel_perturb_ctr(
3499 &self,
3500 x: &CudaSlice<f32>,
3501 y: &mut CudaSlice<f32>,
3502 n: usize,
3503 seed: u64,
3504 ctr: &CudaSlice<u32>,
3505 temp: f32,
3506 ) -> Result<(), Box<dyn std::error::Error>> {
3507 let f = self.func("gumbel_perturb_ctr_f32");
3508 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3509 let cfg = LaunchConfig {
3510 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3511 block_dim: (256, 1, 1),
3512 shared_mem_bytes: 0,
3513 };
3514 let __s_b = self.gpu.stream();
3515 let mut b = __s_b.launch_builder(&f);
3516 b.arg(x)
3517 .arg(&mut *y)
3518 .arg(&ni)
3519 .arg(&slo)
3520 .arg(&shi)
3521 .arg(ctr)
3522 .arg(&temp);
3523 unsafe {
3524 b.launch(cfg)?;
3525 }
3526 Ok(())
3527 }
3528
3529 pub fn softmax_gather(
3533 &self,
3534 x: &CudaSlice<f32>,
3535 row_stride: usize,
3536 ids: &CudaSlice<u32>,
3537 rows: &CudaSlice<i32>,
3538 out: &mut CudaSlice<f32>,
3539 n: usize,
3540 npair: usize,
3541 temp: f32,
3542 ) -> Result<(), Box<dyn std::error::Error>> {
3543 let f = self.func("softmax_gather_f32");
3544 let (ni, rs) = (n as i32, row_stride as i64);
3545 let np = npair as i32;
3546 let cfg = LaunchConfig {
3547 grid_dim: (npair as u32, 1, 1),
3548 block_dim: (256, 1, 1),
3549 shared_mem_bytes: 0,
3550 };
3551 let __s_b = self.gpu.stream();
3552 let mut b = __s_b.launch_builder(&f);
3553 b.arg(x)
3554 .arg(&rs)
3555 .arg(ids)
3556 .arg(rows)
3557 .arg(&mut *out)
3558 .arg(&ni)
3559 .arg(&np)
3560 .arg(&temp);
3561 unsafe {
3562 b.launch(cfg)?;
3563 }
3564 Ok(())
3565 }
3566
3567 pub fn residual_sample(
3571 &self,
3572 p: &CudaSlice<f32>,
3573 q: Option<&CudaSlice<f32>>,
3574 n: usize,
3575 temp: f32,
3576 seed: u64,
3577 stream_pos: u32,
3578 out_tok: &mut CudaSlice<u32>,
3579 ) -> Result<(), Box<dyn std::error::Error>> {
3580 let f = self.func("residual_sample_f32");
3581 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3582 let nth = 1024u32;
3583 let cfg = LaunchConfig {
3584 grid_dim: (1, 1, 1),
3585 block_dim: (nth, 1, 1),
3586 shared_mem_bytes: 0,
3587 };
3588 let has_q: i32 = q.is_some() as i32;
3589 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3591 let mut b = __s_b.launch_builder(&f);
3592 b.arg(p)
3593 .arg(qbuf)
3594 .arg(&has_q)
3595 .arg(&ni)
3596 .arg(&temp)
3597 .arg(&slo)
3598 .arg(&shi)
3599 .arg(&stream_pos)
3600 .arg(&mut *out_tok);
3601 unsafe {
3602 b.launch(cfg)?;
3603 }
3604 Ok(())
3605 }
3606
3607 pub fn with_moe_cache<R>(
3612 &self,
3613 max_block_bytes: usize,
3614 f: impl FnOnce(
3615 &mut crate::moe_cache::MoeSlotCache,
3616 &Engine,
3617 ) -> Result<R, Box<dyn std::error::Error>>,
3618 ) -> Result<R, Box<dyn std::error::Error>> {
3619 let mut guard = self.moe_cache.lock().unwrap();
3620 if guard.is_none() {
3621 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3622 }
3623 let cache = guard.as_mut().unwrap();
3624 f(cache, self)
3625 }
3626
3627 pub fn freeze_moe_cache(&self) {
3630 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3631 cache.freeze();
3632 }
3633 }
3634
3635 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3638 self.moe_cache
3639 .lock()
3640 .unwrap()
3641 .as_ref()
3642 .map(crate::moe_cache::MoeSlotCache::export_residency)
3643 }
3644
3645 pub(crate) fn moe_cache_frozen(&self) -> bool {
3646 self.moe_cache
3647 .lock()
3648 .unwrap()
3649 .as_ref()
3650 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3651 }
3652
3653 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3660 crate::cpu_experts::configured()
3661 && self.moe_cache_frozen()
3662 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3663 }
3664
3665 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3667 assert!(
3668 self.moe_cache.lock().unwrap().is_none(),
3669 "MoE cache layout configured after cache construction"
3670 );
3671 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3672 }
3673
3674 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3675 self.moe_cache_layout.lock().unwrap().clone()
3676 }
3677
3678 pub fn moe_cache_enabled() -> bool {
3680 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3681 }
3682
3683 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3686 let guard = self.moe_cache.lock().unwrap();
3687 guard
3688 .as_ref()
3689 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3690 }
3691
3692 pub fn cpu_expert_stats(
3696 &self,
3697 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3698 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3699 }
3700
3701 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3704 crate::cpu_experts::predictor_stats()
3705 }
3706
3707 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3708 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3709 }
3710
3711 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3714 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3715 }
3716
3717 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3720 let guard = self.moe_cache.lock().unwrap();
3721 guard
3722 .as_ref()
3723 .and_then(|cache| cache.pread_stats())
3724 .map(|stats| {
3725 (
3726 stats.reads,
3727 stats.bytes,
3728 stats.read_errors,
3729 stats.short_reads,
3730 stats.fallbacks,
3731 stats.buffer_waits,
3732 stats.ring_full,
3733 )
3734 })
3735 }
3736
3737 pub fn spill_config_fallbacks(&self) -> u64 {
3739 crate::spill_pread::config_fallbacks()
3740 }
3741
3742 pub fn moe_cache_reset_counters(&self) {
3744 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3745 c.reset_counters();
3746 }
3747 }
3748
3749 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3750 Ok(self.gpu.stream().clone_htod(v)?)
3751 }
3752
3753 pub fn htod_bytes_padded(
3757 &self,
3758 v: &[u8],
3759 pad: usize,
3760 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3761 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3762 {
3763 let mut view = d.slice_mut(0..v.len());
3764 self.gpu.stream().memcpy_htod(v, &mut view)?;
3765 }
3766 Ok(d)
3767 }
3768
3769 pub fn copy_into(
3771 &self,
3772 dst: &mut CudaSlice<f32>,
3773 off: usize,
3774 src: &CudaSlice<f32>,
3775 len: usize,
3776 ) -> Result<(), Box<dyn std::error::Error>> {
3777 let mut view = dst.slice_mut(off..off + len);
3778 self.gpu
3779 .stream()
3780 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3781 Ok(())
3782 }
3783
3784 pub fn copy_u8_into(
3787 &self,
3788 dst: &mut CudaSlice<u8>,
3789 off: usize,
3790 src: &CudaSlice<u8>,
3791 len: usize,
3792 ) -> Result<(), Box<dyn std::error::Error>> {
3793 let mut view = dst.slice_mut(off..off + len);
3794 self.gpu
3795 .stream()
3796 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3797 Ok(())
3798 }
3799
3800 pub fn copy_u8_range_into(
3802 &self,
3803 dst: &mut CudaSlice<u8>,
3804 dst_off: usize,
3805 src: &CudaSlice<u8>,
3806 src_off: usize,
3807 len: usize,
3808 ) -> Result<(), Box<dyn std::error::Error>> {
3809 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3810 self.gpu
3811 .stream()
3812 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3813 Ok(())
3814 }
3815
3816 pub fn prepare_kv_append(
3820 &self,
3821 kv: &mut crate::cache::KvLayer,
3822 retain_from: usize,
3823 append_rows: usize,
3824 ) -> Result<usize, Box<dyn std::error::Error>> {
3825 let Some(plan) = kv
3826 .ring
3827 .as_ref()
3828 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3829 .transpose()?
3830 else {
3831 return Ok(kv.len);
3832 };
3833 match plan {
3834 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3835 crate::cache::KvRingAppend::Rebase {
3836 src_row,
3837 keep_rows,
3838 new_base,
3839 write_row,
3840 } => {
3841 if keep_rows > 0 {
3842 let k_len = keep_rows * kv.k_tok_bytes;
3843 let v_len = keep_rows * kv.v_tok_bytes;
3844 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3845 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3846 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3847 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3848 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3849 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3850 }
3851 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3852 Ok(write_row)
3853 }
3854 }
3855 }
3856
3857 pub fn htod_u8_into(
3860 &self,
3861 dst: &mut CudaSlice<u8>,
3862 off: usize,
3863 src: &[u8],
3864 ) -> Result<(), Box<dyn std::error::Error>> {
3865 let mut view = dst.slice_mut(off..off + src.len());
3866 self.gpu.stream().memcpy_htod(src, &mut view)?;
3867 Ok(())
3868 }
3869
3870 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3871 b.slice(0..len)
3872 }
3873
3874 pub fn view_u8_range<'a>(
3877 &self,
3878 b: &'a CudaSlice<u8>,
3879 start: usize,
3880 end: usize,
3881 ) -> cudarc::driver::CudaView<'a, u8> {
3882 b.slice(start..end)
3883 }
3884 pub fn view_u8<'a>(
3885 &self,
3886 b: &'a CudaSlice<u8>,
3887 len: usize,
3888 ) -> cudarc::driver::CudaView<'a, u8> {
3889 b.slice(0..len)
3890 }
3891
3892 pub fn append_kv_quantized(
3896 &self,
3897 k_row: &CudaSlice<f32>,
3898 v_row: &CudaSlice<f32>,
3899 kc: &mut CudaSlice<u8>,
3900 vc: &mut CudaSlice<u8>,
3901 t: usize,
3902 kv_dim_k: usize,
3903 kv_dim_v: usize,
3904 k_tok_bytes: usize,
3905 v_tok_bytes: usize,
3906 g: bool,
3907 ) -> Result<(), Box<dyn std::error::Error>> {
3908 let f = if g {
3909 self.func_g("append_quantize_kv_q8_0_q5_1")
3910 } else {
3911 self.func("append_quantize_kv_q8_0_q5_1")
3912 };
3913 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3914 let cfg = LaunchConfig {
3915 grid_dim: (nblk, 1, 1),
3916 block_dim: (32, 1, 1),
3917 shared_mem_bytes: 0,
3918 };
3919 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3920 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3921 let __s_b = self.gpu.stream();
3922 let mut b = __s_b.launch_builder(&f);
3923 b.arg(k_row)
3924 .arg(v_row)
3925 .arg(kc)
3926 .arg(vc)
3927 .arg(&ti)
3928 .arg(&kdk)
3929 .arg(&kdv)
3930 .arg(&ktb)
3931 .arg(&vtb);
3932 unsafe {
3933 b.launch(cfg)?;
3934 }
3935 Ok(())
3936 }
3937
3938 pub fn append_kv_quantized_dc(
3942 &self,
3943 k_row: &CudaSlice<f32>,
3944 v_row: &CudaSlice<f32>,
3945 kc: &mut CudaSlice<u8>,
3946 vc: &mut CudaSlice<u8>,
3947 t_dev: &CudaSlice<i32>,
3948 kv_dim_k: usize,
3949 kv_dim_v: usize,
3950 k_tok_bytes: usize,
3951 v_tok_bytes: usize,
3952 g: bool,
3953 ) -> Result<(), Box<dyn std::error::Error>> {
3954 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3955 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3956 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3957 if Self::pdl_on() && Self::pdl_wb_on() {
3959 use cudarc::driver::{DevicePtr, DevicePtrMut};
3960 let s = &self.gpu.stream();
3961 let (pk, _g0) = k_row.device_ptr(s);
3962 let (pv, _g1) = v_row.device_ptr(s);
3963 let (pkc, _g2) = kc.device_ptr_mut(s);
3964 let (pvc, _g3) = vc.device_ptr_mut(s);
3965 let (pt, _g4) = t_dev.device_ptr(s);
3966 let mut ps = [
3967 &pk as *const _ as *mut std::ffi::c_void,
3968 &pv as *const _ as *mut _,
3969 &pkc as *const _ as *mut _,
3970 &pvc as *const _ as *mut _,
3971 &pt as *const _ as *mut _,
3972 &kdk as *const _ as *mut _,
3973 &kdv as *const _ as *mut _,
3974 &ktb as *const _ as *mut _,
3975 &vtb as *const _ as *mut _,
3976 ];
3977 unsafe {
3978 self.launch_pdl_flash(
3979 g,
3980 "append_quantize_kv_q8_0_q5_1_dc",
3981 (nblk, 1, 1),
3982 (32, 1, 1),
3983 0,
3984 &mut ps,
3985 )?;
3986 }
3987 return Ok(());
3988 }
3989 let f = if g {
3990 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3991 } else {
3992 self.func("append_quantize_kv_q8_0_q5_1_dc")
3993 };
3994 let cfg = LaunchConfig {
3995 grid_dim: (nblk, 1, 1),
3996 block_dim: (32, 1, 1),
3997 shared_mem_bytes: 0,
3998 };
3999 let __s_b = self.gpu.stream();
4000 let mut b = __s_b.launch_builder(&f);
4001 b.arg(k_row)
4002 .arg(v_row)
4003 .arg(kc)
4004 .arg(vc)
4005 .arg(t_dev)
4006 .arg(&kdk)
4007 .arg(&kdv)
4008 .arg(&ktb)
4009 .arg(&vtb);
4010 unsafe {
4011 b.launch(cfg)?;
4012 }
4013 Ok(())
4014 }
4015
4016 #[allow(clippy::too_many_arguments)]
4023 pub fn append_kv_quantized_rows(
4024 &self,
4025 k_rows: &CudaSlice<f32>,
4026 v_rows: &CudaSlice<f32>,
4027 kc: &mut CudaSlice<u8>,
4028 vc: &mut CudaSlice<u8>,
4029 t0: usize,
4030 t: usize,
4031 kv_dim_k: usize,
4032 kv_dim_v: usize,
4033 k_tok_bytes: usize,
4034 v_tok_bytes: usize,
4035 g: bool,
4036 ) -> Result<(), Box<dyn std::error::Error>> {
4037 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4038 for i in 0..t {
4039 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4040 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4041 self.append_kv_quantized_view(
4042 &k_row,
4043 &v_row,
4044 kc,
4045 vc,
4046 t0 + i,
4047 kv_dim_k,
4048 kv_dim_v,
4049 k_tok_bytes,
4050 v_tok_bytes,
4051 g,
4052 )?;
4053 }
4054 return Ok(());
4055 }
4056 let f = if g {
4057 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4058 } else {
4059 self.func("append_quantize_kv_q8_0_q5_1_rows")
4060 };
4061 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4062 let cfg = LaunchConfig {
4063 grid_dim: (nblk, t as u32, 1),
4064 block_dim: (32, 1, 1),
4065 shared_mem_bytes: 0,
4066 };
4067 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4068 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4069 let __s_b = self.gpu.stream();
4070 let mut b = __s_b.launch_builder(&f);
4071 b.arg(k_rows)
4072 .arg(v_rows)
4073 .arg(kc)
4074 .arg(vc)
4075 .arg(&t0i)
4076 .arg(&kdk)
4077 .arg(&kdv)
4078 .arg(&ktb)
4079 .arg(&vtb);
4080 unsafe {
4081 b.launch(cfg)?;
4082 }
4083 Ok(())
4084 }
4085
4086 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4090 let f = self.func("inc_i32");
4091 let cfg = LaunchConfig {
4092 grid_dim: (1, 1, 1),
4093 block_dim: (1, 1, 1),
4094 shared_mem_bytes: 0,
4095 };
4096 let __s_b = self.gpu.stream();
4097 let mut b = __s_b.launch_builder(&f);
4098 b.arg(p);
4099 unsafe {
4100 b.launch(cfg)?;
4101 }
4102 Ok(())
4103 }
4104
4105 pub fn append_kv_quantized_view(
4108 &self,
4109 k_row: &cudarc::driver::CudaView<f32>,
4110 v_row: &cudarc::driver::CudaView<f32>,
4111 kc: &mut CudaSlice<u8>,
4112 vc: &mut CudaSlice<u8>,
4113 t: usize,
4114 kv_dim_k: usize,
4115 kv_dim_v: usize,
4116 k_tok_bytes: usize,
4117 v_tok_bytes: usize,
4118 g: bool,
4119 ) -> Result<(), Box<dyn std::error::Error>> {
4120 let stream = self.gpu.stream();
4121 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4122 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4123 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4124 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4125 let f = if g {
4126 self.func_g("append_quantize_kv_q8_0_q5_1")
4127 } else {
4128 self.func("append_quantize_kv_q8_0_q5_1")
4129 };
4130 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4131 let cfg = LaunchConfig {
4132 grid_dim: (nblk, 1, 1),
4133 block_dim: (32, 1, 1),
4134 shared_mem_bytes: 0,
4135 };
4136 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4137 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4138 let mut b = stream.launch_builder(&f);
4139 b.arg(k_row)
4140 .arg(v_row)
4141 .arg(kc)
4142 .arg(vc)
4143 .arg(&ti)
4144 .arg(&kdk)
4145 .arg(&kdv)
4146 .arg(&ktb)
4147 .arg(&vtb);
4148 unsafe {
4149 b.launch(cfg)?;
4150 }
4151 Ok(())
4152 }
4153
4154 pub fn copy_view_into(
4157 &self,
4158 dst: &mut CudaSlice<f32>,
4159 off: usize,
4160 src: &cudarc::driver::CudaView<f32>,
4161 len: usize,
4162 ) -> Result<(), Box<dyn std::error::Error>> {
4163 let mut view = dst.slice_mut(off..off + len);
4164 self.gpu
4165 .stream()
4166 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4167 Ok(())
4168 }
4169
4170 pub fn clone_dtod(
4174 &self,
4175 src: &CudaSlice<f32>,
4176 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4177 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4178 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4179 Ok(dst)
4180 }
4181
4182 pub fn dtod_copy_view(
4185 &self,
4186 src: &cudarc::driver::CudaView<f32>,
4187 dst: &mut CudaSlice<f32>,
4188 ) -> Result<(), Box<dyn std::error::Error>> {
4189 self.gpu.stream().memcpy_dtod(src, dst)?;
4190 Ok(())
4191 }
4192
4193 pub fn dtod_copy_view_i8(
4195 &self,
4196 src: &cudarc::driver::CudaView<i8>,
4197 dst: &mut CudaSlice<i8>,
4198 ) -> Result<(), Box<dyn std::error::Error>> {
4199 self.gpu.stream().memcpy_dtod(src, dst)?;
4200 Ok(())
4201 }
4202
4203 pub fn dtod_copy_into(
4205 &self,
4206 src: &CudaSlice<f32>,
4207 dst: &mut CudaSlice<f32>,
4208 offset: usize,
4209 ) -> Result<(), Box<dyn std::error::Error>> {
4210 let n = src.len();
4211 let mut dv = dst.slice_mut(offset..offset + n);
4212 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4213 Ok(())
4214 }
4215
4216 pub fn copy_batch_uniform_f32(
4222 &self,
4223 table: &CudaSlice<u64>,
4224 n: usize,
4225 words: usize,
4226 ) -> Result<(), Box<dyn std::error::Error>> {
4227 if n == 0 || words == 0 {
4228 return Ok(());
4229 }
4230 debug_assert!(
4231 table.len() >= 2 * n,
4232 "pointer table must hold n srcs + n dsts"
4233 );
4234 let f = self.func("copy_batch_uniform_f32");
4235 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4238 let (ni, wi) = (n as i32, words as i32);
4239 let cfg = LaunchConfig {
4240 grid_dim: (chunks, n as u32, 1),
4241 block_dim: (256, 1, 1),
4242 shared_mem_bytes: 0,
4243 };
4244 let __s = self.gpu.stream();
4245 let mut b = __s.launch_builder(&f);
4246 b.arg(table).arg(&ni).arg(&wi);
4247 unsafe {
4248 b.launch(cfg)?;
4249 }
4250 Ok(())
4251 }
4252
4253 pub fn htod_u64_into(
4256 &self,
4257 v: &[u64],
4258 dst: &mut CudaSlice<u64>,
4259 ) -> Result<(), Box<dyn std::error::Error>> {
4260 let mut view = dst.slice_mut(0..v.len());
4261 self.gpu.stream().memcpy_htod(v, &mut view)?;
4262 Ok(())
4263 }
4264
4265 pub fn copy_indirect_src_f32(
4270 &self,
4271 src_entry: &cudarc::driver::CudaView<u64>,
4272 dst: &mut CudaSlice<f32>,
4273 dst_off: usize,
4274 words: usize,
4275 ) -> Result<(), Box<dyn std::error::Error>> {
4276 let f = self.func("copy_indirect_src_f32");
4277 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4278 let wi = words as i32;
4279 let cfg = LaunchConfig {
4280 grid_dim: (chunks, 1, 1),
4281 block_dim: (256, 1, 1),
4282 shared_mem_bytes: 0,
4283 };
4284 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4285 let __s = self.gpu.stream();
4286 let mut b = __s.launch_builder(&f);
4287 b.arg(src_entry).arg(&mut dv).arg(&wi);
4288 unsafe {
4289 b.launch(cfg)?;
4290 }
4291 Ok(())
4292 }
4293
4294 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4296 self.alloc_uninit::<i8>(n)
4297 }
4298
4299 pub fn qmatvec(
4301 &self,
4302 w: &CudaSlice<u8>,
4303 x: &CudaSlice<f32>,
4304 m: usize,
4305 in_f: usize,
4306 out_f: usize,
4307 qtype: i32,
4308 row_bytes: usize,
4309 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4310 let f = self.func("qmatvec_f32");
4311 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4313 grid_dim: (out_f as u32, m as u32, 1),
4314 block_dim: (256, 1, 1),
4315 shared_mem_bytes: 0,
4316 };
4317 let (inf, outf, mi, qt, rb) =
4318 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4319 let __s_b = self.gpu.stream();
4320 let mut b = __s_b.launch_builder(&f);
4321 b.arg(w)
4322 .arg(x)
4323 .arg(&mut y)
4324 .arg(&inf)
4325 .arg(&outf)
4326 .arg(&mi)
4327 .arg(&qt)
4328 .arg(&rb);
4329 unsafe {
4330 b.launch(cfg)?;
4331 }
4332 Ok(y)
4333 }
4334
4335 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4337 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4338 self.keep_if_capturing(&s);
4339 Ok(s)
4340 }
4341
4342 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4346 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4347 self.keep_if_capturing(&s);
4348 Ok(s)
4349 }
4350
4351 pub fn memset_zeros_view(
4354 &self,
4355 dst: &mut cudarc::driver::CudaViewMut<f32>,
4356 ) -> Result<(), Box<dyn std::error::Error>> {
4357 self.gpu.stream().memset_zeros(dst)?;
4358 Ok(())
4359 }
4360
4361 pub fn stage_expert(
4367 &self,
4368 host_bytes: &[u8],
4369 scratch: &mut CudaSlice<u8>,
4370 off: usize,
4371 ) -> Result<(), Box<dyn std::error::Error>> {
4372 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4375 }
4376
4377 pub fn moe_router_topk(
4383 &self,
4384 logits: &CudaSlice<f32>,
4385 t: usize,
4386 n_expert: usize,
4387 n_used: usize,
4388 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4389 let f = self.func("moe_router_topk_f32");
4390 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 {
4393 grid_dim: (t as u32, 1, 1),
4394 block_dim: (n_expert as u32, 1, 1),
4395 shared_mem_bytes: 0,
4396 };
4397 let (ne, nu) = (n_expert as i32, n_used as i32);
4398 let __s_b = self.gpu.stream();
4399 let mut b = __s_b.launch_builder(&f);
4400 b.arg(logits)
4401 .arg(&mut sel_idx)
4402 .arg(&mut sel_w)
4403 .arg(&ne)
4404 .arg(&nu);
4405 unsafe {
4406 b.launch(cfg)?;
4407 }
4408 Ok((sel_idx, sel_w))
4409 }
4410
4411 pub fn moe_router_topk_scaled(
4414 &self,
4415 logits: &CudaSlice<f32>,
4416 t: usize,
4417 n_expert: usize,
4418 n_used: usize,
4419 ex_scale: &CudaSlice<f32>,
4420 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4421 let f = self.func("moe_router_topk_scaled_f32");
4426 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4427 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4428 let cfg = LaunchConfig {
4429 grid_dim: (t as u32, 1, 1),
4430 block_dim: (n_expert as u32, 1, 1),
4431 shared_mem_bytes: 0,
4432 };
4433 let (ne, nu) = (n_expert as i32, n_used as i32);
4434 let __s_b = self.gpu.stream();
4435 let mut b = __s_b.launch_builder(&f);
4436 b.arg(logits)
4437 .arg(&mut sel_idx)
4438 .arg(&mut sel_w)
4439 .arg(&ne)
4440 .arg(&nu)
4441 .arg(ex_scale);
4442 unsafe {
4443 b.launch(cfg)?;
4444 }
4445 Ok((sel_idx, sel_w))
4446 }
4447
4448 pub fn moe_router_topk_host(
4456 &self,
4457 logits: &CudaSlice<f32>,
4458 t: usize,
4459 n_expert: usize,
4460 n_used: usize,
4461 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4462 let f = self.func("moe_router_topk_f32");
4463 let n = t * n_used;
4464 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4465 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4466 let cfg = LaunchConfig {
4467 grid_dim: (t as u32, 1, 1),
4468 block_dim: (n_expert as u32, 1, 1),
4469 shared_mem_bytes: 0,
4470 };
4471 let (ne, nu) = (n_expert as i32, n_used as i32);
4472 let __s_b = self.gpu.stream();
4473 let mut b = __s_b.launch_builder(&f);
4474 b.arg(logits)
4475 .arg(&mut sel_idx)
4476 .arg(&mut sel_w)
4477 .arg(&ne)
4478 .arg(&nu);
4479 unsafe {
4480 b.launch(cfg)?;
4481 }
4482 let bytes = n * 8;
4484 let mut guard = self.router_stage.lock().unwrap();
4485 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4486 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4487 }
4488 let stage = guard.as_mut().unwrap();
4489 let (si, sw) = unsafe {
4490 (
4491 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4492 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4493 )
4494 };
4495 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()))
4499 }
4500
4501 #[allow(clippy::too_many_arguments)]
4505 pub fn moe_router_sigmoid_topk(
4506 &self,
4507 logits: &CudaSlice<f32>,
4508 t: usize,
4509 n_expert: usize,
4510 n_used: usize,
4511 active_count: usize,
4512 correction_bias: &CudaSlice<f32>,
4513 active: &CudaSlice<u8>,
4514 scaling_factor: f32,
4515 route_norm: bool,
4516 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4517 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4518 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4519 return Err(format!(
4520 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4521 )
4522 .into());
4523 }
4524 if logits.len() < t * n_expert
4525 || correction_bias.len() != n_expert
4526 || active.len() != n_expert
4527 {
4528 return Err(format!(
4529 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4530 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4531 ).into());
4532 }
4533 let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4534 self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4537 } else if crate::sig_expf_dev_on() {
4538 self.func("moe_router_sigmoid_topk_f32_dexp")
4539 } else if crate::topk_fast_on() {
4540 self.func("moe_router_sigmoid_topk_f32_fast")
4541 } else {
4542 self.func("moe_router_sigmoid_topk_f32")
4543 };
4544 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4545 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4546 let threads = n_expert.div_ceil(32) * 32;
4547 let cfg = LaunchConfig {
4548 grid_dim: (t as u32, 1, 1),
4549 block_dim: (threads as u32, 1, 1),
4550 shared_mem_bytes: 0,
4551 };
4552 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4553 let __s_b = self.gpu.stream();
4554 let mut b = __s_b.launch_builder(&f);
4555 b.arg(logits)
4556 .arg(correction_bias)
4557 .arg(active)
4558 .arg(&mut sel_idx)
4559 .arg(&mut sel_w)
4560 .arg(&ne)
4561 .arg(&nu)
4562 .arg(&scaling_factor)
4563 .arg(&rn);
4564 unsafe {
4565 b.launch(cfg)?;
4566 }
4567 Ok((sel_idx, sel_w))
4568 }
4569
4570 #[allow(clippy::too_many_arguments)]
4573 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
4577 if ptr == 0 {
4578 return Err("ring_flag_raw: unarmed flag".into());
4579 }
4580 let f = self.func("memra_ring_flag");
4581 let cfg = LaunchConfig {
4582 grid_dim: (1, 1, 1),
4583 block_dim: (32, 1, 1),
4584 shared_mem_bytes: 0,
4585 };
4586 let __s_b = self.gpu.stream();
4587 let mut b = __s_b.launch_builder(&f);
4588 b.arg(&ptr).arg(&value);
4589 unsafe {
4590 b.launch(cfg)?;
4591 }
4592 Ok(())
4593 }
4594
4595 pub fn moe_sel_w_mirror(
4598 &self,
4599 sel_src: &CudaSlice<i32>,
4600 w_src: &CudaSlice<f32>,
4601 sel_dst: &mut CudaSlice<i32>,
4602 w_dst: &mut CudaSlice<f32>,
4603 n: usize,
4604 ) -> Result<(), Box<dyn std::error::Error>> {
4605 if n == 0
4606 || n > 32
4607 || sel_src.len() < n
4608 || w_src.len() < n
4609 || sel_dst.len() < n
4610 || w_dst.len() < n
4611 {
4612 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
4613 }
4614 let f = self.func("moe_sel_w_mirror");
4615 let cfg = LaunchConfig {
4616 grid_dim: (1, 1, 1),
4617 block_dim: (32, 1, 1),
4618 shared_mem_bytes: 0,
4619 };
4620 let ni = n as i32;
4621 let __s_b = self.gpu.stream();
4622 let mut b = __s_b.launch_builder(&f);
4623 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
4624 unsafe {
4625 b.launch(cfg)?;
4626 }
4627 Ok(())
4628 }
4629
4630 pub fn moe_router_sigmoid_topk_into(
4631 &self,
4632 logits: &CudaSlice<f32>,
4633 t: usize,
4634 n_expert: usize,
4635 n_used: usize,
4636 active_count: usize,
4637 correction_bias: &CudaSlice<f32>,
4638 active: &CudaSlice<u8>,
4639 scaling_factor: f32,
4640 route_norm: bool,
4641 sel_idx: &mut CudaSlice<i32>,
4642 sel_w: &mut CudaSlice<f32>,
4643 ) -> Result<(), Box<dyn std::error::Error>> {
4644 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4645 if n_expert == 0
4646 || n_expert > 1024
4647 || n_used == 0
4648 || n_used > 32 || n_used > n_expert
4650 || logits.len() < t * n_expert
4651 || correction_bias.len() != n_expert
4652 || active.len() != n_expert
4653 || sel_idx.len() < t * n_used
4654 || sel_w.len() < t * n_used
4655 {
4656 return Err("sigmoid router _into geometry mismatch".into());
4657 }
4658 let f = if crate::sig_expf_dev_on() {
4659 self.func("moe_router_sigmoid_topk_f32_dexp")
4660 } else if crate::topk_fast_on() {
4661 self.func("moe_router_sigmoid_topk_f32_fast")
4662 } else {
4663 self.func("moe_router_sigmoid_topk_f32")
4664 };
4665 let threads = n_expert.div_ceil(32) * 32;
4666 let cfg = LaunchConfig {
4667 grid_dim: (t as u32, 1, 1),
4668 block_dim: (threads as u32, 1, 1),
4669 shared_mem_bytes: 0,
4670 };
4671 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4672 let __s_b = self.gpu.stream();
4673 let mut b = __s_b.launch_builder(&f);
4674 b.arg(logits)
4675 .arg(correction_bias)
4676 .arg(active)
4677 .arg(&mut *sel_idx)
4678 .arg(&mut *sel_w)
4679 .arg(&ne)
4680 .arg(&nu)
4681 .arg(&scaling_factor)
4682 .arg(&rn);
4683 unsafe {
4684 b.launch(cfg)?;
4685 }
4686 Ok(())
4687 }
4688
4689 #[allow(clippy::too_many_arguments)]
4692 pub fn moe_router_sigmoid_topk_host(
4693 &self,
4694 logits: &CudaSlice<f32>,
4695 t: usize,
4696 n_expert: usize,
4697 n_used: usize,
4698 active_count: usize,
4699 correction_bias: &CudaSlice<f32>,
4700 active: &CudaSlice<u8>,
4701 scaling_factor: f32,
4702 route_norm: bool,
4703 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4704 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4705 logits,
4706 t,
4707 n_expert,
4708 n_used,
4709 active_count,
4710 correction_bias,
4711 active,
4712 scaling_factor,
4713 route_norm,
4714 )?;
4715 let n = t * n_used;
4716 let bytes = n * 8;
4717 let mut guard = self.router_stage.lock().unwrap();
4718 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4719 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4720 }
4721 let stage = guard.as_mut().unwrap();
4722 let (si, sw) = unsafe {
4723 (
4724 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4725 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4726 )
4727 };
4728 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4729 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4730 self.gpu.stream().synchronize()?;
4731 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4732 }
4733
4734 pub fn stage_expert_async(
4738 &self,
4739 host_bytes: &[u8],
4740 scratch: &mut CudaSlice<u8>,
4741 off: usize,
4742 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4743 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4744 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4745 Ok(self.copy_stream.record_event(None)?)
4746 }
4747
4748 pub fn compute_wait(
4750 &self,
4751 ev: &cudarc::driver::CudaEvent,
4752 ) -> Result<(), Box<dyn std::error::Error>> {
4753 self.gpu.stream().wait(ev)?;
4754 Ok(())
4755 }
4756
4757 pub fn qmatvec_view(
4762 &self,
4763 w: &CudaSlice<u8>,
4764 range: std::ops::Range<usize>,
4765 x: &cudarc::driver::CudaView<f32>,
4766 m: usize,
4767 in_f: usize,
4768 out_f: usize,
4769 qtype: i32,
4770 row_bytes: usize,
4771 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4772 let f = self.func("qmatvec_f32");
4773 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4776 grid_dim: (out_f as u32, m as u32, 1),
4777 block_dim: (256, 1, 1),
4778 shared_mem_bytes: 0,
4779 };
4780 let (inf, outf, mi, qt, rb) =
4781 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4782 let __s_b = self.gpu.stream();
4783 let mut b = __s_b.launch_builder(&f);
4784 b.arg(&wv)
4785 .arg(x)
4786 .arg(&mut y)
4787 .arg(&inf)
4788 .arg(&outf)
4789 .arg(&mi)
4790 .arg(&qt)
4791 .arg(&rb);
4792 unsafe {
4793 b.launch(cfg)?;
4794 }
4795 Ok(y)
4796 }
4797
4798 #[allow(clippy::too_many_arguments)]
4805 pub fn moe_gate_up_silu8_q8(
4809 &self,
4810 gp: WPtr8,
4811 up: WPtr8,
4812 aq: &CudaSlice<i8>,
4813 ad: &CudaSlice<f32>,
4814 in_f: usize,
4815 n_ff: usize,
4816 n_used: usize,
4817 qt_g: i32,
4818 qt_u: i32,
4819 rb_g: usize,
4820 rb_u: usize,
4821 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4822 let f = self.func("moe_gate_up_silu8_q8");
4823 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4824 let cfg = LaunchConfig {
4825 grid_dim: (n_ff as u32, n_used as u32, 1),
4826 block_dim: (32, 1, 1),
4827 shared_mem_bytes: 0,
4828 };
4829 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4830 let __s_b = self.gpu.stream();
4831 let mut b = __s_b.launch_builder(&f);
4832 b.arg(&gp)
4833 .arg(&up)
4834 .arg(aq)
4835 .arg(ad)
4836 .arg(&mut act)
4837 .arg(&inf)
4838 .arg(&nff)
4839 .arg(&qt_g)
4840 .arg(&qt_u)
4841 .arg(&rbg)
4842 .arg(&rbu);
4843 unsafe {
4844 b.launch(cfg)?;
4845 }
4846 Ok(act)
4847 }
4848
4849 #[allow(clippy::too_many_arguments)]
4850 pub fn moe_down8_fma_q8(
4851 &self,
4852 dp: WPtr8,
4853 w: F32x8,
4854 aq2: &CudaSlice<i8>,
4855 ad2: &CudaSlice<f32>,
4856 dst: &mut cudarc::driver::CudaViewMut<f32>,
4857 in_f: usize,
4858 out_f: usize,
4859 n_used: usize,
4860 qt: i32,
4861 rb: usize,
4862 ) -> Result<(), Box<dyn std::error::Error>> {
4863 let f = self.func("moe_down8_fma_q8");
4864 let cfg = LaunchConfig {
4865 grid_dim: (out_f as u32, 1, 1),
4866 block_dim: (32, 1, 1),
4867 shared_mem_bytes: 0,
4868 };
4869 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4870 let __s_b = self.gpu.stream();
4871 let mut b = __s_b.launch_builder(&f);
4872 b.arg(&dp)
4873 .arg(&w)
4874 .arg(aq2)
4875 .arg(ad2)
4876 .arg(dst)
4877 .arg(&inf)
4878 .arg(&outf)
4879 .arg(&nu)
4880 .arg(&qt)
4881 .arg(&rbi);
4882 unsafe {
4883 b.launch(cfg)?;
4884 }
4885 Ok(())
4886 }
4887
4888 pub fn qmatvec_expert_q8(
4890 &self,
4891 w: &CudaSlice<u8>,
4892 range: std::ops::Range<usize>,
4893 aq: &CudaSlice<i8>,
4894 ad: &CudaSlice<f32>,
4895 m: usize,
4896 in_f: usize,
4897 out_f: usize,
4898 qtype: i32,
4899 row_bytes: usize,
4900 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4901 let f = self.func("qmatvec_expert_q8");
4902 let wv = w.slice(range);
4903 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4904 const ROWS: u32 = 4; let cfg = LaunchConfig {
4906 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4907 block_dim: (32, ROWS, 1),
4908 shared_mem_bytes: 0,
4909 };
4910 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4911 let __s_b = self.gpu.stream();
4912 let mut b = __s_b.launch_builder(&f);
4913 b.arg(&wv)
4914 .arg(aq)
4915 .arg(ad)
4916 .arg(&mut y)
4917 .arg(&inf)
4918 .arg(&outf)
4919 .arg(&mi)
4920 .arg(&qtype)
4921 .arg(&rbi);
4922 unsafe {
4923 b.launch(cfg)?;
4924 }
4925 Ok(y)
4926 }
4927
4928 pub fn moe_gate_up_silu8(
4929 &self,
4930 gp: WPtr8,
4931 up: WPtr8,
4932 x: &cudarc::driver::CudaView<f32>,
4933 in_f: usize,
4934 n_ff: usize,
4935 n_used: usize,
4936 qt_g: i32,
4937 qt_u: i32,
4938 rb_g: usize,
4939 rb_u: usize,
4940 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4941 let f = self.func("moe_gate_up_silu8_f32");
4942 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4944 grid_dim: (n_ff as u32, n_used as u32, 1),
4945 block_dim: (256, 1, 1),
4946 shared_mem_bytes: 0,
4947 };
4948 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4949 let __s_b = self.gpu.stream();
4950 let mut b = __s_b.launch_builder(&f);
4951 b.arg(&gp)
4952 .arg(&up)
4953 .arg(x)
4954 .arg(&mut act)
4955 .arg(&inf)
4956 .arg(&nff)
4957 .arg(&qt_g)
4958 .arg(&qt_u)
4959 .arg(&rbg)
4960 .arg(&rbu);
4961 unsafe {
4962 b.launch(cfg)?;
4963 }
4964 Ok(act)
4965 }
4966
4967 #[allow(clippy::too_many_arguments)]
4973 pub fn moe_down8_fma_into(
4974 &self,
4975 dp: WPtr8,
4976 w: F32x8,
4977 act: &CudaSlice<f32>,
4978 dst: &mut cudarc::driver::CudaViewMut<f32>,
4979 in_f: usize,
4980 out_f: usize,
4981 n_used: usize,
4982 qt: i32,
4983 rb: usize,
4984 ) -> Result<(), Box<dyn std::error::Error>> {
4985 let f = self.func("moe_down8_fma_f32");
4986 let cfg = LaunchConfig {
4987 grid_dim: (out_f as u32, 1, 1),
4988 block_dim: (256, 1, 1),
4989 shared_mem_bytes: 0,
4990 };
4991 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4992 let __s_b = self.gpu.stream();
4993 let mut b = __s_b.launch_builder(&f);
4994 b.arg(&dp)
4995 .arg(&w)
4996 .arg(act)
4997 .arg(dst)
4998 .arg(&inf)
4999 .arg(&outf)
5000 .arg(&nu)
5001 .arg(&qt)
5002 .arg(&rbv);
5003 unsafe {
5004 b.launch(cfg)?;
5005 }
5006 Ok(())
5007 }
5008
5009 #[allow(clippy::too_many_arguments)]
5014 #[allow(clippy::too_many_arguments)]
5029 #[allow(clippy::too_many_arguments)]
5031 pub fn moe_pairs_matvec_q8(
5032 &self,
5033 table: &CudaSlice<u64>,
5034 proj: i32,
5035 pair_tok: &CudaSlice<i32>,
5036 pair_ex: &CudaSlice<i32>,
5037 aq: &CudaSlice<i8>,
5038 ad: &CudaSlice<f32>,
5039 in_f: usize,
5040 out_f: usize,
5041 n_expert: usize,
5042 n_pairs: usize,
5043 qtype: i32,
5044 row_bytes: usize,
5045 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5046 let f = self.func("moe_pairs_matvec_q8");
5047 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5048 const ROWS: u32 = 4;
5049 let cfg = LaunchConfig {
5050 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5051 block_dim: (32, ROWS, 1),
5052 shared_mem_bytes: 0,
5053 };
5054 let (inf, outf, ne, np, rbi) = (
5055 in_f as i32,
5056 out_f as i32,
5057 n_expert as i32,
5058 n_pairs as i32,
5059 row_bytes as i64,
5060 );
5061 let __s_b = self.gpu.stream();
5062 let mut b = __s_b.launch_builder(&f);
5063 b.arg(table)
5064 .arg(&proj)
5065 .arg(pair_tok)
5066 .arg(pair_ex)
5067 .arg(aq)
5068 .arg(ad)
5069 .arg(&mut y)
5070 .arg(&inf)
5071 .arg(&outf)
5072 .arg(&ne)
5073 .arg(&np)
5074 .arg(&qtype)
5075 .arg(&rbi);
5076 unsafe {
5077 b.launch(cfg)?;
5078 }
5079 Ok(y)
5080 }
5081
5082 #[allow(clippy::too_many_arguments)]
5084 pub fn moe_pairs_matvec_q8_em(
5085 &self,
5086 table: &CudaSlice<u64>,
5087 proj: i32,
5088 ex_ids: &CudaSlice<i32>,
5089 ex_off: &CudaSlice<i32>,
5090 ex_pairs: &CudaSlice<i32>,
5091 pair_tok: &CudaSlice<i32>,
5092 aq: &CudaSlice<i8>,
5093 ad: &CudaSlice<f32>,
5094 in_f: usize,
5095 out_f: usize,
5096 n_expert: usize,
5097 n_active: usize,
5098 n_pairs: usize,
5099 qtype: i32,
5100 row_bytes: usize,
5101 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5102 let f = self.func("moe_pairs_matvec_q8_em");
5103 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5104 const ROWS: u32 = 4;
5105 let cfg = LaunchConfig {
5106 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5107 block_dim: (32, ROWS, 1),
5108 shared_mem_bytes: 0,
5109 };
5110 let (inf, outf, ne, na, rbi) = (
5111 in_f as i32,
5112 out_f as i32,
5113 n_expert as i32,
5114 n_active as i32,
5115 row_bytes as i64,
5116 );
5117 let __s_b = self.gpu.stream();
5118 let mut b = __s_b.launch_builder(&f);
5119 b.arg(table)
5120 .arg(&proj)
5121 .arg(ex_ids)
5122 .arg(ex_off)
5123 .arg(ex_pairs)
5124 .arg(pair_tok)
5125 .arg(aq)
5126 .arg(ad)
5127 .arg(&mut y)
5128 .arg(&inf)
5129 .arg(&outf)
5130 .arg(&ne)
5131 .arg(&na)
5132 .arg(&qtype)
5133 .arg(&rbi);
5134 unsafe {
5135 b.launch(cfg)?;
5136 }
5137 Ok(y)
5138 }
5139
5140 #[allow(clippy::too_many_arguments)]
5143 pub fn moe_pairs_matvec_q8_dec(
5144 &self,
5145 table: &CudaSlice<u64>,
5146 proj: i32,
5147 ex_ids: &CudaSlice<i32>,
5148 ex_off: &CudaSlice<i32>,
5149 ex_pairs: &CudaSlice<i32>,
5150 pair_tok: &CudaSlice<i32>,
5151 aq: &CudaSlice<i8>,
5152 ad: &CudaSlice<f32>,
5153 in_f: usize,
5154 out_f: usize,
5155 n_expert: usize,
5156 n_active: usize,
5157 n_pairs: usize,
5158 qtype: i32,
5159 row_bytes: usize,
5160 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5161 let f = self.func("moe_pairs_matvec_q8_dec");
5162 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5163 const ROWS: u32 = 4;
5164 let cfg = LaunchConfig {
5165 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5166 block_dim: (32, ROWS, 1),
5167 shared_mem_bytes: 0,
5168 };
5169 let (inf, outf, ne, na, rbi) = (
5170 in_f as i32,
5171 out_f as i32,
5172 n_expert as i32,
5173 n_active as i32,
5174 row_bytes as i64,
5175 );
5176 let __s_b = self.gpu.stream();
5177 let mut b = __s_b.launch_builder(&f);
5178 b.arg(table)
5179 .arg(&proj)
5180 .arg(ex_ids)
5181 .arg(ex_off)
5182 .arg(ex_pairs)
5183 .arg(pair_tok)
5184 .arg(aq)
5185 .arg(ad)
5186 .arg(&mut y)
5187 .arg(&inf)
5188 .arg(&outf)
5189 .arg(&ne)
5190 .arg(&na)
5191 .arg(&qtype)
5192 .arg(&rbi);
5193 unsafe {
5194 b.launch(cfg)?;
5195 }
5196 Ok(y)
5197 }
5198
5199 pub fn moe_pairs_gelu_mul(
5200 &self,
5201 gate: &CudaSlice<f32>,
5202 up: &CudaSlice<f32>,
5203 n: usize,
5204 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5205 let f = self.func("moe_pairs_gelu_mul");
5206 let mut act = self.alloc_uninit::<f32>(n)?;
5207 let cfg = LaunchConfig::for_num_elems(n as u32);
5208 let nl = n as i64;
5209 let __s_b = self.gpu.stream();
5210 let mut b = __s_b.launch_builder(&f);
5211 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5212 unsafe {
5213 b.launch(cfg)?;
5214 }
5215 Ok(act)
5216 }
5217
5218 pub fn moe_pairs_silu_mul(
5219 &self,
5220 gate: &CudaSlice<f32>,
5221 up: &CudaSlice<f32>,
5222 n: usize,
5223 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5224 let f = self.func("moe_pairs_silu_mul");
5225 let mut act = self.alloc_uninit::<f32>(n)?;
5226 let cfg = LaunchConfig::for_num_elems(n as u32);
5227 let nl = n as i64;
5228 let __s_b = self.gpu.stream();
5229 let mut b = __s_b.launch_builder(&f);
5230 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5231 unsafe {
5232 b.launch(cfg)?;
5233 }
5234 Ok(act)
5235 }
5236
5237 #[allow(clippy::too_many_arguments)]
5238 pub fn moe_pairs_scatter(
5239 &self,
5240 y_down: &CudaSlice<f32>,
5241 pair_w: &CudaSlice<f32>,
5242 tok_pair_off: &CudaSlice<i32>,
5243 tok_pair_ids: &CudaSlice<i32>,
5244 moe_out: &mut CudaSlice<f32>,
5245 t: usize,
5246 n_embd: usize,
5247 ) -> Result<(), Box<dyn std::error::Error>> {
5248 let f = self.func("moe_pairs_scatter");
5249 let cfg = LaunchConfig {
5250 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5251 block_dim: (256, 1, 1),
5252 shared_mem_bytes: 0,
5253 };
5254 let ne = n_embd as i32;
5255 let __s_b = self.gpu.stream();
5256 let mut b = __s_b.launch_builder(&f);
5257 b.arg(y_down)
5258 .arg(pair_w)
5259 .arg(tok_pair_off)
5260 .arg(tok_pair_ids)
5261 .arg(moe_out)
5262 .arg(&ne);
5263 unsafe {
5264 b.launch(cfg)?;
5265 }
5266 Ok(())
5267 }
5268
5269 #[allow(clippy::too_many_arguments)]
5273 pub fn moe_gate_up_gelu8_dev_q8(
5274 &self,
5275 table: &CudaSlice<u64>,
5276 sel: &cudarc::driver::CudaView<i32>,
5277 aq: &CudaSlice<i8>,
5278 ad: &CudaSlice<f32>,
5279 in_f: usize,
5280 n_ff: usize,
5281 n_used: usize,
5282 n_expert: usize,
5283 qt_g: i32,
5284 qt_u: i32,
5285 rb_g: usize,
5286 rb_u: usize,
5287 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5288 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5289 let (inf, nff, ne, rbg, rbu) = (
5290 in_f as i32,
5291 n_ff as i32,
5292 n_expert as i32,
5293 rb_g as i64,
5294 rb_u as i64,
5295 );
5296 let f = self.func("moe_gate_up_gelu8_dev_q8");
5297 let cfg = LaunchConfig {
5298 grid_dim: (n_ff as u32, n_used as u32, 1),
5299 block_dim: (32, 1, 1),
5300 shared_mem_bytes: 0,
5301 };
5302 let __s_b = self.gpu.stream();
5303 let mut b = __s_b.launch_builder(&f);
5304 b.arg(table)
5305 .arg(sel)
5306 .arg(aq)
5307 .arg(ad)
5308 .arg(&mut act)
5309 .arg(&inf)
5310 .arg(&nff)
5311 .arg(&ne)
5312 .arg(&qt_g)
5313 .arg(&qt_u)
5314 .arg(&rbg)
5315 .arg(&rbu);
5316 unsafe {
5317 b.launch(cfg)?;
5318 }
5319 Ok(act)
5320 }
5321
5322 #[allow(clippy::too_many_arguments)]
5324 pub fn moe_gate_up_gelu8_dev_q8_rows(
5325 &self,
5326 table: &CudaSlice<u64>,
5327 sel: &CudaSlice<i32>,
5328 aq: &CudaSlice<i8>,
5329 ad: &CudaSlice<f32>,
5330 t: usize,
5331 in_f: usize,
5332 n_ff: usize,
5333 n_used: usize,
5334 n_expert: usize,
5335 qt_g: i32,
5336 qt_u: i32,
5337 rb_g: usize,
5338 rb_u: usize,
5339 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5340 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5341 let (inf, nff, ne, rbg, rbu, nu) = (
5342 in_f as i32,
5343 n_ff as i32,
5344 n_expert as i32,
5345 rb_g as i64,
5346 rb_u as i64,
5347 n_used as i32,
5348 );
5349 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5350 let cfg = LaunchConfig {
5351 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5352 block_dim: (32, 1, 1),
5353 shared_mem_bytes: 0,
5354 };
5355 let __s_b = self.gpu.stream();
5356 let mut b = __s_b.launch_builder(&f);
5357 b.arg(table)
5358 .arg(sel)
5359 .arg(aq)
5360 .arg(ad)
5361 .arg(&mut act)
5362 .arg(&inf)
5363 .arg(&nff)
5364 .arg(&ne)
5365 .arg(&qt_g)
5366 .arg(&qt_u)
5367 .arg(&rbg)
5368 .arg(&rbu)
5369 .arg(&nu);
5370 unsafe {
5371 b.launch(cfg)?;
5372 }
5373 Ok(act)
5374 }
5375
5376 #[allow(clippy::too_many_arguments)]
5378 pub fn moe_gate_up_gelu8_dev_q8_csr(
5379 &self,
5380 table: &CudaSlice<u64>,
5381 sel: &CudaSlice<i32>,
5382 aq: &CudaSlice<i8>,
5383 ad: &CudaSlice<f32>,
5384 n_pairs: usize,
5385 in_f: usize,
5386 n_ff: usize,
5387 n_used: usize,
5388 n_expert: usize,
5389 qt_g: i32,
5390 qt_u: i32,
5391 rb_g: usize,
5392 rb_u: usize,
5393 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5394 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5395 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5396 in_f as i32,
5397 n_ff as i32,
5398 n_expert as i32,
5399 rb_g as i64,
5400 rb_u as i64,
5401 n_used as i32,
5402 n_pairs as i32,
5403 );
5404 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5405 let cfg = LaunchConfig {
5406 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5407 block_dim: (32, 1, 1),
5408 shared_mem_bytes: 0,
5409 };
5410 let __s_b = self.gpu.stream();
5411 let mut b = __s_b.launch_builder(&f);
5412 b.arg(table)
5413 .arg(sel)
5414 .arg(aq)
5415 .arg(ad)
5416 .arg(&mut act)
5417 .arg(&inf)
5418 .arg(&nff)
5419 .arg(&ne)
5420 .arg(&qt_g)
5421 .arg(&qt_u)
5422 .arg(&rbg)
5423 .arg(&rbu)
5424 .arg(&nu)
5425 .arg(&npi);
5426 unsafe {
5427 b.launch(cfg)?;
5428 }
5429 Ok(act)
5430 }
5431
5432 #[allow(clippy::too_many_arguments)]
5434 pub fn moe_down8_fma_dev_q8_rows_g(
5435 &self,
5436 table: &CudaSlice<u64>,
5437 sel: &CudaSlice<i32>,
5438 w: &CudaSlice<f32>,
5439 aq2: &CudaSlice<i8>,
5440 ad2: &CudaSlice<f32>,
5441 dst: &mut CudaSlice<f32>,
5442 t: usize,
5443 in_f: usize,
5444 out_f: usize,
5445 n_used: usize,
5446 n_expert: usize,
5447 qt: i32,
5448 rb: usize,
5449 ) -> Result<(), Box<dyn std::error::Error>> {
5450 let (inf, outf, nu, ne, rbi) = (
5451 in_f as i32,
5452 out_f as i32,
5453 n_used as i32,
5454 n_expert as i32,
5455 rb as i64,
5456 );
5457 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5461 let f = self.func(if step_b1_w8 {
5462 "moe_down8_fma_dev_q8_rows_w8"
5463 } else {
5464 "moe_down8_fma_dev_q8_rows_g"
5465 });
5466 let cfg = LaunchConfig {
5467 grid_dim: (out_f as u32, 1, t as u32),
5468 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5469 shared_mem_bytes: 0,
5470 };
5471 let __s_b = self.gpu.stream();
5472 let mut b = __s_b.launch_builder(&f);
5473 b.arg(table)
5474 .arg(sel)
5475 .arg(w)
5476 .arg(aq2)
5477 .arg(ad2)
5478 .arg(dst)
5479 .arg(&inf)
5480 .arg(&outf)
5481 .arg(&nu)
5482 .arg(&ne)
5483 .arg(&qt)
5484 .arg(&rbi);
5485 unsafe {
5486 b.launch(cfg)?;
5487 }
5488 Ok(())
5489 }
5490
5491 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5495 let (out_f, in_f) = (2048usize, 2816usize);
5496 let nblk = in_f / 32;
5497 let mut seed = 0x9E3779B97F4A7C15u64;
5498 let mut rng = move || {
5499 seed = seed
5500 .wrapping_mul(6364136223846793005)
5501 .wrapping_add(1442695040888963407);
5502 (seed >> 33) as u8
5503 };
5504 let mut w = vec![0u8; out_f * nblk * 18];
5505 for b in w.iter_mut() {
5506 *b = rng();
5507 }
5508 for r in 0..out_f {
5509 for g in 0..nblk {
5510 let off = (r * nblk + g) * 18;
5511 w[off] = 0x00;
5512 w[off + 1] = 0x2C; }
5514 }
5515 let qplane = out_f * nblk * 16;
5516 let mut wrp = vec![0u8; w.len()];
5517 for r in 0..out_f {
5518 for g in 0..nblk {
5519 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5520 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5521 .copy_from_slice(&src[0..2]);
5522 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5523 }
5524 }
5525 let w_d = self.htod_bytes(&w)?;
5526 let wrp_d = self.htod_bytes(&wrp)?;
5527 let mut aq = vec![0i8; m * in_f];
5528 for v in aq.iter_mut() {
5529 *v = rng() as i8;
5530 }
5531 let aq_d = self.htod_i8(&aq)?;
5532 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5533 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5534 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5535 const RPB: u32 = 4;
5536 let cfg = LaunchConfig {
5537 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5538 block_dim: (32, RPB, 1),
5539 shared_mem_bytes: 0,
5540 };
5541 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5542 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5543 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5544 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5545 {
5546 let __s_b = self.gpu.stream();
5547 let mut b = __s_b.launch_builder(&fb);
5548 b.arg(&w_d)
5549 .arg(&aq_d)
5550 .arg(&ad_d)
5551 .arg(&mut y0)
5552 .arg(&inf)
5553 .arg(&outf)
5554 .arg(&mi)
5555 .arg(&rb);
5556 unsafe {
5557 b.launch(cfg)?;
5558 }
5559 let __s_b = self.gpu.stream();
5560 let mut b = __s_b.launch_builder(&fr);
5561 b.arg(&wrp_d)
5562 .arg(&aq_d)
5563 .arg(&ad_d)
5564 .arg(&mut y1)
5565 .arg(&inf)
5566 .arg(&outf)
5567 .arg(&mi)
5568 .arg(&qp);
5569 unsafe {
5570 b.launch(cfg)?;
5571 }
5572 }
5573 self.gpu.stream().synchronize()?;
5574 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5575 let nd = h0
5576 .iter()
5577 .zip(&h1)
5578 .filter(|(a, b)| a.to_bits() != b.to_bits())
5579 .count();
5580 if nd != 0 {
5581 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5582 }
5583 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5584 self.gpu.stream().synchronize()?;
5585 let t0 = std::time::Instant::now();
5586 for _ in 0..500 {
5587 if rp {
5588 let __s_b = self.gpu.stream();
5589 let mut b = __s_b.launch_builder(&fr);
5590 b.arg(&wrp_d)
5591 .arg(&aq_d)
5592 .arg(&ad_d)
5593 .arg(&mut y1)
5594 .arg(&inf)
5595 .arg(&outf)
5596 .arg(&mi)
5597 .arg(&qp);
5598 unsafe {
5599 b.launch(cfg)?;
5600 }
5601 } else {
5602 let __s_b = self.gpu.stream();
5603 let mut b = __s_b.launch_builder(&fb);
5604 b.arg(&w_d)
5605 .arg(&aq_d)
5606 .arg(&ad_d)
5607 .arg(&mut y0)
5608 .arg(&inf)
5609 .arg(&outf)
5610 .arg(&mi)
5611 .arg(&rb);
5612 unsafe {
5613 b.launch(cfg)?;
5614 }
5615 }
5616 }
5617 self.gpu.stream().synchronize()?;
5618 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5619 };
5620 let _ = time(false)?;
5621 let _ = time(true)?; Ok((time(false)?, time(true)?))
5623 }
5624
5625 pub fn build_q4_rp4(
5630 &self,
5631 t: &mut crate::model::GpuTensor,
5632 ) -> Result<(), Box<dyn std::error::Error>> {
5633 use crate::model::GpuTensor;
5634 let GpuTensor::Quant {
5635 bytes,
5636 qtype,
5637 row_bytes,
5638 ne,
5639 rp4,
5640 ..
5641 } = t
5642 else {
5643 return Ok(());
5644 };
5645 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5646 return Ok(());
5647 }
5648 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5649 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5650 return Ok(());
5651 }
5652 let nblk = in_f / 32;
5653 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5654 let f = self.func("q4_0_split_rp_build");
5655 let n = (out_f * nblk) as i32;
5656 let cfg = LaunchConfig {
5657 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5658 block_dim: (256, 1, 1),
5659 shared_mem_bytes: 0,
5660 };
5661 let (of, nb) = (out_f as i32, nblk as i32);
5662 let _ = n;
5663 let __s_b = self.gpu.stream();
5664 let mut b = __s_b.launch_builder(&f);
5665 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5666 unsafe {
5667 b.launch(cfg)?;
5668 }
5669 *rp4 = Some(dst);
5670 Ok(())
5671 }
5672
5673 pub fn build_q8_rp4(
5678 &self,
5679 t: &mut crate::model::GpuTensor,
5680 ) -> Result<(), Box<dyn std::error::Error>> {
5681 use crate::model::GpuTensor;
5682 let GpuTensor::Quant {
5683 bytes,
5684 qtype,
5685 row_bytes,
5686 ne,
5687 rp4,
5688 ..
5689 } = t
5690 else {
5691 return Ok(());
5692 };
5693 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5694 return Ok(());
5695 }
5696 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5697 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5698 return Ok(());
5699 }
5700 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5701 Ok(())
5702 }
5703
5704 pub fn build_q8_rp4_raw(
5707 &self,
5708 bytes: &CudaSlice<u8>,
5709 in_f: usize,
5710 out_f: usize,
5711 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5712 assert!(in_f % 32 == 0);
5713 let nblk = in_f / 32;
5714 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5715 let f = self.func("q8_0_split_rp_build");
5716 let cfg = LaunchConfig {
5717 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5718 block_dim: (256, 1, 1),
5719 shared_mem_bytes: 0,
5720 };
5721 let (of, nb) = (out_f as i32, nblk as i32);
5722 let __s_b = self.gpu.stream();
5723 let mut b = __s_b.launch_builder(&f);
5724 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5725 unsafe {
5726 b.launch(cfg)?;
5727 }
5728 Ok(dst)
5729 }
5730
5731 pub fn build_q4k_rp4(
5739 &self,
5740 t: &mut crate::model::GpuTensor,
5741 ) -> Result<(), Box<dyn std::error::Error>> {
5742 use crate::model::GpuTensor;
5743 let GpuTensor::Quant {
5744 bytes,
5745 qtype,
5746 row_bytes,
5747 ne,
5748 rp4,
5749 ..
5750 } = t
5751 else {
5752 return Ok(());
5753 };
5754 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5755 return Ok(());
5756 }
5757 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5758 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5759 return Ok(());
5760 }
5761 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5762 Ok(())
5763 }
5764
5765 pub fn build_q6k_rp4(
5766 &self,
5767 t: &mut crate::model::GpuTensor,
5768 ) -> Result<(), Box<dyn std::error::Error>> {
5769 use crate::model::GpuTensor;
5770 let GpuTensor::Quant {
5771 bytes,
5772 qtype,
5773 row_bytes,
5774 ne,
5775 rp4,
5776 ..
5777 } = t
5778 else {
5779 return Ok(());
5780 };
5781 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5782 return Ok(());
5783 }
5784 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5785 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5786 return Ok(());
5787 }
5788 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5789 Ok(())
5790 }
5791
5792 pub fn build_kq_rp4_raw(
5794 &self,
5795 bytes: &CudaSlice<u8>,
5796 in_f: usize,
5797 out_f: usize,
5798 qtype: i32,
5799 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5800 assert!(in_f % 256 == 0);
5801 let nsbk = in_f / 256;
5802 let (sb_bytes, kname) = match qtype {
5803 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5804 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5805 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5806 };
5807 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5808 let f = self.func(kname);
5809 let cfg = LaunchConfig {
5810 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5811 block_dim: (256, 1, 1),
5812 shared_mem_bytes: 0,
5813 };
5814 let (of, nb) = (out_f as i32, nsbk as i32);
5815 let __s_b = self.gpu.stream();
5816 let mut b = __s_b.launch_builder(&f);
5817 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5818 unsafe {
5819 b.launch(cfg)?;
5820 }
5821 Ok(dst)
5822 }
5823
5824 pub fn kqrp_enabled() -> bool {
5828 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5829 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5830 Ok("0") => false,
5831 Ok(_) => true,
5832 Err(_) => cfg!(memra_hopper_mma),
5833 })
5834 }
5835
5836 pub fn build_q4_rp_swap(
5842 &self,
5843 t: &mut crate::model::GpuTensor,
5844 ) -> Result<bool, Box<dyn std::error::Error>> {
5845 use crate::model::GpuTensor;
5846 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5856 return Ok(false);
5857 }
5858 self.build_q4_rp4(t)?;
5859 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5861 return Ok(false);
5862 };
5863 match rp4.take() {
5864 Some(split) => {
5865 *bytes = split; *rp = true;
5867 Ok(true)
5868 }
5869 None => Ok(false),
5870 }
5871 }
5872
5873 pub fn q4rp_enabled() -> bool {
5875 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5876 *ON.get_or_init(|| {
5877 std::env::var("MEMRA_Q4RP")
5878 .map(|v| v != "0")
5879 .unwrap_or(true)
5880 })
5881 }
5882
5883 pub fn copy_rows_strided(
5886 &self,
5887 src: &CudaSlice<f32>,
5888 dst: &mut CudaSlice<f32>,
5889 row_elems: usize,
5890 n_rows: usize,
5891 src_stride: usize,
5892 src_off: usize,
5893 ) -> Result<(), Box<dyn std::error::Error>> {
5894 let f = self.func("copy_rows_strided_f32");
5895 let cfg = LaunchConfig {
5896 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5897 block_dim: (256, 1, 1),
5898 shared_mem_bytes: 0,
5899 };
5900 let (re, nr) = (row_elems as i32, n_rows as i32);
5901 let (st, off) = (src_stride as i64, src_off as i64);
5902 let __s_b = self.gpu.stream();
5903 let mut b = __s_b.launch_builder(&f);
5904 b.arg(src)
5905 .arg(&mut *dst)
5906 .arg(&re)
5907 .arg(&nr)
5908 .arg(&st)
5909 .arg(&off);
5910 unsafe {
5911 b.launch(cfg)?;
5912 }
5913 Ok(())
5914 }
5915
5916 pub fn place_rows_strided(
5922 &self,
5923 src: &CudaSlice<f32>,
5924 dst: &mut CudaSlice<f32>,
5925 row_elems: usize,
5926 n_rows: usize,
5927 dst_stride: usize,
5928 dst_off: usize,
5929 ) -> Result<(), Box<dyn std::error::Error>> {
5930 if row_elems == 0 || n_rows == 0 {
5931 return Err("strided row placement requires nonzero rows and row width".into());
5932 }
5933 let src_len = n_rows
5934 .checked_mul(row_elems)
5935 .ok_or("strided row placement source size overflow")?;
5936 let dst_len = n_rows
5937 .checked_sub(1)
5938 .and_then(|rows| rows.checked_mul(dst_stride))
5939 .and_then(|base| base.checked_add(dst_off))
5940 .and_then(|base| base.checked_add(row_elems))
5941 .ok_or("strided row placement destination size overflow")?;
5942 let row_end = dst_off
5943 .checked_add(row_elems)
5944 .ok_or("strided row placement row size overflow")?;
5945 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
5946 return Err(format!(
5947 "strided row placement geometry mismatch: src={} need_src={src_len} \
5948 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
5949 dst_stride={dst_stride} dst_off={dst_off}",
5950 src.len(),
5951 dst.len(),
5952 )
5953 .into());
5954 }
5955 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
5956 return Err("strided row placement exceeds CUDA kernel geometry".into());
5957 }
5958 let f = self.func("place_rows_strided_f32");
5959 let cfg = LaunchConfig {
5960 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5961 block_dim: (256, 1, 1),
5962 shared_mem_bytes: 0,
5963 };
5964 let (re, nr) = (row_elems as i32, n_rows as i32);
5965 let (st, off) = (dst_stride as i64, dst_off as i64);
5966 let __s_b = self.gpu.stream();
5967 let mut b = __s_b.launch_builder(&f);
5968 b.arg(src)
5969 .arg(&mut *dst)
5970 .arg(&re)
5971 .arg(&nr)
5972 .arg(&st)
5973 .arg(&off);
5974 unsafe {
5975 b.launch(cfg)?;
5976 }
5977 Ok(())
5978 }
5979
5980 pub fn u32_set_k(
5982 &self,
5983 dst: &mut CudaSlice<u32>,
5984 v: u32,
5985 idx: usize,
5986 ) -> Result<(), Box<dyn std::error::Error>> {
5987 let f = self.func("u32_set_k");
5988 let cfg = LaunchConfig {
5989 grid_dim: (1, 1, 1),
5990 block_dim: (1, 1, 1),
5991 shared_mem_bytes: 0,
5992 };
5993 let ii = idx as i32;
5994 let __s_b = self.gpu.stream();
5995 let mut b = __s_b.launch_builder(&f);
5996 b.arg(dst).arg(&v).arg(&ii);
5997 unsafe {
5998 b.launch(cfg)?;
5999 }
6000 Ok(())
6001 }
6002
6003 pub fn i32_add_k(
6005 &self,
6006 d: &mut CudaSlice<i32>,
6007 v: i32,
6008 ) -> Result<(), Box<dyn std::error::Error>> {
6009 let f = self.func("i32_add_k");
6010 let cfg = LaunchConfig {
6011 grid_dim: (1, 1, 1),
6012 block_dim: (32, 1, 1),
6013 shared_mem_bytes: 0,
6014 };
6015 let __s_b = self.gpu.stream();
6016 let mut b = __s_b.launch_builder(&f);
6017 b.arg(d).arg(&v);
6018 unsafe {
6019 b.launch(cfg)?;
6020 }
6021 Ok(())
6022 }
6023
6024 pub fn i32_iota_from(
6026 &self,
6027 ctr: &CudaSlice<i32>,
6028 dst: &mut CudaSlice<i32>,
6029 n: usize,
6030 ) -> Result<(), Box<dyn std::error::Error>> {
6031 let f = self.func("i32_iota_from");
6032 let cfg = LaunchConfig::for_num_elems(n as u32);
6033 let ni = n as i32;
6034 let __s_b = self.gpu.stream();
6035 let mut b = __s_b.launch_builder(&f);
6036 b.arg(ctr).arg(dst).arg(&ni);
6037 unsafe {
6038 b.launch(cfg)?;
6039 }
6040 Ok(())
6041 }
6042
6043 pub fn u32_map_k(
6045 &self,
6046 buf: &mut CudaSlice<u32>,
6047 map: &CudaSlice<u32>,
6048 idx: usize,
6049 ) -> Result<(), Box<dyn std::error::Error>> {
6050 let f = self.func("u32_map_k");
6051 let cfg = LaunchConfig {
6052 grid_dim: (1, 1, 1),
6053 block_dim: (1, 1, 1),
6054 shared_mem_bytes: 0,
6055 };
6056 let ii = idx as i32;
6057 let __s_b = self.gpu.stream();
6058 let mut b = __s_b.launch_builder(&f);
6059 b.arg(buf).arg(map).arg(&ii);
6060 unsafe {
6061 b.launch(cfg)?;
6062 }
6063 Ok(())
6064 }
6065
6066 #[allow(clippy::too_many_arguments)]
6068 pub fn u32_pack2(
6069 &self,
6070 a: &CudaSlice<u32>,
6071 off_a: usize,
6072 n1: usize,
6073 b_in: &CudaSlice<u32>,
6074 n2: usize,
6075 out: &mut CudaSlice<u32>,
6076 ) -> Result<(), Box<dyn std::error::Error>> {
6077 let f = self.func("u32_pack2");
6078 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6079 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6080 let __s_b = self.gpu.stream();
6081 let mut b = __s_b.launch_builder(&f);
6082 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6083 unsafe {
6084 b.launch(cfg)?;
6085 }
6086 Ok(())
6087 }
6088
6089 pub fn moe_w_exscale(
6091 &self,
6092 w: &mut CudaSlice<f32>,
6093 sel: &CudaSlice<i32>,
6094 s: &CudaSlice<f32>,
6095 n: usize,
6096 ) -> Result<(), Box<dyn std::error::Error>> {
6097 let f = self.func("moe_w_exscale");
6098 let cfg = LaunchConfig::for_num_elems(n as u32);
6099 let ni = n as i32;
6100 let __s_b = self.gpu.stream();
6101 let mut b = __s_b.launch_builder(&f);
6102 b.arg(w).arg(sel).arg(s).arg(&ni);
6103 unsafe {
6104 b.launch(cfg)?;
6105 }
6106 Ok(())
6107 }
6108
6109 pub fn moe_w_scale_by_expert(
6112 &self,
6113 w: &mut CudaSlice<f32>,
6114 sel: &CudaSlice<i32>,
6115 macros: &CudaSlice<f32>,
6116 n_expert: usize,
6117 n: usize,
6118 ) -> Result<(), Box<dyn std::error::Error>> {
6119 let f = self.func("moe_w_scale_by_expert");
6120 let cfg = LaunchConfig {
6121 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6122 block_dim: (64, 1, 1),
6123 shared_mem_bytes: 0,
6124 };
6125 let (ne, nn) = (n_expert as i32, n as i32);
6126 let __s_b = self.gpu.stream();
6127 let mut b = __s_b.launch_builder(&f);
6128 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6129 unsafe {
6130 b.launch(cfg)?;
6131 }
6132 Ok(())
6133 }
6134
6135 pub fn moe_gate_up_silu8_dev_q8(
6136 &self,
6137 table: &CudaSlice<u64>,
6138 sel: &cudarc::driver::CudaView<i32>,
6139 aq: &CudaSlice<i8>,
6140 ad: &CudaSlice<f32>,
6141 in_f: usize,
6142 n_ff: usize,
6143 n_used: usize,
6144 n_expert: usize,
6145 qt_g: i32,
6146 qt_u: i32,
6147 rb_g: usize,
6148 rb_u: usize,
6149 macros: &CudaSlice<f32>,
6150 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6151 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6152 let (mode, wpb) = GU.get_or_init(|| {
6153 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6154 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6155 .ok()
6156 .and_then(|v| v.parse().ok())
6157 .unwrap_or(4u32)
6158 .clamp(1, 16);
6159 (mode, wpb)
6160 });
6161 let (mode, wpb) = (mode.as_str(), *wpb);
6162 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6163 let (inf, nff, ne, rbg, rbu) = (
6164 in_f as i32,
6165 n_ff as i32,
6166 n_expert as i32,
6167 rb_g as i64,
6168 rb_u as i64,
6169 );
6170 let (f, cfg) = match mode {
6171 "1" | "2" | "4" => {
6172 let rpw: u32 = mode.parse().unwrap();
6173 let f = self.func(match rpw {
6174 1 => "moe_gate_up_silu8_dev_q8_r1",
6175 2 => "moe_gate_up_silu8_dev_q8_r2",
6176 _ => "moe_gate_up_silu8_dev_q8_r4",
6177 });
6178 let rows_per_block = (rpw * wpb) as usize;
6179 let gx = n_ff.div_ceil(rows_per_block) as u32;
6180 (
6181 f,
6182 LaunchConfig {
6183 grid_dim: (gx, n_used as u32, 1),
6184 block_dim: (32, wpb, 1),
6185 shared_mem_bytes: 0,
6186 },
6187 )
6188 }
6189 "j8" if n_used <= 32 => (
6190 self.func("moe_gate_up_silu8_dev_q8_j8"),
6191 LaunchConfig {
6192 grid_dim: (n_ff as u32, 1, 1),
6193 block_dim: (32, n_used as u32, 1),
6194 shared_mem_bytes: 0,
6195 },
6196 ),
6197 "vsm2" => {
6199 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6200 let sh = (rb_g + rb_u) as u32;
6201 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6202 f.set_attribute(
6203 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6204 sh as i32,
6205 )?;
6206 (
6207 f,
6208 LaunchConfig {
6209 grid_dim: (n_ff as u32, n_used as u32, 1),
6210 block_dim: (32, 1, 1),
6211 shared_mem_bytes: sh,
6212 },
6213 )
6214 }
6215 "vsm" => {
6216 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6217 let sh = (rb_g + rb_u) as u32;
6218 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6219 f.set_attribute(
6220 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6221 sh as i32,
6222 )?;
6223 (
6224 f,
6225 LaunchConfig {
6226 grid_dim: (n_ff as u32, n_used as u32, 1),
6227 block_dim: (32, 1, 1),
6228 shared_mem_bytes: sh,
6229 },
6230 )
6231 }
6232 "sg" => (
6233 self.func("moe_gate_up_silu8_dev_q8_sg"),
6234 LaunchConfig {
6235 grid_dim: (n_ff as u32, n_used as u32, 1),
6236 block_dim: (32, 1, 1),
6237 shared_mem_bytes: 0,
6238 },
6239 ),
6240 "j8sg" if n_used <= 32 => (
6241 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6242 LaunchConfig {
6243 grid_dim: (n_ff as u32, 1, 1),
6244 block_dim: (32, n_used as u32, 1),
6245 shared_mem_bytes: 0,
6246 },
6247 ),
6248 "u64" if in_f == 2048 => (
6249 self.func("moe_gate_up_silu8_dev_q8_u64"),
6250 LaunchConfig {
6251 grid_dim: (n_ff as u32, n_used as u32, 1),
6252 block_dim: (32, 1, 1),
6253 shared_mem_bytes: 0,
6254 },
6255 ),
6256 "gs4" if in_f == 2048 => (
6257 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6258 LaunchConfig {
6259 grid_dim: (n_ff as u32, n_used as u32, 1),
6260 block_dim: (32, 4, 1),
6261 shared_mem_bytes: 0,
6262 },
6263 ),
6264 "v" | "" => (
6266 self.func("moe_gate_up_silu8_dev_q8_v"),
6267 LaunchConfig {
6268 grid_dim: (n_ff as u32, n_used as u32, 1),
6269 block_dim: (32, 1, 1),
6270 shared_mem_bytes: 0,
6271 },
6272 ),
6273 "s2" => (
6274 self.func("moe_gate_up_silu8_dev_q8_s2"),
6275 LaunchConfig {
6276 grid_dim: (n_ff as u32, n_used as u32, 1),
6277 block_dim: (32, 2, 1),
6278 shared_mem_bytes: 0,
6279 },
6280 ),
6281 "s2z" => {
6282 let rz = wpb.min(16); (
6284 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6285 LaunchConfig {
6286 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6287 block_dim: (32, 2, rz),
6288 shared_mem_bytes: 0,
6289 },
6290 )
6291 }
6292 _ => (
6293 self.func("moe_gate_up_silu8_dev_q8"),
6294 LaunchConfig {
6295 grid_dim: (n_ff as u32, n_used as u32, 1),
6296 block_dim: (32, 1, 1),
6297 shared_mem_bytes: 0,
6298 },
6299 ),
6300 };
6301 let __s_b = self.gpu.stream();
6302 let mut b = __s_b.launch_builder(&f);
6303 b.arg(table)
6304 .arg(sel)
6305 .arg(aq)
6306 .arg(ad)
6307 .arg(&mut act)
6308 .arg(&inf)
6309 .arg(&nff)
6310 .arg(&ne)
6311 .arg(&qt_g)
6312 .arg(&qt_u)
6313 .arg(&rbg)
6314 .arg(&rbu)
6315 .arg(macros);
6316 unsafe {
6317 b.launch(cfg)?;
6318 }
6319 Ok(act)
6320 }
6321
6322 #[allow(clippy::too_many_arguments)]
6323 pub fn moe_down8_fma_dev_q8(
6324 &self,
6325 table: &CudaSlice<u64>,
6326 sel: &cudarc::driver::CudaView<i32>,
6327 w: &cudarc::driver::CudaView<f32>,
6328 aq2: &CudaSlice<i8>,
6329 ad2: &CudaSlice<f32>,
6330 dst: &mut cudarc::driver::CudaViewMut<f32>,
6331 in_f: usize,
6332 out_f: usize,
6333 n_used: usize,
6334 n_expert: usize,
6335 qt: i32,
6336 rb: usize,
6337 ) -> Result<(), Box<dyn std::error::Error>> {
6338 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6339 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6340 let (inf, outf, nu, ne, rbi) = (
6341 in_f as i32,
6342 out_f as i32,
6343 n_used as i32,
6344 n_expert as i32,
6345 rb as i64,
6346 );
6347 let (f, cfg) = match mode.as_str() {
6350 m @ ("1" | "2" | "4") if n_used <= 8 => {
6351 let rpw: usize = m.parse().unwrap();
6352 let f = self.func(match rpw {
6353 1 => "moe_down8_fma_dev_q8_w8r1",
6354 2 => "moe_down8_fma_dev_q8_w8r2",
6355 _ => "moe_down8_fma_dev_q8_w8r4",
6356 });
6357 (
6358 f,
6359 LaunchConfig {
6360 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6361 block_dim: (32, n_used as u32, 1),
6362 shared_mem_bytes: 0,
6363 },
6364 )
6365 }
6366 "h2" if in_f == 512 => (
6367 self.func("moe_down8_fma_dev_q8_h2"),
6368 LaunchConfig {
6369 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6370 block_dim: (32, 1, 1),
6371 shared_mem_bytes: 0,
6372 },
6373 ),
6374 "" if in_f == 704 && n_used <= 8 => (
6377 self.func("moe_down8_fma_dev_q8_w8r2"),
6378 LaunchConfig {
6379 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6380 block_dim: (32, n_used as u32, 1),
6381 shared_mem_bytes: 0,
6382 },
6383 ),
6384 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6388 self.func("moe_down8_fma_dev_q8_w8h2v"),
6389 LaunchConfig {
6390 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6391 block_dim: (32, n_used as u32, 1),
6392 shared_mem_bytes: 0,
6393 },
6394 ),
6395 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6396 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6397 LaunchConfig {
6398 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6399 block_dim: (32, n_used as u32, 1),
6400 shared_mem_bytes: 0,
6401 },
6402 ),
6403 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6404 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6405 LaunchConfig {
6406 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6407 block_dim: (32, n_used as u32, 1),
6408 shared_mem_bytes: 0,
6409 },
6410 ),
6411 "w8h2" if in_f == 512 && n_used <= 8 => (
6412 self.func("moe_down8_fma_dev_q8_w8h2"),
6413 LaunchConfig {
6414 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6415 block_dim: (32, n_used as u32, 1),
6416 shared_mem_bytes: 0,
6417 },
6418 ),
6419 _ => (
6420 self.func("moe_down8_fma_dev_q8"),
6421 LaunchConfig {
6422 grid_dim: (out_f as u32, 1, 1),
6423 block_dim: (32, 1, 1),
6424 shared_mem_bytes: 0,
6425 },
6426 ),
6427 };
6428 let __s_b = self.gpu.stream();
6429 let mut b = __s_b.launch_builder(&f);
6430 b.arg(table)
6431 .arg(sel)
6432 .arg(w)
6433 .arg(aq2)
6434 .arg(ad2)
6435 .arg(dst)
6436 .arg(&inf)
6437 .arg(&outf)
6438 .arg(&nu)
6439 .arg(&ne)
6440 .arg(&qt)
6441 .arg(&rbi);
6442 unsafe {
6443 b.launch(cfg)?;
6444 }
6445 Ok(())
6446 }
6447
6448 #[allow(clippy::too_many_arguments)]
6455 pub fn moe_gate_up_silu8_dev_q8_rows(
6456 &self,
6457 table: &CudaSlice<u64>,
6458 sel: &CudaSlice<i32>,
6459 aq: &CudaSlice<i8>,
6460 ad: &CudaSlice<f32>,
6461 t: usize,
6462 in_f: usize,
6463 n_ff: usize,
6464 n_used: usize,
6465 n_expert: usize,
6466 qt_g: i32,
6467 qt_u: i32,
6468 rb_g: usize,
6469 rb_u: usize,
6470 macros: &CudaSlice<f32>,
6471 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6472 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6473 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6474 let cfg = LaunchConfig {
6475 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6476 block_dim: (32, 1, 1),
6477 shared_mem_bytes: 0,
6478 };
6479 let (inf, nff, ne, nu, rbg, rbu) = (
6480 in_f as i32,
6481 n_ff as i32,
6482 n_expert as i32,
6483 n_used as i32,
6484 rb_g as i64,
6485 rb_u as i64,
6486 );
6487 let __s_b = self.gpu.stream();
6488 let mut b = __s_b.launch_builder(&f);
6489 b.arg(table)
6490 .arg(sel)
6491 .arg(aq)
6492 .arg(ad)
6493 .arg(&mut act)
6494 .arg(&inf)
6495 .arg(&nff)
6496 .arg(&ne)
6497 .arg(&qt_g)
6498 .arg(&qt_u)
6499 .arg(&rbg)
6500 .arg(&rbu)
6501 .arg(&nu)
6502 .arg(macros);
6503 unsafe {
6504 b.launch(cfg)?;
6505 }
6506 Ok(act)
6507 }
6508
6509 #[allow(clippy::too_many_arguments)]
6514 pub fn moe_down8_fma_dev_q8_rows(
6515 &self,
6516 table: &CudaSlice<u64>,
6517 sel: &CudaSlice<i32>,
6518 w: &CudaSlice<f32>,
6519 aq2: &CudaSlice<i8>,
6520 ad2: &CudaSlice<f32>,
6521 dst: &mut CudaSlice<f32>,
6522 t: usize,
6523 in_f: usize,
6524 out_f: usize,
6525 n_used: usize,
6526 n_expert: usize,
6527 qt: i32,
6528 rb: usize,
6529 ) -> Result<(), Box<dyn std::error::Error>> {
6530 assert!(
6531 in_f == 512 && n_used <= 8,
6532 "down rows twin is w8h2v shape-gated"
6533 );
6534 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6535 let cfg = LaunchConfig {
6536 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6537 block_dim: (32, n_used as u32, 1),
6538 shared_mem_bytes: 0,
6539 };
6540 let (inf, outf, nu, ne, rbi) = (
6541 in_f as i32,
6542 out_f as i32,
6543 n_used as i32,
6544 n_expert as i32,
6545 rb as i64,
6546 );
6547 let __s_b = self.gpu.stream();
6548 let mut b = __s_b.launch_builder(&f);
6549 b.arg(table)
6550 .arg(sel)
6551 .arg(w)
6552 .arg(aq2)
6553 .arg(ad2)
6554 .arg(dst)
6555 .arg(&inf)
6556 .arg(&outf)
6557 .arg(&nu)
6558 .arg(&ne)
6559 .arg(&qt)
6560 .arg(&rbi);
6561 unsafe {
6562 b.launch(cfg)?;
6563 }
6564 Ok(())
6565 }
6566
6567 #[allow(clippy::too_many_arguments)]
6571 pub fn moe_gate_up_silu8_dev_q8_csr(
6572 &self,
6573 table: &CudaSlice<u64>,
6574 sel: &CudaSlice<i32>,
6575 aq: &CudaSlice<i8>,
6576 ad: &CudaSlice<f32>,
6577 n_pairs: usize,
6578 in_f: usize,
6579 n_ff: usize,
6580 n_used: usize,
6581 n_expert: usize,
6582 qt_g: i32,
6583 qt_u: i32,
6584 rb_g: usize,
6585 rb_u: usize,
6586 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6587 let f = if qt_g == crate::QT_NVFP4 {
6590 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6591 } else {
6592 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6593 };
6594 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6595 let cfg = LaunchConfig {
6596 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6597 block_dim: (32, 1, 1),
6598 shared_mem_bytes: 0,
6599 };
6600 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6601 in_f as i32,
6602 n_ff as i32,
6603 n_expert as i32,
6604 n_used as i32,
6605 n_pairs as i32,
6606 rb_g as i64,
6607 rb_u as i64,
6608 );
6609 let __s_b = self.gpu.stream();
6610 let mut b = __s_b.launch_builder(&f);
6611 b.arg(table)
6612 .arg(sel)
6613 .arg(aq)
6614 .arg(ad)
6615 .arg(&mut act)
6616 .arg(&inf)
6617 .arg(&nff)
6618 .arg(&ne)
6619 .arg(&qt_g)
6620 .arg(&qt_u)
6621 .arg(&rbg)
6622 .arg(&rbu)
6623 .arg(&nu)
6624 .arg(&npi);
6625 unsafe {
6626 b.launch(cfg)?;
6627 }
6628 Ok(act)
6629 }
6630
6631 #[allow(clippy::too_many_arguments)]
6635 pub fn moe_down8_fma_dev_q8_variant(
6636 &self,
6637 variant: &str,
6638 table: &CudaSlice<u64>,
6639 sel: &cudarc::driver::CudaView<i32>,
6640 w: &cudarc::driver::CudaView<f32>,
6641 aq2: &CudaSlice<i8>,
6642 ad2: &CudaSlice<f32>,
6643 dst: &mut cudarc::driver::CudaViewMut<f32>,
6644 in_f: usize,
6645 out_f: usize,
6646 n_used: usize,
6647 n_expert: usize,
6648 qt: i32,
6649 rb: usize,
6650 ) -> Result<(), Box<dyn std::error::Error>> {
6651 let (inf, outf, nu, ne, rbi) = (
6652 in_f as i32,
6653 out_f as i32,
6654 n_used as i32,
6655 n_expert as i32,
6656 rb as i64,
6657 );
6658 let (f, cfg) = match variant {
6659 "w8h2" | "w8h2v" => (
6660 self.func(if variant == "w8h2" {
6661 "moe_down8_fma_dev_q8_w8h2"
6662 } else {
6663 "moe_down8_fma_dev_q8_w8h2v"
6664 }),
6665 LaunchConfig {
6666 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6667 block_dim: (32, n_used as u32, 1),
6668 shared_mem_bytes: 0,
6669 },
6670 ),
6671 "w8h2r2" | "w8h2r2v" => (
6672 self.func(if variant == "w8h2r2" {
6673 "moe_down8_fma_dev_q8_w8h2r2"
6674 } else {
6675 "moe_down8_fma_dev_q8_w8h2r2v"
6676 }),
6677 LaunchConfig {
6678 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6679 block_dim: (32, n_used as u32, 1),
6680 shared_mem_bytes: 0,
6681 },
6682 ),
6683 _ => (
6684 self.func("moe_down8_fma_dev_q8"),
6685 LaunchConfig {
6686 grid_dim: (out_f as u32, 1, 1),
6687 block_dim: (32, 1, 1),
6688 shared_mem_bytes: 0,
6689 },
6690 ),
6691 };
6692 let __s_b = self.gpu.stream();
6693 let mut b = __s_b.launch_builder(&f);
6694 b.arg(table)
6695 .arg(sel)
6696 .arg(w)
6697 .arg(aq2)
6698 .arg(ad2)
6699 .arg(dst)
6700 .arg(&inf)
6701 .arg(&outf)
6702 .arg(&nu)
6703 .arg(&ne)
6704 .arg(&qt)
6705 .arg(&rbi);
6706 unsafe {
6707 b.launch(cfg)?;
6708 }
6709 Ok(())
6710 }
6711
6712 #[allow(clippy::too_many_arguments)]
6714 pub fn moe_gate_up_silu8_dev_q8_variant(
6715 &self,
6716 variant: &str,
6717 table: &CudaSlice<u64>,
6718 sel: &cudarc::driver::CudaView<i32>,
6719 aq: &CudaSlice<i8>,
6720 ad: &CudaSlice<f32>,
6721 in_f: usize,
6722 n_ff: usize,
6723 n_used: usize,
6724 n_expert: usize,
6725 qt_g: i32,
6726 qt_u: i32,
6727 rb_g: usize,
6728 rb_u: usize,
6729 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6730 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6731 let (inf, nff, ne, rbg, rbu) = (
6732 in_f as i32,
6733 n_ff as i32,
6734 n_expert as i32,
6735 rb_g as i64,
6736 rb_u as i64,
6737 );
6738 let f = self.func(if variant == "v" {
6739 "moe_gate_up_silu8_dev_q8_v"
6740 } else {
6741 "moe_gate_up_silu8_dev_q8"
6742 });
6743 let cfg = LaunchConfig {
6744 grid_dim: (n_ff as u32, n_used as u32, 1),
6745 block_dim: (32, 1, 1),
6746 shared_mem_bytes: 0,
6747 };
6748 let __s_b = self.gpu.stream();
6749 let mut b = __s_b.launch_builder(&f);
6750 b.arg(table)
6751 .arg(sel)
6752 .arg(aq)
6753 .arg(ad)
6754 .arg(&mut act)
6755 .arg(&inf)
6756 .arg(&nff)
6757 .arg(&ne)
6758 .arg(&qt_g)
6759 .arg(&qt_u)
6760 .arg(&rbg)
6761 .arg(&rbu);
6762 unsafe {
6763 b.launch(cfg)?;
6764 }
6765 Ok(act)
6766 }
6767
6768 pub fn moe_gate_up_silu8_dev(
6769 &self,
6770 table: &CudaSlice<u64>,
6771 sel: &cudarc::driver::CudaView<i32>,
6772 x: &cudarc::driver::CudaView<f32>,
6773 in_f: usize,
6774 n_ff: usize,
6775 n_used: usize,
6776 n_expert: usize,
6777 qt_g: i32,
6778 qt_u: i32,
6779 rb_g: usize,
6780 rb_u: usize,
6781 macros: &CudaSlice<f32>,
6782 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6783 let f = self.func("moe_gate_up_silu8_dev");
6784 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6786 grid_dim: (n_ff as u32, n_used as u32, 1),
6787 block_dim: (256, 1, 1),
6788 shared_mem_bytes: 0,
6789 };
6790 let (inf, nff, ne, rbg, rbu) = (
6791 in_f as i32,
6792 n_ff as i32,
6793 n_expert as i32,
6794 rb_g as i64,
6795 rb_u as i64,
6796 );
6797 let __s_b = self.gpu.stream();
6798 let mut b = __s_b.launch_builder(&f);
6799 b.arg(table)
6800 .arg(sel)
6801 .arg(x)
6802 .arg(&mut act)
6803 .arg(&inf)
6804 .arg(&nff)
6805 .arg(&ne)
6806 .arg(&qt_g)
6807 .arg(&qt_u)
6808 .arg(&rbg)
6809 .arg(&rbu)
6810 .arg(macros);
6811 unsafe {
6812 b.launch(cfg)?;
6813 }
6814 Ok(act)
6815 }
6816
6817 #[allow(clippy::too_many_arguments)]
6820 pub fn moe_down8_fma_dev(
6821 &self,
6822 table: &CudaSlice<u64>,
6823 sel: &cudarc::driver::CudaView<i32>,
6824 w: &cudarc::driver::CudaView<f32>,
6825 act: &CudaSlice<f32>,
6826 dst: &mut cudarc::driver::CudaViewMut<f32>,
6827 in_f: usize,
6828 out_f: usize,
6829 n_used: usize,
6830 n_expert: usize,
6831 qt: i32,
6832 rb: usize,
6833 ) -> Result<(), Box<dyn std::error::Error>> {
6834 let f = self.func("moe_down8_fma_dev");
6835 let cfg = LaunchConfig {
6836 grid_dim: (out_f as u32, 1, 1),
6837 block_dim: (256, 1, 1),
6838 shared_mem_bytes: 0,
6839 };
6840 let (inf, outf, nu, ne, rbv) = (
6841 in_f as i32,
6842 out_f as i32,
6843 n_used as i32,
6844 n_expert as i32,
6845 rb as i64,
6846 );
6847 let __s_b = self.gpu.stream();
6848 let mut b = __s_b.launch_builder(&f);
6849 b.arg(table)
6850 .arg(sel)
6851 .arg(w)
6852 .arg(act)
6853 .arg(dst)
6854 .arg(&inf)
6855 .arg(&outf)
6856 .arg(&nu)
6857 .arg(&ne)
6858 .arg(&qt)
6859 .arg(&rbv);
6860 unsafe {
6861 b.launch(cfg)?;
6862 }
6863 Ok(())
6864 }
6865
6866 pub fn axpy_into(
6868 &self,
6869 src: &CudaSlice<f32>,
6870 alpha: f32,
6871 dst: &mut cudarc::driver::CudaViewMut<f32>,
6872 n: usize,
6873 ) -> Result<(), Box<dyn std::error::Error>> {
6874 let f = self.func("axpy_f32");
6875 let cfg = LaunchConfig::for_num_elems(n as u32);
6876 let (a, ni) = (alpha, n as i32);
6877 let __s_b = self.gpu.stream();
6878 let mut b = __s_b.launch_builder(&f);
6879 b.arg(src).arg(dst).arg(&a).arg(&ni);
6880 unsafe {
6881 b.launch(cfg)?;
6882 }
6883 Ok(())
6884 }
6885
6886 pub fn axpy_host_into(
6888 &self,
6889 src: &cudarc::driver::CudaView<'_, f32>,
6890 alpha: f32,
6891 dst: &mut cudarc::driver::CudaViewMut<f32>,
6892 n: usize,
6893 ) -> Result<(), Box<dyn std::error::Error>> {
6894 let f = self.func("axpy_host_f32");
6895 let cfg = LaunchConfig::for_num_elems(n as u32);
6896 let (a, ni) = (alpha, n as i32);
6897 let __s_b = self.gpu.stream();
6898 let mut b = __s_b.launch_builder(&f);
6899 b.arg(src).arg(dst).arg(&a).arg(&ni);
6900 unsafe {
6901 b.launch(cfg)?;
6902 }
6903 Ok(())
6904 }
6905
6906 pub fn add_scaled_rows(
6908 &self,
6909 src: &CudaSlice<f32>,
6910 scale: &CudaSlice<f32>,
6911 dst: &mut CudaSlice<f32>,
6912 ncols: usize,
6913 nrows: usize,
6914 ) -> Result<(), Box<dyn std::error::Error>> {
6915 let f = self.func("add_scaled_rows_f32");
6916 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6917 let (nc, nr) = (ncols as i32, nrows as i32);
6918 let __s_b = self.gpu.stream();
6919 let mut b = __s_b.launch_builder(&f);
6920 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6921 unsafe {
6922 b.launch(cfg)?;
6923 }
6924 Ok(())
6925 }
6926
6927 pub fn gather_rows(
6931 &self,
6932 src: &CudaSlice<f32>,
6933 idx: &CudaSlice<i32>,
6934 dst: &mut CudaSlice<f32>,
6935 ncols: usize,
6936 m_e: usize,
6937 ) -> Result<(), Box<dyn std::error::Error>> {
6938 let f = self.func("gather_rows_f32");
6939 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6940 let (nc, me) = (ncols as i32, m_e as i32);
6941 let __s_b = self.gpu.stream();
6942 let mut b = __s_b.launch_builder(&f);
6943 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6944 unsafe {
6945 b.launch(cfg)?;
6946 }
6947 Ok(())
6948 }
6949
6950 pub fn scatter_slot(
6955 &self,
6956 src: &CudaSlice<f32>,
6957 tok_idx: &CudaSlice<i32>,
6958 slot_idx: &CudaSlice<i32>,
6959 weight: &CudaSlice<f32>,
6960 dst: &mut CudaSlice<f32>,
6961 wbuf: &mut CudaSlice<f32>,
6962 ncols: usize,
6963 n_used: usize,
6964 m_e: usize,
6965 ) -> Result<(), Box<dyn std::error::Error>> {
6966 let f = self.func("scatter_add_slot_f32");
6967 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6968 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6969 let __s_b = self.gpu.stream();
6970 let mut b = __s_b.launch_builder(&f);
6971 b.arg(src)
6972 .arg(tok_idx)
6973 .arg(slot_idx)
6974 .arg(weight)
6975 .arg(dst)
6976 .arg(wbuf)
6977 .arg(&nc)
6978 .arg(&nu)
6979 .arg(&me);
6980 unsafe {
6981 b.launch(cfg)?;
6982 }
6983 Ok(())
6984 }
6985
6986 pub fn reduce_slots(
6990 &self,
6991 slots: &CudaSlice<f32>,
6992 wbuf: &CudaSlice<f32>,
6993 dst: &mut CudaSlice<f32>,
6994 ncols: usize,
6995 n_used: usize,
6996 t: usize,
6997 ) -> Result<(), Box<dyn std::error::Error>> {
6998 let f = self.func("reduce_slots_f32");
6999 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7000 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7001 let __s_b = self.gpu.stream();
7002 let mut b = __s_b.launch_builder(&f);
7003 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7004 unsafe {
7005 b.launch(cfg)?;
7006 }
7007 Ok(())
7008 }
7009
7010 pub fn reduce_slots_host(
7015 &self,
7016 slots: &CudaSlice<f32>,
7017 wbuf: &CudaSlice<f32>,
7018 dst: &mut CudaSlice<f32>,
7019 ncols: usize,
7020 n_used: usize,
7021 t: usize,
7022 ) -> Result<(), Box<dyn std::error::Error>> {
7023 let f = self.func("reduce_slots_host_f32");
7024 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7025 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7026 let __s_b = self.gpu.stream();
7027 let mut b = __s_b.launch_builder(&f);
7028 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7029 unsafe {
7030 b.launch(cfg)?;
7031 }
7032 Ok(())
7033 }
7034
7035 pub fn quantize_q8_1_view(
7042 &self,
7043 x: &cudarc::driver::CudaView<f32>,
7044 m: usize,
7045 in_f: usize,
7046 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7047 let f = self.func("quantize_q8_1");
7048 let nblk = in_f / 32;
7049 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7050 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7051 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7052 let (inf, mi) = (in_f as i32, m as i32);
7053 let __s_b = self.gpu.stream();
7054 let mut b = __s_b.launch_builder(&f);
7055 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7056 unsafe {
7057 b.launch(cfg)?;
7058 }
7059 Ok((q, d))
7060 }
7061
7062 pub fn quantize_q8_1(
7063 &self,
7064 x: &CudaSlice<f32>,
7065 m: usize,
7066 in_f: usize,
7067 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7068 let nblk = in_f / 32;
7069 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);
7073 let (inf, mi) = (in_f as i32, m as i32);
7074 if Self::pdl_on() && Self::pdl_wb_on() {
7075 {
7076 use cudarc::driver::{DevicePtr, DevicePtrMut};
7077 let s = &self.gpu.stream();
7078 let (px, _g0) = x.device_ptr(s);
7079 let (pq, _g1) = q.device_ptr_mut(s);
7080 let (pd, _g2) = d.device_ptr_mut(s);
7081 let mut ps = [
7082 &px as *const _ as *mut std::ffi::c_void,
7083 &pq as *const _ as *mut _,
7084 &pd as *const _ as *mut _,
7085 &inf as *const _ as *mut _,
7086 &mi as *const _ as *mut _,
7087 ];
7088 unsafe {
7089 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7090 }
7091 }
7092 return Ok((q, d));
7093 }
7094 let f = self.func("quantize_q8_1");
7095 let __s_b = self.gpu.stream();
7096 let mut b = __s_b.launch_builder(&f);
7097 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7098 unsafe {
7099 b.launch(cfg)?;
7100 }
7101 Ok((q, d))
7102 }
7103
7104 pub fn quantize_fp4_act(
7108 &self,
7109 x: &CudaSlice<f32>,
7110 m: usize,
7111 in_f: usize,
7112 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7113 let f = self.func("quantize_fp4_act");
7114 let nb16 = in_f / 16;
7115 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);
7118 let (inf, mi) = (in_f as i32, m as i32);
7119 let __s_b = self.gpu.stream();
7120 let mut b = __s_b.launch_builder(&f);
7121 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7122 unsafe {
7123 b.launch(cfg)?;
7124 }
7125 Ok((aq4, ad4))
7126 }
7127
7128 pub fn qmatvec_gemm_nvfp4_fp4(
7133 &self,
7134 bytes: &CudaSlice<u8>,
7135 x: &CudaSlice<f32>,
7136 m: usize,
7137 in_f: usize,
7138 out_f: usize,
7139 row_bytes: usize,
7140 scale: f32,
7141 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7142 assert!(
7143 in_f % 64 == 0,
7144 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7145 );
7146 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7147 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7148 if scale != 1.0 {
7149 self.scale_inplace(&mut y, scale, m * out_f)?;
7150 }
7151 Ok(y)
7152 }
7153
7154 fn fp4_gemm_launch(
7157 &self,
7158 bytes: &CudaSlice<u8>,
7159 aq4: &CudaSlice<u32>,
7160 ad4: &CudaSlice<u8>,
7161 m: usize,
7162 in_f: usize,
7163 out_f: usize,
7164 row_bytes: usize,
7165 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7166 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7167 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7169 const BN: u32 = 256;
7170 let cfg = LaunchConfig {
7171 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7172 block_dim: (32, 4, 1),
7173 shared_mem_bytes: 0,
7174 };
7175 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7176 let __s_b = self.gpu.stream();
7177 let mut b = __s_b.launch_builder(&f);
7178 b.arg(bytes)
7179 .arg(aq4)
7180 .arg(ad4)
7181 .arg(&mut y)
7182 .arg(&inf)
7183 .arg(&outf)
7184 .arg(&mi)
7185 .arg(&rb);
7186 unsafe {
7187 b.launch(cfg)?;
7188 }
7189 Ok(y)
7190 }
7191
7192 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7194 &self,
7195 bytes: &CudaSlice<u8>,
7196 x: &CudaSlice<f32>,
7197 m: usize,
7198 in_f: usize,
7199 out_f: usize,
7200 row_bytes: usize,
7201 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7202 assert!(
7203 in_f % 64 == 0,
7204 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7205 );
7206 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7207 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7208 }
7209
7210 pub fn qmatvec_q8_0_fast(
7212 &self,
7213 w: &CudaSlice<u8>,
7214 x: &CudaSlice<f32>,
7215 m: usize,
7216 in_f: usize,
7217 out_f: usize,
7218 row_bytes: usize,
7219 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7220 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7221 let f = self.func("qmatvec_q8_0_dp4a");
7222 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7224 grid_dim: (out_f as u32, m as u32, 1),
7225 block_dim: (128, 1, 1),
7226 shared_mem_bytes: 0,
7227 };
7228 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7229 let __s_b = self.gpu.stream();
7230 let mut b = __s_b.launch_builder(&f);
7231 b.arg(w)
7232 .arg(&aq)
7233 .arg(&ad)
7234 .arg(&mut y)
7235 .arg(&inf)
7236 .arg(&outf)
7237 .arg(&mi)
7238 .arg(&rb);
7239 unsafe {
7240 b.launch(cfg)?;
7241 }
7242 Ok(y)
7243 }
7244
7245 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7248 &self,
7249 w: &CudaSlice<u8>,
7250 x: &CudaSlice<f32>,
7251 m: usize,
7252 in_f: usize,
7253 out_f: usize,
7254 row_bytes: usize,
7255 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7256 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7257 let f = self.func("qmatvec_q4_K_dp4a");
7258 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7260 grid_dim: (out_f as u32, m as u32, 1),
7261 block_dim: (128, 1, 1),
7262 shared_mem_bytes: 0,
7263 };
7264 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7265 let __s_b = self.gpu.stream();
7266 let mut b = __s_b.launch_builder(&f);
7267 b.arg(w)
7268 .arg(&aq)
7269 .arg(&ad)
7270 .arg(&mut y)
7271 .arg(&inf)
7272 .arg(&outf)
7273 .arg(&mi)
7274 .arg(&rb);
7275 unsafe {
7276 b.launch(cfg)?;
7277 }
7278 Ok(y)
7279 }
7280
7281 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7284 &self,
7285 w: &CudaSlice<u8>,
7286 x: &CudaSlice<f32>,
7287 m: usize,
7288 in_f: usize,
7289 out_f: usize,
7290 row_bytes: usize,
7291 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7292 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7293 let f = self.func("qmatvec_q6_K_dp4a");
7294 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7296 grid_dim: (out_f as u32, m as u32, 1),
7297 block_dim: (128, 1, 1),
7298 shared_mem_bytes: 0,
7299 };
7300 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7301 let __s_b = self.gpu.stream();
7302 let mut b = __s_b.launch_builder(&f);
7303 b.arg(w)
7304 .arg(&aq)
7305 .arg(&ad)
7306 .arg(&mut y)
7307 .arg(&inf)
7308 .arg(&outf)
7309 .arg(&mi)
7310 .arg(&rb);
7311 unsafe {
7312 b.launch(cfg)?;
7313 }
7314 Ok(y)
7315 }
7316
7317 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7320 &self,
7321 w: &CudaSlice<u8>,
7322 x: &CudaSlice<f32>,
7323 m: usize,
7324 in_f: usize,
7325 out_f: usize,
7326 row_bytes: usize,
7327 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7328 self.qmatvec_dp4a_named(
7329 "qmatvec_q5_K_dp4a",
7330 &w.slice(0..w.len()),
7331 x,
7332 m,
7333 in_f,
7334 out_f,
7335 row_bytes,
7336 )
7337 }
7338 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7341 &self,
7342 w: &CudaSlice<u8>,
7343 x: &CudaSlice<f32>,
7344 m: usize,
7345 in_f: usize,
7346 out_f: usize,
7347 row_bytes: usize,
7348 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7349 self.qmatvec_dp4a_named(
7350 "qmatvec_q3_K_dp4a",
7351 &w.slice(0..w.len()),
7352 x,
7353 m,
7354 in_f,
7355 out_f,
7356 row_bytes,
7357 )
7358 }
7359 pub fn qmatvec_nvfp4_fast_rp(
7361 &self,
7362 w: &CudaSlice<u8>,
7363 x: &CudaSlice<f32>,
7364 m: usize,
7365 in_f: usize,
7366 out_f: usize,
7367 row_bytes: usize,
7368 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7369 assert!(
7370 in_f % 64 == 0,
7371 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7372 );
7373 self.qmatvec_dp4a_named(
7374 "qmatvec_nvfp4_dp4a_rp",
7375 &w.slice(0..w.len()),
7376 x,
7377 m,
7378 in_f,
7379 out_f,
7380 row_bytes,
7381 )
7382 }
7383 pub fn qmatvec_nvfp4_fast(
7385 &self,
7386 w: &cudarc::driver::CudaView<'_, u8>,
7387 x: &CudaSlice<f32>,
7388 m: usize,
7389 in_f: usize,
7390 out_f: usize,
7391 row_bytes: usize,
7392 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7393 assert!(
7396 in_f % 64 == 0,
7397 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7398 );
7399 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7400 }
7401 pub fn qmatvec_nvfp4_fast_v2(
7404 &self,
7405 w: &cudarc::driver::CudaView<'_, u8>,
7406 x: &CudaSlice<f32>,
7407 m: usize,
7408 in_f: usize,
7409 out_f: usize,
7410 row_bytes: usize,
7411 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7412 assert!(
7413 in_f % 64 == 0,
7414 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7415 );
7416 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7417 }
7418 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7421 &self,
7422 w: &CudaSlice<u8>,
7423 x: &CudaSlice<f32>,
7424 m: usize,
7425 in_f: usize,
7426 out_f: usize,
7427 row_bytes: usize,
7428 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7429 self.qmatvec_dp4a_named(
7430 "qmatvec_iq4_XS_dp4a",
7431 &w.slice(0..w.len()),
7432 x,
7433 m,
7434 in_f,
7435 out_f,
7436 row_bytes,
7437 )
7438 }
7439
7440 fn qmatvec_dp4a_named(
7442 &self,
7443 name: &str,
7444 w: &cudarc::driver::CudaView<'_, u8>,
7445 x: &CudaSlice<f32>,
7446 m: usize,
7447 in_f: usize,
7448 out_f: usize,
7449 row_bytes: usize,
7450 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7451 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7452 let f = self.func(name);
7453 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7455 grid_dim: (out_f as u32, m as u32, 1),
7456 block_dim: (128, 1, 1),
7457 shared_mem_bytes: 0,
7458 };
7459 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7460 let __s_b = self.gpu.stream();
7461 let mut b = __s_b.launch_builder(&f);
7462 b.arg(w)
7463 .arg(&aq)
7464 .arg(&ad)
7465 .arg(&mut y)
7466 .arg(&inf)
7467 .arg(&outf)
7468 .arg(&mi)
7469 .arg(&rb);
7470 unsafe {
7471 b.launch(cfg)?;
7472 }
7473 Ok(y)
7474 }
7475
7476 #[allow(clippy::too_many_arguments)]
7482 pub fn qmatvec_nvfp4_fast_prequant_into(
7483 &self,
7484 w: &CudaSlice<u8>,
7485 aq: &CudaSlice<i8>,
7486 ad: &CudaSlice<f32>,
7487 y: &mut CudaSlice<f32>,
7488 m: usize,
7489 in_f: usize,
7490 out_f: usize,
7491 row_bytes: usize,
7492 ) -> Result<(), Box<dyn std::error::Error>> {
7493 assert!(
7494 in_f % 64 == 0,
7495 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7496 );
7497 if y.len() < m * out_f {
7498 return Err(format!(
7499 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7500 y.len()
7501 )
7502 .into());
7503 }
7504 let f = self.func("qmatvec_nvfp4_dp4a");
7505 let cfg = LaunchConfig {
7506 grid_dim: (out_f as u32, m as u32, 1),
7507 block_dim: (128, 1, 1),
7508 shared_mem_bytes: 0,
7509 };
7510 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7511 let __s_b = self.gpu.stream();
7512 let mut b = __s_b.launch_builder(&f);
7513 b.arg(w)
7514 .arg(aq)
7515 .arg(ad)
7516 .arg(y)
7517 .arg(&inf)
7518 .arg(&outf)
7519 .arg(&mi)
7520 .arg(&rb);
7521 unsafe {
7522 b.launch(cfg)?;
7523 }
7524 Ok(())
7525 }
7526
7527 #[allow(clippy::too_many_arguments)]
7530 pub fn matvec_f32_qkv_into(
7531 &self,
7532 wq: &CudaSlice<f32>,
7533 wk: &CudaSlice<f32>,
7534 wv: &CudaSlice<f32>,
7535 wg: &CudaSlice<f32>,
7536 x: &CudaSlice<f32>,
7537 yq: &mut CudaSlice<f32>,
7538 yk: &mut CudaSlice<f32>,
7539 yv: &mut CudaSlice<f32>,
7540 yg: &mut CudaSlice<f32>,
7541 in_f: usize,
7542 out_q: usize,
7543 out_kv: usize,
7544 out_g: usize,
7545 ) -> Result<(), Box<dyn std::error::Error>> {
7546 if in_f % 4 != 0
7547 || wq.len() != out_q * in_f
7548 || wk.len() != out_kv * in_f
7549 || wv.len() != out_kv * in_f
7550 || wg.len() < out_g * in_f
7551 || x.len() < in_f
7552 || yq.len() < out_q
7553 || yk.len() < out_kv
7554 || yv.len() < out_kv
7555 || (out_g > 0 && yg.len() < out_g)
7556 {
7557 return Err(format!(
7558 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7559 wq={} wk={} wv={} wg={}",
7560 wq.len(),
7561 wk.len(),
7562 wv.len(),
7563 wg.len()
7564 )
7565 .into());
7566 }
7567 let f = self.func("matvec_f32_qkv");
7568 let cfg = LaunchConfig {
7569 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7570 block_dim: (128, 1, 1),
7571 shared_mem_bytes: 0,
7572 };
7573 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7574 let __s_b = self.gpu.stream();
7575 let mut b = __s_b.launch_builder(&f);
7576 b.arg(wq)
7577 .arg(wk)
7578 .arg(wv)
7579 .arg(wg)
7580 .arg(x)
7581 .arg(yq)
7582 .arg(yk)
7583 .arg(yv)
7584 .arg(yg)
7585 .arg(&inf)
7586 .arg(&oq)
7587 .arg(&okv)
7588 .arg(&og);
7589 unsafe {
7590 b.launch(cfg)?;
7591 }
7592 Ok(())
7593 }
7594
7595 #[allow(clippy::too_many_arguments)]
7601 #[allow(clippy::too_many_arguments)]
7604 pub fn qmatvec_nvfp4_sel_gu_into(
7605 &self,
7606 gate_bank: &CudaSlice<u8>,
7607 up_bank: &CudaSlice<u8>,
7608 sel: &CudaSlice<i32>,
7609 aq: &CudaSlice<i8>,
7610 ad: &CudaSlice<f32>,
7611 yg: &mut CudaSlice<f32>,
7612 yu: &mut CudaSlice<f32>,
7613 n_sel: usize,
7614 in_f: usize,
7615 out_f: usize,
7616 row_bytes: usize,
7617 expert_stride: usize,
7618 ) -> Result<(), Box<dyn std::error::Error>> {
7619 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7620 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7621 return Err("NVFP4 gu sel geometry".into());
7622 }
7623 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
7626 let rpw = *RPW.get_or_init(|| {
7627 std::env::var("MEMRA_SEL_GU_RPW")
7628 .ok()
7629 .and_then(|v| v.parse().ok())
7630 .filter(|r| *r == 2 || *r == 4)
7631 .unwrap_or(1)
7632 });
7633 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
7634 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7637 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
7638 let f = self.func(match (wpr, rpw) {
7639 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
7640 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
7641 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
7642 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
7643 });
7644 let cfg = LaunchConfig {
7645 grid_dim: if wpr {
7646 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
7647 } else if rpw == 1 {
7648 ((2 * out_f) as u32, n_sel as u32, 1)
7649 } else {
7650 ((out_f / rpw) as u32, n_sel as u32, 1)
7651 },
7652 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
7653 shared_mem_bytes: 0,
7654 };
7655 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7656 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7657 let (ars, adrs) = (0i64, 0i64);
7658 let __s_b = self.gpu.stream();
7659 let mut b = __s_b.launch_builder(&f);
7660 b.arg(gate_bank)
7661 .arg(up_bank)
7662 .arg(sel)
7663 .arg(aq)
7664 .arg(ad)
7665 .arg(yg)
7666 .arg(yu)
7667 .arg(&inf)
7668 .arg(&outf)
7669 .arg(&ns)
7670 .arg(&rb)
7671 .arg(&es)
7672 .arg(&ars)
7673 .arg(&adrs);
7674 unsafe {
7675 b.launch(cfg)?;
7676 }
7677 Ok(())
7678 }
7679
7680 #[allow(clippy::too_many_arguments)]
7688 pub fn qmatvec_nvfp4_sel_down8_into(
7689 &self,
7690 bank: &CudaSlice<u8>,
7691 sel: &CudaSlice<i32>,
7692 aq: &CudaSlice<i8>,
7693 ad: &CudaSlice<f32>,
7694 route_w: &CudaSlice<f32>,
7695 md: &CudaSlice<f32>,
7696 dst: &mut CudaSlice<f32>,
7697 n_sel: usize,
7698 in_f: usize,
7699 out_f: usize,
7700 row_bytes: usize,
7701 expert_stride: usize,
7702 act_row_stride: usize,
7703 ad_row_stride: usize,
7704 ) -> Result<(), Box<dyn std::error::Error>> {
7705 if in_f % 64 != 0
7706 || n_sel == 0
7707 || n_sel > 8
7708 || (in_f >> 5) > 32
7709 || dst.len() < out_f
7710 || sel.len() < n_sel
7711 || route_w.len() < n_sel
7712 {
7713 return Err(format!(
7714 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
7715 dst.len()
7716 )
7717 .into());
7718 }
7719 if !crate::tp::nvfp4_bank_v2_on() {
7720 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
7721 }
7722 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
7723 let cfg = LaunchConfig {
7724 grid_dim: (out_f as u32, 1, 1),
7725 block_dim: (32, n_sel as u32, 1),
7726 shared_mem_bytes: 0,
7727 };
7728 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7729 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7730 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7731 let __s_b = self.gpu.stream();
7732 let mut b = __s_b.launch_builder(&f);
7733 b.arg(bank)
7734 .arg(sel)
7735 .arg(aq)
7736 .arg(ad)
7737 .arg(route_w)
7738 .arg(md)
7739 .arg(dst)
7740 .arg(&inf)
7741 .arg(&outf)
7742 .arg(&ns)
7743 .arg(&rb)
7744 .arg(&es)
7745 .arg(&ars)
7746 .arg(&adrs);
7747 unsafe {
7748 b.launch(cfg)?;
7749 }
7750 Ok(())
7751 }
7752
7753 pub fn qmatvec_nvfp4_sel_into(
7754 &self,
7755 bank: &CudaSlice<u8>,
7756 sel: &CudaSlice<i32>,
7757 aq: &CudaSlice<i8>,
7758 ad: &CudaSlice<f32>,
7759 y: &mut CudaSlice<f32>,
7760 n_sel: usize,
7761 in_f: usize,
7762 out_f: usize,
7763 row_bytes: usize,
7764 expert_stride: usize,
7765 act_row_stride: usize,
7766 ad_row_stride: usize,
7767 ) -> Result<(), Box<dyn std::error::Error>> {
7768 assert!(
7769 in_f % 64 == 0,
7770 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7771 );
7772 if y.len() < n_sel * out_f || sel.len() < n_sel {
7773 return Err(format!(
7774 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7775 y.len(),
7776 sel.len()
7777 )
7778 .into());
7779 }
7780 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7787 let mode = *MR.get_or_init(|| {
7788 if crate::tp::nvfp4_bank_v2_on() {
7789 3
7790 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7791 2
7792 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7793 1
7794 } else {
7795 0
7796 }
7797 });
7798 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7799 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7803 let v2s = mode == 3
7804 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7805 && row_bytes % 16 == 0
7806 && in_f <= 4096;
7807 let f = match (mode, v2s) {
7808 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7809 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7810 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7811 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7812 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7813 };
7814 let nsb = in_f >> 5;
7819 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7820 32
7821 } else if mode == 1 {
7822 512
7823 } else {
7824 128
7825 };
7826 let cfg = LaunchConfig {
7827 grid_dim: (
7828 if v2s {
7829 (out_f as u32).div_ceil(8)
7830 } else {
7831 match mode {
7832 2 => (out_f as u32).div_ceil(16),
7833 1 => (out_f as u32).div_ceil(4),
7834 _ => out_f as u32,
7835 }
7836 },
7837 n_sel as u32,
7838 1,
7839 ),
7840 block_dim: (fit_block, 1, 1),
7841 shared_mem_bytes: 0,
7842 };
7843 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7844 let (rb, es, ars, adrs) = (
7845 row_bytes as i64,
7846 expert_stride as i64,
7847 act_row_stride as i64,
7848 ad_row_stride as i64,
7849 );
7850 let __s_b = self.gpu.stream();
7851 let mut b = __s_b.launch_builder(&f);
7852 b.arg(bank)
7853 .arg(sel)
7854 .arg(aq)
7855 .arg(ad)
7856 .arg(y)
7857 .arg(&inf)
7858 .arg(&outf)
7859 .arg(&ns)
7860 .arg(&rb)
7861 .arg(&es)
7862 .arg(&ars)
7863 .arg(&adrs);
7864 unsafe {
7865 b.launch(cfg)?;
7866 }
7867 Ok(())
7868 }
7869
7870 #[allow(clippy::too_many_arguments)]
7875 pub fn silu_mul_scaled_q8_1_sel_into(
7876 &self,
7877 gate: &CudaSlice<f32>,
7878 up: &CudaSlice<f32>,
7879 gmac: &CudaSlice<f32>,
7880 umac: &CudaSlice<f32>,
7881 sel: &CudaSlice<i32>,
7882 limit: Option<f32>,
7883 out_q: &mut CudaSlice<i8>,
7884 out_d: &mut CudaSlice<f32>,
7885 n_per: usize,
7886 n_sel: usize,
7887 ) -> Result<(), Box<dyn std::error::Error>> {
7888 let n = n_per * n_sel;
7889 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
7890 return Err(format!(
7891 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
7892 out_q.len(),
7893 out_d.len()
7894 )
7895 .into());
7896 }
7897 if let Some(limit) = limit {
7898 if limit <= 1e-6 {
7899 return Err(format!(
7900 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
7901 )
7902 .into());
7903 }
7904 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
7905 let cfg = LaunchConfig::for_num_elems(n as u32);
7906 let (np, ns) = (n_per as i32, n_sel as i32);
7907 let __s_b = self.gpu.stream();
7908 let mut b = __s_b.launch_builder(&f);
7909 b.arg(gate)
7910 .arg(up)
7911 .arg(gmac)
7912 .arg(umac)
7913 .arg(sel)
7914 .arg(&limit)
7915 .arg(out_q)
7916 .arg(out_d)
7917 .arg(&np)
7918 .arg(&ns);
7919 unsafe {
7920 b.launch(cfg)?;
7921 }
7922 return Ok(());
7923 }
7924 let f = self.func("silu_mul_scaled_q8_1_sel");
7925 let cfg = LaunchConfig::for_num_elems(n as u32);
7926 let (np, ns) = (n_per as i32, n_sel as i32);
7927 let __s_b = self.gpu.stream();
7928 let mut b = __s_b.launch_builder(&f);
7929 b.arg(gate)
7930 .arg(up)
7931 .arg(gmac)
7932 .arg(umac)
7933 .arg(sel)
7934 .arg(out_q)
7935 .arg(out_d)
7936 .arg(&np)
7937 .arg(&ns);
7938 unsafe {
7939 b.launch(cfg)?;
7940 }
7941 Ok(())
7942 }
7943
7944 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7945 Ok(self.gpu.stream().clone_htod(v)?)
7946 }
7947 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7948 Ok(self.gpu.stream().clone_htod(v)?)
7949 }
7950 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7952 Ok(self.gpu.stream().clone_htod(v)?)
7953 }
7954 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7955 Ok(self.gpu.stream().clone_htod(v)?)
7956 }
7957 pub fn dtoh_view(
7959 &self,
7960 d: &cudarc::driver::CudaView<f32>,
7961 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7962 let v = self.gpu.stream().clone_dtoh(d)?;
7963 self.gpu.stream().synchronize()?;
7964 Ok(v)
7965 }
7966 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7967 let v = self.gpu.stream().clone_dtoh(d)?;
7968 self.gpu.stream().synchronize()?;
7969 Ok(v)
7970 }
7971 pub fn dtoh_pair(
7975 &self,
7976 a: &CudaSlice<f32>,
7977 b: &CudaSlice<f32>,
7978 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7979 let av = self.gpu.stream().clone_dtoh(a)?;
7980 let bv = self.gpu.stream().clone_dtoh(b)?;
7981 self.gpu.stream().synchronize()?;
7982 Ok((av, bv))
7983 }
7984 pub fn dtoh_pair_views(
7987 &self,
7988 a: &cudarc::driver::CudaView<f32>,
7989 b: &cudarc::driver::CudaView<f32>,
7990 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7991 let av = self.gpu.stream().clone_dtoh(a)?;
7992 let bv = self.gpu.stream().clone_dtoh(b)?;
7993 self.gpu.stream().synchronize()?;
7994 Ok((av, bv))
7995 }
7996 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7998 let v = self.gpu.stream().clone_dtoh(d)?;
7999 self.gpu.stream().synchronize()?;
8000 Ok(v)
8001 }
8002 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8004 let v = self.gpu.stream().clone_dtoh(d)?;
8005 self.gpu.stream().synchronize()?;
8006 Ok(v)
8007 }
8008 pub fn dtoh_u8_view(
8009 &self,
8010 d: &cudarc::driver::CudaView<u8>,
8011 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8012 let v = self.gpu.stream().clone_dtoh(d)?;
8013 self.gpu.stream().synchronize()?;
8014 Ok(v)
8015 }
8016 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8017 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8018 self.keep_if_capturing(&s);
8019 Ok(s)
8020 }
8021
8022 pub fn prob_of_token_device(
8031 &self,
8032 logits: &CudaSlice<f32>,
8033 tok: &CudaSlice<u32>,
8034 n_vocab: usize,
8035 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8036 let nb = ARGMAX_NB;
8037 let mut part = self.alloc_uninit::<f32>(nb)?;
8038 let mut p = self.alloc_uninit::<f32>(1)?;
8039 let f1 = self.func("prob_of_token_partial_f32");
8040 let cfg1 = LaunchConfig {
8041 grid_dim: (nb as u32, 1, 1),
8042 block_dim: (256, 1, 1),
8043 shared_mem_bytes: 0,
8044 };
8045 let nv = n_vocab as i32;
8046 let __s_b1 = self.gpu.stream();
8047 let mut b1 = __s_b1.launch_builder(&f1);
8048 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8049 unsafe {
8050 b1.launch(cfg1)?;
8051 }
8052 let f2 = self.func("prob_of_token_final_f32");
8053 let cfg2 = LaunchConfig {
8054 grid_dim: (1, 1, 1),
8055 block_dim: (256, 1, 1),
8056 shared_mem_bytes: 0,
8057 };
8058 let nbi = nb as i32;
8059 let __s_b2 = self.gpu.stream();
8060 let mut b2 = __s_b2.launch_builder(&f2);
8061 b2.arg(&part).arg(&mut p).arg(&nbi);
8062 unsafe {
8063 b2.launch(cfg2)?;
8064 }
8065 Ok(p)
8066 }
8067
8068 pub fn prob_of_token_device_col(
8075 &self,
8076 logits: &CudaSlice<f32>,
8077 tok_all: &CudaSlice<u32>,
8078 tok_idx: usize,
8079 p_out: &mut CudaSlice<f32>,
8080 p_idx: usize,
8081 n_vocab: usize,
8082 ) -> Result<(), Box<dyn std::error::Error>> {
8083 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8084 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8085 let nb = ARGMAX_NB;
8086 let mut part = self.alloc_uninit::<f32>(nb)?;
8087 let f1 = self.func("prob_of_token_partial_f32");
8088 let cfg1 = LaunchConfig {
8089 grid_dim: (nb as u32, 1, 1),
8090 block_dim: (256, 1, 1),
8091 shared_mem_bytes: 0,
8092 };
8093 let nv = n_vocab as i32;
8094 let __s_b1 = self.gpu.stream();
8095 let mut b1 = __s_b1.launch_builder(&f1);
8096 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8097 unsafe {
8098 b1.launch(cfg1)?;
8099 }
8100 let f2 = self.func("prob_of_token_final_f32");
8101 let cfg2 = LaunchConfig {
8102 grid_dim: (1, 1, 1),
8103 block_dim: (256, 1, 1),
8104 shared_mem_bytes: 0,
8105 };
8106 let nbi = nb as i32;
8107 let __s_b2 = self.gpu.stream();
8108 let mut b2 = __s_b2.launch_builder(&f2);
8109 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8110 unsafe {
8111 b2.launch(cfg2)?;
8112 }
8113 Ok(())
8114 }
8115
8116 pub fn prob_of_token_device_into(
8117 &self,
8118 logits: &CudaSlice<f32>,
8119 tok: &CudaSlice<u32>,
8120 p_out: &mut CudaSlice<f32>,
8121 n_vocab: usize,
8122 ) -> Result<(), Box<dyn std::error::Error>> {
8123 let nb = ARGMAX_NB;
8124 let mut part = self.alloc_uninit::<f32>(nb)?;
8125 let f1 = self.func("prob_of_token_partial_f32");
8126 let cfg1 = LaunchConfig {
8127 grid_dim: (nb as u32, 1, 1),
8128 block_dim: (256, 1, 1),
8129 shared_mem_bytes: 0,
8130 };
8131 let nv = n_vocab as i32;
8132 let __s_b1 = self.gpu.stream();
8133 let mut b1 = __s_b1.launch_builder(&f1);
8134 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8135 unsafe {
8136 b1.launch(cfg1)?;
8137 }
8138 let f2 = self.func("prob_of_token_final_f32");
8139 let cfg2 = LaunchConfig {
8140 grid_dim: (1, 1, 1),
8141 block_dim: (256, 1, 1),
8142 shared_mem_bytes: 0,
8143 };
8144 let nbi = nb as i32;
8145 let __s_b2 = self.gpu.stream();
8146 let mut b2 = __s_b2.launch_builder(&f2);
8147 b2.arg(&part).arg(p_out).arg(&nbi);
8148 unsafe {
8149 b2.launch(cfg2)?;
8150 }
8151 Ok(())
8152 }
8153
8154 pub fn u32_hist_append(
8157 &self,
8158 tok: &CudaSlice<u32>,
8159 hist: &mut CudaSlice<u32>,
8160 idx: &mut CudaSlice<i32>,
8161 ) -> Result<(), Box<dyn std::error::Error>> {
8162 let f = self.func("u32_hist_append");
8163 let cfg = LaunchConfig {
8164 grid_dim: (1, 1, 1),
8165 block_dim: (32, 1, 1),
8166 shared_mem_bytes: 0,
8167 };
8168 let __s_b = self.gpu.stream();
8169 let mut b = __s_b.launch_builder(&f);
8170 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8171 unsafe {
8172 b.launch(cfg)?;
8173 }
8174 Ok(())
8175 }
8176
8177 pub fn argmax_token_device(
8178 &self,
8179 logits: &CudaSlice<f32>,
8180 n_vocab: usize,
8181 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8182 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8183 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8184 Ok(tok)
8185 }
8186 pub fn argmax_token_device_into(
8193 &self,
8194 logits: &CudaSlice<f32>,
8195 tok: &mut CudaSlice<u32>,
8196 n_vocab: usize,
8197 ) -> Result<(), Box<dyn std::error::Error>> {
8198 let nb = ARGMAX_NB;
8199 let f1 = self.func("argmax_partial_f32");
8200 let f2 = self.func("argmax_final_f32");
8201 let mut guard = self.argmax_partials.lock().unwrap();
8202 if guard.is_none() {
8203 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8206 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8207 *guard = Some((pv, pi));
8208 }
8209 let (part_v, part_i) = guard.as_mut().unwrap();
8210 let nv = n_vocab as i32;
8211 let nbi = nb as i32;
8212 let cfg1 = LaunchConfig {
8214 grid_dim: (nb as u32, 1, 1),
8215 block_dim: (256, 1, 1),
8216 shared_mem_bytes: 0,
8217 };
8218 let __s_b1 = self.gpu.stream();
8219 let mut b1 = __s_b1.launch_builder(&f1);
8220 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8221 unsafe {
8222 b1.launch(cfg1)?;
8223 }
8224 let cfg2 = LaunchConfig {
8226 grid_dim: (1, 1, 1),
8227 block_dim: (256, 1, 1),
8228 shared_mem_bytes: 0,
8229 };
8230 let __s_b2 = self.gpu.stream();
8231 let mut b2 = __s_b2.launch_builder(&f2);
8232 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8233 unsafe {
8234 b2.launch(cfg2)?;
8235 }
8236 Ok(())
8237 }
8238 pub fn argmax_token_device_col(
8244 &self,
8245 logits: &CudaSlice<f32>,
8246 col: usize,
8247 n_vocab: usize,
8248 toks: &mut CudaSlice<u32>,
8249 out_idx: usize,
8250 ) -> Result<(), Box<dyn std::error::Error>> {
8251 let nb = ARGMAX_NB;
8252 let f1 = self.func("argmax_partial_f32");
8253 let f2 = self.func("argmax_final_f32");
8254 let mut guard = self.argmax_partials.lock().unwrap();
8255 if guard.is_none() {
8256 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8257 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8258 *guard = Some((pv, pi));
8259 }
8260 let (part_v, part_i) = guard.as_mut().unwrap();
8261 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8262 let nv = n_vocab as i32;
8263 let nbi = nb as i32;
8264 let cfg1 = LaunchConfig {
8265 grid_dim: (nb as u32, 1, 1),
8266 block_dim: (256, 1, 1),
8267 shared_mem_bytes: 0,
8268 };
8269 let __s_b1 = self.gpu.stream();
8270 let mut b1 = __s_b1.launch_builder(&f1);
8271 b1.arg(&col_view)
8272 .arg(&mut *part_v)
8273 .arg(&mut *part_i)
8274 .arg(&nv);
8275 unsafe {
8276 b1.launch(cfg1)?;
8277 }
8278 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8279 let cfg2 = LaunchConfig {
8280 grid_dim: (1, 1, 1),
8281 block_dim: (256, 1, 1),
8282 shared_mem_bytes: 0,
8283 };
8284 let __s_b2 = self.gpu.stream();
8285 let mut b2 = __s_b2.launch_builder(&f2);
8286 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8287 unsafe {
8288 b2.launch(cfg2)?;
8289 }
8290 Ok(())
8291 }
8292 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8294 Ok(self.gpu.stream().clone_htod(v)?)
8295 }
8296 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8297 let v = self.gpu.stream().clone_dtoh(d)?;
8298 self.gpu.stream().synchronize()?;
8299 Ok(v)
8300 }
8301
8302 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8303 let v = self.gpu.stream().clone_dtoh(d)?;
8304 self.gpu.stream().synchronize()?;
8305 Ok(v)
8306 }
8307 pub fn htod_u32_into(
8311 &self,
8312 dst: &mut CudaSlice<u32>,
8313 src: &[u32],
8314 ) -> Result<(), Box<dyn std::error::Error>> {
8315 let mut view = dst.slice_mut(0..src.len());
8316 self.gpu.stream().memcpy_htod(src, &mut view)?;
8317 Ok(())
8318 }
8319
8320 pub fn htod_i32_into(
8323 &self,
8324 dst: &mut CudaSlice<i32>,
8325 src: &[i32],
8326 ) -> Result<(), Box<dyn std::error::Error>> {
8327 let mut view = dst.slice_mut(0..src.len());
8328 self.gpu.stream().memcpy_htod(src, &mut view)?;
8329 Ok(())
8330 }
8331
8332 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8333 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8334 self.keep_if_capturing(&s);
8335 Ok(s)
8336 }
8337 pub fn embed_gather_device_into(
8340 &self,
8341 embd: &CudaSlice<u8>,
8342 token_d: &CudaSlice<u32>,
8343 x_out: &mut CudaSlice<f32>,
8344 n_embd: usize,
8345 qtype: i32,
8346 row_bytes: usize,
8347 ) -> Result<(), Box<dyn std::error::Error>> {
8348 let f = self.func("embed_gather_u32");
8349 let cfg = LaunchConfig {
8350 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8351 block_dim: (256, 1, 1),
8352 shared_mem_bytes: 0,
8353 };
8354 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8355 let __s_b = self.gpu.stream();
8356 let mut b = __s_b.launch_builder(&f);
8357 b.arg(embd)
8358 .arg(token_d)
8359 .arg(x_out)
8360 .arg(&ne)
8361 .arg(&qt)
8362 .arg(&rb);
8363 unsafe {
8364 b.launch(cfg)?;
8365 }
8366 Ok(())
8367 }
8368 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8370 let v = self.gpu.stream().clone_dtoh(d)?;
8371 self.gpu.stream().synchronize()?;
8372 Ok(v[0])
8373 }
8374 pub fn i32_set_k(
8381 &self,
8382 dst: &mut CudaSlice<i32>,
8383 v: i32,
8384 ) -> Result<(), Box<dyn std::error::Error>> {
8385 let f = self.func("i32_set_k");
8386 let cfg = LaunchConfig {
8387 grid_dim: (1, 1, 1),
8388 block_dim: (1, 1, 1),
8389 shared_mem_bytes: 0,
8390 };
8391 let idx = 0i32;
8392 let __s_b = self.gpu.stream();
8393 let mut b = __s_b.launch_builder(&f);
8394 b.arg(dst).arg(&v).arg(&idx);
8395 unsafe {
8396 b.launch(cfg)?;
8397 }
8398 Ok(())
8399 }
8400
8401 pub fn set_i32_one(
8402 &self,
8403 d: &mut CudaSlice<i32>,
8404 v: i32,
8405 ) -> Result<(), Box<dyn std::error::Error>> {
8406 self.gpu.stream().memcpy_htod(&[v], d)?;
8407 Ok(())
8408 }
8409 pub fn set_u32_one(
8412 &self,
8413 d: &mut CudaSlice<u32>,
8414 v: u32,
8415 ) -> Result<(), Box<dyn std::error::Error>> {
8416 self.gpu.stream().memcpy_htod(&[v], d)?;
8417 Ok(())
8418 }
8419 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8421 let v = self.gpu.stream().clone_dtoh(d)?;
8422 self.gpu.stream().synchronize()?;
8423 Ok(v[0])
8424 }
8425 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8427 Ok(self.gpu.stream().clone_htod(bytes)?)
8428 }
8429 pub fn embed_gather_device(
8433 &self,
8434 embd: &CudaSlice<u8>,
8435 token_d: &CudaSlice<u32>,
8436 n_embd: usize,
8437 qtype: i32,
8438 row_bytes: usize,
8439 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8440 let f = self.func("embed_gather_u32");
8441 let mut x = self.alloc_uninit::<f32>(n_embd)?;
8442 let cfg = LaunchConfig {
8443 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8444 block_dim: (256, 1, 1),
8445 shared_mem_bytes: 0,
8446 };
8447 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8448 let __s_b = self.gpu.stream();
8449 let mut b = __s_b.launch_builder(&f);
8450 b.arg(embd)
8451 .arg(token_d)
8452 .arg(&mut x)
8453 .arg(&ne)
8454 .arg(&qt)
8455 .arg(&rb);
8456 unsafe {
8457 b.launch(cfg)?;
8458 }
8459 Ok(x)
8460 }
8461
8462 pub fn embed_gather_device_t(
8466 &self,
8467 embd: &CudaSlice<u8>,
8468 tokens: &[u32],
8469 n_embd: usize,
8470 qtype: i32,
8471 row_bytes: usize,
8472 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8473 let t = tokens.len();
8474 let tok_d = self.gpu.stream().clone_htod(tokens)?;
8475 let f = self.func("embed_gather_u32_t");
8476 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8477 let cfg = LaunchConfig {
8478 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8479 block_dim: (256, 1, 1),
8480 shared_mem_bytes: 0,
8481 };
8482 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8483 let __s_b = self.gpu.stream();
8484 let mut b = __s_b.launch_builder(&f);
8485 b.arg(embd)
8486 .arg(&tok_d)
8487 .arg(&mut x)
8488 .arg(&ne)
8489 .arg(&qt)
8490 .arg(&rb)
8491 .arg(&ti);
8492 unsafe {
8493 b.launch(cfg)?;
8494 }
8495 Ok(x)
8496 }
8497
8498 pub fn embed_gather_device_tv(
8503 &self,
8504 embd: &CudaSlice<u8>,
8505 tok_v: &cudarc::driver::CudaView<u32>,
8506 t: usize,
8507 n_embd: usize,
8508 qtype: i32,
8509 row_bytes: usize,
8510 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8511 let f = self.func("embed_gather_u32_t");
8512 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8513 let cfg = LaunchConfig {
8514 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8515 block_dim: (256, 1, 1),
8516 shared_mem_bytes: 0,
8517 };
8518 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8519 let __s_b = self.gpu.stream();
8520 let mut b = __s_b.launch_builder(&f);
8521 b.arg(embd)
8522 .arg(tok_v)
8523 .arg(&mut x)
8524 .arg(&ne)
8525 .arg(&qt)
8526 .arg(&rb)
8527 .arg(&ti);
8528 unsafe {
8529 b.launch(cfg)?;
8530 }
8531 Ok(x)
8532 }
8533
8534 pub fn embed_gather_device_td(
8535 &self,
8536 embd: &CudaSlice<u8>,
8537 tok_d: &CudaSlice<u32>,
8538 t: usize,
8539 n_embd: usize,
8540 qtype: i32,
8541 row_bytes: usize,
8542 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8543 let f = self.func("embed_gather_u32_t");
8544 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8545 let cfg = LaunchConfig {
8546 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8547 block_dim: (256, 1, 1),
8548 shared_mem_bytes: 0,
8549 };
8550 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8551 let __s_b = self.gpu.stream();
8552 let mut b = __s_b.launch_builder(&f);
8553 b.arg(embd)
8554 .arg(tok_d)
8555 .arg(&mut x)
8556 .arg(&ne)
8557 .arg(&qt)
8558 .arg(&rb)
8559 .arg(&ti);
8560 unsafe {
8561 b.launch(cfg)?;
8562 }
8563 Ok(x)
8564 }
8565
8566 #[inline]
8572 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8574 if self
8575 .capture_keep_on
8576 .load(std::sync::atomic::Ordering::Relaxed)
8577 {
8578 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8579 }
8580 }
8581
8582 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8583 &self,
8584 n: usize,
8585 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8586 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8587 {
8591 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8592 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8593 use cudarc::driver::DevicePtrMut;
8595 let n_bytes = s.len() * std::mem::size_of::<T>();
8596 let stream = self.gpu.stream();
8597 let (p_, _g) = s.device_ptr_mut(&stream);
8598 unsafe {
8599 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8600 .result()?;
8601 }
8602 }
8603 }
8604 self.keep_if_capturing(&s);
8605 Ok(s)
8606 }
8607
8608 pub fn uninit_q8_pair(
8613 &self,
8614 n: usize,
8615 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8616 Ok((
8617 self.alloc_uninit::<i8>(n)?,
8618 self.alloc_uninit::<f32>(n / 32)?,
8619 ))
8620 }
8621
8622 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8623 self.alloc_uninit::<f32>(n)
8624 }
8625
8626 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8628 self.alloc_uninit::<i8>(n)
8629 }
8630
8631 #[allow(clippy::too_many_arguments)]
8635 pub fn rms_norm3(
8636 &self,
8637 x: &CudaSlice<f32>,
8638 w0: &CudaSlice<f32>,
8639 w1: &CudaSlice<f32>,
8640 w2: &CudaSlice<f32>,
8641 d0: &mut CudaSlice<f32>,
8642 d1: &mut CudaSlice<f32>,
8643 d2: &mut CudaSlice<f32>,
8644 ncols: usize,
8645 nrows: usize,
8646 eps: f32,
8647 ) -> Result<(), Box<dyn std::error::Error>> {
8648 let f = self.func("rms_norm3_f32");
8649 let cfg = LaunchConfig {
8650 grid_dim: (nrows as u32, 1, 1),
8651 block_dim: (rms_block(), 1, 1),
8652 shared_mem_bytes: 0,
8653 };
8654 let (nc, e) = (ncols as i32, eps);
8655 let __s_b = self.gpu.stream();
8656 let mut b = __s_b.launch_builder(&f);
8657 b.arg(x)
8658 .arg(w0)
8659 .arg(w1)
8660 .arg(w2)
8661 .arg(d0)
8662 .arg(d1)
8663 .arg(d2)
8664 .arg(&nc)
8665 .arg(&e);
8666 unsafe {
8667 b.launch(cfg)?;
8668 }
8669 Ok(())
8670 }
8671
8672 #[allow(clippy::too_many_arguments)]
8674 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8677 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8678 *WARP_ON.get_or_init(|| {
8679 std::env::var("MEMRA_QKVNORM_W")
8680 .map(|v| v != "0")
8681 .unwrap_or(true)
8682 }) && ncols % 4 == 0
8683 && rows >= 64
8684 }
8685
8686 #[allow(clippy::too_many_arguments)]
8689 pub fn rms_norm_qkv_w4b(
8690 &self,
8691 q: &CudaSlice<f32>,
8692 k: &CudaSlice<f32>,
8693 v: &CudaSlice<f32>,
8694 wq: &CudaSlice<f32>,
8695 wk: &CudaSlice<f32>,
8696 wv: &CudaSlice<f32>,
8697 dq: &mut CudaSlice<f32>,
8698 dk: &mut CudaSlice<f32>,
8699 dv: &mut CudaSlice<f32>,
8700 dvb: &mut CudaSlice<u8>,
8701 ncols: usize,
8702 rq: usize,
8703 rk: usize,
8704 eps: f32,
8705 vf16: bool,
8706 ) -> Result<(), Box<dyn std::error::Error>> {
8707 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8708 let f = self.func("rms_norm_qkv_w4b_f32");
8709 let rows = (rq + 2 * rk) as u32;
8710 let cfg = LaunchConfig {
8711 grid_dim: (rows.div_ceil(8), 1, 1),
8712 block_dim: (256, 1, 1),
8713 shared_mem_bytes: 0,
8714 };
8715 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8716 let vf = vf16 as i32;
8717 let __s_b = self.gpu.stream();
8718 let mut b = __s_b.launch_builder(&f);
8719 b.arg(q)
8720 .arg(k)
8721 .arg(v)
8722 .arg(wq)
8723 .arg(wk)
8724 .arg(wv)
8725 .arg(dq)
8726 .arg(dk)
8727 .arg(dv)
8728 .arg(&mut *dvb)
8729 .arg(&nc)
8730 .arg(&rqi)
8731 .arg(&rki)
8732 .arg(&rvi)
8733 .arg(&e)
8734 .arg(&vf);
8735 unsafe {
8736 b.launch(cfg)?;
8737 }
8738 Ok(())
8739 }
8740
8741 pub fn rms_norm_qkv(
8742 &self,
8743 q: &CudaSlice<f32>,
8744 k: &CudaSlice<f32>,
8745 v: &CudaSlice<f32>,
8746 wq: &CudaSlice<f32>,
8747 wk: &CudaSlice<f32>,
8748 wv: &CudaSlice<f32>,
8749 dq: &mut CudaSlice<f32>,
8750 dk: &mut CudaSlice<f32>,
8751 dv: &mut CudaSlice<f32>,
8752 ncols: usize,
8753 rq: usize,
8754 rk: usize,
8755 eps: f32,
8756 ) -> Result<(), Box<dyn std::error::Error>> {
8757 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8761 let warp_on = *WARP_ON.get_or_init(|| {
8762 std::env::var("MEMRA_QKVNORM_W")
8763 .map(|v| v != "0")
8764 .unwrap_or(true)
8765 });
8766 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8769 let f = self.func("rms_norm_qkv_w4_f32");
8770 let rows = (rq + 2 * rk) as u32;
8771 let cfg = LaunchConfig {
8772 grid_dim: (rows.div_ceil(8), 1, 1),
8773 block_dim: (256, 1, 1),
8774 shared_mem_bytes: 0,
8775 };
8776 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8777 let __s_b = self.gpu.stream();
8778 let mut b = __s_b.launch_builder(&f);
8779 b.arg(q)
8780 .arg(k)
8781 .arg(v)
8782 .arg(wq)
8783 .arg(wk)
8784 .arg(wv)
8785 .arg(dq)
8786 .arg(dk)
8787 .arg(dv)
8788 .arg(&nc)
8789 .arg(&rqi)
8790 .arg(&rki)
8791 .arg(&rvi)
8792 .arg(&e);
8793 unsafe {
8794 b.launch(cfg)?;
8795 }
8796 return Ok(());
8797 }
8798 let f = self.func("rms_norm_qkv_f32");
8799 let grid = (rq + 2 * rk) as u32;
8800 let cfg = LaunchConfig {
8801 grid_dim: (grid, 1, 1),
8802 block_dim: (rms_block(), 1, 1),
8803 shared_mem_bytes: 0,
8804 };
8805 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8806 let __s_b = self.gpu.stream();
8807 let mut b = __s_b.launch_builder(&f);
8808 b.arg(q)
8809 .arg(k)
8810 .arg(v)
8811 .arg(wq)
8812 .arg(wk)
8813 .arg(wv)
8814 .arg(dq)
8815 .arg(dk)
8816 .arg(dv)
8817 .arg(&nc)
8818 .arg(&rqi)
8819 .arg(&rki)
8820 .arg(&e);
8821 unsafe {
8822 b.launch(cfg)?;
8823 }
8824 Ok(())
8825 }
8826
8827 #[allow(clippy::too_many_arguments)]
8829 pub fn rms_norm2x(
8830 &self,
8831 a: &CudaSlice<f32>,
8832 bb: &CudaSlice<f32>,
8833 wa: &CudaSlice<f32>,
8834 wb: &CudaSlice<f32>,
8835 da: &mut CudaSlice<f32>,
8836 db: &mut CudaSlice<f32>,
8837 ncols: usize,
8838 nrows: usize,
8839 eps: f32,
8840 ) -> Result<(), Box<dyn std::error::Error>> {
8841 let f = self.func("rms_norm2x_f32");
8842 let cfg = LaunchConfig {
8843 grid_dim: (2 * nrows as u32, 1, 1),
8844 block_dim: (rms_block(), 1, 1),
8845 shared_mem_bytes: 0,
8846 };
8847 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
8848 let __s_b = self.gpu.stream();
8849 let mut b = __s_b.launch_builder(&f);
8850 b.arg(a)
8851 .arg(bb)
8852 .arg(wa)
8853 .arg(wb)
8854 .arg(da)
8855 .arg(db)
8856 .arg(&nc)
8857 .arg(&nr)
8858 .arg(&e);
8859 unsafe {
8860 b.launch(cfg)?;
8861 }
8862 Ok(())
8863 }
8864
8865 pub fn softcap(
8867 &self,
8868 y: &mut CudaSlice<f32>,
8869 cap: f32,
8870 n: usize,
8871 ) -> Result<(), Box<dyn std::error::Error>> {
8872 let f = self.func("softcap_f32");
8873 let cfg = LaunchConfig::for_num_elems(n as u32);
8874 let ni = n as i32;
8875 let __s_b = self.gpu.stream();
8876 let mut b = __s_b.launch_builder(&f);
8877 b.arg(y).arg(&cap).arg(&ni);
8878 unsafe {
8879 b.launch(cfg)?;
8880 }
8881 Ok(())
8882 }
8883
8884 pub fn mask_ids_rows(
8887 &self,
8888 y: &mut CudaSlice<f32>,
8889 ids: &CudaSlice<i32>,
8890 n_ids: usize,
8891 n_vocab: usize,
8892 t: usize,
8893 ) -> Result<(), Box<dyn std::error::Error>> {
8894 let f = self.func("mask_ids_rows_f32");
8895 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
8896 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
8897 let __s_b = self.gpu.stream();
8898 let mut b = __s_b.launch_builder(&f);
8899 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
8900 unsafe {
8901 b.launch(cfg)?;
8902 }
8903 Ok(())
8904 }
8905
8906 #[allow(clippy::too_many_arguments)]
8908 pub fn add_scale_rms_norm(
8909 &self,
8910 a: &CudaSlice<f32>,
8911 b_in: &CudaSlice<f32>,
8912 c: f32,
8913 w: &CudaSlice<f32>,
8914 res: &mut CudaSlice<f32>,
8915 dst: &mut CudaSlice<f32>,
8916 ncols: usize,
8917 nrows: usize,
8918 eps: f32,
8919 ) -> Result<(), Box<dyn std::error::Error>> {
8920 let f = self.func("add_scale_rms_norm_f32");
8921 let cfg = LaunchConfig {
8922 grid_dim: (nrows as u32, 1, 1),
8923 block_dim: (rms_block(), 1, 1),
8924 shared_mem_bytes: 0,
8925 };
8926 let (nc, e2) = (ncols as i32, eps);
8927 let __s_b = self.gpu.stream();
8928 let mut b = __s_b.launch_builder(&f);
8929 b.arg(a)
8930 .arg(b_in)
8931 .arg(&c)
8932 .arg(w)
8933 .arg(res)
8934 .arg(dst)
8935 .arg(&nc)
8936 .arg(&e2);
8937 unsafe {
8938 b.launch(cfg)?;
8939 }
8940 Ok(())
8941 }
8942
8943 #[allow(clippy::too_many_arguments)]
8946 pub fn add_scale_rms_norm_q8_1(
8947 &self,
8948 a: &CudaSlice<f32>,
8949 b_in: &CudaSlice<f32>,
8950 c: f32,
8951 w: &CudaSlice<f32>,
8952 res: &mut CudaSlice<f32>,
8953 ncols: usize,
8954 nrows: usize,
8955 eps: f32,
8956 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8957 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8958 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8959 let (nc, e2) = (ncols as i32, eps);
8960 if Self::pdl_on() && Self::pdl_wb_on() {
8961 {
8962 use cudarc::driver::{DevicePtr, DevicePtrMut};
8963 let s = &self.gpu.stream();
8964 let (pa, _g0) = a.device_ptr(s);
8965 let (pb, _g1) = b_in.device_ptr(s);
8966 let (pw, _g2) = w.device_ptr(s);
8967 let (pr, _g3) = res.device_ptr_mut(s);
8968 let (pq, _g4) = out_q.device_ptr_mut(s);
8969 let (pd, _g5) = out_d.device_ptr_mut(s);
8970 let mut ps = [
8971 &pa as *const _ as *mut std::ffi::c_void,
8972 &pb as *const _ as *mut _,
8973 &c as *const _ as *mut _,
8974 &pw as *const _ as *mut _,
8975 &pr as *const _ as *mut _,
8976 &pq as *const _ as *mut _,
8977 &pd as *const _ as *mut _,
8978 &nc as *const _ as *mut _,
8979 &e2 as *const _ as *mut _,
8980 ];
8981 unsafe {
8982 self.launch_pdl(
8983 "add_scale_rms_norm_q8_1",
8984 (nrows as u32, 1, 1),
8985 (rms_block(), 1, 1),
8986 &mut ps,
8987 )?;
8988 }
8989 }
8990 return Ok((out_q, out_d));
8991 }
8992 let f = self.func("add_scale_rms_norm_q8_1");
8993 let cfg = LaunchConfig {
8994 grid_dim: (nrows as u32, 1, 1),
8995 block_dim: (rms_block(), 1, 1),
8996 shared_mem_bytes: 0,
8997 };
8998 let __s_b = self.gpu.stream();
8999 let mut b = __s_b.launch_builder(&f);
9000 b.arg(a)
9001 .arg(b_in)
9002 .arg(&c)
9003 .arg(w)
9004 .arg(res)
9005 .arg(&mut out_q)
9006 .arg(&mut out_d)
9007 .arg(&nc)
9008 .arg(&e2);
9009 unsafe {
9010 b.launch(cfg)?;
9011 }
9012 Ok((out_q, out_d))
9013 }
9014
9015 #[allow(clippy::too_many_arguments)]
9017 pub fn add_scale_rms_norm_q8_1_into(
9018 &self,
9019 a: &CudaSlice<f32>,
9020 b_in: &CudaSlice<f32>,
9021 c: f32,
9022 w: &CudaSlice<f32>,
9023 res: &mut CudaSlice<f32>,
9024 ncols: usize,
9025 nrows: usize,
9026 eps: f32,
9027 out_q: &mut CudaSlice<i8>,
9028 out_d: &mut CudaSlice<f32>,
9029 ) -> Result<(), Box<dyn std::error::Error>> {
9030 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9031 let (nc, e2) = (ncols as i32, eps);
9032 if Self::pdl_on() && Self::pdl_wb_on() {
9033 use cudarc::driver::{DevicePtr, DevicePtrMut};
9034 let s = &self.gpu.stream();
9035 let (pa, _g0) = a.device_ptr(s);
9036 let (pb, _g1) = b_in.device_ptr(s);
9037 let (pw, _g2) = w.device_ptr(s);
9038 let (pr, _g3) = res.device_ptr_mut(s);
9039 let (pq, _g4) = out_q.device_ptr_mut(s);
9040 let (pd, _g5) = out_d.device_ptr_mut(s);
9041 let mut ps = [
9042 &pa as *const _ as *mut std::ffi::c_void,
9043 &pb as *const _ as *mut _,
9044 &c as *const _ as *mut _,
9045 &pw as *const _ as *mut _,
9046 &pr as *const _ as *mut _,
9047 &pq as *const _ as *mut _,
9048 &pd as *const _ as *mut _,
9049 &nc as *const _ as *mut _,
9050 &e2 as *const _ as *mut _,
9051 ];
9052 unsafe {
9053 self.launch_pdl(
9054 "add_scale_rms_norm_q8_1",
9055 (nrows as u32, 1, 1),
9056 (rms_block(), 1, 1),
9057 &mut ps,
9058 )?;
9059 }
9060 return Ok(());
9061 }
9062 let f = self.func("add_scale_rms_norm_q8_1");
9063 let cfg = LaunchConfig {
9064 grid_dim: (nrows as u32, 1, 1),
9065 block_dim: (rms_block(), 1, 1),
9066 shared_mem_bytes: 0,
9067 };
9068 let __s_b = self.gpu.stream();
9069 let mut b = __s_b.launch_builder(&f);
9070 b.arg(a)
9071 .arg(b_in)
9072 .arg(&c)
9073 .arg(w)
9074 .arg(res)
9075 .arg(&mut *out_q)
9076 .arg(&mut *out_d)
9077 .arg(&nc)
9078 .arg(&e2);
9079 unsafe {
9080 b.launch(cfg)?;
9081 }
9082 Ok(())
9083 }
9084
9085 #[allow(clippy::too_many_arguments)]
9088 pub fn rms_pre_add_scale_rms_norm_q8_1(
9089 &self,
9090 a: &CudaSlice<f32>,
9091 wa: &CudaSlice<f32>,
9092 b_in: &CudaSlice<f32>,
9093 c: f32,
9094 w: &CudaSlice<f32>,
9095 res: &mut CudaSlice<f32>,
9096 ncols: usize,
9097 nrows: usize,
9098 eps: f32,
9099 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9100 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9101 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9102 let (nc, e2) = (ncols as i32, eps);
9103 if Self::pdl_on() {
9104 {
9105 use cudarc::driver::{DevicePtr, DevicePtrMut};
9106 let s = &self.gpu.stream();
9107 let (pa, _g0) = a.device_ptr(s);
9108 let (pwa, _g1) = wa.device_ptr(s);
9109 let (pb, _g2) = b_in.device_ptr(s);
9110 let (pw, _g3) = w.device_ptr(s);
9111 let (pr, _g4) = res.device_ptr_mut(s);
9112 let (pq, _g5) = out_q.device_ptr_mut(s);
9113 let (pd, _g6) = out_d.device_ptr_mut(s);
9114 let mut ps = [
9115 &pa as *const _ as *mut std::ffi::c_void,
9116 &pwa as *const _ as *mut _,
9117 &pb as *const _ as *mut _,
9118 &c as *const _ as *mut _,
9119 &pw as *const _ as *mut _,
9120 &pr as *const _ as *mut _,
9121 &pq as *const _ as *mut _,
9122 &pd as *const _ as *mut _,
9123 &nc as *const _ as *mut _,
9124 &e2 as *const _ as *mut _,
9125 ];
9126 unsafe {
9127 self.launch_pdl(
9128 "rms_pre_add_scale_rms_norm_q8_1",
9129 (nrows as u32, 1, 1),
9130 (rms_block(), 1, 1),
9131 &mut ps,
9132 )?;
9133 }
9134 }
9135 return Ok((out_q, out_d));
9136 }
9137 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9138 let cfg = LaunchConfig {
9139 grid_dim: (nrows as u32, 1, 1),
9140 block_dim: (rms_block(), 1, 1),
9141 shared_mem_bytes: 0,
9142 };
9143 let __s_b = self.gpu.stream();
9144 let mut b = __s_b.launch_builder(&f);
9145 b.arg(a)
9146 .arg(wa)
9147 .arg(b_in)
9148 .arg(&c)
9149 .arg(w)
9150 .arg(res)
9151 .arg(&mut out_q)
9152 .arg(&mut out_d)
9153 .arg(&nc)
9154 .arg(&e2);
9155 unsafe {
9156 b.launch(cfg)?;
9157 }
9158 Ok((out_q, out_d))
9159 }
9160
9161 pub fn gelu_tanh_mul_q8_1(
9164 &self,
9165 gate: &CudaSlice<f32>,
9166 up: &cudarc::driver::CudaView<f32>,
9167 act: &mut CudaSlice<f32>,
9168 ncols: usize,
9169 nrows: usize,
9170 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9171 debug_assert!(ncols % 128 == 0);
9172 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9173 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9174 let nc = ncols as i32;
9175 if Self::pdl_on() {
9176 {
9177 use cudarc::driver::{DevicePtr, DevicePtrMut};
9178 let s = &self.gpu.stream();
9179 let (pg, _g0) = gate.device_ptr(s);
9180 let (pu, _g1) = up.device_ptr(s);
9181 let (pact, _g2) = act.device_ptr_mut(s);
9182 let (pq, _g3) = out_q.device_ptr_mut(s);
9183 let (pd, _g4) = out_d.device_ptr_mut(s);
9184 let mut ps = [
9185 &pg as *const _ as *mut std::ffi::c_void,
9186 &pu as *const _ as *mut _,
9187 &pact as *const _ as *mut _,
9188 &pq as *const _ as *mut _,
9189 &pd as *const _ as *mut _,
9190 &nc as *const _ as *mut _,
9191 ];
9192 unsafe {
9193 self.launch_pdl(
9194 "gelu_tanh_mul_q8_1",
9195 (nrows as u32, 1, 1),
9196 (rms_block(), 1, 1),
9197 &mut ps,
9198 )?;
9199 }
9200 }
9201 return Ok((out_q, out_d));
9202 }
9203 let f = self.func("gelu_tanh_mul_q8_1");
9204 let cfg = LaunchConfig {
9205 grid_dim: (nrows as u32, 1, 1),
9206 block_dim: (rms_block(), 1, 1),
9207 shared_mem_bytes: 0,
9208 };
9209 let __s_b = self.gpu.stream();
9210 let mut b = __s_b.launch_builder(&f);
9211 b.arg(gate)
9212 .arg(up)
9213 .arg(act)
9214 .arg(&mut out_q)
9215 .arg(&mut out_d)
9216 .arg(&nc);
9217 unsafe {
9218 b.launch(cfg)?;
9219 }
9220 Ok((out_q, out_d))
9221 }
9222
9223 #[allow(clippy::too_many_arguments)]
9225 pub fn gelu_tanh_mul_q8_1_into(
9226 &self,
9227 gate: &CudaSlice<f32>,
9228 up: &cudarc::driver::CudaView<f32>,
9229 act: &mut CudaSlice<f32>,
9230 ncols: usize,
9231 nrows: usize,
9232 out_q: &mut CudaSlice<i8>,
9233 out_d: &mut CudaSlice<f32>,
9234 ) -> Result<(), Box<dyn std::error::Error>> {
9235 debug_assert!(ncols % 128 == 0);
9236 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9237 let nc = ncols as i32;
9238 if Self::pdl_on() {
9239 use cudarc::driver::{DevicePtr, DevicePtrMut};
9240 let s = &self.gpu.stream();
9241 let (pg, _g0) = gate.device_ptr(s);
9242 let (pu, _g1) = up.device_ptr(s);
9243 let (pact, _g2) = act.device_ptr_mut(s);
9244 let (pq, _g3) = out_q.device_ptr_mut(s);
9245 let (pd, _g4) = out_d.device_ptr_mut(s);
9246 let mut ps = [
9247 &pg as *const _ as *mut std::ffi::c_void,
9248 &pu as *const _ as *mut _,
9249 &pact as *const _ as *mut _,
9250 &pq as *const _ as *mut _,
9251 &pd as *const _ as *mut _,
9252 &nc as *const _ as *mut _,
9253 ];
9254 unsafe {
9255 self.launch_pdl(
9256 "gelu_tanh_mul_q8_1",
9257 (nrows as u32, 1, 1),
9258 (rms_block(), 1, 1),
9259 &mut ps,
9260 )?;
9261 }
9262 return Ok(());
9263 }
9264 let f = self.func("gelu_tanh_mul_q8_1");
9265 let cfg = LaunchConfig {
9266 grid_dim: (nrows as u32, 1, 1),
9267 block_dim: (rms_block(), 1, 1),
9268 shared_mem_bytes: 0,
9269 };
9270 let __s_b = self.gpu.stream();
9271 let mut b = __s_b.launch_builder(&f);
9272 b.arg(gate)
9273 .arg(up)
9274 .arg(&mut *act)
9275 .arg(&mut *out_q)
9276 .arg(&mut *out_d)
9277 .arg(&nc);
9278 unsafe {
9279 b.launch(cfg)?;
9280 }
9281 Ok(())
9282 }
9283
9284 #[allow(clippy::too_many_arguments)]
9286 pub fn add_rms_norm3_q8z(
9287 &self,
9288 a: &CudaSlice<f32>,
9289 b_in: &CudaSlice<f32>,
9290 w0: &CudaSlice<f32>,
9291 w1: &CudaSlice<f32>,
9292 w2: &CudaSlice<f32>,
9293 res: &mut CudaSlice<f32>,
9294 out1: &mut CudaSlice<f32>,
9295 ncols: usize,
9296 nrows: usize,
9297 eps: f32,
9298 ) -> Result<
9299 (
9300 (CudaSlice<i8>, CudaSlice<f32>),
9301 (CudaSlice<i8>, CudaSlice<f32>),
9302 ),
9303 Box<dyn std::error::Error>,
9304 > {
9305 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9306 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9307 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9308 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9309 let f = self.func("add_rms_norm3_q8z_f32");
9310 let cfg = LaunchConfig {
9311 grid_dim: (nrows as u32, 1, 1),
9312 block_dim: (rms_block(), 1, 1),
9313 shared_mem_bytes: 0,
9314 };
9315 let (nc, e2) = (ncols as i32, eps);
9316 let __s_b = self.gpu.stream();
9317 let mut b = __s_b.launch_builder(&f);
9318 b.arg(a)
9319 .arg(b_in)
9320 .arg(w0)
9321 .arg(w1)
9322 .arg(w2)
9323 .arg(res)
9324 .arg(&mut q0)
9325 .arg(&mut d0)
9326 .arg(out1)
9327 .arg(&mut q2)
9328 .arg(&mut d2)
9329 .arg(&nc)
9330 .arg(&e2);
9331 unsafe {
9332 b.launch(cfg)?;
9333 }
9334 Ok(((q0, d0), (q2, d2)))
9335 }
9336
9337 #[allow(clippy::too_many_arguments)]
9339 pub fn add_rms_norm3(
9340 &self,
9341 a: &CudaSlice<f32>,
9342 b_in: &CudaSlice<f32>,
9343 w0: &CudaSlice<f32>,
9344 w1: &CudaSlice<f32>,
9345 w2: &CudaSlice<f32>,
9346 res: &mut CudaSlice<f32>,
9347 d0: &mut CudaSlice<f32>,
9348 d1: &mut CudaSlice<f32>,
9349 d2: &mut CudaSlice<f32>,
9350 ncols: usize,
9351 nrows: usize,
9352 eps: f32,
9353 ) -> Result<(), Box<dyn std::error::Error>> {
9354 let f = self.func("add_rms_norm3_f32");
9355 let cfg = LaunchConfig {
9356 grid_dim: (nrows as u32, 1, 1),
9357 block_dim: (rms_block(), 1, 1),
9358 shared_mem_bytes: 0,
9359 };
9360 let (nc, e2) = (ncols as i32, eps);
9361 let __s_b = self.gpu.stream();
9362 let mut b = __s_b.launch_builder(&f);
9363 b.arg(a)
9364 .arg(b_in)
9365 .arg(w0)
9366 .arg(w1)
9367 .arg(w2)
9368 .arg(res)
9369 .arg(d0)
9370 .arg(d1)
9371 .arg(d2)
9372 .arg(&nc)
9373 .arg(&e2);
9374 unsafe {
9375 b.launch(cfg)?;
9376 }
9377 Ok(())
9378 }
9379
9380 pub fn add_scale(
9382 &self,
9383 a: &CudaSlice<f32>,
9384 b_in: &CudaSlice<f32>,
9385 c: f32,
9386 dst: &mut CudaSlice<f32>,
9387 n: usize,
9388 ) -> Result<(), Box<dyn std::error::Error>> {
9389 let f = self.func("add_scale_f32");
9390 let cfg = LaunchConfig::for_num_elems(n as u32);
9391 let ni = n as i32;
9392 let __s_b = self.gpu.stream();
9393 let mut b = __s_b.launch_builder(&f);
9394 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9395 unsafe {
9396 b.launch(cfg)?;
9397 }
9398 Ok(())
9399 }
9400
9401 pub fn layer_norm_bias(
9403 &self,
9404 x: &CudaSlice<f32>,
9405 w: &CudaSlice<f32>,
9406 b: &CudaSlice<f32>,
9407 dst: &mut CudaSlice<f32>,
9408 ncols: usize,
9409 nrows: usize,
9410 eps: f32,
9411 ) -> Result<(), Box<dyn std::error::Error>> {
9412 let f = self.func("layer_norm_bias_f32");
9413 let (nc, e) = (ncols as i32, eps);
9414 let cfg = LaunchConfig {
9415 grid_dim: (nrows as u32, 1, 1),
9416 block_dim: (256, 1, 1),
9417 shared_mem_bytes: 0,
9418 };
9419 let __s_b = self.gpu.stream();
9420 let mut lb = __s_b.launch_builder(&f);
9421 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9422 unsafe {
9423 lb.launch(cfg)?;
9424 }
9425 Ok(())
9426 }
9427
9428 pub fn gelu_tanh(
9430 &self,
9431 x: &CudaSlice<f32>,
9432 dst: &mut CudaSlice<f32>,
9433 n: usize,
9434 ) -> Result<(), Box<dyn std::error::Error>> {
9435 let f = self.func("gelu_tanh_f32");
9436 let ni = n as i64;
9437 let cfg = LaunchConfig {
9438 grid_dim: (n.div_ceil(256) as u32, 1, 1),
9439 block_dim: (256, 1, 1),
9440 shared_mem_bytes: 0,
9441 };
9442 let __s_b = self.gpu.stream();
9443 let mut lb = __s_b.launch_builder(&f);
9444 lb.arg(x).arg(&mut *dst).arg(&ni);
9445 unsafe {
9446 lb.launch(cfg)?;
9447 }
9448 Ok(())
9449 }
9450
9451 pub fn row_softmax(
9453 &self,
9454 x: &mut CudaSlice<f32>,
9455 ncols: usize,
9456 nrows: usize,
9457 ) -> Result<(), Box<dyn std::error::Error>> {
9458 let f = self.func("row_softmax_f32");
9459 let nc = ncols as i32;
9460 let cfg = LaunchConfig {
9461 grid_dim: (nrows as u32, 1, 1),
9462 block_dim: (256, 1, 1),
9463 shared_mem_bytes: 0,
9464 };
9465 let __s_b = self.gpu.stream();
9466 let mut lb = __s_b.launch_builder(&f);
9467 lb.arg(&mut *x).arg(&nc);
9468 unsafe {
9469 lb.launch(cfg)?;
9470 }
9471 Ok(())
9472 }
9473
9474 pub fn rms_norm(
9475 &self,
9476 x: &CudaSlice<f32>,
9477 w: &CudaSlice<f32>,
9478 dst: &mut CudaSlice<f32>,
9479 ncols: usize,
9480 nrows: usize,
9481 eps: f32,
9482 ) -> Result<(), Box<dyn std::error::Error>> {
9483 let (nc, e) = (ncols as i32, eps);
9484 let kname = if Self::norm_ilp_on() {
9485 "rms_norm_f32_v2"
9486 } else {
9487 "rms_norm_f32"
9488 };
9489 if Self::pdl_on() && Self::pdl_wb_on() {
9490 use cudarc::driver::{DevicePtr, DevicePtrMut};
9491 let s = &self.gpu.stream();
9492 let (px, _g0) = x.device_ptr(s);
9493 let (pw, _g1) = w.device_ptr(s);
9494 let (pd, _g2) = dst.device_ptr_mut(s);
9495 let mut ps = [
9496 &px as *const _ as *mut std::ffi::c_void,
9497 &pw as *const _ as *mut _,
9498 &pd as *const _ as *mut _,
9499 &nc as *const _ as *mut _,
9500 &e as *const _ as *mut _,
9501 ];
9502 unsafe {
9503 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9504 }
9505 return Ok(());
9506 }
9507 let f = self.func(kname);
9508 let cfg = LaunchConfig {
9509 grid_dim: (nrows as u32, 1, 1),
9510 block_dim: (rms_block(), 1, 1),
9511 shared_mem_bytes: 0,
9512 };
9513 let __s_b = self.gpu.stream();
9514 let mut b = __s_b.launch_builder(&f);
9515 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9516 unsafe {
9517 b.launch(cfg)?;
9518 }
9519 Ok(())
9520 }
9521
9522 pub fn rms_norm_decode(
9530 &self,
9531 x: &CudaSlice<f32>,
9532 w: &CudaSlice<f32>,
9533 dst: &mut CudaSlice<f32>,
9534 ncols: usize,
9535 nrows: usize,
9536 eps: f32,
9537 ) -> Result<(), Box<dyn std::error::Error>> {
9538 let f = self.func(if Self::norm_ilp_on() {
9539 "rms_norm_f32_v2"
9540 } else {
9541 "rms_norm_f32"
9542 });
9543 let cfg = LaunchConfig {
9544 grid_dim: (nrows as u32, 1, 1),
9545 block_dim: (1024, 1, 1),
9546 shared_mem_bytes: 0,
9547 };
9548 let (nc, e) = (ncols as i32, eps);
9549 let __s_b = self.gpu.stream();
9550 let mut b = __s_b.launch_builder(&f);
9551 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9552 unsafe {
9553 b.launch(cfg)?;
9554 }
9555 Ok(())
9556 }
9557
9558 pub fn rms_norm_q8_1(
9562 &self,
9563 x: &CudaSlice<f32>,
9564 w: &CudaSlice<f32>,
9565 ncols: usize,
9566 nrows: usize,
9567 eps: f32,
9568 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9569 let nblk = ncols / 32;
9570 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9571 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9572 let (nc, e) = (ncols as i32, eps);
9573 if Self::pdl_on() {
9574 {
9575 use cudarc::driver::{DevicePtr, DevicePtrMut};
9576 let s = &self.gpu.stream();
9577 let (px, _g0) = x.device_ptr(s);
9578 let (pw, _g1) = w.device_ptr(s);
9579 let (pq, _g2) = q.device_ptr_mut(s);
9580 let (pd, _g3) = d.device_ptr_mut(s);
9581 let mut ps = [
9582 &px as *const _ as *mut std::ffi::c_void,
9583 &pw as *const _ as *mut _,
9584 &pq as *const _ as *mut _,
9585 &pd as *const _ as *mut _,
9586 &nc as *const _ as *mut _,
9587 &e as *const _ as *mut _,
9588 ];
9589 unsafe {
9590 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9591 }
9592 }
9593 return Ok((q, d));
9594 }
9595 let f = self.func("rms_norm_q8_1");
9596 let cfg = LaunchConfig {
9599 grid_dim: (nrows as u32, 1, 1),
9600 block_dim: (1024, 1, 1),
9601 shared_mem_bytes: 0,
9602 };
9603 let __s_b = self.gpu.stream();
9604 let mut b = __s_b.launch_builder(&f);
9605 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9606 unsafe {
9607 b.launch(cfg)?;
9608 }
9609 Ok((q, d))
9610 }
9611
9612 pub fn rms_norm_q8_1_into(
9615 &self,
9616 x: &CudaSlice<f32>,
9617 w: &CudaSlice<f32>,
9618 ncols: usize,
9619 nrows: usize,
9620 eps: f32,
9621 q: &mut CudaSlice<i8>,
9622 d: &mut CudaSlice<f32>,
9623 ) -> Result<(), Box<dyn std::error::Error>> {
9624 let nblk = ncols / 32;
9625 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9626 let (nc, e) = (ncols as i32, eps);
9627 if Self::pdl_on() {
9628 use cudarc::driver::{DevicePtr, DevicePtrMut};
9629 let s = &self.gpu.stream();
9630 let (px, _g0) = x.device_ptr(s);
9631 let (pw, _g1) = w.device_ptr(s);
9632 let (pq, _g2) = q.device_ptr_mut(s);
9633 let (pd, _g3) = d.device_ptr_mut(s);
9634 let mut ps = [
9635 &px as *const _ as *mut std::ffi::c_void,
9636 &pw as *const _ as *mut _,
9637 &pq as *const _ as *mut _,
9638 &pd as *const _ as *mut _,
9639 &nc as *const _ as *mut _,
9640 &e as *const _ as *mut _,
9641 ];
9642 unsafe {
9643 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9644 }
9645 return Ok(());
9646 }
9647 let f = self.func("rms_norm_q8_1");
9648 let cfg = LaunchConfig {
9649 grid_dim: (nrows as u32, 1, 1),
9650 block_dim: (1024, 1, 1),
9651 shared_mem_bytes: 0,
9652 };
9653 let __s_b = self.gpu.stream();
9654 let mut b = __s_b.launch_builder(&f);
9655 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9656 unsafe {
9657 b.launch(cfg)?;
9658 }
9659 Ok(())
9660 }
9661
9662 pub fn quantize_q8_1_into(
9664 &self,
9665 x: &CudaSlice<f32>,
9666 m: usize,
9667 in_f: usize,
9668 q: &mut CudaSlice<i8>,
9669 d: &mut CudaSlice<f32>,
9670 ) -> Result<(), Box<dyn std::error::Error>> {
9671 let nblk = in_f / 32;
9672 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9673 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9674 let (inf, mi) = (in_f as i32, m as i32);
9675 if Self::pdl_on() && Self::pdl_wb_on() {
9676 use cudarc::driver::{DevicePtr, DevicePtrMut};
9677 let s = &self.gpu.stream();
9678 let (px, _g0) = x.device_ptr(s);
9679 let (pq, _g1) = q.device_ptr_mut(s);
9680 let (pd, _g2) = d.device_ptr_mut(s);
9681 let mut ps = [
9682 &px as *const _ as *mut std::ffi::c_void,
9683 &pq as *const _ as *mut _,
9684 &pd as *const _ as *mut _,
9685 &inf as *const _ as *mut _,
9686 &mi as *const _ as *mut _,
9687 ];
9688 unsafe {
9689 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9690 }
9691 return Ok(());
9692 }
9693 let f = self.func("quantize_q8_1");
9694 let __s_b = self.gpu.stream();
9695 let mut b = __s_b.launch_builder(&f);
9696 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9697 unsafe {
9698 b.launch(cfg)?;
9699 }
9700 Ok(())
9701 }
9702
9703 pub fn add_rms_norm_q8_1(
9707 &self,
9708 a: &CudaSlice<f32>,
9709 b_in: &CudaSlice<f32>,
9710 w: &CudaSlice<f32>,
9711 res: &mut CudaSlice<f32>,
9712 ncols: usize,
9713 nrows: usize,
9714 eps: f32,
9715 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9716 let nblk = ncols / 32;
9717 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9718 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9719 let f = self.func("add_rms_norm_q8_1");
9720 let cfg = LaunchConfig {
9722 grid_dim: (nrows as u32, 1, 1),
9723 block_dim: (1024, 1, 1),
9724 shared_mem_bytes: 0,
9725 };
9726 let (nc, e) = (ncols as i32, eps);
9727 let __s_bld = self.gpu.stream();
9728 let mut bld = __s_bld.launch_builder(&f);
9729 bld.arg(a)
9730 .arg(b_in)
9731 .arg(w)
9732 .arg(res)
9733 .arg(&mut q)
9734 .arg(&mut d)
9735 .arg(&nc)
9736 .arg(&e);
9737 unsafe {
9738 bld.launch(cfg)?;
9739 }
9740 Ok((q, d))
9741 }
9742
9743 #[allow(clippy::too_many_arguments)]
9749 pub fn join_add_rms_norm_raw(
9750 &self,
9751 a0_raw: u64,
9752 a1_raw: u64,
9753 x: &CudaSlice<f32>,
9754 w: &CudaSlice<f32>,
9755 res: &mut CudaSlice<f32>,
9756 dst: &mut CudaSlice<f32>,
9757 ncols: usize,
9758 eps: f32,
9759 ) -> Result<(), Box<dyn std::error::Error>> {
9760 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9761 return Err("join_add_rms_norm geometry".into());
9762 }
9763 let f = self.func("join_add_rms_norm_f32");
9764 let cfg = LaunchConfig {
9765 grid_dim: (1, 1, 1),
9766 block_dim: (rms_block(), 1, 1),
9767 shared_mem_bytes: 0,
9768 };
9769 let (nc, e) = (ncols as i32, eps);
9770 let __s_b = self.gpu.stream();
9771 let mut b = __s_b.launch_builder(&f);
9772 b.arg(&a0_raw)
9773 .arg(&a1_raw)
9774 .arg(x)
9775 .arg(w)
9776 .arg(&mut *res)
9777 .arg(&mut *dst)
9778 .arg(&nc)
9779 .arg(&e);
9780 unsafe {
9781 b.launch(cfg)?;
9782 }
9783 Ok(())
9784 }
9785
9786 pub fn add_rms_norm(
9787 &self,
9788 a: &CudaSlice<f32>,
9789 b: &CudaSlice<f32>,
9790 w: &CudaSlice<f32>,
9791 res: &mut CudaSlice<f32>,
9792 dst: &mut CudaSlice<f32>,
9793 ncols: usize,
9794 nrows: usize,
9795 eps: f32,
9796 ) -> Result<(), Box<dyn std::error::Error>> {
9797 let (nc, e) = (ncols as i32, eps);
9798 let kname = if Self::norm_ilp_on() {
9799 "add_rms_norm_f32_v2"
9800 } else {
9801 "add_rms_norm_f32"
9802 };
9803 if Self::pdl_on() && Self::pdl_wb_on() {
9804 use cudarc::driver::{DevicePtr, DevicePtrMut};
9805 let s = &self.gpu.stream();
9806 let (pa, _g0) = a.device_ptr(s);
9807 let (pb, _g1) = b.device_ptr(s);
9808 let (pw, _g2) = w.device_ptr(s);
9809 let (pr, _g3) = res.device_ptr_mut(s);
9810 let (pd, _g4) = dst.device_ptr_mut(s);
9811 let mut ps = [
9812 &pa as *const _ as *mut std::ffi::c_void,
9813 &pb as *const _ as *mut _,
9814 &pw as *const _ as *mut _,
9815 &pr as *const _ as *mut _,
9816 &pd as *const _ as *mut _,
9817 &nc as *const _ as *mut _,
9818 &e as *const _ as *mut _,
9819 ];
9820 unsafe {
9821 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9822 }
9823 return Ok(());
9824 }
9825 let f = self.func(kname);
9826 let cfg = LaunchConfig {
9827 grid_dim: (nrows as u32, 1, 1),
9828 block_dim: (rms_block(), 1, 1),
9829 shared_mem_bytes: 0,
9830 };
9831 let __s_b2 = self.gpu.stream();
9832 let mut b2 = __s_b2.launch_builder(&f);
9833 b2.arg(a)
9834 .arg(b)
9835 .arg(w)
9836 .arg(&mut *res)
9837 .arg(&mut *dst)
9838 .arg(&nc)
9839 .arg(&e);
9840 unsafe {
9841 b2.launch(cfg)?;
9842 }
9843 Ok(())
9844 }
9845
9846 #[allow(clippy::too_many_arguments)]
9849 pub fn rms_pre_add_rms_norm(
9850 &self,
9851 a: &CudaSlice<f32>,
9852 wa: &CudaSlice<f32>,
9853 b: &CudaSlice<f32>,
9854 w: &CudaSlice<f32>,
9855 res: &mut CudaSlice<f32>,
9856 dst: &mut CudaSlice<f32>,
9857 ncols: usize,
9858 nrows: usize,
9859 eps: f32,
9860 ) -> Result<(), Box<dyn std::error::Error>> {
9861 let f = self.func("rms_pre_add_rms_norm_f32");
9862 let cfg = LaunchConfig {
9863 grid_dim: (nrows as u32, 1, 1),
9864 block_dim: (rms_block(), 1, 1),
9865 shared_mem_bytes: 0,
9866 };
9867 let (nc, e) = (ncols as i32, eps);
9868 let __s_b2 = self.gpu.stream();
9869 let mut b2 = __s_b2.launch_builder(&f);
9870 b2.arg(a)
9871 .arg(wa)
9872 .arg(b)
9873 .arg(w)
9874 .arg(&mut *res)
9875 .arg(&mut *dst)
9876 .arg(&nc)
9877 .arg(&e);
9878 unsafe {
9879 b2.launch(cfg)?;
9880 }
9881 Ok(())
9882 }
9883
9884 #[allow(clippy::too_many_arguments)]
9886 pub fn rms_pre_add_rms_norm_q8z(
9887 &self,
9888 a: &CudaSlice<f32>,
9889 wa: &CudaSlice<f32>,
9890 b: &CudaSlice<f32>,
9891 w: &CudaSlice<f32>,
9892 res: &mut CudaSlice<f32>,
9893 dst: &mut CudaSlice<f32>,
9894 ncols: usize,
9895 nrows: usize,
9896 eps: f32,
9897 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9898 debug_assert!(ncols % 128 == 0);
9899 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9900 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9901 let (nc, e) = (ncols as i32, eps);
9902 if Self::pdl_on() {
9903 {
9904 use cudarc::driver::{DevicePtr, DevicePtrMut};
9905 let s = &self.gpu.stream();
9906 let (pa, _g0) = a.device_ptr(s);
9907 let (pwa, _g1) = wa.device_ptr(s);
9908 let (pb, _g2) = b.device_ptr(s);
9909 let (pw, _g3) = w.device_ptr(s);
9910 let (pr, _g4) = res.device_ptr_mut(s);
9911 let (pdst, _g5) = dst.device_ptr_mut(s);
9912 let (pq, _g6) = out_q.device_ptr_mut(s);
9913 let (pd, _g7) = out_d.device_ptr_mut(s);
9914 let mut ps = [
9915 &pa as *const _ as *mut std::ffi::c_void,
9916 &pwa as *const _ as *mut _,
9917 &pb as *const _ as *mut _,
9918 &pw as *const _ as *mut _,
9919 &pr as *const _ as *mut _,
9920 &pdst as *const _ as *mut _,
9921 &pq as *const _ as *mut _,
9922 &pd as *const _ as *mut _,
9923 &nc as *const _ as *mut _,
9924 &e as *const _ as *mut _,
9925 ];
9926 unsafe {
9927 self.launch_pdl(
9928 "rms_pre_add_rms_norm_q8z_f32",
9929 (nrows as u32, 1, 1),
9930 (rms_block(), 1, 1),
9931 &mut ps,
9932 )?;
9933 }
9934 }
9935 return Ok((out_q, out_d));
9936 }
9937 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9938 let cfg = LaunchConfig {
9939 grid_dim: (nrows as u32, 1, 1),
9940 block_dim: (rms_block(), 1, 1),
9941 shared_mem_bytes: 0,
9942 };
9943 let __s_b2 = self.gpu.stream();
9944 let mut b2 = __s_b2.launch_builder(&f);
9945 b2.arg(a)
9946 .arg(wa)
9947 .arg(b)
9948 .arg(w)
9949 .arg(&mut *res)
9950 .arg(&mut *dst)
9951 .arg(&mut out_q)
9952 .arg(&mut out_d)
9953 .arg(&nc)
9954 .arg(&e);
9955 unsafe {
9956 b2.launch(cfg)?;
9957 }
9958 Ok((out_q, out_d))
9959 }
9960
9961 #[allow(clippy::too_many_arguments)]
9965 pub fn rms_pre_add_rms_norm_q8z_into(
9966 &self,
9967 a: &CudaSlice<f32>,
9968 wa: &CudaSlice<f32>,
9969 b: &CudaSlice<f32>,
9970 w: &CudaSlice<f32>,
9971 res: &mut CudaSlice<f32>,
9972 dst: &mut CudaSlice<f32>,
9973 ncols: usize,
9974 nrows: usize,
9975 eps: f32,
9976 out_q: &mut CudaSlice<i8>,
9977 out_d: &mut CudaSlice<f32>,
9978 ) -> Result<(), Box<dyn std::error::Error>> {
9979 debug_assert!(ncols % 128 == 0);
9980 let (nc, e) = (ncols as i32, eps);
9981 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9982 let cfg = LaunchConfig {
9983 grid_dim: (nrows as u32, 1, 1),
9984 block_dim: (rms_block(), 1, 1),
9985 shared_mem_bytes: 0,
9986 };
9987 let __s_b = self.gpu.stream();
9988 let mut b2 = __s_b.launch_builder(&f);
9989 b2.arg(a)
9990 .arg(wa)
9991 .arg(b)
9992 .arg(w)
9993 .arg(&mut *res)
9994 .arg(&mut *dst)
9995 .arg(&mut *out_q)
9996 .arg(&mut *out_d)
9997 .arg(&nc)
9998 .arg(&e);
9999 unsafe {
10000 b2.launch(cfg)?;
10001 }
10002 Ok(())
10003 }
10004
10005 #[allow(clippy::too_many_arguments)]
10008 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10009 &self,
10010 a: &CudaSlice<f32>,
10011 wa: &CudaSlice<f32>,
10012 b_in: &CudaSlice<f32>,
10013 c: f32,
10014 w: &CudaSlice<f32>,
10015 res: &mut CudaSlice<f32>,
10016 ncols: usize,
10017 nrows: usize,
10018 eps: f32,
10019 out_q: &mut CudaSlice<i8>,
10020 out_d: &mut CudaSlice<f32>,
10021 ) -> Result<(), Box<dyn std::error::Error>> {
10022 debug_assert!(ncols % 128 == 0);
10023 let (nc, e2) = (ncols as i32, eps);
10024 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10025 let cfg = LaunchConfig {
10026 grid_dim: (nrows as u32, 1, 1),
10027 block_dim: (rms_block(), 1, 1),
10028 shared_mem_bytes: 0,
10029 };
10030 let __s_b = self.gpu.stream();
10031 let mut b2 = __s_b.launch_builder(&f);
10032 b2.arg(a)
10033 .arg(wa)
10034 .arg(b_in)
10035 .arg(&c)
10036 .arg(w)
10037 .arg(&mut *res)
10038 .arg(&mut *out_q)
10039 .arg(&mut *out_d)
10040 .arg(&nc)
10041 .arg(&e2);
10042 unsafe {
10043 b2.launch(cfg)?;
10044 }
10045 Ok(())
10046 }
10047
10048 pub fn g4_pnfold_on() -> bool {
10056 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10057 *ON.get_or_init(|| {
10058 std::env::var("MEMRA_G4_PNFOLD")
10059 .map(|v| v != "0")
10060 .unwrap_or(true)
10061 })
10062 }
10063
10064 pub fn build_q4_out_concat3(
10068 &self,
10069 w0: &crate::model::GpuTensor,
10070 w1: &crate::model::GpuTensor,
10071 w2: &crate::model::GpuTensor,
10072 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10073 use crate::model::GpuTensor;
10074 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10075 match w {
10076 GpuTensor::Quant {
10077 qtype,
10078 row_bytes,
10079 rp,
10080 ..
10081 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10082 _ => None,
10083 }
10084 };
10085 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10086 else {
10087 return Ok(None);
10088 };
10089 if rb0 != rb1
10090 || rb0 != rb2
10091 || w0.in_features() != w1.in_features()
10092 || w0.in_features() != w2.in_features()
10093 {
10094 return Ok(None);
10095 }
10096 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10097 match w {
10098 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10099 _ => unreachable!(),
10100 }
10101 }
10102 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10103 let total = rb0 * (o0 + o1 + o2);
10104 let mut cat = self.alloc_u8(total)?;
10105 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10106 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10107 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10108 Ok(Some(GpuTensor::Quant {
10109 bytes: cat,
10110 qtype: QT_Q4_0,
10111 row_bytes: rb0,
10112 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10113 scale: 1.0,
10114 rp: false,
10115 #[cfg(memra_cutlass)]
10116 cutlass: None,
10117 fp8: None,
10118 blk: None,
10119 rp4: None,
10120 f16: None,
10121 }))
10122 }
10123
10124 fn full_width_rope_only(
10142 kernel: &str,
10143 n_rot: usize,
10144 head_dim: usize,
10145 ) -> Result<(), Box<dyn std::error::Error>> {
10146 if n_rot == head_dim {
10147 return Ok(());
10148 }
10149 Err(format!(
10150 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10151 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10152 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10153 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10154 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10155 )
10156 .into())
10157 }
10158
10159 #[allow(clippy::too_many_arguments)]
10163 pub fn rms_norm_qkv_rope_cat(
10164 &self,
10165 qkv: &CudaSlice<f32>,
10166 wq: &CudaSlice<f32>,
10167 wk: &CudaSlice<f32>,
10168 wv: &CudaSlice<f32>,
10169 q: &mut CudaSlice<f32>,
10170 k: &mut CudaSlice<f32>,
10171 v: &mut CudaSlice<f32>,
10172 head_dim: usize,
10173 n_rot: usize,
10174 rq: usize,
10175 rk: usize,
10176 pos: &CudaSlice<i32>,
10177 nh_q: usize,
10178 nh_k: usize,
10179 base: f32,
10180 freq_scale: f32,
10181 ff: Option<&CudaSlice<f32>>,
10182 eps: f32,
10183 ) -> Result<(), Box<dyn std::error::Error>> {
10184 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10185 let rows = rq + rk + rk;
10186 let theta_scale = base.powf(-2.0 / head_dim as f32);
10187 let (nc, rqi, rki, nhq, nhk) = (
10188 head_dim as i32,
10189 rq as i32,
10190 rk as i32,
10191 nh_q as i32,
10192 nh_k as i32,
10193 );
10194 if Self::pdl_on() {
10195 use cudarc::driver::{DevicePtr, DevicePtrMut};
10196 let s = &self.gpu.stream();
10197 let (pqkv, _g0) = qkv.device_ptr(s);
10198 let (pwq, _g1) = wq.device_ptr(s);
10199 let (pwk, _g2) = wk.device_ptr(s);
10200 let (pwv, _g3) = wv.device_ptr(s);
10201 let (pq, _g4) = q.device_ptr_mut(s);
10202 let (pk, _g5) = k.device_ptr_mut(s);
10203 let (pv, _g6) = v.device_ptr_mut(s);
10204 let (ppos, _g7) = pos.device_ptr(s);
10205 let (pff, _g8) = match ff {
10206 Some(t) => {
10207 let (p, g) = t.device_ptr(s);
10208 (p, Some(g))
10209 }
10210 None => (0, None),
10211 };
10212 let mut ps = [
10213 &pqkv as *const _ as *mut std::ffi::c_void,
10214 &pwq as *const _ as *mut _,
10215 &pwk as *const _ as *mut _,
10216 &pwv as *const _ as *mut _,
10217 &pq as *const _ as *mut _,
10218 &pk as *const _ as *mut _,
10219 &pv as *const _ as *mut _,
10220 &nc as *const _ as *mut _,
10221 &rqi as *const _ as *mut _,
10222 &rki as *const _ as *mut _,
10223 &ppos as *const _ as *mut _,
10224 &nhq as *const _ as *mut _,
10225 &nhk as *const _ as *mut _,
10226 &theta_scale as *const _ as *mut _,
10227 &freq_scale as *const _ as *mut _,
10228 &pff as *const _ as *mut _,
10229 &eps as *const _ as *mut _,
10230 ];
10231 unsafe {
10232 self.launch_pdl(
10233 "rms_norm_qkv_rope_cat_f32",
10234 (rows as u32, 1, 1),
10235 (rms_block(), 1, 1),
10236 &mut ps,
10237 )?;
10238 }
10239 return Ok(());
10240 }
10241 let f = self.func("rms_norm_qkv_rope_cat_f32");
10242 let cfg = LaunchConfig {
10243 grid_dim: (rows as u32, 1, 1),
10244 block_dim: (rms_block(), 1, 1),
10245 shared_mem_bytes: 0,
10246 };
10247 let __s_b = self.gpu.stream();
10248 let mut b = __s_b.launch_builder(&f);
10249 match ff {
10250 Some(t) => {
10251 b.arg(qkv)
10252 .arg(wq)
10253 .arg(wk)
10254 .arg(wv)
10255 .arg(&mut *q)
10256 .arg(&mut *k)
10257 .arg(&mut *v)
10258 .arg(&nc)
10259 .arg(&rqi)
10260 .arg(&rki)
10261 .arg(pos)
10262 .arg(&nhq)
10263 .arg(&nhk)
10264 .arg(&theta_scale)
10265 .arg(&freq_scale)
10266 .arg(t)
10267 .arg(&eps);
10268 unsafe {
10269 b.launch(cfg)?;
10270 }
10271 }
10272 None => {
10273 let null: u64 = 0;
10274 b.arg(qkv)
10275 .arg(wq)
10276 .arg(wk)
10277 .arg(wv)
10278 .arg(&mut *q)
10279 .arg(&mut *k)
10280 .arg(&mut *v)
10281 .arg(&nc)
10282 .arg(&rqi)
10283 .arg(&rki)
10284 .arg(pos)
10285 .arg(&nhq)
10286 .arg(&nhk)
10287 .arg(&theta_scale)
10288 .arg(&freq_scale)
10289 .arg(&null)
10290 .arg(&eps);
10291 unsafe {
10292 b.launch(cfg)?;
10293 }
10294 }
10295 }
10296 Ok(())
10297 }
10298
10299 #[allow(clippy::too_many_arguments)]
10303 pub fn rms_norm_qkv_rope(
10304 &self,
10305 q0: &CudaSlice<f32>,
10306 k0: &CudaSlice<f32>,
10307 v0: &CudaSlice<f32>,
10308 wq: &CudaSlice<f32>,
10309 wk: &CudaSlice<f32>,
10310 wv: &CudaSlice<f32>,
10311 q: &mut CudaSlice<f32>,
10312 k: &mut CudaSlice<f32>,
10313 v: &mut CudaSlice<f32>,
10314 head_dim: usize,
10315 n_rot: usize,
10316 rq: usize,
10317 rk: usize,
10318 pos: &CudaSlice<i32>,
10319 nh_q: usize,
10320 nh_k: usize,
10321 base: f32,
10322 freq_scale: f32,
10323 ff: Option<&CudaSlice<f32>>,
10324 eps: f32,
10325 ) -> Result<(), Box<dyn std::error::Error>> {
10326 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10327 let f = self.func("rms_norm_qkv_rope_f32");
10328 let rows = rq + rk + rk; let cfg = LaunchConfig {
10330 grid_dim: (rows as u32, 1, 1),
10331 block_dim: (rms_block(), 1, 1),
10332 shared_mem_bytes: 0,
10333 };
10334 let theta_scale = base.powf(-2.0 / head_dim as f32);
10335 let (nc, rqi, rki, nhq, nhk) = (
10336 head_dim as i32,
10337 rq as i32,
10338 rk as i32,
10339 nh_q as i32,
10340 nh_k as i32,
10341 );
10342 let __s_b = self.gpu.stream();
10343 let mut b = __s_b.launch_builder(&f);
10344 match ff {
10345 Some(t) => {
10346 b.arg(q0)
10347 .arg(k0)
10348 .arg(v0)
10349 .arg(wq)
10350 .arg(wk)
10351 .arg(wv)
10352 .arg(&mut *q)
10353 .arg(&mut *k)
10354 .arg(&mut *v)
10355 .arg(&nc)
10356 .arg(&rqi)
10357 .arg(&rki)
10358 .arg(pos)
10359 .arg(&nhq)
10360 .arg(&nhk)
10361 .arg(&theta_scale)
10362 .arg(&freq_scale)
10363 .arg(t)
10364 .arg(&eps);
10365 unsafe {
10366 b.launch(cfg)?;
10367 }
10368 }
10369 None => {
10370 let null: u64 = 0;
10371 b.arg(q0)
10372 .arg(k0)
10373 .arg(v0)
10374 .arg(wq)
10375 .arg(wk)
10376 .arg(wv)
10377 .arg(&mut *q)
10378 .arg(&mut *k)
10379 .arg(&mut *v)
10380 .arg(&nc)
10381 .arg(&rqi)
10382 .arg(&rki)
10383 .arg(pos)
10384 .arg(&nhq)
10385 .arg(&nhk)
10386 .arg(&theta_scale)
10387 .arg(&freq_scale)
10388 .arg(&null)
10389 .arg(&eps);
10390 unsafe {
10391 b.launch(cfg)?;
10392 }
10393 }
10394 }
10395 Ok(())
10396 }
10397
10398 #[allow(clippy::too_many_arguments)]
10404 pub fn rms_norm_qkv_rope_append_dc(
10405 &self,
10406 q0: &CudaSlice<f32>,
10407 k0: &CudaSlice<f32>,
10408 v0: &CudaSlice<f32>,
10409 wq: &CudaSlice<f32>,
10410 wk: &CudaSlice<f32>,
10411 wv: &CudaSlice<f32>,
10412 q: &mut CudaSlice<f32>,
10413 k: &mut CudaSlice<f32>,
10414 v: &mut CudaSlice<f32>,
10415 head_dim: usize,
10416 n_rot: usize,
10417 rq: usize,
10418 rk: usize,
10419 pos: &CudaSlice<i32>,
10420 nh_q: usize,
10421 nh_k: usize,
10422 base: f32,
10423 freq_scale: f32,
10424 ff: Option<&CudaSlice<f32>>,
10425 eps: f32,
10426 kc: &mut CudaSlice<u8>,
10427 vc: &mut CudaSlice<u8>,
10428 t_dev: &CudaSlice<i32>,
10429 k_tok_bytes: usize,
10430 v_tok_bytes: usize,
10431 g: bool,
10432 ) -> Result<(), Box<dyn std::error::Error>> {
10433 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10434 let rows = rq + rk + rk;
10435 let theta_scale = base.powf(-2.0 / head_dim as f32);
10436 let (nc, rqi, rki, nhq, nhk) = (
10437 head_dim as i32,
10438 rq as i32,
10439 rk as i32,
10440 nh_q as i32,
10441 nh_k as i32,
10442 );
10443 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10444 if Self::pdl_on() && Self::pdl_wb_on() {
10445 use cudarc::driver::{DevicePtr, DevicePtrMut};
10446 let s = &self.gpu.stream();
10447 let (p0, _a0) = q0.device_ptr(s);
10448 let (p1, _a1) = k0.device_ptr(s);
10449 let (p2, _a2) = v0.device_ptr(s);
10450 let (pwq, _a3) = wq.device_ptr(s);
10451 let (pwk, _a4) = wk.device_ptr(s);
10452 let (pwv, _a5) = wv.device_ptr(s);
10453 let (pq, _a6) = q.device_ptr_mut(s);
10454 let (pk, _a7) = k.device_ptr_mut(s);
10455 let (pv, _a8) = v.device_ptr_mut(s);
10456 let (pp, _a9) = pos.device_ptr(s);
10457 let pff: u64 = match ff {
10458 Some(t) => {
10459 let (p, _gg) = t.device_ptr(s);
10460 p as u64
10461 }
10462 None => 0,
10463 };
10464 let (pkc, _a10) = kc.device_ptr_mut(s);
10465 let (pvc, _a11) = vc.device_ptr_mut(s);
10466 let (pt, _a12) = t_dev.device_ptr(s);
10467 let mut ps = [
10468 &p0 as *const _ as *mut std::ffi::c_void,
10469 &p1 as *const _ as *mut _,
10470 &p2 as *const _ as *mut _,
10471 &pwq as *const _ as *mut _,
10472 &pwk as *const _ as *mut _,
10473 &pwv as *const _ as *mut _,
10474 &pq as *const _ as *mut _,
10475 &pk as *const _ as *mut _,
10476 &pv as *const _ as *mut _,
10477 &nc as *const _ as *mut _,
10478 &rqi as *const _ as *mut _,
10479 &rki as *const _ as *mut _,
10480 &pp as *const _ as *mut _,
10481 &nhq as *const _ as *mut _,
10482 &nhk as *const _ as *mut _,
10483 &theta_scale as *const _ as *mut _,
10484 &freq_scale as *const _ as *mut _,
10485 &pff as *const _ as *mut _,
10486 &eps as *const _ as *mut _,
10487 &pkc as *const _ as *mut _,
10488 &pvc as *const _ as *mut _,
10489 &pt as *const _ as *mut _,
10490 &ktb as *const _ as *mut _,
10491 &vtb as *const _ as *mut _,
10492 ];
10493 unsafe {
10494 self.launch_pdl_flash(
10495 g,
10496 "rms_norm_qkv_rope_append_dc_f32",
10497 (rows as u32, 1, 1),
10498 (rms_block(), 1, 1),
10499 0,
10500 &mut ps,
10501 )?;
10502 }
10503 return Ok(());
10504 }
10505 let f = if g {
10506 self.func_g("rms_norm_qkv_rope_append_dc_f32")
10507 } else {
10508 self.func("rms_norm_qkv_rope_append_dc_f32")
10509 };
10510 let cfg = LaunchConfig {
10511 grid_dim: (rows as u32, 1, 1),
10512 block_dim: (rms_block(), 1, 1),
10513 shared_mem_bytes: 0,
10514 };
10515 let __s_b = self.gpu.stream();
10516 let mut b = __s_b.launch_builder(&f);
10517 match ff {
10518 Some(t) => {
10519 b.arg(q0)
10520 .arg(k0)
10521 .arg(v0)
10522 .arg(wq)
10523 .arg(wk)
10524 .arg(wv)
10525 .arg(&mut *q)
10526 .arg(&mut *k)
10527 .arg(&mut *v)
10528 .arg(&nc)
10529 .arg(&rqi)
10530 .arg(&rki)
10531 .arg(pos)
10532 .arg(&nhq)
10533 .arg(&nhk)
10534 .arg(&theta_scale)
10535 .arg(&freq_scale)
10536 .arg(t)
10537 .arg(&eps)
10538 .arg(&mut *kc)
10539 .arg(&mut *vc)
10540 .arg(t_dev)
10541 .arg(&ktb)
10542 .arg(&vtb);
10543 unsafe {
10544 b.launch(cfg)?;
10545 }
10546 }
10547 None => {
10548 let null: u64 = 0;
10549 b.arg(q0)
10550 .arg(k0)
10551 .arg(v0)
10552 .arg(wq)
10553 .arg(wk)
10554 .arg(wv)
10555 .arg(&mut *q)
10556 .arg(&mut *k)
10557 .arg(&mut *v)
10558 .arg(&nc)
10559 .arg(&rqi)
10560 .arg(&rki)
10561 .arg(pos)
10562 .arg(&nhq)
10563 .arg(&nhk)
10564 .arg(&theta_scale)
10565 .arg(&freq_scale)
10566 .arg(&null)
10567 .arg(&eps)
10568 .arg(&mut *kc)
10569 .arg(&mut *vc)
10570 .arg(t_dev)
10571 .arg(&ktb)
10572 .arg(&vtb);
10573 unsafe {
10574 b.launch(cfg)?;
10575 }
10576 }
10577 }
10578 Ok(())
10579 }
10580
10581 #[allow(clippy::too_many_arguments)]
10589 pub fn rms_norm_qkv_rope_append(
10590 &self,
10591 q0: &CudaSlice<f32>,
10592 k0: &CudaSlice<f32>,
10593 v0: &CudaSlice<f32>,
10594 wq: &CudaSlice<f32>,
10595 wk: &CudaSlice<f32>,
10596 wv: &CudaSlice<f32>,
10597 q: &mut CudaSlice<f32>,
10598 k: &mut CudaSlice<f32>,
10599 v: &mut CudaSlice<f32>,
10600 head_dim: usize,
10601 n_rot: usize,
10602 rq: usize,
10603 rk: usize,
10604 pos: &CudaSlice<i32>,
10605 nh_q: usize,
10606 nh_k: usize,
10607 base: f32,
10608 freq_scale: f32,
10609 ff: Option<&CudaSlice<f32>>,
10610 eps: f32,
10611 kc: &mut CudaSlice<u8>,
10612 vc: &mut CudaSlice<u8>,
10613 t: usize,
10614 k_tok_bytes: usize,
10615 v_tok_bytes: usize,
10616 g: bool,
10617 ) -> Result<(), Box<dyn std::error::Error>> {
10618 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10619 let rows = rq + rk + rk;
10620 let theta_scale = base.powf(-2.0 / head_dim as f32);
10621 let (nc, rqi, rki, nhq, nhk) = (
10622 head_dim as i32,
10623 rq as i32,
10624 rk as i32,
10625 nh_q as i32,
10626 nh_k as i32,
10627 );
10628 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10629 let ti = t as i32;
10630 if Self::pdl_on() && Self::pdl_wb_on() {
10631 use cudarc::driver::{DevicePtr, DevicePtrMut};
10632 let s = &self.gpu.stream();
10633 let (p0, _a0) = q0.device_ptr(s);
10634 let (p1, _a1) = k0.device_ptr(s);
10635 let (p2, _a2) = v0.device_ptr(s);
10636 let (pwq, _a3) = wq.device_ptr(s);
10637 let (pwk, _a4) = wk.device_ptr(s);
10638 let (pwv, _a5) = wv.device_ptr(s);
10639 let (pq, _a6) = q.device_ptr_mut(s);
10640 let (pk, _a7) = k.device_ptr_mut(s);
10641 let (pv, _a8) = v.device_ptr_mut(s);
10642 let (pp, _a9) = pos.device_ptr(s);
10643 let pff: u64 = match ff {
10644 Some(t) => {
10645 let (p, _gg) = t.device_ptr(s);
10646 p as u64
10647 }
10648 None => 0,
10649 };
10650 let (pkc, _a10) = kc.device_ptr_mut(s);
10651 let (pvc, _a11) = vc.device_ptr_mut(s);
10652 let mut ps = [
10653 &p0 as *const _ as *mut std::ffi::c_void,
10654 &p1 as *const _ as *mut _,
10655 &p2 as *const _ as *mut _,
10656 &pwq as *const _ as *mut _,
10657 &pwk as *const _ as *mut _,
10658 &pwv as *const _ as *mut _,
10659 &pq as *const _ as *mut _,
10660 &pk as *const _ as *mut _,
10661 &pv as *const _ as *mut _,
10662 &nc as *const _ as *mut _,
10663 &rqi as *const _ as *mut _,
10664 &rki as *const _ as *mut _,
10665 &pp as *const _ as *mut _,
10666 &nhq as *const _ as *mut _,
10667 &nhk as *const _ as *mut _,
10668 &theta_scale as *const _ as *mut _,
10669 &freq_scale as *const _ as *mut _,
10670 &pff as *const _ as *mut _,
10671 &eps as *const _ as *mut _,
10672 &pkc as *const _ as *mut _,
10673 &pvc as *const _ as *mut _,
10674 &ti as *const _ as *mut _,
10675 &ktb as *const _ as *mut _,
10676 &vtb as *const _ as *mut _,
10677 ];
10678 unsafe {
10679 self.launch_pdl_flash(
10680 g,
10681 "rms_norm_qkv_rope_append_f32",
10682 (rows as u32, 1, 1),
10683 (rms_block(), 1, 1),
10684 0,
10685 &mut ps,
10686 )?;
10687 }
10688 return Ok(());
10689 }
10690 let f = if g {
10691 self.func_g("rms_norm_qkv_rope_append_f32")
10692 } else {
10693 self.func("rms_norm_qkv_rope_append_f32")
10694 };
10695 let cfg = LaunchConfig {
10696 grid_dim: (rows as u32, 1, 1),
10697 block_dim: (rms_block(), 1, 1),
10698 shared_mem_bytes: 0,
10699 };
10700 let __s_b = self.gpu.stream();
10701 let mut b = __s_b.launch_builder(&f);
10702 let null: u64 = 0;
10703 b.arg(q0)
10704 .arg(k0)
10705 .arg(v0)
10706 .arg(wq)
10707 .arg(wk)
10708 .arg(wv)
10709 .arg(&mut *q)
10710 .arg(&mut *k)
10711 .arg(&mut *v)
10712 .arg(&nc)
10713 .arg(&rqi)
10714 .arg(&rki)
10715 .arg(pos)
10716 .arg(&nhq)
10717 .arg(&nhk)
10718 .arg(&theta_scale)
10719 .arg(&freq_scale);
10720 match ff {
10721 Some(t) => {
10722 b.arg(t);
10723 }
10724 None => {
10725 b.arg(&null);
10726 }
10727 }
10728 b.arg(&eps)
10729 .arg(&mut *kc)
10730 .arg(&mut *vc)
10731 .arg(&ti)
10732 .arg(&ktb)
10733 .arg(&vtb);
10734 unsafe {
10735 b.launch(cfg)?;
10736 }
10737 Ok(())
10738 }
10739
10740 pub fn add_q8_1(
10741 &self,
10742 a: &CudaSlice<f32>,
10743 b: &CudaSlice<f32>,
10744 res: &mut CudaSlice<f32>,
10745 ncols: usize,
10746 nrows: usize,
10747 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10748 debug_assert!(ncols % 128 == 0);
10749 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10750 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10751 let f = self.func("add_q8_1_f32");
10752 let cfg = LaunchConfig {
10753 grid_dim: (nrows as u32, 1, 1),
10754 block_dim: (rms_block(), 1, 1),
10755 shared_mem_bytes: 0,
10756 };
10757 let nc = ncols as i32;
10758 let __s_b2 = self.gpu.stream();
10759 let mut b2 = __s_b2.launch_builder(&f);
10760 b2.arg(a)
10761 .arg(b)
10762 .arg(&mut *res)
10763 .arg(&mut out_q)
10764 .arg(&mut out_d)
10765 .arg(&nc);
10766 unsafe {
10767 b2.launch(cfg)?;
10768 }
10769 Ok((out_q, out_d))
10770 }
10771
10772 pub fn rms_pre_add_q8_1(
10776 &self,
10777 a: &CudaSlice<f32>,
10778 wa: &CudaSlice<f32>,
10779 b: &CudaSlice<f32>,
10780 res: &mut CudaSlice<f32>,
10781 ncols: usize,
10782 nrows: usize,
10783 eps: f32,
10784 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10785 debug_assert!(ncols % 128 == 0);
10786 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10787 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10788 let f = self.func("rms_pre_add_q8_1_f32");
10789 let cfg = LaunchConfig {
10790 grid_dim: (nrows as u32, 1, 1),
10791 block_dim: (rms_block(), 1, 1),
10792 shared_mem_bytes: 0,
10793 };
10794 let (nc, ep) = (ncols as i32, eps);
10795 let __s_b2 = self.gpu.stream();
10796 let mut b2 = __s_b2.launch_builder(&f);
10797 b2.arg(a)
10798 .arg(wa)
10799 .arg(b)
10800 .arg(&mut *res)
10801 .arg(&mut out_q)
10802 .arg(&mut out_d)
10803 .arg(&nc)
10804 .arg(&ep);
10805 unsafe {
10806 b2.launch(cfg)?;
10807 }
10808 Ok((out_q, out_d))
10809 }
10810
10811 pub fn l2_v2_on(ncols: usize) -> bool {
10815 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10816 }
10817
10818 pub fn l2_norm_pp(
10819 &self,
10820 x: &CudaSlice<f32>,
10821 dst: &mut CudaSlice<f32>,
10822 dst16: Option<&mut CudaSlice<u8>>,
10823 ncols: usize,
10824 nrows: usize,
10825 eps: f32,
10826 ) -> Result<(), Box<dyn std::error::Error>> {
10827 if Self::l2_v2_on(ncols) {
10828 let f = self.func("l2_norm_pp_v2_f32");
10829 let rows_per_block = 8u32; let cfg = LaunchConfig {
10831 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10832 block_dim: (256, 1, 1),
10833 shared_mem_bytes: 0,
10834 };
10835 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10836 let d16: u64 = match dst16 {
10838 Some(d) => self.addr_u8(d),
10839 None => 0,
10840 };
10841 let __s_b = self.gpu.stream();
10842 let mut b = __s_b.launch_builder(&f);
10843 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
10844 unsafe {
10845 b.launch(cfg)?;
10846 }
10847 return Ok(());
10848 }
10849 self.l2_norm(x, dst, ncols, nrows, eps)
10850 }
10851
10852 pub fn l2_norm(
10853 &self,
10854 x: &CudaSlice<f32>,
10855 dst: &mut CudaSlice<f32>,
10856 ncols: usize,
10857 nrows: usize,
10858 eps: f32,
10859 ) -> Result<(), Box<dyn std::error::Error>> {
10860 let f = self.func("l2_norm_f32");
10861 let cfg = LaunchConfig {
10862 grid_dim: (nrows as u32, 1, 1),
10863 block_dim: (256, 1, 1),
10864 shared_mem_bytes: 0,
10865 };
10866 let (nc, e) = (ncols as i32, eps);
10867 let __s_b = self.gpu.stream();
10868 let mut b = __s_b.launch_builder(&f);
10869 b.arg(x).arg(dst).arg(&nc).arg(&e);
10870 unsafe {
10871 b.launch(cfg)?;
10872 }
10873 Ok(())
10874 }
10875
10876 pub fn l2_norm_decode(
10882 &self,
10883 x: &CudaSlice<f32>,
10884 dst: &mut CudaSlice<f32>,
10885 ncols: usize,
10886 nrows: usize,
10887 eps: f32,
10888 ) -> Result<(), Box<dyn std::error::Error>> {
10889 let f = self.func("l2_norm_f32");
10890 let cfg = LaunchConfig {
10891 grid_dim: (nrows as u32, 1, 1),
10892 block_dim: (32, 1, 1),
10893 shared_mem_bytes: 0,
10894 };
10895 let (nc, e) = (ncols as i32, eps);
10896 let __s_b = self.gpu.stream();
10897 let mut b = __s_b.launch_builder(&f);
10898 b.arg(x).arg(dst).arg(&nc).arg(&e);
10899 unsafe {
10900 b.launch(cfg)?;
10901 }
10902 Ok(())
10903 }
10904
10905 pub fn rope_neox(
10907 &self,
10908 x: &mut CudaSlice<f32>,
10909 pos: &CudaSlice<i32>,
10910 head_dim: usize,
10911 n_dims: usize,
10912 n_heads: usize,
10913 n_tokens: usize,
10914 freq_base: f32,
10915 freq_scale: f32,
10916 ) -> Result<(), Box<dyn std::error::Error>> {
10917 let f = self.func("rope_neox_f32");
10918 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10919 let grid = (n_heads * n_tokens) as u32;
10920 let cfg = LaunchConfig {
10921 grid_dim: (grid, 1, 1),
10922 block_dim: ((head_dim / 2) as u32, 1, 1),
10923 shared_mem_bytes: 0,
10924 };
10925 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10926 let __s_b = self.gpu.stream();
10927 let mut b = __s_b.launch_builder(&f);
10928 b.arg(x)
10929 .arg(pos)
10930 .arg(&hd)
10931 .arg(&nd)
10932 .arg(&nh)
10933 .arg(&theta_scale)
10934 .arg(&freq_scale);
10935 unsafe {
10936 b.launch(cfg)?;
10937 }
10938 Ok(())
10939 }
10940
10941 pub fn rope_neox_ff(
10943 &self,
10944 x: &mut CudaSlice<f32>,
10945 pos: &CudaSlice<i32>,
10946 head_dim: usize,
10947 n_dims: usize,
10948 n_heads: usize,
10949 n_tokens: usize,
10950 freq_base: f32,
10951 freq_scale: f32,
10952 ff: &CudaSlice<f32>,
10953 ) -> Result<(), Box<dyn std::error::Error>> {
10954 let f = self.func("rope_neox_ff_f32");
10955 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10956 let grid = (n_heads * n_tokens) as u32;
10957 let cfg = LaunchConfig {
10958 grid_dim: (grid, 1, 1),
10959 block_dim: ((head_dim / 2) as u32, 1, 1),
10960 shared_mem_bytes: 0,
10961 };
10962 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10963 let __s_b = self.gpu.stream();
10964 let mut b = __s_b.launch_builder(&f);
10965 b.arg(x)
10966 .arg(pos)
10967 .arg(&hd)
10968 .arg(&nd)
10969 .arg(&nh)
10970 .arg(&theta_scale)
10971 .arg(&freq_scale)
10972 .arg(ff);
10973 unsafe {
10974 b.launch(cfg)?;
10975 }
10976 Ok(())
10977 }
10978
10979 #[allow(clippy::too_many_arguments)]
10981 pub fn rope_neox2(
10982 &self,
10983 q: &mut CudaSlice<f32>,
10984 k: &mut CudaSlice<f32>,
10985 pos: &CudaSlice<i32>,
10986 head_dim: usize,
10987 n_dims: usize,
10988 nh_q: usize,
10989 nh_k: usize,
10990 n_tokens: usize,
10991 freq_base: f32,
10992 freq_scale: f32,
10993 ff: Option<&CudaSlice<f32>>,
10994 ) -> Result<(), Box<dyn std::error::Error>> {
10995 let f = self.func("rope_neox2_f32");
10996 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10997 let grid = ((nh_q + nh_k) * n_tokens) as u32;
10998 let cfg = LaunchConfig {
10999 grid_dim: (grid, 1, 1),
11000 block_dim: ((head_dim / 2) as u32, 1, 1),
11001 shared_mem_bytes: 0,
11002 };
11003 let (hd, nd, nq, nk, nt) = (
11004 head_dim as i32,
11005 n_dims as i32,
11006 nh_q as i32,
11007 nh_k as i32,
11008 n_tokens as i32,
11009 );
11010 let __s_b = self.gpu.stream();
11011 let mut b = __s_b.launch_builder(&f);
11012 b.arg(q)
11013 .arg(k)
11014 .arg(pos)
11015 .arg(&hd)
11016 .arg(&nd)
11017 .arg(&nq)
11018 .arg(&nk)
11019 .arg(&nt)
11020 .arg(&theta_scale)
11021 .arg(&freq_scale);
11022 match ff {
11023 Some(ffv) => {
11024 b.arg(ffv);
11025 unsafe {
11026 b.launch(cfg)?;
11027 }
11028 }
11029 None => {
11030 let null: u64 = 0;
11031 b.arg(&null);
11032 unsafe {
11033 b.launch(cfg)?;
11034 }
11035 }
11036 }
11037 Ok(())
11038 }
11039
11040 pub fn gelu_tanh_mul(
11042 &self,
11043 gate: &CudaSlice<f32>,
11044 up: &CudaSlice<f32>,
11045 dst: &mut CudaSlice<f32>,
11046 n: usize,
11047 ) -> Result<(), Box<dyn std::error::Error>> {
11048 let f = self.func("gelu_tanh_mul_f32");
11049 let cfg = LaunchConfig::for_num_elems(n as u32);
11050 let ni = n as i32;
11051 let __s_b = self.gpu.stream();
11052 let mut b = __s_b.launch_builder(&f);
11053 b.arg(gate).arg(up).arg(dst).arg(&ni);
11054 unsafe {
11055 b.launch(cfg)?;
11056 }
11057 Ok(())
11058 }
11059
11060 pub fn silu_mul(
11061 &self,
11062 gate: &CudaSlice<f32>,
11063 up: &CudaSlice<f32>,
11064 dst: &mut CudaSlice<f32>,
11065 n: usize,
11066 ) -> Result<(), Box<dyn std::error::Error>> {
11067 let f = self.func("silu_mul_f32");
11068 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11070 let ni = n as i32;
11071 let __s_b = self.gpu.stream();
11072 let mut b = __s_b.launch_builder(&f);
11073 b.arg(gate).arg(up).arg(dst).arg(&ni);
11074 unsafe {
11075 b.launch(cfg)?;
11076 }
11077 Ok(())
11078 }
11079
11080 pub fn silu_mul_host_expf(
11082 &self,
11083 gate: &CudaSlice<f32>,
11084 up: &CudaSlice<f32>,
11085 dst: &mut CudaSlice<f32>,
11086 n: usize,
11087 ) -> Result<(), Box<dyn std::error::Error>> {
11088 let f = self.func("silu_mul_host_expf_f32");
11089 let cfg = LaunchConfig::for_num_elems(n as u32);
11090 let ni = n as i32;
11091 let __s_b = self.gpu.stream();
11092 let mut b = __s_b.launch_builder(&f);
11093 b.arg(gate).arg(up).arg(dst).arg(&ni);
11094 unsafe {
11095 b.launch(cfg)?;
11096 }
11097 Ok(())
11098 }
11099
11100 pub fn silu_clamped_mul_host_expf(
11102 &self,
11103 gate: &CudaSlice<f32>,
11104 up: &CudaSlice<f32>,
11105 limit: f32,
11106 dst: &mut CudaSlice<f32>,
11107 n: usize,
11108 ) -> Result<(), Box<dyn std::error::Error>> {
11109 if !limit.is_finite() || limit <= 0.0 {
11110 return Err(
11111 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11112 );
11113 }
11114 let f = self.func("silu_clamped_mul_host_expf_f32");
11115 let cfg = LaunchConfig::for_num_elems(n as u32);
11116 let ni = n as i32;
11117 let __s_b = self.gpu.stream();
11118 let mut b = __s_b.launch_builder(&f);
11119 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11120 unsafe {
11121 b.launch(cfg)?;
11122 }
11123 Ok(())
11124 }
11125
11126 pub fn silu_mul_f16out(
11129 &self,
11130 gate: &CudaSlice<f32>,
11131 up: &CudaSlice<f32>,
11132 dst: &mut CudaSlice<f32>,
11133 dst16: &mut CudaSlice<u8>,
11134 n: usize,
11135 ) -> Result<(), Box<dyn std::error::Error>> {
11136 let f = self.func("silu_mul_f16out_f32");
11137 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11138 let ni = n as i32;
11139 let __s_b = self.gpu.stream();
11140 let mut b = __s_b.launch_builder(&f);
11141 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11142 unsafe {
11143 b.launch(cfg)?;
11144 }
11145 Ok(())
11146 }
11147
11148 pub fn silu_mul_scaled(
11155 &self,
11156 gate: &CudaSlice<f32>,
11157 up: &CudaSlice<f32>,
11158 gs: f32,
11159 us: f32,
11160 dst: &mut CudaSlice<f32>,
11161 n: usize,
11162 ) -> Result<(), Box<dyn std::error::Error>> {
11163 let f = self.func("silu_mul_scaled_f32");
11164 let cfg = LaunchConfig::for_num_elems(n as u32);
11165 let ni = n as i32;
11166 let (gsf, usf) = (gs, us);
11167 let __s_b = self.gpu.stream();
11168 let mut b = __s_b.launch_builder(&f);
11169 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11170 unsafe {
11171 b.launch(cfg)?;
11172 }
11173 Ok(())
11174 }
11175
11176 #[allow(clippy::too_many_arguments)]
11180 pub fn swigluoai_mul_scaled(
11181 &self,
11182 gate: &CudaSlice<f32>,
11183 up: &CudaSlice<f32>,
11184 gs: f32,
11185 us: f32,
11186 alpha: f32,
11187 limit: f32,
11188 dst: &mut CudaSlice<f32>,
11189 n: usize,
11190 ) -> Result<(), Box<dyn std::error::Error>> {
11191 let f = self.func("swigluoai_mul_scaled_f32");
11192 let cfg = LaunchConfig::for_num_elems(n as u32);
11193 let ni = n as i32;
11194 let __s_b = self.gpu.stream();
11195 let mut b = __s_b.launch_builder(&f);
11196 b.arg(gate)
11197 .arg(up)
11198 .arg(&gs)
11199 .arg(&us)
11200 .arg(&alpha)
11201 .arg(&limit)
11202 .arg(dst)
11203 .arg(&ni);
11204 unsafe {
11205 b.launch(cfg)?;
11206 }
11207 Ok(())
11208 }
11209
11210 pub fn silu_mul_scaled_q8_1(
11218 &self,
11219 gate: &CudaSlice<f32>,
11220 up: &CudaSlice<f32>,
11221 gs: f32,
11222 us: f32,
11223 n: usize,
11224 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11225 let f = self.func("silu_mul_scaled_q8_1");
11226 let nblk = n / 32;
11227 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);
11231 let (gsf, usf, ni) = (gs, us, n as i32);
11232 let __s_b = self.gpu.stream();
11233 let mut b = __s_b.launch_builder(&f);
11234 b.arg(gate)
11235 .arg(up)
11236 .arg(&gsf)
11237 .arg(&usf)
11238 .arg(&mut aq)
11239 .arg(&mut ad)
11240 .arg(&ni);
11241 unsafe {
11242 b.launch(cfg)?;
11243 }
11244 Ok((aq, ad))
11245 }
11246
11247 pub fn add(
11248 &self,
11249 a: &CudaSlice<f32>,
11250 b_in: &CudaSlice<f32>,
11251 dst: &mut CudaSlice<f32>,
11252 n: usize,
11253 ) -> Result<(), Box<dyn std::error::Error>> {
11254 let f = self.func("add_f32");
11255 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11257 let ni = n as i32;
11258 let __s_bld = self.gpu.stream();
11259 let mut bld = __s_bld.launch_builder(&f);
11260 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11261 unsafe {
11262 bld.launch(cfg)?;
11263 }
11264 Ok(())
11265 }
11266
11267 pub fn mul(
11268 &self,
11269 a: &CudaSlice<f32>,
11270 b_in: &CudaSlice<f32>,
11271 dst: &mut CudaSlice<f32>,
11272 n: usize,
11273 ) -> Result<(), Box<dyn std::error::Error>> {
11274 let f = self.func("mul_f32");
11275 let cfg = LaunchConfig::for_num_elems(n as u32);
11276 let ni = n as i32;
11277 let __s_bld = self.gpu.stream();
11278 let mut bld = __s_bld.launch_builder(&f);
11279 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11280 unsafe {
11281 bld.launch(cfg)?;
11282 }
11283 Ok(())
11284 }
11285
11286 pub fn matmul(
11289 &self,
11290 w: &crate::model::GpuTensor,
11291 x: &CudaSlice<f32>,
11292 m: usize,
11293 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11294 use crate::model::GpuTensor;
11295 let in_f = w.in_features();
11296 let out_f = w.out_features();
11297 #[allow(non_snake_case)]
11305 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11308 usize::MAX
11309 } else {
11310 16usize
11311 };
11312
11313 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11338 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11339 return Ok(y);
11340 }
11341 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11348 return Ok(y);
11349 }
11350 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11353 return Ok(y);
11354 }
11355 }
11356 if let GpuTensor::Quant { qtype, .. } = w {
11371 if *qtype == QT_F8_E4M3_BLK {
11372 if m >= GEMM_M_THRESHOLD {
11373 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11374 return Ok(y);
11375 }
11376 }
11377 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11378 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11379 return Ok(y);
11380 }
11381 }
11382 }
11383 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11384 return self.qmatvec_mmq(w, x, m);
11385 }
11386 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11387 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11388 return self.qmatvec_gemm(w, &aq, &ad, m);
11389 }
11390 if m >= GEMM_M_THRESHOLD {
11393 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11394 return Ok(y);
11395 }
11396 }
11397 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11401 if m == 1 && fast {
11406 if let GpuTensor::Quant {
11407 bytes,
11408 qtype,
11409 row_bytes,
11410 rp,
11411 rp4,
11412 scale,
11413 ..
11414 } = w
11415 {
11416 if self.mmvq_supports(*qtype) {
11417 let (bytes, rp) = match rp4 {
11421 Some(m4) => (m4, true),
11422 None => (bytes, *rp),
11423 };
11424 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11425 return self.qmatvec_mmvq(
11426 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11427 );
11428 }
11429 }
11430 }
11431 if (2..=16).contains(&m)
11447 && fast
11448 && std::env::var("MEMRA_NO_BATCHED").is_err()
11449 && (m <= 4 || Self::b8_enabled())
11450 {
11451 let m_ok = m <= 8
11461 || matches!(w, GpuTensor::Quant { qtype, .. }
11462 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11463 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11464 if m_ok {
11465 if let GpuTensor::Quant {
11466 bytes,
11467 qtype,
11468 row_bytes,
11469 rp,
11470 rp4,
11471 ..
11472 } = w
11473 {
11474 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11475 let (bytes, rp) = match rp4 {
11476 Some(m4) => (m4, true),
11477 None => (bytes, *rp),
11478 };
11479 let mcols = Self::batched_mcols(m);
11480 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11481 let mut y = self.qmatvec_mmvq_batched(
11482 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11483 )?;
11484 if let GpuTensor::Quant { scale, .. } = w {
11485 if *scale != 1.0 {
11486 self.scale_inplace(&mut y, *scale, m * out_f)?;
11487 }
11488 }
11489 return Ok(y);
11490 }
11491 }
11492 }
11493 }
11494 if fast {
11500 if let GpuTensor::Quant {
11501 bytes,
11502 qtype,
11503 row_bytes,
11504 scale,
11505 ..
11506 } = w
11507 {
11508 if *qtype == QT_F8_E4M3 {
11509 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11510 return self.qmatvec_mmvq(
11511 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11512 );
11513 }
11514 }
11515 }
11516 let mut y = match w {
11517 GpuTensor::Quant {
11518 bytes,
11519 qtype,
11520 row_bytes,
11521 ..
11522 } if fast && *qtype == QT_Q8_0 => {
11523 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11524 }
11525 GpuTensor::Quant {
11526 bytes,
11527 qtype,
11528 row_bytes,
11529 ..
11530 } if fast && *qtype == QT_Q4_K => {
11531 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11532 }
11533 GpuTensor::Quant {
11534 bytes,
11535 qtype,
11536 row_bytes,
11537 ..
11538 } if fast && *qtype == QT_Q6_K => {
11539 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11540 }
11541 GpuTensor::Quant {
11542 bytes,
11543 qtype,
11544 row_bytes,
11545 ..
11546 } if fast && *qtype == QT_Q5_K => {
11547 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11548 }
11549 GpuTensor::Quant {
11550 bytes,
11551 qtype,
11552 row_bytes,
11553 ..
11554 } if fast && *qtype == QT_Q3_K => {
11555 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11556 }
11557 GpuTensor::Quant {
11558 bytes,
11559 qtype,
11560 row_bytes,
11561 rp,
11562 ..
11563 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11564 if *rp {
11565 "qmatvec_nvfp4_dp4a_rp"
11566 } else {
11567 "qmatvec_nvfp4_dp4a"
11568 },
11569 &bytes.slice(0..bytes.len()),
11570 x,
11571 m,
11572 in_f,
11573 out_f,
11574 *row_bytes,
11575 )?,
11576 GpuTensor::Quant {
11580 bytes,
11581 qtype,
11582 row_bytes,
11583 ..
11584 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11585 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11586 }
11587 GpuTensor::Quant {
11592 bytes,
11593 qtype,
11594 row_bytes,
11595 rp,
11596 ..
11597 } =>
11598 {
11601 self.qmatvec(
11602 bytes,
11603 x,
11604 m,
11605 in_f,
11606 out_f,
11607 if *rp && *qtype == QT_NVFP4 {
11608 QT_NVFP4_RP
11609 } else {
11610 *qtype
11611 },
11612 *row_bytes,
11613 )?
11614 }
11615 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11616 GpuTensor::FloatBf16 { data, .. } => {
11619 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11620 }
11621 };
11622 if let GpuTensor::Quant { scale, .. } = w {
11624 if *scale != 1.0 {
11625 self.scale_inplace(&mut y, *scale, m * out_f)?;
11626 }
11627 }
11628 Ok(y)
11629 }
11630
11631 pub fn stage_a_raw_needed() -> bool {
11641 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11642 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11643 }
11644
11645 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11648 use crate::model::GpuTensor;
11649 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11650 return false;
11651 }
11652 match w {
11653 GpuTensor::Quant { qtype, .. } => {
11660 matches!(
11661 *qtype,
11662 QT_Q8_0
11663 | QT_Q4_K
11664 | QT_Q6_K
11665 | QT_Q5_K
11666 | QT_Q3_K
11667 | QT_NVFP4
11668 | QT_F8_E4M3
11669 | QT_F8_E4M3_BLK
11670 | QT_Q4_0
11671 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11672 }
11673 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11674 }
11675 }
11676
11677 pub fn matmul_pre(
11682 &self,
11683 w: &crate::model::GpuTensor,
11684 aq: &CudaSlice<i8>,
11685 ad: &CudaSlice<f32>,
11686 x_fallback: &CudaSlice<f32>,
11687 m: usize,
11688 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11689 use crate::model::GpuTensor;
11690 let x_raw_ok = x_fallback.len() >= m * w.in_features();
11696 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11699 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11700 return Ok(y);
11701 }
11702 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11705 return Ok(y);
11706 }
11707 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11709 return Ok(y);
11710 }
11711 }
11712 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11718 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11719 return Ok(y);
11720 }
11721 }
11722 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11723 return Ok(y);
11724 }
11725 if m >= 16
11730 && w.out_features() >= 128
11731 && self.mmq_supports(w)
11732 && !self.verify_exact_on()
11733 && x_raw_ok
11734 {
11735 return self.qmatvec_mmq(w, x_fallback, m);
11736 }
11737 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11740 if let Some(y) =
11741 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11742 {
11743 return Ok(y);
11744 }
11745 }
11746 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11749 return self.qmatvec_gemm(w, aq, ad, m);
11750 }
11751 if !self.uses_q8_1_fast(w) {
11770 if !x_raw_ok {
11771 return Err(format!(
11772 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11773 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11774 dropped the f32, so there is nothing to fall back to — pass the real f32 \
11775 activation (see Engine::rms_norm_decode, which is bit-identical to \
11776 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11777 x_fallback.len(),
11778 m,
11779 w.in_features(),
11780 m * w.in_features()
11781 )
11782 .into());
11783 }
11784 return self.matmul(w, x_fallback, m);
11785 }
11786 let in_f = w.in_features();
11787 let out_f = w.out_features();
11788 let (bytes, qtype, row_bytes, scale, rp) = match w {
11789 GpuTensor::Quant {
11790 bytes,
11791 qtype,
11792 row_bytes,
11793 scale,
11794 rp,
11795 ..
11796 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11797 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11798 };
11799 let (mbytes, mrp) = match w {
11802 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11803 _ => (bytes, rp),
11804 };
11805 if m == 1 && self.mmvq_supports(qtype) {
11809 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11810 }
11811 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11824 && std::env::var("MEMRA_NO_BATCHED").is_err()
11825 && (m <= 4 || Self::b8_enabled())
11826 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11830 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11831 {
11832 let mcols = Self::batched_mcols(m);
11833 return self.qmatvec_mmvq_batched(
11834 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11835 );
11836 }
11837 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
11843 let (b2, r2) = if qtype == QT_Q4_0 {
11844 (mbytes, mrp)
11845 } else {
11846 (bytes, rp)
11847 };
11848 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
11849 }
11850 let name = match qtype {
11851 QT_Q8_0 => "qmatvec_q8_0_dp4a",
11852 QT_Q4_K => "qmatvec_q4_K_dp4a",
11853 QT_Q6_K => "qmatvec_q6_K_dp4a",
11854 QT_Q5_K => "qmatvec_q5_K_dp4a",
11855 QT_Q3_K => "qmatvec_q3_K_dp4a",
11856 QT_NVFP4 => {
11857 if rp {
11858 "qmatvec_nvfp4_dp4a_rp"
11859 } else {
11860 "qmatvec_nvfp4_dp4a"
11861 }
11862 }
11863 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
11864 _ => unreachable!(),
11865 };
11866 let f = self.func(name);
11867 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
11869 grid_dim: (out_f as u32, m as u32, 1),
11870 block_dim: (128, 1, 1),
11871 shared_mem_bytes: 0,
11872 };
11873 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11874 let __s_b = self.gpu.stream();
11875 let mut b = __s_b.launch_builder(&f);
11876 b.arg(bytes)
11877 .arg(aq)
11878 .arg(ad)
11879 .arg(&mut y)
11880 .arg(&inf)
11881 .arg(&outf)
11882 .arg(&mi)
11883 .arg(&rb);
11884 unsafe {
11885 b.launch(cfg)?;
11886 }
11887 if scale != 1.0 {
11888 self.scale_inplace(&mut y, scale, m * out_f)?;
11889 }
11890 Ok(y)
11891 }
11892
11893 pub fn matmul_decode_exact(
11901 &self,
11902 w: &crate::model::GpuTensor,
11903 x: &CudaSlice<f32>,
11904 m: usize,
11905 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11906 use crate::model::GpuTensor;
11907 if let GpuTensor::Float { data, .. } = w {
11915 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
11916 }
11917 if let GpuTensor::FloatBf16 { data, .. } = w {
11920 let (in_f, out_f) = (w.in_features(), w.out_features());
11921 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
11922 }
11923 if !self.uses_q8_1_fast(w) {
11924 return self.matmul(w, x, m);
11925 }
11926 let in_f = w.in_features();
11927 let out_f = w.out_features();
11928 let (bytes, qtype, row_bytes, scale, rp) = match w {
11929 GpuTensor::Quant {
11930 bytes,
11931 qtype,
11932 row_bytes,
11933 scale,
11934 rp,
11935 ..
11936 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11937 _ => return self.matmul(w, x, m),
11938 };
11939 let (bytes, rp) = match w {
11942 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11943 _ => (bytes, rp),
11944 };
11945 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11946 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11950 return Ok(y);
11951 }
11952 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11961 && std::env::var("MEMRA_NO_BATCHED").is_err()
11962 && (m <= 4 || Self::b8_enabled())
11963 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
11966 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
11967 {
11968 let mcols = Self::batched_mcols(m);
11969 return self.qmatvec_mmvq_batched(
11970 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11971 );
11972 }
11973 if self.mmvq_supports(qtype) {
11974 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11977 }
11978 self.matmul_pre(w, &aq, &ad, x, m)
11981 }
11982
11983 pub fn matmul_decode_exact_pre(
11993 &self,
11994 w: &crate::model::GpuTensor,
11995 aq: &CudaSlice<i8>,
11996 ad: &CudaSlice<f32>,
11997 m: usize,
11998 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11999 use crate::model::GpuTensor;
12000 debug_assert!(
12001 self.uses_q8_1_fast(w),
12002 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12003 );
12004 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12006 return Ok(y);
12007 }
12008 let in_f = w.in_features();
12009 let out_f = w.out_features();
12010 let (bytes, qtype, row_bytes, scale, rp) = match w {
12011 GpuTensor::Quant {
12012 bytes,
12013 qtype,
12014 row_bytes,
12015 scale,
12016 rp,
12017 ..
12018 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12019 _ => {
12020 return Err(
12021 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12022 );
12023 }
12024 };
12025 let (bytes, rp) = match w {
12027 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12028 _ => (bytes, rp),
12029 };
12030 if (2..=16).contains(&m)
12032 && self.batched_supports(qtype)
12033 && self.mmvq_supports(qtype)
12034 && std::env::var("MEMRA_NO_BATCHED").is_err()
12035 && (m <= 4 || Self::b8_enabled())
12036 && (m <= 8
12037 || qtype == QT_Q4_0
12038 || qtype == QT_Q6_K
12039 || qtype == QT_F8_E4M3
12040 || qtype == QT_NVFP4
12041 || qtype == QT_Q4_K
12042 || qtype == QT_Q5_K
12043 || qtype == QT_Q8_0)
12044 {
12045 let mcols = Self::batched_mcols(m);
12046 return self.qmatvec_mmvq_batched(
12047 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12048 );
12049 }
12050 if self.mmvq_supports(qtype) {
12051 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12052 }
12053 let x0 = self.zeros(0)?;
12056 self.matmul_pre(w, aq, ad, &x0, m)
12057 }
12058
12059 pub fn matmul_decode_exact_dual_pre(
12068 &self,
12069 w0: &crate::model::GpuTensor,
12070 w1: &crate::model::GpuTensor,
12071 aq: &CudaSlice<i8>,
12072 ad: &CudaSlice<f32>,
12073 m: usize,
12074 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12075 {
12076 use crate::model::GpuTensor;
12077 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12078 let on = *ON.get_or_init(|| {
12079 std::env::var("MEMRA_SPEC_DUAL_T")
12080 .map(|v| v != "0")
12081 .unwrap_or(true)
12082 });
12083 if !on
12084 || !(2..=7).contains(&m)
12085 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12086 || !self.uses_q8_1_fast(w0)
12087 || !self.uses_q8_1_fast(w1)
12088 {
12089 return Ok(None);
12090 }
12091 if !self.mmvq_supports(QT_NVFP4) {
12096 return Ok(None);
12097 }
12098 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12099 if w1.in_features() != in_f || w1.out_features() != out_f {
12100 return Ok(None);
12101 }
12102 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12103 (
12104 GpuTensor::Quant {
12105 bytes: b0,
12106 qtype: q0,
12107 row_bytes: rb0,
12108 scale: s0,
12109 rp: rp0,
12110 rp4: None,
12111 ..
12112 },
12113 GpuTensor::Quant {
12114 bytes: b1,
12115 qtype: q1,
12116 row_bytes: rb1,
12117 scale: s1,
12118 rp: rp1,
12119 rp4: None,
12120 ..
12121 },
12122 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12123 (b0, b1, *rb0, *s0, *s1, *rp0)
12124 }
12125 _ => return Ok(None),
12126 };
12127 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12130 {
12131 return Ok(None);
12132 }
12133 let (y0, y1) =
12134 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12135 Ok(Some(((y0, s0), (y1, s1))))
12136 }
12137
12138 pub fn matmul_decode_exact_group4_pre(
12150 &self,
12151 ws: [&crate::model::GpuTensor; 4],
12152 aq: &CudaSlice<i8>,
12153 ad: &CudaSlice<f32>,
12154 m: usize,
12155 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12156 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12157 let on = *ON.get_or_init(|| {
12158 std::env::var("MEMRA_TK_GDN_GROUP")
12159 .map(|v| v != "0")
12160 .unwrap_or(true)
12161 });
12162 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12163 }
12164
12165 pub fn matmul_decode_exact_group3_pre(
12170 &self,
12171 ws: [&crate::model::GpuTensor; 3],
12172 aq: &CudaSlice<i8>,
12173 ad: &CudaSlice<f32>,
12174 m: usize,
12175 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12176 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12177 let on = *ON.get_or_init(|| {
12178 std::env::var("MEMRA_TK_FA_GROUP")
12179 .map(|v| v != "0")
12180 .unwrap_or(true)
12181 });
12182 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12183 }
12184
12185 fn matmul_decode_exact_group_pre(
12189 &self,
12190 ws: &[&crate::model::GpuTensor],
12191 aq: &CudaSlice<i8>,
12192 ad: &CudaSlice<f32>,
12193 m: usize,
12194 on: bool,
12195 tag: &'static str,
12196 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12197 use crate::model::GpuTensor;
12198 if !on
12199 || !(2..=16).contains(&m)
12200 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12201 || (m > 4 && !Self::b8_enabled())
12202 || !self.mmvq_supports(QT_NVFP4)
12203 || !self.batched_supports(QT_NVFP4)
12204 {
12205 return Ok(None);
12206 }
12207 let in_f = ws[0].in_features();
12208 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12209 for w in ws {
12210 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12211 return Ok(None);
12212 }
12213 match w {
12214 GpuTensor::Quant {
12215 bytes,
12216 qtype,
12217 scale,
12218 rp: true,
12219 rp4: None,
12220 ..
12221 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12222 parts.push((bytes, w.out_features(), *scale));
12223 }
12224 _ => return Ok(None),
12225 }
12226 }
12227 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12229 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12230 let mcols = if (5..=7).contains(&m) && b567 {
12231 m
12232 } else {
12233 Self::batched_mcols(m)
12234 };
12235 let kname: &'static str = match mcols {
12236 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12237 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12238 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12239 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12240 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12241 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12242 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12243 _ => return Ok(None),
12244 };
12245 if std::env::var("MEMRA_DEBUG").is_ok() {
12248 use std::sync::Mutex;
12249 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12250 let mut seen = SEEN.lock().unwrap();
12251 if !seen.contains(&tag) {
12252 seen.push(tag);
12253 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12254 }
12255 }
12256 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12259 let three = parts.len() == 3;
12260 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12261 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12262 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12263 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12266 let cfg = LaunchConfig {
12267 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12268 block_dim: (32, ROWS_PER_BLOCK, 1),
12269 shared_mem_bytes: 0,
12270 };
12271 let (inf, mi) = (in_f as i32, m as i32);
12272 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12273 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12274 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12275 let s3 = if three { 1.0f32 } else { parts[3].2 };
12276 let w3 = if three { parts[0].0 } else { parts[3].0 };
12277 let f = self.func(kname);
12278 let __s_b = self.gpu.stream();
12279 let mut b = __s_b.launch_builder(&f);
12280 b.arg(parts[0].0)
12281 .arg(parts[1].0)
12282 .arg(parts[2].0)
12283 .arg(w3)
12284 .arg(aq)
12285 .arg(ad)
12286 .arg(&mut y0)
12287 .arg(&mut y1)
12288 .arg(&mut y2)
12289 .arg(&mut y3)
12290 .arg(&inf)
12291 .arg(&n0)
12292 .arg(&n1)
12293 .arg(&n2)
12294 .arg(&n3)
12295 .arg(&mi)
12296 .arg(&s0)
12297 .arg(&s1)
12298 .arg(&s2)
12299 .arg(&s3);
12300 unsafe {
12301 b.launch(cfg)?;
12302 }
12303 Ok(Some(if three {
12304 vec![y0, y1, y2]
12305 } else {
12306 vec![y0, y1, y2, y3]
12307 }))
12308 }
12309
12310 pub fn matmul_decode_exact_dual(
12326 &self,
12327 w0: &crate::model::GpuTensor,
12328 w1: &crate::model::GpuTensor,
12329 x: &CudaSlice<f32>,
12330 m: usize,
12331 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12332 use crate::model::GpuTensor;
12333 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12334 let on = *ON.get_or_init(|| {
12335 std::env::var("MEMRA_SPEC_DUAL_T")
12336 .map(|v| v != "0")
12337 .unwrap_or(true)
12338 });
12339 if !on
12340 || !(2..=4).contains(&m)
12341 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12342 || !self.uses_q8_1_fast(w0)
12343 || !self.uses_q8_1_fast(w1)
12344 {
12345 return Ok(None);
12346 }
12347 if !self.mmvq_supports(QT_NVFP4) {
12352 return Ok(None);
12353 }
12354 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12355 if w1.in_features() != in_f || w1.out_features() != out_f {
12356 return Ok(None);
12357 }
12358 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12359 (
12360 GpuTensor::Quant {
12361 bytes: b0,
12362 qtype: q0,
12363 row_bytes: rb0,
12364 scale: s0,
12365 rp: rp0,
12366 rp4: None,
12367 ..
12368 },
12369 GpuTensor::Quant {
12370 bytes: b1,
12371 qtype: q1,
12372 row_bytes: rb1,
12373 scale: s1,
12374 rp: rp1,
12375 rp4: None,
12376 ..
12377 },
12378 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12379 (b0, b1, *rb0, *s0, *s1, *rp0)
12380 }
12381 _ => return Ok(None),
12382 };
12383 if std::env::var("MEMRA_DEBUG").is_ok() {
12386 static ONCE: std::sync::Once = std::sync::Once::new();
12387 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12388 }
12389 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12390 let (y0, y1) =
12391 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12392 let mut y0 = y0;
12393 let mut y1 = y1;
12394 if s0 != 1.0 {
12395 self.scale_inplace(&mut y0, s0, m * out_f)?;
12396 }
12397 if s1 != 1.0 {
12398 self.scale_inplace(&mut y1, s1, m * out_f)?;
12399 }
12400 Ok(Some((y0, y1)))
12401 }
12402
12403 #[allow(clippy::too_many_arguments)]
12408 pub fn qmatvec_batched_dual_raw(
12409 &self,
12410 b0: &CudaSlice<u8>,
12411 b1: &CudaSlice<u8>,
12412 aq: &CudaSlice<i8>,
12413 ad: &CudaSlice<f32>,
12414 m: usize,
12415 in_f: usize,
12416 out_f: usize,
12417 row_bytes: usize,
12418 rp: bool,
12419 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12420 const ROWS_PER_BLOCK: u32 = 4;
12421 let mcols = Self::batched_mcols(m);
12422 let tiny_rp1 = rp
12425 && mcols == 4
12426 && out_f <= 128
12427 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12428 let (name, rows_per_block) = if tiny_rp1 {
12429 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12430 } else {
12431 match (mcols, rp, m) {
12432 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12433 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12434 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12435 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12436 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12437 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12438 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12439 _ => {
12440 return Err(
12441 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12442 );
12443 }
12444 }
12445 };
12446 let f = self.func(name);
12447 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12448 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12449 let cfg = LaunchConfig {
12450 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12451 block_dim: (32, ROWS_PER_BLOCK, 1),
12452 shared_mem_bytes: 0,
12453 };
12454 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12455 let __s_b = self.gpu.stream();
12456 let mut b = __s_b.launch_builder(&f);
12457 b.arg(b0)
12458 .arg(b1)
12459 .arg(aq)
12460 .arg(ad)
12461 .arg(&mut y0)
12462 .arg(&mut y1)
12463 .arg(&inf)
12464 .arg(&outf)
12465 .arg(&mi)
12466 .arg(&rb);
12467 unsafe {
12468 b.launch(cfg)?;
12469 }
12470 Ok((y0, y1))
12471 }
12472
12473 pub fn matmul_pre_dual_noscale(
12485 &self,
12486 w0: &crate::model::GpuTensor,
12487 w1: &crate::model::GpuTensor,
12488 aq: &CudaSlice<i8>,
12489 ad: &CudaSlice<f32>,
12490 m: usize,
12491 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12492 {
12493 use crate::model::GpuTensor;
12494 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12495 return Ok(None);
12496 }
12497 if !self.mmvq_supports(QT_NVFP4) {
12507 return Ok(None);
12508 }
12509 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12510 if w1.in_features() != in_f || w1.out_features() != out_f {
12511 return Ok(None);
12512 }
12513 let no_mirror =
12526 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12527 if self.q8_ffn_fuse2_on()
12528 && no_mirror(w0)
12529 && no_mirror(w1)
12530 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12531 {
12532 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12533 return Ok(Some(((y0, 1.0), (y1, 1.0))));
12534 }
12535 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12545 let (y0, y1) =
12546 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12547 return Ok(Some(((y0, p0.3), (y1, p1.3))));
12548 }
12549 let (b0, q0, rb0, s0, rp0) = match w0 {
12550 GpuTensor::Quant {
12551 bytes,
12552 qtype,
12553 row_bytes,
12554 scale,
12555 rp,
12556 ..
12557 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12558 _ => return Ok(None),
12559 };
12560 let (b1, q1, rb1, s1, rp1) = match w1 {
12561 GpuTensor::Quant {
12562 bytes,
12563 qtype,
12564 row_bytes,
12565 scale,
12566 rp,
12567 ..
12568 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12569 _ => return Ok(None),
12570 };
12571 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12572 return Ok(None);
12573 }
12574 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12576 let rows_per_block = ROWS_PER_BLOCK * RPW;
12577 let f = self.func(if rp0 {
12578 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12579 } else {
12580 "qmatvec_nvfp4_mmvq_dual_mr2"
12581 });
12582 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12583 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12584 let cfg = LaunchConfig {
12585 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12586 block_dim: (32, ROWS_PER_BLOCK, 1),
12587 shared_mem_bytes: 0,
12588 };
12589 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12590 let one = 1.0f32;
12593 let __s_b = self.gpu.stream();
12594 let mut b = __s_b.launch_builder(&f);
12595 b.arg(b0)
12596 .arg(b1)
12597 .arg(aq)
12598 .arg(ad)
12599 .arg(&mut y0)
12600 .arg(&mut y1)
12601 .arg(&inf)
12602 .arg(&outf)
12603 .arg(&mi)
12604 .arg(&rb)
12605 .arg(&one)
12606 .arg(&one);
12607 unsafe {
12608 b.launch(cfg)?;
12609 }
12610 Ok(Some(((y0, s0), (y1, s1))))
12611 }
12612
12613 #[allow(clippy::too_many_arguments)]
12621 pub fn matmul_nvfp4_fused3(
12622 &self,
12623 w0: &crate::model::GpuTensor,
12624 w1: &crate::model::GpuTensor,
12625 w2: &crate::model::GpuTensor,
12626 aq: &CudaSlice<i8>,
12627 ad: &CudaSlice<f32>,
12628 m: usize,
12629 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12630 {
12631 use crate::model::GpuTensor;
12632 if !self.mmvq_supports(QT_NVFP4)
12639 || !self.uses_q8_1_fast(w0)
12640 || !self.uses_q8_1_fast(w1)
12641 || !self.uses_q8_1_fast(w2)
12642 {
12643 return Ok(None);
12644 }
12645 if (9..=16).contains(&m) {
12648 return Ok(
12649 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12650 Some(mut ys) => {
12651 let y2 = ys.pop().unwrap();
12652 let y1 = ys.pop().unwrap();
12653 let y0 = ys.pop().unwrap();
12654 Some((y0, y1, y2))
12655 }
12656 None => None,
12657 },
12658 );
12659 }
12660 if !(1..=8).contains(&m) {
12661 return Ok(None);
12662 }
12663 if m > 1 {
12664 let in_f = w0.in_features();
12665 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12666 || !self.batched_supports(QT_NVFP4)
12667 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12668 || (m > 4 && !Self::b8_enabled())
12669 || in_f % 512 != 0
12670 || in_f / 64 > 272
12671 {
12672 return Ok(None);
12673 }
12674 }
12675 let unpack = |w: &crate::model::GpuTensor| match w {
12676 GpuTensor::Quant {
12677 bytes,
12678 qtype,
12679 scale,
12680 rp,
12681 ..
12682 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12683 _ => None,
12684 };
12685 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12686 return Ok(None);
12687 };
12688 let in_f = w0.in_features();
12689 if w1.in_features() != in_f || w2.in_features() != in_f {
12690 return Ok(None);
12691 }
12692 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12693 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12695 let rows_pb = ROWS_PER_BLOCK * RPW;
12696 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12697 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12698 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12699 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12700 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12701 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12704 if m > 1 {
12705 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12707 return Ok(None);
12708 }
12709 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12710 let cfg = LaunchConfig {
12711 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12712 block_dim: (32, ROWS_PER_BLOCK, 1),
12713 shared_mem_bytes: 0,
12714 };
12715 let __s_b = self.gpu.stream();
12716 let mut b = __s_b.launch_builder(&f);
12717 b.arg(b0)
12718 .arg(b1)
12719 .arg(b2)
12720 .arg(aq)
12721 .arg(ad)
12722 .arg(&mut y0)
12723 .arg(&mut y1)
12724 .arg(&mut y2)
12725 .arg(&inf)
12726 .arg(&oi0)
12727 .arg(&oi1)
12728 .arg(&oi2)
12729 .arg(&mi);
12730 unsafe {
12731 b.launch(cfg)?;
12732 }
12733 return Ok(Some((y0, y1, y2)));
12734 }
12735 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12736 let cfg = LaunchConfig {
12737 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12738 block_dim: (32, ROWS_PER_BLOCK, 1),
12739 shared_mem_bytes: 0,
12740 };
12741 let __s_b = self.gpu.stream();
12742 let mut b = __s_b.launch_builder(&f);
12743 b.arg(b0)
12744 .arg(b1)
12745 .arg(b2)
12746 .arg(aq)
12747 .arg(ad)
12748 .arg(&mut y0)
12749 .arg(&mut y1)
12750 .arg(&mut y2)
12751 .arg(&inf)
12752 .arg(&oi0)
12753 .arg(&oi1)
12754 .arg(&oi2)
12755 .arg(&mi)
12756 .arg(&p0.1)
12757 .arg(&p1.1)
12758 .arg(&p2.1);
12759 unsafe {
12760 b.launch(cfg)?;
12761 }
12762 Ok(Some((y0, y1, y2)))
12763 }
12764
12765 pub fn matmul_nvfp4_fused2(
12774 &self,
12775 w0: &crate::model::GpuTensor,
12776 w1: &crate::model::GpuTensor,
12777 aq: &CudaSlice<i8>,
12778 ad: &CudaSlice<f32>,
12779 m: usize,
12780 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12781 use crate::model::GpuTensor;
12782 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12783 let off =
12784 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12785 if off
12788 || m != 1
12789 || !self.mmvq_supports(QT_NVFP4)
12790 || !self.uses_q8_1_fast(w0)
12791 || !self.uses_q8_1_fast(w1)
12792 {
12793 return Ok(None);
12794 }
12795 let unpack = |w: &crate::model::GpuTensor| match w {
12796 GpuTensor::Quant {
12797 bytes,
12798 qtype,
12799 scale,
12800 rp,
12801 ..
12802 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12803 _ => None,
12804 };
12805 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12806 return Ok(None);
12807 };
12808 let in_f = w0.in_features();
12809 if w1.in_features() != in_f {
12810 return Ok(None);
12811 }
12812 let (o0, o1) = (w0.out_features(), w1.out_features());
12813 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12815 let rows_pb = ROWS_PER_BLOCK * RPW;
12816 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12817 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12818 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12819 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12820 let cfg = LaunchConfig {
12821 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12822 block_dim: (32, ROWS_PER_BLOCK, 1),
12823 shared_mem_bytes: 0,
12824 };
12825 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12826 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12829 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12832 {
12833 use cudarc::driver::{DevicePtr, DevicePtrMut};
12834 let s = &self.gpu.stream();
12835 let (pw0, _g0) = b0.device_ptr(s);
12836 let (pw1, _g1) = b1.device_ptr(s);
12837 let (paq, _g2) = aq.device_ptr(s);
12838 let (pad, _g3) = ad.device_ptr(s);
12839 let (py0, _g4) = y0.device_ptr_mut(s);
12840 let (py1, _g5) = y1.device_ptr_mut(s);
12841 let (s0, s1) = (p0.1, p1.1);
12842 let mut ps = [
12843 &pw0 as *const _ as *mut std::ffi::c_void,
12844 &pw1 as *const _ as *mut _,
12845 &paq as *const _ as *mut _,
12846 &pad as *const _ as *mut _,
12847 &py0 as *const _ as *mut _,
12848 &py1 as *const _ as *mut _,
12849 &inf as *const _ as *mut _,
12850 &oi0 as *const _ as *mut _,
12851 &oi1 as *const _ as *mut _,
12852 &mi as *const _ as *mut _,
12853 &s0 as *const _ as *mut _,
12854 &s1 as *const _ as *mut _,
12855 ];
12856 unsafe {
12857 self.launch_pdl(
12858 "qmatvec_nvfp4_mmvq_fused2_rp",
12859 cfg.grid_dim,
12860 cfg.block_dim,
12861 &mut ps,
12862 )?;
12863 }
12864 }
12865 return Ok(Some((y0, y1)));
12866 }
12867 let __s_b = self.gpu.stream();
12868 let mut b = __s_b.launch_builder(&f);
12869 b.arg(b0)
12870 .arg(b1)
12871 .arg(aq)
12872 .arg(ad)
12873 .arg(&mut y0)
12874 .arg(&mut y1)
12875 .arg(&inf)
12876 .arg(&oi0)
12877 .arg(&oi1)
12878 .arg(&mi)
12879 .arg(&p0.1)
12880 .arg(&p1.1);
12881 unsafe {
12882 b.launch(cfg)?;
12883 }
12884 Ok(Some((y0, y1)))
12885 }
12886
12887 pub fn matmul_nvfp4_fused2_into(
12892 &self,
12893 w0: &crate::model::GpuTensor,
12894 w1: &crate::model::GpuTensor,
12895 aq: &CudaSlice<i8>,
12896 ad: &CudaSlice<f32>,
12897 y0: &mut CudaSlice<f32>,
12898 y1: &mut CudaSlice<f32>,
12899 ) -> Result<bool, Box<dyn std::error::Error>> {
12900 use crate::model::GpuTensor;
12901 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12902 let off =
12903 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12904 if off
12905 || !self.mmvq_supports(QT_NVFP4)
12906 || !self.uses_q8_1_fast(w0)
12907 || !self.uses_q8_1_fast(w1)
12908 {
12909 return Ok(false);
12910 }
12911 let unpack = |w: &crate::model::GpuTensor| match w {
12912 GpuTensor::Quant {
12913 bytes,
12914 qtype,
12915 scale,
12916 rp,
12917 ..
12918 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12919 _ => None,
12920 };
12921 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12922 return Ok(false);
12923 };
12924 let in_f = w0.in_features();
12925 if w1.in_features() != in_f {
12926 return Ok(false);
12927 }
12928 let (o0, o1) = (w0.out_features(), w1.out_features());
12929 if y0.len() < o0 || y1.len() < o1 {
12930 return Ok(false);
12931 }
12932 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12934 let rows_pb = ROWS_PER_BLOCK * RPW;
12935 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12936 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12937 let cfg = LaunchConfig {
12938 grid_dim: (nb(o0) + nb(o1), 1, 1),
12939 block_dim: (32, ROWS_PER_BLOCK, 1),
12940 shared_mem_bytes: 0,
12941 };
12942 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
12943 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12946 let __s_b = self.gpu.stream();
12947 let mut b = __s_b.launch_builder(&f);
12948 b.arg(b0)
12949 .arg(b1)
12950 .arg(aq)
12951 .arg(ad)
12952 .arg(&mut *y0)
12953 .arg(&mut *y1)
12954 .arg(&inf)
12955 .arg(&oi0)
12956 .arg(&oi1)
12957 .arg(&mi)
12958 .arg(&p0.1)
12959 .arg(&p1.1);
12960 unsafe {
12961 b.launch(cfg)?;
12962 }
12963 Ok(true)
12964 }
12965
12966 #[allow(clippy::type_complexity)]
12971 pub fn matmul_nvfp4_fused4(
12972 &self,
12973 w0: &crate::model::GpuTensor,
12974 w1: &crate::model::GpuTensor,
12975 w2: &crate::model::GpuTensor,
12976 w3: &crate::model::GpuTensor,
12977 aq: &CudaSlice<i8>,
12978 ad: &CudaSlice<f32>,
12979 m: usize,
12980 ) -> Result<
12981 Option<(
12982 CudaSlice<f32>,
12983 CudaSlice<f32>,
12984 CudaSlice<f32>,
12985 CudaSlice<f32>,
12986 )>,
12987 Box<dyn std::error::Error>,
12988 > {
12989 use crate::model::GpuTensor;
12990 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
12997 || !self.mmvq_supports(QT_NVFP4)
12998 || !self.uses_q8_1_fast(w0)
12999 || !self.uses_q8_1_fast(w1)
13000 || !self.uses_q8_1_fast(w2)
13001 || !self.uses_q8_1_fast(w3)
13002 {
13003 return Ok(None);
13004 }
13005 if (9..=16).contains(&m) {
13010 return Ok(
13011 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13012 Some(mut ys) => {
13013 let y3 = ys.pop().unwrap();
13014 let y2 = ys.pop().unwrap();
13015 let y1 = ys.pop().unwrap();
13016 let y0 = ys.pop().unwrap();
13017 Some((y0, y1, y2, y3))
13018 }
13019 None => None,
13020 },
13021 );
13022 }
13023 if !(1..=8).contains(&m) {
13024 return Ok(None);
13025 }
13026 if m > 1 {
13027 let in_f = w0.in_features();
13030 if !self.batched_supports(QT_NVFP4)
13031 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13032 || (m > 4 && !Self::b8_enabled())
13033 || in_f % 512 != 0
13034 || in_f / 64 > 272
13035 {
13036 return Ok(None);
13037 }
13038 }
13039 let unpack = |w: &crate::model::GpuTensor| match w {
13040 GpuTensor::Quant {
13041 bytes,
13042 qtype,
13043 scale,
13044 rp,
13045 ..
13046 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13047 _ => None,
13048 };
13049 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13050 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13051 else {
13052 return Ok(None);
13053 };
13054 let in_f = w0.in_features();
13055 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13056 return Ok(None);
13057 }
13058 let (o0, o1, o2, o3) = (
13059 w0.out_features(),
13060 w1.out_features(),
13061 w2.out_features(),
13062 w3.out_features(),
13063 );
13064 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13066 let rows_pb = ROWS_PER_BLOCK * RPW;
13067 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13068 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13069 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13070 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13071 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13072 let (inf, oi0, oi1, oi2, oi3, mi) = (
13073 in_f as i32,
13074 o0 as i32,
13075 o1 as i32,
13076 o2 as i32,
13077 o3 as i32,
13078 m as i32,
13079 );
13080 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13083 if m > 1 {
13084 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13087 return Ok(None);
13088 }
13089 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13090 let cfg = LaunchConfig {
13091 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13092 block_dim: (32, ROWS_PER_BLOCK, 1),
13093 shared_mem_bytes: 0,
13094 };
13095 let __s_b = self.gpu.stream();
13096 let mut b = __s_b.launch_builder(&f);
13097 b.arg(b0)
13098 .arg(b1)
13099 .arg(b2)
13100 .arg(b3)
13101 .arg(aq)
13102 .arg(ad)
13103 .arg(&mut y0)
13104 .arg(&mut y1)
13105 .arg(&mut y2)
13106 .arg(&mut y3)
13107 .arg(&inf)
13108 .arg(&oi0)
13109 .arg(&oi1)
13110 .arg(&oi2)
13111 .arg(&oi3)
13112 .arg(&mi);
13113 unsafe {
13114 b.launch(cfg)?;
13115 }
13116 return Ok(Some((y0, y1, y2, y3)));
13117 }
13118 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13119 let cfg = LaunchConfig {
13120 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13121 block_dim: (32, ROWS_PER_BLOCK, 1),
13122 shared_mem_bytes: 0,
13123 };
13124 let __s_b = self.gpu.stream();
13125 let mut b = __s_b.launch_builder(&f);
13126 b.arg(b0)
13127 .arg(b1)
13128 .arg(b2)
13129 .arg(b3)
13130 .arg(aq)
13131 .arg(ad)
13132 .arg(&mut y0)
13133 .arg(&mut y1)
13134 .arg(&mut y2)
13135 .arg(&mut y3)
13136 .arg(&inf)
13137 .arg(&oi0)
13138 .arg(&oi1)
13139 .arg(&oi2)
13140 .arg(&oi3)
13141 .arg(&mi)
13142 .arg(&p0.1)
13143 .arg(&p1.1)
13144 .arg(&p2.1)
13145 .arg(&p3.1);
13146 unsafe {
13147 b.launch(cfg)?;
13148 }
13149 Ok(Some((y0, y1, y2, y3)))
13150 }
13151
13152 pub fn matmul_q8_fused2(
13160 &self,
13161 w0: &crate::model::GpuTensor,
13162 w1: &crate::model::GpuTensor,
13163 aq: &CudaSlice<i8>,
13164 ad: &CudaSlice<f32>,
13165 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13166 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13172 return Ok(Some(self.e4m3_fused2_core(
13173 p0.0,
13174 p1.0,
13175 aq,
13176 ad,
13177 w0.in_features(),
13178 p0.1,
13179 p1.1,
13180 p0.2,
13181 p0.3,
13182 p1.3,
13183 )?));
13184 }
13185 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13186 return Ok(None);
13187 };
13188 Ok(Some(self.q8_fused2_core(
13189 p0.0,
13190 p1.0,
13191 aq,
13192 ad,
13193 w0.in_features(),
13194 p0.1,
13195 p1.1,
13196 p0.2,
13197 )?))
13198 }
13199
13200 #[allow(clippy::too_many_arguments)]
13201 fn q8_fused2_core(
13202 &self,
13203 b0: &CudaSlice<u8>,
13204 b1: &CudaSlice<u8>,
13205 aq: &CudaSlice<i8>,
13206 ad: &CudaSlice<f32>,
13207 in_f: usize,
13208 out0: usize,
13209 out1: usize,
13210 row_bytes: usize,
13211 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13212 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13214 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13215 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13216 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13217 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13218 let cfg = LaunchConfig {
13219 grid_dim: (nb0 + nb1, 1, 1),
13220 block_dim: (32, ROWS_PER_BLOCK, 1),
13221 shared_mem_bytes: 0,
13222 };
13223 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13224 let __s_b = self.gpu.stream();
13225 let mut b = __s_b.launch_builder(&f);
13226 b.arg(b0)
13227 .arg(b1)
13228 .arg(aq)
13229 .arg(ad)
13230 .arg(&mut y0)
13231 .arg(&mut y1)
13232 .arg(&inf)
13233 .arg(&o0)
13234 .arg(&o1)
13235 .arg(&rbl);
13236 unsafe {
13237 b.launch(cfg)?;
13238 }
13239 Ok((y0, y1))
13240 }
13241
13242 pub fn matmul_q8_fused2_x(
13248 &self,
13249 w0: &crate::model::GpuTensor,
13250 w1: &crate::model::GpuTensor,
13251 x: &CudaSlice<f32>,
13252 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13253 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13254 return Ok(None);
13255 }
13256 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13257 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13258 return Ok(Some(self.e4m3_fused2_core(
13259 p0.0,
13260 p1.0,
13261 &aq,
13262 &ad,
13263 w0.in_features(),
13264 p0.1,
13265 p1.1,
13266 p0.2,
13267 p0.3,
13268 p1.3,
13269 )?));
13270 }
13271 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13272 return Ok(None);
13273 };
13274 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13275 Ok(Some(self.q8_fused2_core(
13276 p0.0,
13277 p1.0,
13278 &aq,
13279 &ad,
13280 w0.in_features(),
13281 p0.1,
13282 p1.1,
13283 p0.2,
13284 )?))
13285 }
13286
13287 #[allow(clippy::too_many_arguments)]
13290 pub fn qmatvec_q8_fused2_raw(
13291 &self,
13292 b0: &CudaSlice<u8>,
13293 b1: &CudaSlice<u8>,
13294 x: &CudaSlice<f32>,
13295 in_f: usize,
13296 out0: usize,
13297 out1: usize,
13298 row_bytes: usize,
13299 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13300 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13301 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13302 }
13303
13304 pub fn matmul_q4_fused3(
13310 &self,
13311 w0: &crate::model::GpuTensor,
13312 w1: &crate::model::GpuTensor,
13313 w2: &crate::model::GpuTensor,
13314 aq: &CudaSlice<i8>,
13315 ad: &CudaSlice<f32>,
13316 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13317 {
13318 use crate::model::GpuTensor;
13319 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13320 match w {
13321 GpuTensor::Quant {
13322 qtype, row_bytes, ..
13323 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13324 _ => None,
13325 }
13326 };
13327 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13328 return Ok(None);
13329 };
13330 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13331 return Ok(None);
13332 }
13333 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13337 match w {
13338 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13339 Some(m) => (m, true),
13340 None => (bytes, *rp),
13341 },
13342 _ => unreachable!(),
13343 }
13344 }
13345 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13346 if rp0 != rp1 || rp1 != rp2 {
13347 return Ok(None);
13348 }
13349 let rp = rp0;
13350 let rpb: u32 = 4;
13351 let mr1 = rp && Self::q40_mr1_on();
13355 let nb = |o: usize| {
13356 if mr1 {
13357 (o as u32).div_ceil(rpb)
13358 } else {
13359 (o as u32).div_ceil(2).div_ceil(rpb)
13360 }
13361 };
13362 let grid = nb(o0) + nb(o1) + nb(o2);
13363 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13364 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13365 let mut y2 = self.alloc_uninit::<f32>(o2)?;
13366 let f = self.func(if mr1 {
13367 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13368 } else if rp {
13369 "qmatvec_q4_0_mmvq_fused3_rp"
13370 } else {
13371 "qmatvec_q4_0_mmvq_fused3"
13372 });
13373 let cfg = LaunchConfig {
13374 grid_dim: (grid, 1, 1),
13375 block_dim: (32, rpb, 1),
13376 shared_mem_bytes: 0,
13377 };
13378 let inf = w0.in_features() as i32;
13379 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13380 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13381 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13384 {
13385 use cudarc::driver::{DevicePtr, DevicePtrMut};
13386 let s = &self.gpu.stream();
13387 let (p0, _g0) = b0.device_ptr(s);
13388 let (p1, _g1) = b1.device_ptr(s);
13389 let (p2, _g2) = b2.device_ptr(s);
13390 let (paq, _g3) = aq.device_ptr(s);
13391 let (pad, _g4) = ad.device_ptr(s);
13392 let (py0, _g5) = y0.device_ptr_mut(s);
13393 let (py1, _g6) = y1.device_ptr_mut(s);
13394 let (py2, _g7) = y2.device_ptr_mut(s);
13395 let mut ps = [
13396 &p0 as *const _ as *mut std::ffi::c_void,
13397 &p1 as *const _ as *mut _,
13398 &p2 as *const _ as *mut _,
13399 &paq as *const _ as *mut _,
13400 &pad as *const _ as *mut _,
13401 &py0 as *const _ as *mut _,
13402 &py1 as *const _ as *mut _,
13403 &py2 as *const _ as *mut _,
13404 &inf as *const _ as *mut _,
13405 &oo0 as *const _ as *mut _,
13406 &oo1 as *const _ as *mut _,
13407 &oo2 as *const _ as *mut _,
13408 &r0 as *const _ as *mut _,
13409 &r1 as *const _ as *mut _,
13410 &r2 as *const _ as *mut _,
13411 ];
13412 unsafe {
13413 self.launch_pdl(
13414 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13415 (grid, 1, 1),
13416 (32, rpb, 1),
13417 &mut ps,
13418 )?;
13419 }
13420 }
13421 return Ok(Some((y0, y1, y2)));
13422 }
13423 let __s_b = self.gpu.stream();
13424 let mut b = __s_b.launch_builder(&f);
13425 b.arg(b0)
13426 .arg(b1)
13427 .arg(b2)
13428 .arg(aq)
13429 .arg(ad)
13430 .arg(&mut y0)
13431 .arg(&mut y1)
13432 .arg(&mut y2)
13433 .arg(&inf)
13434 .arg(&oo0)
13435 .arg(&oo1)
13436 .arg(&oo2)
13437 .arg(&r0)
13438 .arg(&r1)
13439 .arg(&r2);
13440 unsafe {
13441 b.launch(cfg)?;
13442 }
13443 Ok(Some((y0, y1, y2)))
13444 }
13445
13446 #[allow(clippy::too_many_arguments)]
13449 pub fn matmul_q4_fused3_into(
13450 &self,
13451 w0: &crate::model::GpuTensor,
13452 w1: &crate::model::GpuTensor,
13453 w2: &crate::model::GpuTensor,
13454 aq: &CudaSlice<i8>,
13455 ad: &CudaSlice<f32>,
13456 y0: &mut CudaSlice<f32>,
13457 y1: &mut CudaSlice<f32>,
13458 y2: &mut CudaSlice<f32>,
13459 ) -> Result<bool, Box<dyn std::error::Error>> {
13460 use crate::model::GpuTensor;
13461 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13462 match w {
13463 GpuTensor::Quant {
13464 qtype, row_bytes, ..
13465 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13466 _ => None,
13467 }
13468 };
13469 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13470 return Ok(false);
13471 };
13472 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13473 return Ok(false);
13474 }
13475 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13476 match w {
13477 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13478 Some(m) => (m, true),
13479 None => (bytes, *rp),
13480 },
13481 _ => unreachable!(),
13482 }
13483 }
13484 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13485 if rp0 != rp1 || rp1 != rp2 {
13486 return Ok(false);
13487 }
13488 let rp = rp0;
13489 let rpb: u32 = 4;
13490 let mr1 = rp && Self::q40_mr1_on();
13491 let nb = |o: usize| {
13492 if mr1 {
13493 (o as u32).div_ceil(rpb)
13494 } else {
13495 (o as u32).div_ceil(2).div_ceil(rpb)
13496 }
13497 };
13498 let grid = nb(o0) + nb(o1) + nb(o2);
13499 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13500 let f = self.func(if mr1 {
13501 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13502 } else if rp {
13503 "qmatvec_q4_0_mmvq_fused3_rp"
13504 } else {
13505 "qmatvec_q4_0_mmvq_fused3"
13506 });
13507 let cfg = LaunchConfig {
13508 grid_dim: (grid, 1, 1),
13509 block_dim: (32, rpb, 1),
13510 shared_mem_bytes: 0,
13511 };
13512 let inf = w0.in_features() as i32;
13513 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13514 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13515 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13517 use cudarc::driver::{DevicePtr, DevicePtrMut};
13518 let s = &self.gpu.stream();
13519 let (p0, _g0) = b0.device_ptr(s);
13520 let (p1, _g1) = b1.device_ptr(s);
13521 let (p2, _g2) = b2.device_ptr(s);
13522 let (paq, _g3) = aq.device_ptr(s);
13523 let (pad, _g4) = ad.device_ptr(s);
13524 let (py0, _g5) = y0.device_ptr_mut(s);
13525 let (py1, _g6) = y1.device_ptr_mut(s);
13526 let (py2, _g7) = y2.device_ptr_mut(s);
13527 let mut ps = [
13528 &p0 as *const _ as *mut std::ffi::c_void,
13529 &p1 as *const _ as *mut _,
13530 &p2 as *const _ as *mut _,
13531 &paq as *const _ as *mut _,
13532 &pad as *const _ as *mut _,
13533 &py0 as *const _ as *mut _,
13534 &py1 as *const _ as *mut _,
13535 &py2 as *const _ as *mut _,
13536 &inf as *const _ as *mut _,
13537 &oo0 as *const _ as *mut _,
13538 &oo1 as *const _ as *mut _,
13539 &oo2 as *const _ as *mut _,
13540 &r0 as *const _ as *mut _,
13541 &r1 as *const _ as *mut _,
13542 &r2 as *const _ as *mut _,
13543 ];
13544 unsafe {
13545 self.launch_pdl(
13546 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13547 (grid, 1, 1),
13548 (32, rpb, 1),
13549 &mut ps,
13550 )?;
13551 }
13552 return Ok(true);
13553 }
13554 let __s_b = self.gpu.stream();
13555 let mut b = __s_b.launch_builder(&f);
13556 b.arg(b0)
13557 .arg(b1)
13558 .arg(b2)
13559 .arg(aq)
13560 .arg(ad)
13561 .arg(&mut *y0)
13562 .arg(&mut *y1)
13563 .arg(&mut *y2)
13564 .arg(&inf)
13565 .arg(&oo0)
13566 .arg(&oo1)
13567 .arg(&oo2)
13568 .arg(&r0)
13569 .arg(&r1)
13570 .arg(&r2);
13571 unsafe {
13572 b.launch(cfg)?;
13573 }
13574 Ok(true)
13575 }
13576
13577 pub fn matmul_q4_fused2(
13579 &self,
13580 w0: &crate::model::GpuTensor,
13581 w1: &crate::model::GpuTensor,
13582 aq: &CudaSlice<i8>,
13583 ad: &CudaSlice<f32>,
13584 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13585 use crate::model::GpuTensor;
13586 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13587 match w {
13588 GpuTensor::Quant {
13589 qtype, row_bytes, ..
13590 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13591 _ => None,
13592 }
13593 };
13594 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13595 return Ok(None);
13596 };
13597 if w0.in_features() != w1.in_features() {
13598 return Ok(None);
13599 }
13600 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13602 match w {
13603 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13604 Some(m) => (m, true),
13605 None => (bytes, *rp),
13606 },
13607 _ => unreachable!(),
13608 }
13609 }
13610 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13611 if rp0 != rp1 {
13612 return Ok(None);
13613 }
13614 let rp = rp0;
13615 let rpb: u32 = 4;
13616 let mr1 = rp && Self::q40_mr1_on();
13618 let nb = |o: usize| {
13619 if mr1 {
13620 (o as u32).div_ceil(rpb)
13621 } else {
13622 (o as u32).div_ceil(2).div_ceil(rpb)
13623 }
13624 };
13625 let grid = nb(o0) + nb(o1);
13626 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13627 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13628 let f = self.func(if mr1 {
13629 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13630 } else if rp {
13631 "qmatvec_q4_0_mmvq_fused2_rp"
13632 } else {
13633 "qmatvec_q4_0_mmvq_fused2"
13634 });
13635 let cfg = LaunchConfig {
13636 grid_dim: (grid, 1, 1),
13637 block_dim: (32, rpb, 1),
13638 shared_mem_bytes: 0,
13639 };
13640 let inf = w0.in_features() as i32;
13641 let (oo0, oo1) = (o0 as i32, o1 as i32);
13642 let (r0, r1) = (rb0 as i64, rb1 as i64);
13643 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13645 {
13646 use cudarc::driver::{DevicePtr, DevicePtrMut};
13647 let s = &self.gpu.stream();
13648 let (p0, _g0) = b0.device_ptr(s);
13649 let (p1, _g1) = b1.device_ptr(s);
13650 let (paq, _g2) = aq.device_ptr(s);
13651 let (pad, _g3) = ad.device_ptr(s);
13652 let (py0, _g4) = y0.device_ptr_mut(s);
13653 let (py1, _g5) = y1.device_ptr_mut(s);
13654 let mut ps = [
13655 &p0 as *const _ as *mut std::ffi::c_void,
13656 &p1 as *const _ as *mut _,
13657 &paq as *const _ as *mut _,
13658 &pad as *const _ as *mut _,
13659 &py0 as *const _ as *mut _,
13660 &py1 as *const _ as *mut _,
13661 &inf as *const _ as *mut _,
13662 &oo0 as *const _ as *mut _,
13663 &oo1 as *const _ as *mut _,
13664 &r0 as *const _ as *mut _,
13665 &r1 as *const _ as *mut _,
13666 ];
13667 unsafe {
13668 self.launch_pdl(
13669 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13670 (grid, 1, 1),
13671 (32, rpb, 1),
13672 &mut ps,
13673 )?;
13674 }
13675 }
13676 return Ok(Some((y0, y1)));
13677 }
13678 let __s_b = self.gpu.stream();
13679 let mut b = __s_b.launch_builder(&f);
13680 b.arg(b0)
13681 .arg(b1)
13682 .arg(aq)
13683 .arg(ad)
13684 .arg(&mut y0)
13685 .arg(&mut y1)
13686 .arg(&inf)
13687 .arg(&oo0)
13688 .arg(&oo1)
13689 .arg(&r0)
13690 .arg(&r1);
13691 unsafe {
13692 b.launch(cfg)?;
13693 }
13694 Ok(Some((y0, y1)))
13695 }
13696
13697 pub fn matmul_q4_fused2_into(
13699 &self,
13700 w0: &crate::model::GpuTensor,
13701 w1: &crate::model::GpuTensor,
13702 aq: &CudaSlice<i8>,
13703 ad: &CudaSlice<f32>,
13704 y0: &mut CudaSlice<f32>,
13705 y1: &mut CudaSlice<f32>,
13706 ) -> Result<bool, Box<dyn std::error::Error>> {
13707 use crate::model::GpuTensor;
13708 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13709 match w {
13710 GpuTensor::Quant {
13711 qtype, row_bytes, ..
13712 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13713 _ => None,
13714 }
13715 };
13716 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13717 return Ok(false);
13718 };
13719 if w0.in_features() != w1.in_features() {
13720 return Ok(false);
13721 }
13722 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13723 match w {
13724 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13725 Some(m) => (m, true),
13726 None => (bytes, *rp),
13727 },
13728 _ => unreachable!(),
13729 }
13730 }
13731 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13732 if rp0 != rp1 {
13733 return Ok(false);
13734 }
13735 let rp = rp0;
13736 let rpb: u32 = 4;
13737 let mr1 = rp && Self::q40_mr1_on();
13738 let nb = |o: usize| {
13739 if mr1 {
13740 (o as u32).div_ceil(rpb)
13741 } else {
13742 (o as u32).div_ceil(2).div_ceil(rpb)
13743 }
13744 };
13745 let grid = nb(o0) + nb(o1);
13746 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13747 let f = self.func(if mr1 {
13748 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13749 } else if rp {
13750 "qmatvec_q4_0_mmvq_fused2_rp"
13751 } else {
13752 "qmatvec_q4_0_mmvq_fused2"
13753 });
13754 let cfg = LaunchConfig {
13755 grid_dim: (grid, 1, 1),
13756 block_dim: (32, rpb, 1),
13757 shared_mem_bytes: 0,
13758 };
13759 let inf = w0.in_features() as i32;
13760 let (oo0, oo1) = (o0 as i32, o1 as i32);
13761 let (r0, r1) = (rb0 as i64, rb1 as i64);
13762 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13764 use cudarc::driver::{DevicePtr, DevicePtrMut};
13765 let s = &self.gpu.stream();
13766 let (p0, _g0) = b0.device_ptr(s);
13767 let (p1, _g1) = b1.device_ptr(s);
13768 let (paq, _g2) = aq.device_ptr(s);
13769 let (pad, _g3) = ad.device_ptr(s);
13770 let (py0, _g4) = y0.device_ptr_mut(s);
13771 let (py1, _g5) = y1.device_ptr_mut(s);
13772 let mut ps = [
13773 &p0 as *const _ as *mut std::ffi::c_void,
13774 &p1 as *const _ as *mut _,
13775 &paq as *const _ as *mut _,
13776 &pad as *const _ as *mut _,
13777 &py0 as *const _ as *mut _,
13778 &py1 as *const _ as *mut _,
13779 &inf as *const _ as *mut _,
13780 &oo0 as *const _ as *mut _,
13781 &oo1 as *const _ as *mut _,
13782 &r0 as *const _ as *mut _,
13783 &r1 as *const _ as *mut _,
13784 ];
13785 unsafe {
13786 self.launch_pdl(
13787 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13788 (grid, 1, 1),
13789 (32, rpb, 1),
13790 &mut ps,
13791 )?;
13792 }
13793 return Ok(true);
13794 }
13795 let __s_b = self.gpu.stream();
13796 let mut b = __s_b.launch_builder(&f);
13797 b.arg(b0)
13798 .arg(b1)
13799 .arg(aq)
13800 .arg(ad)
13801 .arg(&mut *y0)
13802 .arg(&mut *y1)
13803 .arg(&inf)
13804 .arg(&oo0)
13805 .arg(&oo1)
13806 .arg(&r0)
13807 .arg(&r1);
13808 unsafe {
13809 b.launch(cfg)?;
13810 }
13811 Ok(true)
13812 }
13813
13814 pub fn matmul_q4_fused2_batched(
13819 &self,
13820 w0: &crate::model::GpuTensor,
13821 w1: &crate::model::GpuTensor,
13822 aq: &CudaSlice<i8>,
13823 ad: &CudaSlice<f32>,
13824 m: usize,
13825 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13826 use crate::model::GpuTensor;
13827 if m < 2 || m > 8 {
13828 return Ok(None);
13829 }
13830 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13831 match w {
13832 GpuTensor::Quant {
13833 qtype, row_bytes, ..
13834 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13835 _ => None,
13836 }
13837 };
13838 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
13839 return Ok(None);
13840 };
13841 if w0.in_features() != w1.in_features() {
13842 return Ok(None);
13843 }
13844 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13845 match w {
13846 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13847 Some(mr) => (mr, true),
13848 None => (bytes, *rp),
13849 },
13850 _ => unreachable!(),
13851 }
13852 }
13853 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13854 if !rp0 || !rp1 {
13855 return Ok(None);
13856 }
13857 let mcols = Self::batched_mcols(m);
13858 let rpb: u32 = 4;
13859 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13860 let grid = nb(o0) + nb(o1);
13861 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13862 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13863 let f = self.func(match mcols {
13864 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
13865 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
13866 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
13867 });
13868 let cfg = LaunchConfig {
13869 grid_dim: (grid, 1, 1),
13870 block_dim: (32, rpb, 1),
13871 shared_mem_bytes: 0,
13872 };
13873 let inf = w0.in_features() as i32;
13874 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
13875 let rb = rb0 as i64;
13876 let __s_b = self.gpu.stream();
13877 let mut b = __s_b.launch_builder(&f);
13878 b.arg(b0)
13879 .arg(b1)
13880 .arg(aq)
13881 .arg(ad)
13882 .arg(&mut y0)
13883 .arg(&mut y1)
13884 .arg(&inf)
13885 .arg(&oo0)
13886 .arg(&oo1)
13887 .arg(&mi)
13888 .arg(&rb);
13889 unsafe {
13890 b.launch(cfg)?;
13891 }
13892 Ok(Some((y0, y1)))
13893 }
13894
13895 #[allow(clippy::too_many_arguments)]
13898 pub fn matmul_q4_fused3_batched(
13899 &self,
13900 w0: &crate::model::GpuTensor,
13901 w1: &crate::model::GpuTensor,
13902 w2: &crate::model::GpuTensor,
13903 aq: &CudaSlice<i8>,
13904 ad: &CudaSlice<f32>,
13905 m: usize,
13906 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13907 {
13908 use crate::model::GpuTensor;
13909 if m < 2 || m > 8 {
13910 return Ok(None);
13911 }
13912 let q4 = |w: &GpuTensor| -> Option<usize> {
13913 match w {
13914 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
13915 _ => None,
13916 }
13917 };
13918 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
13919 return Ok(None);
13920 };
13921 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13922 return Ok(None);
13923 }
13924 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13925 match w {
13926 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13927 Some(mr) => (mr, true),
13928 None => (bytes, *rp),
13929 },
13930 _ => unreachable!(),
13931 }
13932 }
13933 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13934 if !rp0 || !rp1 || !rp2 {
13935 return Ok(None);
13936 }
13937 let mcols = Self::batched_mcols(m);
13938 let rpb: u32 = 4;
13939 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13940 let grid = nb(o0) + nb(o1) + nb(o2);
13941 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13942 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13943 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13944 let f = self.func(match mcols {
13945 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
13946 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
13947 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
13948 });
13949 let cfg = LaunchConfig {
13950 grid_dim: (grid, 1, 1),
13951 block_dim: (32, rpb, 1),
13952 shared_mem_bytes: 0,
13953 };
13954 let inf = w0.in_features() as i32;
13955 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
13956 let rb = 0i64;
13957 let __s_b = self.gpu.stream();
13958 let mut b = __s_b.launch_builder(&f);
13959 b.arg(b0)
13960 .arg(b1)
13961 .arg(b2)
13962 .arg(aq)
13963 .arg(ad)
13964 .arg(&mut y0)
13965 .arg(&mut y1)
13966 .arg(&mut y2)
13967 .arg(&inf)
13968 .arg(&oo0)
13969 .arg(&oo1)
13970 .arg(&oo2)
13971 .arg(&mi)
13972 .arg(&rb);
13973 unsafe {
13974 b.launch(cfg)?;
13975 }
13976 Ok(Some((y0, y1, y2)))
13977 }
13978
13979 pub fn matmul_q8_fused3(
13980 &self,
13981 w0: &crate::model::GpuTensor,
13982 w1: &crate::model::GpuTensor,
13983 w2: &crate::model::GpuTensor,
13984 aq: &CudaSlice<i8>,
13985 ad: &CudaSlice<f32>,
13986 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13987 {
13988 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13991 return Ok(Some(self.e4m3_fused3_core(
13992 p0.0,
13993 p1.0,
13994 p2.0,
13995 aq,
13996 ad,
13997 w0.in_features(),
13998 p0.1,
13999 p1.1,
14000 p2.1,
14001 p0.2,
14002 p0.3,
14003 p1.3,
14004 p2.3,
14005 )?));
14006 }
14007 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14008 return Ok(None);
14009 };
14010 Ok(Some(self.q8_fused3_core(
14011 p0.0,
14012 p1.0,
14013 p2.0,
14014 aq,
14015 ad,
14016 w0.in_features(),
14017 p0.1,
14018 p1.1,
14019 p2.1,
14020 p0.2,
14021 )?))
14022 }
14023
14024 #[allow(clippy::too_many_arguments)]
14025 fn q8_fused3_core(
14026 &self,
14027 b0: &CudaSlice<u8>,
14028 b1: &CudaSlice<u8>,
14029 b2: &CudaSlice<u8>,
14030 aq: &CudaSlice<i8>,
14031 ad: &CudaSlice<f32>,
14032 in_f: usize,
14033 out0: usize,
14034 out1: usize,
14035 out2: usize,
14036 row_bytes: usize,
14037 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14038 const ROWS_PER_BLOCK: u32 = 4;
14039 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14040 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14041 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14042 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14043 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14044 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14045 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14046 let cfg = LaunchConfig {
14047 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14048 block_dim: (32, ROWS_PER_BLOCK, 1),
14049 shared_mem_bytes: 0,
14050 };
14051 let (inf, o0, o1, o2, rbl) = (
14052 in_f as i32,
14053 out0 as i32,
14054 out1 as i32,
14055 out2 as i32,
14056 row_bytes as i64,
14057 );
14058 let __s_b = self.gpu.stream();
14059 let mut b = __s_b.launch_builder(&f);
14060 b.arg(b0)
14061 .arg(b1)
14062 .arg(b2)
14063 .arg(aq)
14064 .arg(ad)
14065 .arg(&mut y0)
14066 .arg(&mut y1)
14067 .arg(&mut y2)
14068 .arg(&inf)
14069 .arg(&o0)
14070 .arg(&o1)
14071 .arg(&o2)
14072 .arg(&rbl);
14073 unsafe {
14074 b.launch(cfg)?;
14075 }
14076 Ok((y0, y1, y2))
14077 }
14078
14079 #[allow(clippy::too_many_arguments)]
14081 pub fn qmatvec_q8_fused3_raw(
14082 &self,
14083 b0: &CudaSlice<u8>,
14084 b1: &CudaSlice<u8>,
14085 b2: &CudaSlice<u8>,
14086 x: &CudaSlice<f32>,
14087 in_f: usize,
14088 out0: usize,
14089 out1: usize,
14090 out2: usize,
14091 row_bytes: usize,
14092 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14093 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14094 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14095 }
14096
14097 pub fn matmul_q8_fused2_t(
14108 &self,
14109 w0: &crate::model::GpuTensor,
14110 w1: &crate::model::GpuTensor,
14111 aq: &CudaSlice<i8>,
14112 ad: &CudaSlice<f32>,
14113 m: usize,
14114 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14115 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14119 return Ok(None);
14120 }
14121 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14124 if m > 4 && !Self::b8_enabled() {
14125 return Ok(None);
14126 }
14127 return Ok(Some(self.e4m3_fused2_t_core(
14128 p0.0,
14129 p1.0,
14130 aq,
14131 ad,
14132 m,
14133 w0.in_features(),
14134 p0.1,
14135 p1.1,
14136 p0.2,
14137 p0.3,
14138 p1.3,
14139 )?));
14140 }
14141 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14142 return Ok(None);
14143 };
14144 Ok(Some(self.q8_fused2_t_core(
14145 p0.0,
14146 p1.0,
14147 aq,
14148 ad,
14149 m,
14150 w0.in_features(),
14151 p0.1,
14152 p1.1,
14153 p0.2,
14154 )?))
14155 }
14156
14157 #[allow(clippy::too_many_arguments)]
14158 fn q8_fused2_t_core(
14159 &self,
14160 b0: &CudaSlice<u8>,
14161 b1: &CudaSlice<u8>,
14162 aq: &CudaSlice<i8>,
14163 ad: &CudaSlice<f32>,
14164 m: usize,
14165 in_f: usize,
14166 out0: usize,
14167 out1: usize,
14168 row_bytes: usize,
14169 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14170 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14172 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14173 let f = self.func(match Self::batched_mcols(m) {
14174 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14175 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14176 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14178 });
14179 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14180 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14181 let cfg = LaunchConfig {
14182 grid_dim: (nb0 + nb1, 1, 1),
14183 block_dim: (32, ROWS_PER_BLOCK, 1),
14184 shared_mem_bytes: 0,
14185 };
14186 let (inf, o0, o1, mi, rbl) = (
14187 in_f as i32,
14188 out0 as i32,
14189 out1 as i32,
14190 m as i32,
14191 row_bytes as i64,
14192 );
14193 let __s_b = self.gpu.stream();
14194 let mut b = __s_b.launch_builder(&f);
14195 b.arg(b0)
14196 .arg(b1)
14197 .arg(aq)
14198 .arg(ad)
14199 .arg(&mut y0)
14200 .arg(&mut y1)
14201 .arg(&inf)
14202 .arg(&o0)
14203 .arg(&o1)
14204 .arg(&mi)
14205 .arg(&rbl);
14206 unsafe {
14207 b.launch(cfg)?;
14208 }
14209 Ok((y0, y1))
14210 }
14211
14212 #[allow(clippy::too_many_arguments)]
14215 pub fn qmatvec_q8_fused2_t_raw(
14216 &self,
14217 b0: &CudaSlice<u8>,
14218 b1: &CudaSlice<u8>,
14219 x: &CudaSlice<f32>,
14220 m: usize,
14221 in_f: usize,
14222 out0: usize,
14223 out1: usize,
14224 row_bytes: usize,
14225 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14226 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14227 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14228 }
14229
14230 #[allow(clippy::too_many_arguments)]
14233 pub fn matmul_q8_fused3_t(
14234 &self,
14235 w0: &crate::model::GpuTensor,
14236 w1: &crate::model::GpuTensor,
14237 w2: &crate::model::GpuTensor,
14238 aq: &CudaSlice<i8>,
14239 ad: &CudaSlice<f32>,
14240 m: usize,
14241 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14242 {
14243 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14244 return Ok(None);
14245 }
14246 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14247 return Ok(Some(self.e4m3_fused3_t_core(
14248 p0.0,
14249 p1.0,
14250 p2.0,
14251 aq,
14252 ad,
14253 m,
14254 w0.in_features(),
14255 p0.1,
14256 p1.1,
14257 p2.1,
14258 p0.2,
14259 p0.3,
14260 p1.3,
14261 p2.3,
14262 )?));
14263 }
14264 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14265 return Ok(None);
14266 };
14267 Ok(Some(self.q8_fused3_t_core(
14268 p0.0,
14269 p1.0,
14270 p2.0,
14271 aq,
14272 ad,
14273 m,
14274 w0.in_features(),
14275 p0.1,
14276 p1.1,
14277 p2.1,
14278 p0.2,
14279 )?))
14280 }
14281
14282 #[allow(clippy::too_many_arguments)]
14283 fn q8_fused3_t_core(
14284 &self,
14285 b0: &CudaSlice<u8>,
14286 b1: &CudaSlice<u8>,
14287 b2: &CudaSlice<u8>,
14288 aq: &CudaSlice<i8>,
14289 ad: &CudaSlice<f32>,
14290 m: usize,
14291 in_f: usize,
14292 out0: usize,
14293 out1: usize,
14294 out2: usize,
14295 row_bytes: usize,
14296 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14297 const ROWS_PER_BLOCK: u32 = 4;
14298 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14299 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14300 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14301 let f = self.func(if Self::batched_mcols(m) == 2 {
14302 "qmatvec_q8_0_mmvq_fused3_b2"
14303 } else {
14304 "qmatvec_q8_0_mmvq_fused3_b4"
14305 });
14306 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14307 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14308 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14309 let cfg = LaunchConfig {
14310 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14311 block_dim: (32, ROWS_PER_BLOCK, 1),
14312 shared_mem_bytes: 0,
14313 };
14314 let (inf, o0, o1, o2, mi, rbl) = (
14315 in_f as i32,
14316 out0 as i32,
14317 out1 as i32,
14318 out2 as i32,
14319 m as i32,
14320 row_bytes as i64,
14321 );
14322 let __s_b = self.gpu.stream();
14323 let mut b = __s_b.launch_builder(&f);
14324 b.arg(b0)
14325 .arg(b1)
14326 .arg(b2)
14327 .arg(aq)
14328 .arg(ad)
14329 .arg(&mut y0)
14330 .arg(&mut y1)
14331 .arg(&mut y2)
14332 .arg(&inf)
14333 .arg(&o0)
14334 .arg(&o1)
14335 .arg(&o2)
14336 .arg(&mi)
14337 .arg(&rbl);
14338 unsafe {
14339 b.launch(cfg)?;
14340 }
14341 Ok((y0, y1, y2))
14342 }
14343
14344 #[allow(clippy::too_many_arguments)]
14346 pub fn qmatvec_q8_fused3_t_raw(
14347 &self,
14348 b0: &CudaSlice<u8>,
14349 b1: &CudaSlice<u8>,
14350 b2: &CudaSlice<u8>,
14351 x: &CudaSlice<f32>,
14352 m: usize,
14353 in_f: usize,
14354 out0: usize,
14355 out1: usize,
14356 out2: usize,
14357 row_bytes: usize,
14358 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14359 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14360 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14361 }
14362
14363 pub fn q8_ffn_fuse2_on(&self) -> bool {
14367 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14368 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14369 }
14370
14371 #[allow(clippy::type_complexity)]
14377 fn q8_fused_params<'w, const N: usize>(
14378 &self,
14379 ws: &[&'w crate::model::GpuTensor; N],
14380 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14381 use crate::model::GpuTensor;
14382 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14383 return None;
14384 }
14385 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14386 return None;
14387 }
14388 let in_f = ws[0].in_features();
14389 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14390 for (i, w) in ws.iter().enumerate() {
14391 match w {
14392 GpuTensor::Quant {
14393 bytes,
14394 qtype,
14395 row_bytes,
14396 scale,
14397 ..
14398 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14399 out[i] = Some((bytes, w.out_features(), *row_bytes))
14400 }
14401 _ => return None,
14402 }
14403 }
14404 Some(out.map(|o| o.unwrap()))
14405 }
14406
14407 pub fn e4m3_dual_on(&self) -> bool {
14410 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14411 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14412 }
14413
14414 #[allow(clippy::type_complexity)]
14426 fn e4m3_fused_params<'w, const N: usize>(
14427 &self,
14428 ws: &[&'w crate::model::GpuTensor; N],
14429 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14430 use crate::model::GpuTensor;
14431 if !self.e4m3_dual_on() {
14432 return None;
14433 }
14434 let in_f = ws[0].in_features();
14435 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14436 for (i, w) in ws.iter().enumerate() {
14437 match w {
14438 GpuTensor::Quant {
14439 bytes,
14440 qtype,
14441 row_bytes,
14442 scale,
14443 rp,
14444 rp4,
14445 ..
14446 } if *qtype == QT_F8_E4M3
14447 && w.in_features() == in_f
14448 && *row_bytes == in_f
14449 && !*rp
14450 && rp4.is_none() =>
14451 {
14452 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14453 }
14454 _ => return None,
14455 }
14456 }
14457 Some(out.map(|o| o.unwrap()))
14458 }
14459
14460 #[allow(clippy::too_many_arguments)]
14464 fn e4m3_fused2_core(
14465 &self,
14466 b0: &CudaSlice<u8>,
14467 b1: &CudaSlice<u8>,
14468 aq: &CudaSlice<i8>,
14469 ad: &CudaSlice<f32>,
14470 in_f: usize,
14471 out0: usize,
14472 out1: usize,
14473 row_bytes: usize,
14474 ws0: f32,
14475 ws1: f32,
14476 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14477 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14479 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14480 let f = self.func("qmatvec_e4m3_mmvq_fused2");
14481 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14482 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14483 let cfg = LaunchConfig {
14484 grid_dim: (nb0 + nb1, 1, 1),
14485 block_dim: (32, ROWS_PER_BLOCK, 1),
14486 shared_mem_bytes: 0,
14487 };
14488 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14489 let __s_b = self.gpu.stream();
14490 let mut b = __s_b.launch_builder(&f);
14491 b.arg(b0)
14492 .arg(b1)
14493 .arg(aq)
14494 .arg(ad)
14495 .arg(&mut y0)
14496 .arg(&mut y1)
14497 .arg(&inf)
14498 .arg(&o0)
14499 .arg(&o1)
14500 .arg(&rbl)
14501 .arg(&ws0)
14502 .arg(&ws1);
14503 unsafe {
14504 b.launch(cfg)?;
14505 }
14506 Ok((y0, y1))
14507 }
14508
14509 #[allow(clippy::too_many_arguments)]
14511 fn e4m3_fused3_core(
14512 &self,
14513 b0: &CudaSlice<u8>,
14514 b1: &CudaSlice<u8>,
14515 b2: &CudaSlice<u8>,
14516 aq: &CudaSlice<i8>,
14517 ad: &CudaSlice<f32>,
14518 in_f: usize,
14519 out0: usize,
14520 out1: usize,
14521 out2: usize,
14522 row_bytes: usize,
14523 ws0: f32,
14524 ws1: f32,
14525 ws2: f32,
14526 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14527 const ROWS_PER_BLOCK: u32 = 4;
14528 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14529 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14530 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14531 let f = self.func("qmatvec_e4m3_mmvq_fused3");
14532 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14533 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14534 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14535 let cfg = LaunchConfig {
14536 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14537 block_dim: (32, ROWS_PER_BLOCK, 1),
14538 shared_mem_bytes: 0,
14539 };
14540 let (inf, o0, o1, o2, rbl) = (
14541 in_f as i32,
14542 out0 as i32,
14543 out1 as i32,
14544 out2 as i32,
14545 row_bytes as i64,
14546 );
14547 let __s_b = self.gpu.stream();
14548 let mut b = __s_b.launch_builder(&f);
14549 b.arg(b0)
14550 .arg(b1)
14551 .arg(b2)
14552 .arg(aq)
14553 .arg(ad)
14554 .arg(&mut y0)
14555 .arg(&mut y1)
14556 .arg(&mut y2)
14557 .arg(&inf)
14558 .arg(&o0)
14559 .arg(&o1)
14560 .arg(&o2)
14561 .arg(&rbl)
14562 .arg(&ws0)
14563 .arg(&ws1)
14564 .arg(&ws2);
14565 unsafe {
14566 b.launch(cfg)?;
14567 }
14568 Ok((y0, y1, y2))
14569 }
14570
14571 #[allow(clippy::too_many_arguments)]
14575 fn e4m3_fused2_t_core(
14576 &self,
14577 b0: &CudaSlice<u8>,
14578 b1: &CudaSlice<u8>,
14579 aq: &CudaSlice<i8>,
14580 ad: &CudaSlice<f32>,
14581 m: usize,
14582 in_f: usize,
14583 out0: usize,
14584 out1: usize,
14585 row_bytes: usize,
14586 ws0: f32,
14587 ws1: f32,
14588 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14589 const ROWS_PER_BLOCK: u32 = 4;
14590 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14591 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14592 let f = self.func(match Self::batched_mcols(m) {
14593 2 => "qmatvec_e4m3_mmvq_fused2_b2",
14594 4 => "qmatvec_e4m3_mmvq_fused2_b4",
14595 _ => "qmatvec_e4m3_mmvq_fused2_b8",
14596 });
14597 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14598 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14599 let cfg = LaunchConfig {
14600 grid_dim: (nb0 + nb1, 1, 1),
14601 block_dim: (32, ROWS_PER_BLOCK, 1),
14602 shared_mem_bytes: 0,
14603 };
14604 let (inf, o0, o1, mi, rbl) = (
14605 in_f as i32,
14606 out0 as i32,
14607 out1 as i32,
14608 m as i32,
14609 row_bytes as i64,
14610 );
14611 let __s_b = self.gpu.stream();
14612 let mut b = __s_b.launch_builder(&f);
14613 b.arg(b0)
14614 .arg(b1)
14615 .arg(aq)
14616 .arg(ad)
14617 .arg(&mut y0)
14618 .arg(&mut y1)
14619 .arg(&inf)
14620 .arg(&o0)
14621 .arg(&o1)
14622 .arg(&mi)
14623 .arg(&rbl);
14624 unsafe {
14625 b.launch(cfg)?;
14626 }
14627 if ws0 != 1.0 {
14628 self.scale_inplace(&mut y0, ws0, m * out0)?;
14629 }
14630 if ws1 != 1.0 {
14631 self.scale_inplace(&mut y1, ws1, m * out1)?;
14632 }
14633 Ok((y0, y1))
14634 }
14635
14636 #[allow(clippy::too_many_arguments)]
14638 fn e4m3_fused3_t_core(
14639 &self,
14640 b0: &CudaSlice<u8>,
14641 b1: &CudaSlice<u8>,
14642 b2: &CudaSlice<u8>,
14643 aq: &CudaSlice<i8>,
14644 ad: &CudaSlice<f32>,
14645 m: usize,
14646 in_f: usize,
14647 out0: usize,
14648 out1: usize,
14649 out2: usize,
14650 row_bytes: usize,
14651 ws0: f32,
14652 ws1: f32,
14653 ws2: f32,
14654 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14655 const ROWS_PER_BLOCK: u32 = 4;
14656 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14657 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14658 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14659 let f = self.func(if Self::batched_mcols(m) == 2 {
14660 "qmatvec_e4m3_mmvq_fused3_b2"
14661 } else {
14662 "qmatvec_e4m3_mmvq_fused3_b4"
14663 });
14664 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14665 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14666 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14667 let cfg = LaunchConfig {
14668 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14669 block_dim: (32, ROWS_PER_BLOCK, 1),
14670 shared_mem_bytes: 0,
14671 };
14672 let (inf, o0, o1, o2, mi, rbl) = (
14673 in_f as i32,
14674 out0 as i32,
14675 out1 as i32,
14676 out2 as i32,
14677 m as i32,
14678 row_bytes as i64,
14679 );
14680 let __s_b = self.gpu.stream();
14681 let mut b = __s_b.launch_builder(&f);
14682 b.arg(b0)
14683 .arg(b1)
14684 .arg(b2)
14685 .arg(aq)
14686 .arg(ad)
14687 .arg(&mut y0)
14688 .arg(&mut y1)
14689 .arg(&mut y2)
14690 .arg(&inf)
14691 .arg(&o0)
14692 .arg(&o1)
14693 .arg(&o2)
14694 .arg(&mi)
14695 .arg(&rbl);
14696 unsafe {
14697 b.launch(cfg)?;
14698 }
14699 if ws0 != 1.0 {
14700 self.scale_inplace(&mut y0, ws0, m * out0)?;
14701 }
14702 if ws1 != 1.0 {
14703 self.scale_inplace(&mut y1, ws1, m * out1)?;
14704 }
14705 if ws2 != 1.0 {
14706 self.scale_inplace(&mut y2, ws2, m * out2)?;
14707 }
14708 Ok((y0, y1, y2))
14709 }
14710
14711 pub fn qmatvec_e4m3_blk_mmvq(
14721 &self,
14722 bytes: &CudaSlice<u8>,
14723 aq: &CudaSlice<i8>,
14724 ad: &CudaSlice<f32>,
14725 scales: &CudaSlice<f32>,
14726 m: usize,
14727 in_f: usize,
14728 out_f: usize,
14729 row_bytes: usize,
14730 scale_cols: usize,
14731 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14732 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
14734 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14735 )?;
14736 Ok(y)
14737 }
14738
14739 #[allow(clippy::too_many_arguments)]
14741 pub fn qmatvec_e4m3_blk_mmvq_into(
14742 &self,
14743 bytes: &CudaSlice<u8>,
14744 aq: &CudaSlice<i8>,
14745 ad: &CudaSlice<f32>,
14746 scales: &CudaSlice<f32>,
14747 m: usize,
14748 in_f: usize,
14749 out_f: usize,
14750 row_bytes: usize,
14751 scale_cols: usize,
14752 y: &mut CudaSlice<f32>,
14753 ) -> Result<(), Box<dyn std::error::Error>> {
14754 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
14756 let cfg = LaunchConfig {
14757 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14758 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
14761 let (inf, outf, mi, rb, sc) = (
14762 in_f as i32,
14763 out_f as i32,
14764 m as i32,
14765 row_bytes as i64,
14766 scale_cols as i32,
14767 );
14768 let __s_b = self.gpu.stream();
14769 let mut b = __s_b.launch_builder(&f);
14770 b.arg(bytes)
14771 .arg(aq)
14772 .arg(ad)
14773 .arg(scales)
14774 .arg(&mut *y)
14775 .arg(&inf)
14776 .arg(&outf)
14777 .arg(&mi)
14778 .arg(&rb)
14779 .arg(&sc);
14780 unsafe {
14781 b.launch(cfg)?;
14782 }
14783 Ok(())
14784 }
14785
14786 #[allow(clippy::too_many_arguments)]
14792 pub fn qmatvec_e4m3_blk_mmvq_batched(
14793 &self,
14794 bytes: &CudaSlice<u8>,
14795 aq: &CudaSlice<i8>,
14796 ad: &CudaSlice<f32>,
14797 scales: &CudaSlice<f32>,
14798 m: usize,
14799 in_f: usize,
14800 out_f: usize,
14801 row_bytes: usize,
14802 scale_cols: usize,
14803 mcols: usize,
14804 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14805 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14807 let name = match mcols {
14808 2 => "qmatvec_e4m3_blk_mmvq_b2",
14809 4 => "qmatvec_e4m3_blk_mmvq_b4",
14810 8 => "qmatvec_e4m3_blk_mmvq_b8",
14811 16 => "qmatvec_e4m3_blk_mmvq_b16",
14812 _ => {
14813 return Err(
14814 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14815 );
14816 }
14817 };
14818 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14819 let f = self.func(name);
14820 let cfg = LaunchConfig {
14821 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14822 block_dim: (32, ROWS_PER_BLOCK, 1),
14823 shared_mem_bytes: 0,
14824 };
14825 let (inf, outf, mi, rb, sc) = (
14826 in_f as i32,
14827 out_f as i32,
14828 m as i32,
14829 row_bytes as i64,
14830 scale_cols as i32,
14831 );
14832 let __s_b = self.gpu.stream();
14833 let mut b = __s_b.launch_builder(&f);
14834 b.arg(bytes)
14835 .arg(aq)
14836 .arg(ad)
14837 .arg(scales)
14838 .arg(&mut y)
14839 .arg(&inf)
14840 .arg(&outf)
14841 .arg(&mi)
14842 .arg(&rb)
14843 .arg(&sc);
14844 unsafe {
14845 b.launch(cfg)?;
14846 }
14847 Ok(y)
14848 }
14849
14850 #[allow(clippy::too_many_arguments)]
14853 pub fn qmatvec_e4m3_blk_batched_raw(
14854 &self,
14855 bytes: &CudaSlice<u8>,
14856 x: &CudaSlice<f32>,
14857 scales: &CudaSlice<f32>,
14858 m: usize,
14859 in_f: usize,
14860 out_f: usize,
14861 row_bytes: usize,
14862 scale_cols: usize,
14863 mcols: usize,
14864 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14865 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14866 self.qmatvec_e4m3_blk_mmvq_batched(
14867 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
14868 )
14869 }
14870
14871 #[allow(clippy::too_many_arguments)]
14874 pub fn qmatvec_e4m3_blk_mmvq_raw(
14875 &self,
14876 bytes: &CudaSlice<u8>,
14877 x: &CudaSlice<f32>,
14878 scales: &CudaSlice<f32>,
14879 m: usize,
14880 in_f: usize,
14881 out_f: usize,
14882 row_bytes: usize,
14883 scale_cols: usize,
14884 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14885 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14886 self.qmatvec_e4m3_blk_mmvq(
14887 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
14888 )
14889 }
14890
14891 #[allow(clippy::too_many_arguments)]
14894 pub fn qmatvec_e4m3_fused2_raw(
14895 &self,
14896 b0: &CudaSlice<u8>,
14897 b1: &CudaSlice<u8>,
14898 x: &CudaSlice<f32>,
14899 in_f: usize,
14900 out0: usize,
14901 out1: usize,
14902 row_bytes: usize,
14903 ws0: f32,
14904 ws1: f32,
14905 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14906 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14907 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
14908 }
14909
14910 #[allow(clippy::too_many_arguments)]
14911 pub fn qmatvec_e4m3_fused3_raw(
14912 &self,
14913 b0: &CudaSlice<u8>,
14914 b1: &CudaSlice<u8>,
14915 b2: &CudaSlice<u8>,
14916 x: &CudaSlice<f32>,
14917 in_f: usize,
14918 out0: usize,
14919 out1: usize,
14920 out2: usize,
14921 row_bytes: usize,
14922 ws0: f32,
14923 ws1: f32,
14924 ws2: f32,
14925 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14926 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14927 self.e4m3_fused3_core(
14928 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14929 )
14930 }
14931
14932 #[allow(clippy::too_many_arguments)]
14933 pub fn qmatvec_e4m3_fused2_t_raw(
14934 &self,
14935 b0: &CudaSlice<u8>,
14936 b1: &CudaSlice<u8>,
14937 x: &CudaSlice<f32>,
14938 m: usize,
14939 in_f: usize,
14940 out0: usize,
14941 out1: usize,
14942 row_bytes: usize,
14943 ws0: f32,
14944 ws1: f32,
14945 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14946 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14947 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
14948 }
14949
14950 #[allow(clippy::too_many_arguments)]
14951 pub fn qmatvec_e4m3_fused3_t_raw(
14952 &self,
14953 b0: &CudaSlice<u8>,
14954 b1: &CudaSlice<u8>,
14955 b2: &CudaSlice<u8>,
14956 x: &CudaSlice<f32>,
14957 m: usize,
14958 in_f: usize,
14959 out0: usize,
14960 out1: usize,
14961 out2: usize,
14962 row_bytes: usize,
14963 ws0: f32,
14964 ws1: f32,
14965 ws2: f32,
14966 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14967 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14968 self.e4m3_fused3_t_core(
14969 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14970 )
14971 }
14972
14973 fn try_e4m3_blk_pre(
14984 &self,
14985 w: &crate::model::GpuTensor,
14986 aq: &CudaSlice<i8>,
14987 ad: &CudaSlice<f32>,
14988 m: usize,
14989 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14990 use crate::model::GpuTensor;
14991 if let GpuTensor::Quant {
14992 bytes,
14993 qtype,
14994 row_bytes,
14995 blk: Some(g),
14996 ..
14997 } = w
14998 {
14999 if *qtype == QT_F8_E4M3_BLK {
15000 if (2..=16).contains(&m)
15006 && std::env::var("MEMRA_NO_BATCHED").is_err()
15007 && (m <= 4 || Self::b8_enabled())
15008 {
15009 let mcols = Self::batched_mcols(m);
15010 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15011 bytes,
15012 aq,
15013 ad,
15014 &g.scales,
15015 m,
15016 w.in_features(),
15017 w.out_features(),
15018 *row_bytes,
15019 g.cols,
15020 mcols,
15021 )?));
15022 }
15023 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15024 bytes,
15025 aq,
15026 ad,
15027 &g.scales,
15028 m,
15029 w.in_features(),
15030 w.out_features(),
15031 *row_bytes,
15032 g.cols,
15033 )?));
15034 }
15035 }
15036 Ok(None)
15037 }
15038
15039 fn try_e4m3_blk_prefill(
15086 &self,
15087 w: &crate::model::GpuTensor,
15088 x: &CudaSlice<f32>,
15089 m: usize,
15090 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15091 use crate::model::GpuTensor;
15092 let GpuTensor::Quant {
15093 bytes,
15094 qtype,
15095 blk: Some(g),
15096 ..
15097 } = w
15098 else {
15099 return Ok(None);
15100 };
15101 if *qtype != QT_F8_E4M3_BLK {
15102 return Ok(None);
15103 }
15104 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15109 return Ok(Some(y));
15110 }
15111 let (in_f, out_f) = (w.in_features(), w.out_features());
15112 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15113 let tmp = GpuTensor::Quant {
15114 bytes: slab,
15115 qtype: QT_Q8_0,
15116 row_bytes: in_f / 32 * 34,
15117 ne: vec![in_f as u64, out_f as u64],
15118 scale: 1.0,
15119 rp: false,
15120 #[cfg(memra_cutlass)]
15121 cutlass: None,
15122 fp8: None,
15123 blk: None,
15124 f16: None,
15125 rp4: None,
15126 };
15127 Ok(Some(self.matmul(&tmp, x, m)?))
15129 }
15130
15131 pub fn matmul_pre_noscale(
15132 &self,
15133 w: &crate::model::GpuTensor,
15134 aq: &CudaSlice<i8>,
15135 ad: &CudaSlice<f32>,
15136 m: usize,
15137 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15138 use crate::model::GpuTensor;
15139 if m == 1 {
15143 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15144 return Ok(Some((y, 1.0)));
15145 }
15146 }
15147 if m != 1 || !self.uses_q8_1_fast(w) {
15149 return Ok(None);
15150 }
15151 let in_f = w.in_features();
15152 let out_f = w.out_features();
15153 let (bytes, qtype, row_bytes, scale, rp) = match w {
15154 GpuTensor::Quant {
15155 bytes,
15156 qtype,
15157 row_bytes,
15158 scale,
15159 rp,
15160 ..
15161 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15162 _ => return Ok(None),
15163 };
15164 if self.mmvq_supports(qtype) {
15166 let (mbytes, mrp) = match w {
15168 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15169 _ => (bytes, rp),
15170 };
15171 let y = self.qmatvec_mmvq(
15172 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15173 )?;
15174 return Ok(Some((y, scale)));
15175 }
15176 let name = match qtype {
15178 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15179 QT_Q4_K => "qmatvec_q4_K_dp4a",
15180 QT_Q6_K => "qmatvec_q6_K_dp4a",
15181 QT_Q5_K => "qmatvec_q5_K_dp4a",
15182 QT_Q3_K => "qmatvec_q3_K_dp4a",
15183 QT_NVFP4 => {
15184 if rp {
15185 "qmatvec_nvfp4_dp4a_rp"
15186 } else {
15187 "qmatvec_nvfp4_dp4a"
15188 }
15189 }
15190 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15191 _ => return Ok(None),
15192 };
15193 let f = self.func(name);
15194 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15195 let cfg = LaunchConfig {
15196 grid_dim: (out_f as u32, m as u32, 1),
15197 block_dim: (128, 1, 1),
15198 shared_mem_bytes: 0,
15199 };
15200 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15201 let __s_b = self.gpu.stream();
15202 let mut b = __s_b.launch_builder(&f);
15203 b.arg(bytes)
15204 .arg(aq)
15205 .arg(ad)
15206 .arg(&mut y)
15207 .arg(&inf)
15208 .arg(&outf)
15209 .arg(&mi)
15210 .arg(&rb);
15211 unsafe {
15212 b.launch(cfg)?;
15213 }
15214 Ok(Some((y, scale)))
15215 }
15216
15217 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15220 if qtype == QT_F8_E4M3 {
15225 return true;
15226 }
15227 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15228 return false;
15229 }
15230 matches!(
15231 qtype,
15232 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15233 )
15234 }
15235
15236 pub fn qmatvec_mmvq(
15241 &self,
15242 bytes: &CudaSlice<u8>,
15243 aq: &CudaSlice<i8>,
15244 ad: &CudaSlice<f32>,
15245 m: usize,
15246 in_f: usize,
15247 out_f: usize,
15248 qtype: i32,
15249 row_bytes: usize,
15250 scale: f32,
15251 rp: bool,
15252 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15253 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15255 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15256 )?;
15257 Ok(y)
15258 }
15259
15260 #[allow(clippy::too_many_arguments)]
15262 pub fn qmatvec_mmvq_into(
15263 &self,
15264 bytes: &CudaSlice<u8>,
15265 aq: &CudaSlice<i8>,
15266 ad: &CudaSlice<f32>,
15267 m: usize,
15268 in_f: usize,
15269 out_f: usize,
15270 qtype: i32,
15271 row_bytes: usize,
15272 scale: f32,
15273 rp: bool,
15274 y: &mut CudaSlice<f32>,
15275 ) -> Result<(), Box<dyn std::error::Error>> {
15276 debug_assert!(y.len() >= m * out_f);
15277 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15283 && rp
15284 && m == 1
15285 && out_f >= 64
15286 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15287 && {
15288 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15289 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15290 }
15291 {
15292 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15293 let cfg = LaunchConfig {
15294 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15295 block_dim: (32, 2, 1),
15296 shared_mem_bytes: 0,
15297 };
15298 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15299 let __s_b = self.gpu.stream();
15300 let mut b = __s_b.launch_builder(&f);
15301 b.arg(bytes)
15302 .arg(aq)
15303 .arg(ad)
15304 .arg(&mut *y)
15305 .arg(&inf)
15306 .arg(&outf)
15307 .arg(&mi)
15308 .arg(&rb);
15309 unsafe {
15310 b.launch(cfg)?;
15311 }
15312 if scale != 1.0 {
15313 self.scale_inplace(y, scale, out_f)?;
15314 }
15315 return Ok(());
15316 }
15317 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15326 2
15327 } else {
15328 1
15329 };
15330 if m == 1 && qtype == QT_Q4_0 {
15335 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15336 mr = *Q40MR.get_or_init(|| {
15339 std::env::var("MEMRA_Q40_MR")
15340 .ok()
15341 .and_then(|v| v.parse().ok())
15342 .unwrap_or(1)
15343 });
15344 }
15345 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15356 let q5_force = q5_mode.as_deref() == Some("2");
15357 let q5_il = qtype == QT_Q5_K
15360 && m == 1
15361 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15362 if q5_il && !q5_force && out_f > 65536 {
15363 mr = 1;
15364 }
15365 if qtype == QT_Q4_0 && rp && mr != 1 {
15368 mr = 2;
15369 }
15370 if qtype == QT_Q8_0 && rp {
15374 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15375 mr = *Q80MR.get_or_init(|| {
15376 std::env::var("MEMRA_Q80_MR")
15377 .ok()
15378 .and_then(|v| v.parse().ok())
15379 .unwrap_or(1)
15380 });
15381 }
15382 let name = match (qtype, mr, rp) {
15383 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15384 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15385 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15386 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15387 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15388 (QT_Q5_K, 2, _) => {
15389 if q5_il {
15390 "qmatvec_q5_K_mmvq_mr2_il"
15391 } else {
15392 "qmatvec_q5_K_mmvq_mr2"
15393 }
15394 }
15395 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15396 (QT_Q8_0, _, true)
15401 if in_f % 1024 == 0 && {
15402 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15403 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15404 } =>
15405 {
15406 "qmatvec_q8_0_mmvq_rpca"
15407 }
15408 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15409 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15410 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15414 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15415 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15416 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15417 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15418 (QT_Q5_K, _, _) => {
15419 if q5_il {
15420 "qmatvec_q5_K_mmvq_il"
15421 } else {
15422 "qmatvec_q5_K_mmvq"
15423 }
15424 }
15425 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15426 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15427 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15428 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15429 };
15430 let f = self.func(name);
15431 let rows_per_block = ROWS_PER_BLOCK * mr;
15433 let cfg = LaunchConfig {
15434 grid_dim: (
15435 (out_f as u32 + rows_per_block - 1) / rows_per_block,
15436 m as u32,
15437 1,
15438 ),
15439 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15442 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15443 let __s_b = self.gpu.stream();
15444 let mut b = __s_b.launch_builder(&f);
15445 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15450 if Self::pdl_on()
15453 && Self::pdl_mmvq_on()
15454 && Self::pdl_nvfp4q8_on()
15455 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15456 {
15457 use cudarc::driver::{DevicePtr, DevicePtrMut};
15458 let s = &self.gpu.stream();
15459 let (pw, _g0) = bytes.device_ptr(s);
15460 let (paq, _g1) = aq.device_ptr(s);
15461 let (pad, _g2) = ad.device_ptr(s);
15462 let (py, _g3) = y.device_ptr_mut(s);
15463 let mut ps = [
15464 &pw as *const _ as *mut std::ffi::c_void,
15465 &paq as *const _ as *mut _,
15466 &pad as *const _ as *mut _,
15467 &py as *const _ as *mut _,
15468 &inf as *const _ as *mut _,
15469 &outf as *const _ as *mut _,
15470 &mi as *const _ as *mut _,
15471 &rb as *const _ as *mut _,
15472 &scale as *const _ as *mut _,
15473 ];
15474 unsafe {
15475 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15476 }
15477 return Ok(());
15478 }
15479 b.arg(bytes)
15480 .arg(aq)
15481 .arg(ad)
15482 .arg(&mut *y)
15483 .arg(&inf)
15484 .arg(&outf)
15485 .arg(&mi)
15486 .arg(&rb)
15487 .arg(&scale);
15488 unsafe {
15489 b.launch(cfg)?;
15490 }
15491 } else if Self::pdl_on()
15492 && Self::pdl_mmvq_on()
15493 && (matches!(
15494 name,
15495 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15496 ) || (Self::pdl_nvfp4q8_on()
15497 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15498 {
15499 {
15503 use cudarc::driver::{DevicePtr, DevicePtrMut};
15504 let s = &self.gpu.stream();
15505 let (pw, _g0) = bytes.device_ptr(s);
15506 let (paq, _g1) = aq.device_ptr(s);
15507 let (pad, _g2) = ad.device_ptr(s);
15508 let (py, _g3) = y.device_ptr_mut(s);
15509 let mut ps = [
15510 &pw as *const _ as *mut std::ffi::c_void,
15511 &paq as *const _ as *mut _,
15512 &pad as *const _ as *mut _,
15513 &py as *const _ as *mut _,
15514 &inf as *const _ as *mut _,
15515 &outf as *const _ as *mut _,
15516 &mi as *const _ as *mut _,
15517 &rb as *const _ as *mut _,
15518 ];
15519 unsafe {
15520 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15521 }
15522 }
15523 if scale != 1.0 {
15524 self.scale_inplace(y, scale, m * out_f)?;
15525 }
15526 } else {
15527 b.arg(bytes)
15528 .arg(aq)
15529 .arg(ad)
15530 .arg(&mut *y)
15531 .arg(&inf)
15532 .arg(&outf)
15533 .arg(&mi)
15534 .arg(&rb);
15535 unsafe {
15536 b.launch(cfg)?;
15537 }
15538 if scale != 1.0 {
15539 self.scale_inplace(y, scale, m * out_f)?;
15540 }
15541 }
15542 Ok(())
15543 }
15544
15545 pub fn qmatvec_mmvq_raw(
15549 &self,
15550 bytes: &CudaSlice<u8>,
15551 x: &CudaSlice<f32>,
15552 m: usize,
15553 in_f: usize,
15554 out_f: usize,
15555 qtype: i32,
15556 row_bytes: usize,
15557 rp: bool,
15558 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15559 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15560 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15561 }
15562
15563 pub fn batched_supports(&self, qtype: i32) -> bool {
15567 matches!(
15568 qtype,
15569 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15570 )
15571 }
15572
15573 pub fn iq_fast_enabled() -> bool {
15581 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15582 *ON.get_or_init(|| {
15583 std::env::var("MEMRA_IQ_FAST")
15584 .map(|v| v != "0")
15585 .unwrap_or(true)
15586 })
15587 }
15588
15589 pub fn b8_enabled() -> bool {
15592 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15593 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15594 }
15595
15596 pub fn batched_mcols(m: usize) -> usize {
15598 if m == 2 {
15599 2
15600 } else if m <= 4 {
15601 4
15602 } else if m <= 8 {
15603 8
15604 } else {
15605 16
15606 }
15607 }
15608
15609 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15614 Some(match (qtype, mcols) {
15615 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15616 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15617 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15618 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15624 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15625 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15626 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15627 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15630 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15631 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15632 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15633 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15636 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15637 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15638 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15639 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15640 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15641 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15642 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15643 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15647 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15648 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15649 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15650 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15654 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15655 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15656 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15657 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15658 _ => return None,
15659 })
15660 }
15661
15662 pub fn sm_count(&self) -> i32 {
15697 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15698 *SMS.get_or_init(|| {
15699 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15700 self.gpu
15701 .ctx
15702 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15703 .unwrap_or(82)
15704 })
15705 }
15706
15707 pub fn batched_variant(
15708 &self,
15709 _m: usize,
15710 in_f: usize,
15711 out_f: usize,
15712 qtype: i32,
15713 row_bytes: usize,
15714 mcols: usize,
15715 rp: bool,
15716 ) -> &'static str {
15717 if qtype == QT_Q8_0 {
15722 return if rp { "rp" } else { "base" };
15723 }
15724 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15725 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15726 Ok("base") => "base",
15727 Ok("pf") => "pf",
15728 Ok("r2") => "r2",
15729 Ok("r2w8") => "r2w8",
15730 Ok("pfr2") => "pfr2",
15731 Ok("ca") => "ca",
15732 Ok("car2") => "car2",
15733 Ok("rp") => "rp",
15736 Ok("rpr2") => "rpr2",
15737 Ok("rpr2w8") => "rpr2w8",
15738 Ok("rpca") => "rpca",
15741 Ok("rpcar2") => "rpcar2",
15742 Ok("rpsc") => "rpsc",
15749 Ok("rpms") => "rpms",
15750 Ok("rpmsc") => "rpmsc",
15751 Ok("rpks") => "rpks",
15752 Ok("rpksc") => "rpksc",
15753 _ => "auto",
15754 });
15755 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15759 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15764 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15765 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15766 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15767 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15768 let sms = *SMS.get_or_init(|| {
15769 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15770 self.gpu
15771 .ctx
15772 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15773 .unwrap_or(82)
15774 });
15775 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15795 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15798 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15799 Ok("base") => "base",
15800 Ok("r2") => "r2",
15801 Ok("r2w8") => "r2w8",
15802 _ => "auto",
15803 });
15804 let variant: &'static str = if qtype == QT_Q4_0 {
15805 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15809 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15810 Ok("base") => "base",
15816 Ok("r2") => "r2",
15817 Ok("ms") => "ms",
15818 Ok("sm") => "sm",
15819 Ok("la") => "la",
15820 _ => "auto",
15821 });
15822 let v = if q40 != "auto" {
15823 q40
15824 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15825 "r2"
15826 } else {
15827 "base"
15828 };
15829 if rp {
15834 match v {
15835 "ms" => "r2ms_rp",
15836 "sm" => "r2sm_rp",
15837 "la" => "r2la_rp",
15838 "r2" => "r2_rp",
15839 _ => "rp",
15840 }
15841 } else if matches!(v, "ms" | "sm" | "la") {
15842 "r2"
15843 } else {
15844 v
15845 }
15846 } else if qtype != QT_NVFP4 && !kq_r2 {
15847 "base"
15848 } else if kq_r2 && rp {
15849 "rp"
15853 } else if kq_r2 {
15854 if kq_bv != "auto" {
15857 if kq_bv == "r2w8" && mcols != 4 {
15858 "r2"
15859 } else {
15860 kq_bv
15861 }
15862 } else if bv != "auto" {
15863 match bv {
15864 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
15865 "r2w8" | "rpr2w8" => {
15866 if mcols != 4 {
15867 "r2"
15868 } else {
15869 "r2w8"
15870 }
15871 }
15872 _ => "base", }
15874 } else {
15875 let blocks = (out_f + 7) / 8;
15876 let waves = blocks as f64 / (7 * sms as usize) as f64;
15877 let filled = blocks >= 4 * sms as usize;
15878 let use_r2 = if qtype == QT_Q4_K {
15879 filled
15880 } else {
15881 waves >= 2.0
15882 };
15883 if use_r2 { "r2" } else { "base" }
15884 }
15885 } else if bv != "auto" {
15886 let v = if bv == "r2w8" && mcols == 2 {
15891 "r2"
15892 } else if bv == "ca" && (!ca_ok || mcols == 8) {
15893 "pf"
15894 } else if bv == "car2" && (!ca_ok || mcols == 8) {
15895 "r2"
15896 } else if bv == "pfr2" && mcols == 8 {
15897 "r2"
15898 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
15899 "rpr2"
15900 }
15901 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
15903 if mcols == 8 { "rpr2w8" } else { "rpr2" }
15904 } else if bv == "rpcar2" && mcols == 2 {
15905 "rpca"
15906 }
15907 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
15910 "rpr2"
15911 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
15912 "rpr2"
15913 } else {
15914 bv
15915 };
15916 if rp {
15917 match v {
15918 "base" | "pf" | "ca" | "rp" => "rp",
15919 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
15920 "r2w8" | "rpr2w8" => {
15921 if mcols == 2 {
15922 "rpr2"
15923 } else {
15924 "rpr2w8"
15925 }
15926 }
15927 other => other, }
15929 } else {
15930 v
15931 }
15932 } else if mcols == 8 {
15933 if rp {
15944 if sc_ok { "rpsc" } else { "rpr2w8" }
15945 } else {
15946 "r2w8"
15947 }
15948 } else if mcols >= 4 {
15949 let blocks = (out_f + 7) / 8;
15953 let r7 = 7 * sms as usize;
15954 let r8 = 8 * sms as usize;
15955 let waves = blocks as f64 / r7 as f64;
15956 let filled = blocks >= 4 * sms as usize;
15957 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
15961 if rp { "rpr2w8" } else { "r2w8" }
15965 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
15966 if rp { "rpr2" } else { "r2" }
15969 } else {
15970 if rp { "rp" } else { "pf" }
15974 }
15975 } else if in_f >= 6144 {
15976 if rp { "rpr2" } else { "r2" }
15980 } else if rp {
15981 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
15986 if sc_ok && waves >= 0.9 && waves <= 1.1 {
15987 "rpsc"
15988 } else {
15989 "rp"
15990 }
15991 } else {
15992 "base"
15993 };
15994 variant
15995 }
15996
15997 pub fn qmatvec_mmvq_batched(
15998 &self,
15999 bytes: &CudaSlice<u8>,
16000 aq: &CudaSlice<i8>,
16001 ad: &CudaSlice<f32>,
16002 m: usize,
16003 in_f: usize,
16004 out_f: usize,
16005 qtype: i32,
16006 row_bytes: usize,
16007 mcols: usize,
16008 scale: f32,
16009 rp: bool,
16010 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16011 const ROWS_PER_BLOCK: u32 = 4;
16012 let forced: Option<&'static str> = {
16017 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16018 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16019 .as_deref()
16020 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16021 };
16022 let variant = match forced {
16023 Some(v) if !rp || v.contains("rp") => v,
16024 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16025 };
16026 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16027 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16028 })?;
16029 let variant = if mcols == 16 {
16033 if rp { "rp" } else { "base" }
16034 } else {
16035 variant
16036 };
16037 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16044 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16045 if b567
16046 && qtype == QT_NVFP4
16047 && rp
16048 && mcols == 8
16049 && (5..=7).contains(&m)
16050 && matches!(variant, "rpsc" | "rpr2w8")
16051 {
16052 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16053 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16055 let cfg = LaunchConfig {
16056 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16057 block_dim: (32, ROWS_PER_BLOCK, 1),
16058 shared_mem_bytes: 0,
16059 };
16060 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16061 let __s_b = self.gpu.stream();
16062 let mut b = __s_b.launch_builder(&f);
16063 b.arg(bytes)
16064 .arg(aq)
16065 .arg(ad)
16066 .arg(&mut y)
16067 .arg(&inf)
16068 .arg(&outf)
16069 .arg(&mi)
16070 .arg(&rb);
16071 unsafe {
16072 b.launch(cfg)?;
16073 }
16074 if scale != 1.0 {
16075 self.scale_inplace(&mut y, scale, m * out_f)?;
16076 }
16077 return Ok(y);
16078 }
16079 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16080 "base" => (base_name.into(), ROWS_PER_BLOCK),
16081 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16082 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16083 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16084 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16088 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16089 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16090 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16091 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16092 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16093 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16094 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16096 debug_assert!(
16097 !rp || name.contains("_rp"),
16098 "rp weight dispatched to a GGUF-layout kernel"
16099 );
16100 let f = self.func(&name);
16101 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16102 let smem = if name.contains("_r2sm_rp") {
16104 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16105 } else {
16106 0
16107 };
16108 let cfg = LaunchConfig {
16109 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16110 block_dim: (32, ROWS_PER_BLOCK, 1),
16111 shared_mem_bytes: smem,
16112 };
16113 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16114 let __s_b = self.gpu.stream();
16115 let mut b = __s_b.launch_builder(&f);
16116 b.arg(bytes)
16117 .arg(aq)
16118 .arg(ad)
16119 .arg(&mut y)
16120 .arg(&inf)
16121 .arg(&outf)
16122 .arg(&mi)
16123 .arg(&rb);
16124 unsafe {
16125 b.launch(cfg)?;
16126 }
16127 if scale != 1.0 {
16128 self.scale_inplace(&mut y, scale, m * out_f)?;
16129 }
16130 Ok(y)
16131 }
16132
16133 pub fn qmatvec_batched_raw(
16137 &self,
16138 bytes: &CudaSlice<u8>,
16139 x: &CudaSlice<f32>,
16140 m: usize,
16141 in_f: usize,
16142 out_f: usize,
16143 qtype: i32,
16144 row_bytes: usize,
16145 mcols: usize,
16146 rp: bool,
16147 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16148 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16149 self.qmatvec_mmvq_batched(
16150 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16151 )
16152 }
16153
16154 pub fn qmatvec_nvfp4_batched_raw(
16156 &self,
16157 bytes: &CudaSlice<u8>,
16158 x: &CudaSlice<f32>,
16159 m: usize,
16160 in_f: usize,
16161 out_f: usize,
16162 row_bytes: usize,
16163 mcols: usize,
16164 rp: bool,
16165 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16166 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16167 }
16168
16169 fn try_fp4_gemm(
16173 &self,
16174 w: &crate::model::GpuTensor,
16175 x: &CudaSlice<f32>,
16176 m: usize,
16177 in_f: usize,
16178 out_f: usize,
16179 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16180 use crate::model::GpuTensor;
16181 if cfg!(memra_portable_cuda) {
16182 return Ok(None);
16183 }
16184 if std::env::var("MEMRA_FP4").is_ok() {
16187 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16188 }
16189 if std::env::var("MEMRA_FP4").is_err() {
16190 return Ok(None);
16191 }
16192 #[cfg(memra_cutlass)]
16201 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16202 if let GpuTensor::Quant {
16203 bytes,
16204 qtype,
16205 scale,
16206 row_bytes,
16207 cutlass,
16208 ..
16209 } = w
16210 {
16211 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16212 if let Some(cw) = cutlass {
16213 let y = self.cutlass_fp4_gemm(
16215 &cw.b_packed,
16216 &cw.sfb_swizzled,
16217 x,
16218 *scale,
16219 m,
16220 out_f,
16221 in_f,
16222 )?;
16223 return Ok(Some(y));
16224 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16225 let (b_packed, sfb_sw) =
16230 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16231 let y =
16232 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16233 return Ok(Some(y));
16234 }
16235 }
16236 }
16237 }
16238 if let GpuTensor::Quant {
16239 bytes,
16240 qtype,
16241 row_bytes,
16242 scale,
16243 rp,
16244 ..
16245 } = w
16246 {
16247 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16250 let y =
16251 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16252 return Ok(Some(y));
16253 }
16254 }
16255 Ok(None)
16256 }
16257
16258 pub fn rms_norm_f16out(
16262 &self,
16263 x: &CudaSlice<f32>,
16264 w: &CudaSlice<f32>,
16265 dst: &mut CudaSlice<f32>,
16266 dst16: &mut CudaSlice<u8>,
16267 ncols: usize,
16268 nrows: usize,
16269 eps: f32,
16270 ) -> Result<(), Box<dyn std::error::Error>> {
16271 let f = self.func("rms_norm_f16out_f32");
16272 let cfg = LaunchConfig {
16273 grid_dim: (nrows as u32, 1, 1),
16274 block_dim: (rms_block(), 1, 1),
16275 shared_mem_bytes: 0,
16276 };
16277 let (nc, e) = (ncols as i32, eps);
16278 let __s_b = self.gpu.stream();
16279 let mut b = __s_b.launch_builder(&f);
16280 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16281 unsafe {
16282 b.launch(cfg)?;
16283 }
16284 Ok(())
16285 }
16286
16287 #[allow(clippy::too_many_arguments)]
16290 pub fn add_rms_norm_f16out(
16291 &self,
16292 a: &CudaSlice<f32>,
16293 b: &CudaSlice<f32>,
16294 w: &CudaSlice<f32>,
16295 res: &mut CudaSlice<f32>,
16296 dst: &mut CudaSlice<f32>,
16297 dst16: &mut CudaSlice<u8>,
16298 ncols: usize,
16299 nrows: usize,
16300 eps: f32,
16301 ) -> Result<(), Box<dyn std::error::Error>> {
16302 let f = self.func("add_rms_norm_f16out_f32");
16303 let cfg = LaunchConfig {
16304 grid_dim: (nrows as u32, 1, 1),
16305 block_dim: (rms_block(), 1, 1),
16306 shared_mem_bytes: 0,
16307 };
16308 let (nc, e) = (ncols as i32, eps);
16309 let __s_lb = self.gpu.stream();
16310 let mut lb = __s_lb.launch_builder(&f);
16311 lb.arg(a)
16312 .arg(b)
16313 .arg(w)
16314 .arg(res)
16315 .arg(dst)
16316 .arg(dst16)
16317 .arg(&nc)
16318 .arg(&e);
16319 unsafe {
16320 lb.launch(cfg)?;
16321 }
16322 Ok(())
16323 }
16324
16325 pub fn matmul_group_xh(
16328 &self,
16329 ws: &[&crate::model::GpuTensor],
16330 x: &CudaSlice<f32>,
16331 xh: &CudaSlice<u8>,
16332 m: usize,
16333 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16334 let mut out = Vec::with_capacity(ws.len());
16335 let in_f = ws[0].in_features();
16336 for w in ws {
16337 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16338 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16339 out.push(y);
16340 continue;
16341 }
16342 }
16343 out.push(self.matmul(w, x, m)?);
16344 }
16345 Ok(out)
16346 }
16347
16348 pub fn gdn_pad_mask(
16351 &self,
16352 beta: &mut CudaSlice<f32>,
16353 g_log: &mut CudaSlice<f32>,
16354 len_d: &CudaSlice<i32>,
16355 h: usize,
16356 t: usize,
16357 ) -> Result<(), Box<dyn std::error::Error>> {
16358 let f = self.func("gdn_pad_mask_f32");
16359 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16360 let (hi, ti) = (h as i32, t as i32);
16361 let __s_b = self.gpu.stream();
16362 let mut b = __s_b.launch_builder(&f);
16363 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16364 unsafe {
16365 b.launch(cfg)?;
16366 }
16367 Ok(())
16368 }
16369
16370 pub fn row_gather_dev(
16373 &self,
16374 src: &CudaSlice<f32>,
16375 dst: &mut CudaSlice<f32>,
16376 len_d: &CudaSlice<i32>,
16377 ncols: usize,
16378 ) -> Result<(), Box<dyn std::error::Error>> {
16379 let f = self.func("row_gather_dev_f32");
16380 let cfg = LaunchConfig::for_num_elems(ncols as u32);
16381 let nc = ncols as i32;
16382 let __s_b = self.gpu.stream();
16383 let mut b = __s_b.launch_builder(&f);
16384 b.arg(src).arg(dst).arg(len_d).arg(&nc);
16385 unsafe {
16386 b.launch(cfg)?;
16387 }
16388 Ok(())
16389 }
16390
16391 pub fn matmul_group(
16398 &self,
16399 ws: &[&crate::model::GpuTensor],
16400 x: &CudaSlice<f32>,
16401 m: usize,
16402 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16403 use crate::model::GpuTensor;
16404 let mut out = Vec::with_capacity(ws.len());
16405 let any_mirror = ws
16406 .iter()
16407 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16408 if m >= 16 && any_mirror && !self.verify_exact_on() {
16409 let in_f = ws[0].in_features();
16410 let xh = self.f16_act(x, m * in_f, in_f)?;
16411 for w in ws {
16412 if w.in_features() == in_f {
16413 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16414 out.push(y);
16415 continue;
16416 }
16417 }
16418 out.push(self.matmul(w, x, m)?);
16419 }
16420 return Ok(out);
16421 }
16422 for w in ws {
16423 out.push(self.matmul(w, x, m)?);
16424 }
16425 Ok(out)
16426 }
16427
16428 pub fn matmul_group_multi(
16435 &self,
16436 ws: &[&crate::model::GpuTensor],
16437 xs: &[&CudaSlice<f32>],
16438 ms: &[usize],
16439 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16440 assert_eq!(xs.len(), ms.len());
16441 let in_f = ws[0].in_features();
16442 let total: usize = ms.iter().sum();
16443 let mut xcat = self.uninit(total * in_f)?;
16444 let mut off = 0usize;
16445 for (x, &m) in xs.iter().zip(ms) {
16446 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16447 off += m;
16448 }
16449 let ys = self.matmul_group(ws, &xcat, total)?;
16450 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16451 for (w, y) in ws.iter().zip(ys) {
16452 let out_f = w.out_features();
16453 let mut off = 0usize;
16454 for (s, &m) in ms.iter().enumerate() {
16455 let mut ys_s = self.uninit(m * out_f)?;
16456 let src = y.slice(off * out_f..(off + m) * out_f);
16457 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16458 out[s].push(ys_s);
16459 off += m;
16460 }
16461 }
16462 Ok(out)
16463 }
16464
16465 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16475 use crate::model::GpuTensor;
16476 if !legacy_quant_gemm_allowed(
16477 cfg!(memra_portable_cuda),
16478 cfg!(memra_hopper_mma),
16479 std::env::var_os("MEMRA_NO_GEMM").is_some(),
16480 ) {
16481 return false;
16482 }
16483 match w {
16484 GpuTensor::Quant { qtype, .. } => {
16485 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16486 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16487 }
16488 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16489 }
16490 }
16491
16492 pub fn qmatvec_gemm(
16499 &self,
16500 w: &crate::model::GpuTensor,
16501 aq: &CudaSlice<i8>,
16502 ad: &CudaSlice<f32>,
16503 m: usize,
16504 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16505 use crate::model::GpuTensor;
16506 let in_f = w.in_features();
16507 let out_f = w.out_features();
16508 let (bytes, qtype, row_bytes, scale, rp) = match w {
16509 GpuTensor::Quant {
16510 bytes,
16511 qtype,
16512 row_bytes,
16513 scale,
16514 rp,
16515 ..
16516 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16517 _ => unreachable!("gemm_supports guaranteed Quant"),
16518 };
16519 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16525 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16526 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16527 if scale != 1.0 {
16528 self.scale_inplace(&mut y, scale, m * out_f)?;
16529 }
16530 return Ok(y);
16531 }
16532 }
16533 let name = match qtype {
16534 QT_Q8_0 => "qmatvec_gemm_q8_0",
16535 QT_Q4_K => "qmatvec_gemm_q4_K",
16536 QT_Q4_0 => {
16537 if rp {
16538 "qmatvec_gemm_q4_0_rp"
16539 } else {
16540 "qmatvec_gemm_q4_0"
16541 }
16542 }
16543 QT_Q5_K => "qmatvec_gemm_q5_K",
16544 QT_Q6_K => "qmatvec_gemm_q6_K",
16545 QT_NVFP4 => {
16546 if rp {
16547 "qmatvec_gemm_nvfp4_rp"
16548 } else {
16549 "qmatvec_gemm_nvfp4"
16550 }
16551 }
16552 _ => unreachable!(),
16553 };
16554 let f = self.func(name);
16555 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);
16560 let k1_tile = if is_k1 {
16562 k1_launch_override().unwrap_or((128, 128, 8))
16563 } else {
16564 (128, 128, 8)
16565 };
16566 let (bm, bn): (u32, u32) = if is_k1 {
16567 (k1_tile.0, k1_tile.1)
16568 } else {
16569 (64, 256)
16570 };
16571 let warps: u32 = if is_k1 {
16572 k1_tile.2
16573 } else {
16574 match qtype {
16575 QT_NVFP4 => 8,
16576 _ => 4,
16577 }
16578 };
16579 let cfg = LaunchConfig {
16580 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16581 block_dim: (32, warps, 1),
16582 shared_mem_bytes: 0,
16583 };
16584 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16585 let __s_b = self.gpu.stream();
16586 let mut b = __s_b.launch_builder(&f);
16587 b.arg(bytes)
16588 .arg(aq)
16589 .arg(ad)
16590 .arg(&mut y)
16591 .arg(&inf)
16592 .arg(&outf)
16593 .arg(&mi)
16594 .arg(&rb);
16595 unsafe {
16596 b.launch(cfg)?;
16597 }
16598 if scale != 1.0 {
16599 self.scale_inplace(&mut y, scale, m * out_f)?;
16600 }
16601 Ok(y)
16602 }
16603
16604 pub fn qmatvec_gemm_raw(
16609 &self,
16610 bytes: &CudaSlice<u8>,
16611 x: &CudaSlice<f32>,
16612 m: usize,
16613 in_f: usize,
16614 out_f: usize,
16615 qtype: i32,
16616 row_bytes: usize,
16617 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16618 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16619 let name = match qtype {
16620 QT_Q8_0 => "qmatvec_gemm_q8_0",
16621 QT_Q4_K => "qmatvec_gemm_q4_K",
16622 QT_Q4_0 => "qmatvec_gemm_q4_0",
16623 QT_Q5_K => "qmatvec_gemm_q5_K",
16624 QT_Q6_K => "qmatvec_gemm_q6_K",
16625 QT_NVFP4 => "qmatvec_gemm_nvfp4",
16626 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16627 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16628 };
16629 let f = self.func(name);
16630 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);
16634 let k1_tile = if is_k1 {
16636 k1_launch_override().unwrap_or((128, 128, 8))
16637 } else {
16638 (128, 128, 8)
16639 };
16640 let (bm, bn): (u32, u32) = if is_k1 {
16641 (k1_tile.0, k1_tile.1)
16642 } else {
16643 (64, 256)
16644 };
16645 let warps: u32 = if is_k1 {
16646 k1_tile.2
16647 } else {
16648 match qtype {
16649 QT_NVFP4 | QT_NVFP4_RP => 8,
16650 _ => 4,
16651 }
16652 };
16653 let cfg = LaunchConfig {
16654 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16655 block_dim: (32, warps, 1),
16656 shared_mem_bytes: 0,
16657 };
16658 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16659 let __s_b = self.gpu.stream();
16660 let mut b = __s_b.launch_builder(&f);
16661 b.arg(bytes)
16662 .arg(&aq)
16663 .arg(&ad)
16664 .arg(&mut y)
16665 .arg(&inf)
16666 .arg(&outf)
16667 .arg(&mi)
16668 .arg(&rb);
16669 unsafe {
16670 b.launch(cfg)?;
16671 }
16672 Ok(y)
16673 }
16674
16675 pub fn qmatvec_gemm_q8_0_wgmma_raw(
16682 &self,
16683 rp4: &CudaSlice<u8>,
16684 aq: &CudaSlice<i8>,
16685 ad: &CudaSlice<f32>,
16686 m: usize,
16687 in_f: usize,
16688 out_f: usize,
16689 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16690 assert!(
16691 out_f % 64 == 0 && in_f % 32 == 0,
16692 "wgmma GEMM needs out_f%64==0, in_f%32==0"
16693 );
16694 let f = self.func("qmatvec_gemm_q8_0_wgmma");
16695 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
16697 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16698 block_dim: (128, 1, 1),
16699 shared_mem_bytes: 0,
16700 };
16701 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16702 let __s_b = self.gpu.stream();
16703 let mut b = __s_b.launch_builder(&f);
16704 b.arg(rp4)
16705 .arg(aq)
16706 .arg(ad)
16707 .arg(&mut y)
16708 .arg(&inf)
16709 .arg(&outf)
16710 .arg(&mi);
16711 unsafe {
16712 b.launch(cfg)?;
16713 }
16714 Ok(y)
16715 }
16716
16717 pub fn scale_inplace(
16719 &self,
16720 y: &mut CudaSlice<f32>,
16721 s: f32,
16722 n: usize,
16723 ) -> Result<(), Box<dyn std::error::Error>> {
16724 let f = self.func("scale_f32");
16725 let cfg = LaunchConfig::for_num_elems(n as u32);
16726 let (sf, ni) = (s, n as i32);
16727 let __s_b = self.gpu.stream();
16728 let mut b = __s_b.launch_builder(&f);
16729 b.arg(y).arg(&sf).arg(&ni);
16730 unsafe {
16731 b.launch(cfg)?;
16732 }
16733 Ok(())
16734 }
16735
16736 pub fn bf16_to_f32(
16741 &self,
16742 data: &cudarc::driver::CudaView<'_, u8>,
16743 n: usize,
16744 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16745 let mut out = self.alloc_uninit::<f32>(n)?;
16746 let f = self.func("bf16_to_f32");
16747 let cfg = LaunchConfig::for_num_elems(n as u32);
16748 let ni = n as i32;
16749 let __s_b = self.gpu.stream();
16750 let mut b = __s_b.launch_builder(&f);
16751 b.arg(data).arg(&mut out).arg(&ni);
16752 unsafe {
16753 b.launch(cfg)?;
16754 }
16755 Ok(out)
16756 }
16757
16758 fn linear_bf16_chunked(
16765 &self,
16766 x: &CudaSlice<f32>,
16767 data: &CudaSlice<u8>,
16768 m: usize,
16769 in_f: usize,
16770 out_f: usize,
16771 exact: bool,
16772 canonical_chunk_rows: Option<usize>,
16773 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16774 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16778 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16779 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16780 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16781 let started = timing.then(std::time::Instant::now);
16782 let result =
16783 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16784 if let Some(started) = started {
16785 use std::sync::atomic::Ordering;
16786 self.stream().synchronize()?;
16787 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16788 + started.elapsed().as_nanos() as u64;
16789 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16790 + (in_f * out_f * 2) as u64;
16791 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16792 if calls % 1024 == 0 {
16793 eprintln!(
16794 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16795 weight_gb={:.2}",
16796 ns as f64 / 1.0e6,
16797 ns as f64 / calls as f64 / 1.0e3,
16798 wb as f64 / 1.0e9,
16799 );
16800 }
16801 }
16802 result
16803 }
16804
16805 pub(crate) fn bf16_mmv_on() -> bool {
16810 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16811 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16812 }
16813
16814 fn matvec_bf16(
16817 &self,
16818 data: &CudaSlice<u8>,
16819 x: &CudaSlice<f32>,
16820 in_f: usize,
16821 out_f: usize,
16822 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16823 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16824 return Err(format!(
16825 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16826 data.len(),
16827 x.len()
16828 )
16829 .into());
16830 }
16831 let mut y = self.alloc_uninit::<f32>(out_f)?;
16832 let f = self.func("matvec_bf16_f32acc");
16833 let cfg = LaunchConfig {
16834 grid_dim: (out_f as u32, 1, 1),
16835 block_dim: (mmv_block(), 1, 1),
16836 shared_mem_bytes: 0,
16837 };
16838 let ini = in_f as i32;
16839 let __s_bld = self.gpu.stream();
16840 let mut bld = __s_bld.launch_builder(&f);
16841 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
16842 unsafe {
16843 bld.launch(cfg)?;
16844 }
16845 Ok(y)
16846 }
16847
16848 #[allow(clippy::too_many_arguments)]
16852 #[allow(clippy::too_many_arguments)]
16857 pub fn qk_norm_rope_append_inc_dcw(
16858 &self,
16859 q_raw: &CudaSlice<f32>,
16860 k_raw: &CudaSlice<f32>,
16861 v_raw: &CudaSlice<f32>,
16862 qw: &CudaSlice<f32>,
16863 kw: &CudaSlice<f32>,
16864 q_out: &mut CudaSlice<f32>,
16865 k_out: &mut CudaSlice<f32>,
16866 pos: &CudaSlice<i32>,
16867 k_plane: &mut CudaSlice<u8>,
16868 v_plane: &mut CudaSlice<u8>,
16869 len_dev: &CudaSlice<i32>,
16872 base_dev: Option<&CudaSlice<i32>>,
16873 done_ctr: &mut CudaSlice<u32>,
16874 kv_dim_k: usize,
16875 kv_dim_v: usize,
16876 k_tok_bytes: usize,
16877 v_tok_bytes: usize,
16878 head_dim: usize,
16879 n_dims: usize,
16880 nh_q: usize,
16881 nh_k: usize,
16882 eps: f32,
16883 freq_base: f32,
16884 freq_scale: f32,
16885 ff: Option<&CudaSlice<f32>>,
16886 ) -> Result<(), Box<dyn std::error::Error>> {
16887 if head_dim != 128
16888 || kv_dim_v != kv_dim_k
16889 || kv_dim_k != nh_k * head_dim
16890 || q_raw.len() < nh_q * head_dim
16891 || k_raw.len() < nh_k * head_dim
16892 || v_raw.len() < kv_dim_v
16893 || q_out.len() < nh_q * head_dim
16894 || k_out.len() < nh_k * head_dim
16895 || pos.is_empty()
16896 || done_ctr.is_empty()
16897 {
16898 return Err(format!(
16899 "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}"
16900 )
16901 .into());
16902 }
16903 let f = self.func("qk_norm_rope_append_inc_dcw");
16904 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16905 let cfg = LaunchConfig {
16906 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16907 block_dim: (128, 1, 1),
16908 shared_mem_bytes: 0,
16909 };
16910 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
16911 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16912 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16913 let null: u64 = 0;
16914 let __s_b = self.gpu.stream();
16915 let mut b = __s_b.launch_builder(&f);
16916 b.arg(q_raw)
16917 .arg(k_raw)
16918 .arg(v_raw)
16919 .arg(qw)
16920 .arg(kw)
16921 .arg(q_out)
16922 .arg(k_out)
16923 .arg(pos)
16924 .arg(&mut *k_plane)
16925 .arg(&mut *v_plane)
16926 .arg(len_dev);
16927 match base_dev {
16928 Some(base) => {
16929 b.arg(base);
16930 }
16931 None => {
16932 b.arg(&null);
16933 }
16934 }
16935 b.arg(&mut *done_ctr)
16936 .arg(&kvk)
16937 .arg(&kvv)
16938 .arg(&ktb)
16939 .arg(&vtb)
16940 .arg(&hd)
16941 .arg(&nd)
16942 .arg(&nq)
16943 .arg(&eps)
16944 .arg(&theta_scale)
16945 .arg(&freq_scale);
16946 match ff {
16947 Some(freqs) => {
16948 b.arg(freqs);
16949 }
16950 None => {
16951 b.arg(&null);
16952 }
16953 }
16954 unsafe {
16955 b.launch(cfg)?;
16956 }
16957 Ok(())
16958 }
16959
16960 pub fn qk_norm_rope_into(
16961 &self,
16962 q_raw: &CudaSlice<f32>,
16963 k_raw: &CudaSlice<f32>,
16964 qw: &CudaSlice<f32>,
16965 kw: &CudaSlice<f32>,
16966 q_out: &mut CudaSlice<f32>,
16967 k_out: &mut CudaSlice<f32>,
16968 pos: &CudaSlice<i32>,
16969 head_dim: usize,
16970 n_dims: usize,
16971 nh_q: usize,
16972 nh_k: usize,
16973 eps: f32,
16974 freq_base: f32,
16975 freq_scale: f32,
16976 ff: Option<&CudaSlice<f32>>,
16977 ) -> Result<(), Box<dyn std::error::Error>> {
16978 if head_dim > 512
16979 || q_raw.len() < nh_q * head_dim
16980 || k_raw.len() < nh_k * head_dim
16981 || q_out.len() < nh_q * head_dim
16982 || k_out.len() < nh_k * head_dim
16983 || qw.len() < head_dim
16984 || kw.len() < head_dim
16985 || pos.is_empty()
16986 {
16987 return Err(format!(
16988 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
16989 )
16990 .into());
16991 }
16992 let f = self.func("qk_norm_rope_f32");
16993 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16994 let cfg = LaunchConfig {
16995 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16996 block_dim: (128, 1, 1),
16997 shared_mem_bytes: 0,
16998 };
16999 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17000 let __s_b = self.gpu.stream();
17001 let mut b = __s_b.launch_builder(&f);
17002 b.arg(q_raw)
17003 .arg(k_raw)
17004 .arg(qw)
17005 .arg(kw)
17006 .arg(q_out)
17007 .arg(k_out)
17008 .arg(pos)
17009 .arg(&hd)
17010 .arg(&nd)
17011 .arg(&nq)
17012 .arg(&eps)
17013 .arg(&theta_scale)
17014 .arg(&freq_scale);
17015 match ff {
17016 Some(ffv) => {
17017 b.arg(ffv);
17018 unsafe {
17019 b.launch(cfg)?;
17020 }
17021 }
17022 None => {
17023 let null: u64 = 0;
17024 b.arg(&null);
17025 unsafe {
17026 b.launch(cfg)?;
17027 }
17028 }
17029 }
17030 Ok(())
17031 }
17032
17033 #[allow(clippy::too_many_arguments)]
17036 pub fn matvec_f32_b4_into(
17037 &self,
17038 w: [&CudaSlice<f32>; 4],
17039 x: &CudaSlice<f32>,
17040 y: &mut CudaSlice<f32>,
17041 block_cols: usize,
17042 out_f: usize,
17043 ) -> Result<(), Box<dyn std::error::Error>> {
17044 if block_cols % 4 != 0
17045 || x.len() < 4 * block_cols
17046 || y.len() < out_f
17047 || w.iter().any(|w| w.len() != out_f * block_cols)
17048 {
17049 return Err(format!(
17050 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17051 x.len()
17052 )
17053 .into());
17054 }
17055 let f = self.func("matvec_f32_b4");
17056 let cfg = LaunchConfig {
17057 grid_dim: (out_f as u32, 1, 1),
17058 block_dim: (128, 1, 1),
17059 shared_mem_bytes: 0,
17060 };
17061 let (bc, of) = (block_cols as i32, out_f as i32);
17062 let __s_b = self.gpu.stream();
17063 let mut b = __s_b.launch_builder(&f);
17064 b.arg(w[0])
17065 .arg(w[1])
17066 .arg(w[2])
17067 .arg(w[3])
17068 .arg(x)
17069 .arg(y)
17070 .arg(&bc)
17071 .arg(&of);
17072 unsafe {
17073 b.launch(cfg)?;
17074 }
17075 Ok(())
17076 }
17077
17078 pub fn axpy_rows_seq_into(
17081 &self,
17082 x: &CudaSlice<f32>,
17083 w: &CudaSlice<f32>,
17084 y: &mut CudaSlice<f32>,
17085 width: usize,
17086 n_rows: usize,
17087 ) -> Result<(), Box<dyn std::error::Error>> {
17088 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17089 return Err(format!(
17090 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17091 x.len(),
17092 w.len(),
17093 y.len()
17094 )
17095 .into());
17096 }
17097 let f = self.func("axpy_rows_seq_f32");
17098 let cfg = LaunchConfig::for_num_elems(width as u32);
17099 let (wi, nr) = (width as i32, n_rows as i32);
17100 let __s_b = self.gpu.stream();
17101 let mut b = __s_b.launch_builder(&f);
17102 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17103 unsafe {
17104 b.launch(cfg)?;
17105 }
17106 Ok(())
17107 }
17108
17109 #[allow(clippy::too_many_arguments)]
17113 pub fn axpy_rows_seq_md_off_into(
17114 &self,
17115 x: &CudaSlice<f32>,
17116 w_route: &CudaSlice<f32>,
17117 md: &CudaSlice<f32>,
17118 sel: &CudaSlice<i32>,
17119 y: &mut CudaSlice<f32>,
17120 width: usize,
17121 n_rows: usize,
17122 row0: usize,
17123 ) -> Result<(), Box<dyn std::error::Error>> {
17124 if x.len() < (row0 + n_rows) * width
17125 || w_route.len() < row0 + n_rows
17126 || sel.len() < row0 + n_rows
17127 || y.len() < width
17128 {
17129 return Err(format!(
17130 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17131 rows={n_rows} row0={row0}",
17132 x.len(),
17133 w_route.len(),
17134 sel.len(),
17135 y.len()
17136 )
17137 .into());
17138 }
17139 let f = self.func("axpy_rows_seq_md_off_f32");
17140 let cfg = LaunchConfig::for_num_elems(width as u32);
17141 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17142 let __s_b = self.gpu.stream();
17143 let mut b = __s_b.launch_builder(&f);
17144 b.arg(x)
17145 .arg(w_route)
17146 .arg(md)
17147 .arg(sel)
17148 .arg(y)
17149 .arg(&wi)
17150 .arg(&nr)
17151 .arg(&r0);
17152 unsafe {
17153 b.launch(cfg)?;
17154 }
17155 Ok(())
17156 }
17157
17158 #[allow(clippy::too_many_arguments)]
17163 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17164 &self,
17165 gate_bank: &CudaSlice<u8>,
17166 up_bank: &CudaSlice<u8>,
17167 sel: &CudaSlice<i32>,
17168 aq: &CudaSlice<i8>,
17169 ad: &CudaSlice<f32>,
17170 yg: &mut CudaSlice<f32>,
17171 yu: &mut CudaSlice<f32>,
17172 n_sel: usize,
17173 n_sel_col: usize,
17174 in_f: usize,
17175 out_f: usize,
17176 row_bytes: usize,
17177 expert_stride: usize,
17178 act_row_stride: usize,
17179 ad_row_stride: usize,
17180 ) -> Result<(), Box<dyn std::error::Error>> {
17181 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17182 if yg.len() < n_sel * out_f
17183 || yu.len() < n_sel * out_f
17184 || sel.len() < n_sel
17185 || n_sel_col == 0
17186 || n_sel % n_sel_col != 0
17187 {
17188 return Err("NVFP4 gu tcol geometry".into());
17189 }
17190 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17191 let cfg = LaunchConfig {
17192 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17193 block_dim: (128, 1, 1),
17194 shared_mem_bytes: 0,
17195 };
17196 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17197 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17198 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17199 let __s_b = self.gpu.stream();
17200 let mut b = __s_b.launch_builder(&f);
17201 b.arg(gate_bank)
17202 .arg(up_bank)
17203 .arg(sel)
17204 .arg(aq)
17205 .arg(ad)
17206 .arg(yg)
17207 .arg(yu)
17208 .arg(&inf)
17209 .arg(&outf)
17210 .arg(&ns)
17211 .arg(&rb)
17212 .arg(&es)
17213 .arg(&ars)
17214 .arg(&adrs)
17215 .arg(&nsc);
17216 unsafe {
17217 b.launch(cfg)?;
17218 }
17219 Ok(())
17220 }
17221
17222 #[allow(clippy::too_many_arguments)]
17225 pub fn axpy_rows_seq_md_into(
17226 &self,
17227 x: &CudaSlice<f32>,
17228 w_route: &CudaSlice<f32>,
17229 md: &CudaSlice<f32>,
17230 sel: &CudaSlice<i32>,
17231 y: &mut CudaSlice<f32>,
17232 width: usize,
17233 n_rows: usize,
17234 ) -> Result<(), Box<dyn std::error::Error>> {
17235 if x.len() < n_rows * width
17236 || w_route.len() < n_rows
17237 || sel.len() < n_rows
17238 || y.len() < width
17239 {
17240 return Err(format!(
17241 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17242 x.len(),
17243 w_route.len(),
17244 sel.len(),
17245 y.len()
17246 )
17247 .into());
17248 }
17249 let f = self.func("axpy_rows_seq_md_f32");
17250 let cfg = LaunchConfig::for_num_elems(width as u32);
17251 let (wi, nr) = (width as i32, n_rows as i32);
17252 let __s_b = self.gpu.stream();
17253 let mut b = __s_b.launch_builder(&f);
17254 b.arg(x)
17255 .arg(w_route)
17256 .arg(md)
17257 .arg(sel)
17258 .arg(y)
17259 .arg(&wi)
17260 .arg(&nr);
17261 unsafe {
17262 b.launch(cfg)?;
17263 }
17264 Ok(())
17265 }
17266
17267 #[allow(clippy::too_many_arguments)]
17269 #[allow(clippy::too_many_arguments)]
17273 pub fn matvec_bf16_qkvg_tcol_into(
17274 &self,
17275 wq: &CudaSlice<u8>,
17276 wk: &CudaSlice<u8>,
17277 wv: &CudaSlice<u8>,
17278 wg: &CudaSlice<u8>,
17279 x_t: &CudaSlice<f32>,
17280 yq: &mut CudaSlice<f32>,
17281 yk: &mut CudaSlice<f32>,
17282 yv: &mut CudaSlice<f32>,
17283 yg: &mut CudaSlice<f32>,
17284 in_f: usize,
17285 out_q: usize,
17286 out_kv: usize,
17287 out_g: usize,
17288 t: usize,
17289 ) -> Result<(), Box<dyn std::error::Error>> {
17290 if t == 0
17291 || t > 8
17292 || in_f % 8 != 0
17293 || x_t.len() < t * in_f
17294 || yq.len() < t * out_q
17295 || yk.len() < t * out_kv
17296 || yv.len() < t * out_kv
17297 || (out_g > 0 && yg.len() < t * out_g)
17298 {
17299 return Err("matvec_bf16_qkvg_tcol geometry".into());
17300 }
17301 let f = self.func("matvec_bf16_qkvg_tcol");
17302 let grid = out_q + 2 * out_kv + out_g;
17303 let cfg = LaunchConfig {
17304 grid_dim: (grid as u32, 1, 1),
17305 block_dim: (mmv_block(), 1, 1),
17306 shared_mem_bytes: 0,
17307 };
17308 let (ini, oq, okv, og, ti) = (
17309 in_f as i32,
17310 out_q as i32,
17311 out_kv as i32,
17312 out_g as i32,
17313 t as i32,
17314 );
17315 let __s_b = self.gpu.stream();
17316 let mut b = __s_b.launch_builder(&f);
17317 b.arg(wq)
17318 .arg(wk)
17319 .arg(wv)
17320 .arg(wg)
17321 .arg(x_t)
17322 .arg(yq)
17323 .arg(yk)
17324 .arg(yv)
17325 .arg(yg)
17326 .arg(&ini)
17327 .arg(&oq)
17328 .arg(&okv)
17329 .arg(&og)
17330 .arg(&ti);
17331 unsafe {
17332 b.launch(cfg)?;
17333 }
17334 Ok(())
17335 }
17336
17337 pub fn matvec_bf16_qkvg_into(
17338 &self,
17339 wq: &CudaSlice<u8>,
17340 wk: &CudaSlice<u8>,
17341 wv: &CudaSlice<u8>,
17342 wg: &CudaSlice<u8>,
17343 x: &CudaSlice<f32>,
17344 yq: &mut CudaSlice<f32>,
17345 yk: &mut CudaSlice<f32>,
17346 yv: &mut CudaSlice<f32>,
17347 yg: &mut CudaSlice<f32>,
17348 in_f: usize,
17349 out_q: usize,
17350 out_kv: usize,
17351 out_g: usize,
17352 ) -> Result<(), Box<dyn std::error::Error>> {
17353 if in_f % 8 != 0
17354 || wq.len() != out_q * in_f * 2
17355 || wk.len() != out_kv * in_f * 2
17356 || wv.len() != out_kv * in_f * 2
17357 || wg.len() < out_g * in_f * 2
17358 || x.len() < in_f
17359 || yq.len() < out_q
17360 || yk.len() < out_kv
17361 || yv.len() < out_kv
17362 || (out_g > 0 && yg.len() < out_g)
17363 {
17364 return Err(format!(
17365 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17366 )
17367 .into());
17368 }
17369 let f = self.func("matvec_bf16_qkvg");
17370 let cfg = LaunchConfig {
17371 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17372 block_dim: (mmv_block(), 1, 1),
17373 shared_mem_bytes: 0,
17374 };
17375 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17376 let __s_b = self.gpu.stream();
17377 let mut b = __s_b.launch_builder(&f);
17378 b.arg(wq)
17379 .arg(wk)
17380 .arg(wv)
17381 .arg(wg)
17382 .arg(x)
17383 .arg(yq)
17384 .arg(yk)
17385 .arg(yv)
17386 .arg(yg)
17387 .arg(&inf)
17388 .arg(&oq)
17389 .arg(&okv)
17390 .arg(&og);
17391 unsafe {
17392 b.launch(cfg)?;
17393 }
17394 Ok(())
17395 }
17396
17397 pub fn matvec_bf16_b4_into(
17399 &self,
17400 w: [&CudaSlice<u8>; 4],
17401 x: &CudaSlice<f32>,
17402 y: &mut CudaSlice<f32>,
17403 block_cols: usize,
17404 out_f: usize,
17405 ) -> Result<(), Box<dyn std::error::Error>> {
17406 if block_cols % 8 != 0
17407 || x.len() < 4 * block_cols
17408 || y.len() < out_f
17409 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17410 {
17411 return Err(format!(
17412 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17413 x.len()
17414 )
17415 .into());
17416 }
17417 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17420 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17421 let f = self.func(if x2 {
17422 "matvec_bf16_b4_x2"
17423 } else {
17424 "matvec_bf16_b4"
17425 });
17426 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17427 let cfg = LaunchConfig {
17428 grid_dim: (grid as u32, 1, 1),
17429 block_dim: (mmv_block(), 1, 1),
17430 shared_mem_bytes: 0,
17431 };
17432 let (bc, of) = (block_cols as i32, out_f as i32);
17433 let __s_b = self.gpu.stream();
17434 let mut b = __s_b.launch_builder(&f);
17435 b.arg(w[0])
17436 .arg(w[1])
17437 .arg(w[2])
17438 .arg(w[3])
17439 .arg(x)
17440 .arg(y)
17441 .arg(&bc)
17442 .arg(&of);
17443 unsafe {
17444 b.launch(cfg)?;
17445 }
17446 Ok(())
17447 }
17448
17449 pub fn matvec_bf16_b4_tcol_into(
17455 &self,
17456 w: [&CudaSlice<u8>; 4],
17457 x_t: &CudaSlice<f32>,
17458 y_t: &mut CudaSlice<f32>,
17459 block_cols: usize,
17460 out_f: usize,
17461 t: usize,
17462 ) -> Result<(), Box<dyn std::error::Error>> {
17463 if block_cols % 8 != 0
17464 || t == 0
17465 || t > 8
17466 || x_t.len() < t * 4 * block_cols
17467 || y_t.len() < t * out_f
17468 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17469 {
17470 return Err(format!(
17471 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17472 x_t.len()
17473 )
17474 .into());
17475 }
17476 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17477 return Err(
17478 "b4 tcol verify is qualified against the plain b4 kernel only \
17479 (MEMRA_B4_X2=1 is a different t=1 program)"
17480 .into(),
17481 );
17482 }
17483 let f = self.func("matvec_bf16_b4_tcol");
17484 let cfg = LaunchConfig {
17485 grid_dim: (out_f as u32, 1, 1),
17486 block_dim: (mmv_block(), 1, 1),
17487 shared_mem_bytes: 0,
17488 };
17489 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17490 let __s_b = self.gpu.stream();
17491 let mut b = __s_b.launch_builder(&f);
17492 b.arg(w[0])
17493 .arg(w[1])
17494 .arg(w[2])
17495 .arg(w[3])
17496 .arg(x_t)
17497 .arg(y_t)
17498 .arg(&bc)
17499 .arg(&of)
17500 .arg(&ti);
17501 unsafe {
17502 b.launch(cfg)?;
17503 }
17504 Ok(())
17505 }
17506
17507 pub fn matvec_bf16_into(
17509 &self,
17510 data: &CudaSlice<u8>,
17511 x: &CudaSlice<f32>,
17512 y: &mut CudaSlice<f32>,
17513 in_f: usize,
17514 out_f: usize,
17515 ) -> Result<(), Box<dyn std::error::Error>> {
17516 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17517 return Err(format!(
17518 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17519 data.len(),
17520 x.len(),
17521 y.len()
17522 )
17523 .into());
17524 }
17525 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17529 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17530 && in_f <= 2048;
17531 if x4 {
17532 let f = self.func("matvec_bf16_f32acc_x4");
17533 let cfg = LaunchConfig {
17534 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17535 block_dim: (mmv_block(), 1, 1),
17536 shared_mem_bytes: 0,
17537 };
17538 let (ini, outi) = (in_f as i32, out_f as i32);
17539 let __s_b = self.gpu.stream();
17540 let mut b = __s_b.launch_builder(&f);
17541 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17542 unsafe {
17543 b.launch(cfg)?;
17544 }
17545 return Ok(());
17546 }
17547 let f = self.func("matvec_bf16_f32acc");
17548 let cfg = LaunchConfig {
17549 grid_dim: (out_f as u32, 1, 1),
17550 block_dim: (mmv_block(), 1, 1),
17551 shared_mem_bytes: 0,
17552 };
17553 let ini = in_f as i32;
17554 let __s_b = self.gpu.stream();
17555 let mut b = __s_b.launch_builder(&f);
17556 b.arg(data).arg(x).arg(y).arg(&ini);
17557 unsafe {
17558 b.launch(cfg)?;
17559 }
17560 Ok(())
17561 }
17562
17563 pub fn matvec_bf16_view_into(
17566 &self,
17567 data: &cudarc::driver::CudaView<'_, u8>,
17568 x: &CudaSlice<f32>,
17569 y: &mut CudaSlice<f32>,
17570 in_f: usize,
17571 out_f: usize,
17572 ) -> Result<(), Box<dyn std::error::Error>> {
17573 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17574 return Err(format!(
17575 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17576 data.len(),
17577 x.len(),
17578 y.len()
17579 )
17580 .into());
17581 }
17582 let f = self.func("matvec_bf16_f32acc");
17583 let cfg = LaunchConfig {
17584 grid_dim: (out_f as u32, 1, 1),
17585 block_dim: (mmv_block(), 1, 1),
17586 shared_mem_bytes: 0,
17587 };
17588 let ini = in_f as i32;
17589 let __s_b = self.gpu.stream();
17590 let mut b = __s_b.launch_builder(&f);
17591 b.arg(data).arg(x).arg(y).arg(&ini);
17592 unsafe {
17593 b.launch(cfg)?;
17594 }
17595 Ok(())
17596 }
17597
17598 pub fn matvec_bf16_raw_out(
17601 &self,
17602 w: &CudaSlice<u8>,
17603 x: &CudaSlice<f32>,
17604 y_raw: u64,
17605 in_f: usize,
17606 out_f: usize,
17607 ) -> Result<(), Box<dyn std::error::Error>> {
17608 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17609 return Err("matvec_bf16_raw_out geometry".into());
17610 }
17611 let f = self.func("matvec_bf16_f32acc");
17612 let cfg = LaunchConfig {
17613 grid_dim: (out_f as u32, 1, 1),
17614 block_dim: (mmv_block(), 1, 1),
17615 shared_mem_bytes: 0,
17616 };
17617 let ini = in_f as i32;
17618 let __s_b = self.gpu.stream();
17619 let mut b = __s_b.launch_builder(&f);
17620 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17621 unsafe {
17622 b.launch(cfg)?;
17623 }
17624 Ok(())
17625 }
17626
17627 pub fn add3_raw(
17631 &self,
17632 a: &CudaSlice<f32>,
17633 b: &CudaSlice<f32>,
17634 sh_raw: u64,
17635 scale_raw: u64,
17636 dst: &mut CudaSlice<f32>,
17637 n: usize,
17638 ) -> Result<(), Box<dyn std::error::Error>> {
17639 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17640 return Err("add3_raw geometry".into());
17641 }
17642 let f = self.func("add3_f32");
17643 let cfg = LaunchConfig {
17644 grid_dim: ((n as u32).div_ceil(256), 1, 1),
17645 block_dim: (256, 1, 1),
17646 shared_mem_bytes: 0,
17647 };
17648 let ni = n as i32;
17649 let __s_b = self.gpu.stream();
17650 let mut bld = __s_b.launch_builder(&f);
17651 bld.arg(a)
17652 .arg(b)
17653 .arg(&sh_raw)
17654 .arg(&scale_raw)
17655 .arg(dst)
17656 .arg(&ni);
17657 unsafe {
17658 bld.launch(cfg)?;
17659 }
17660 Ok(())
17661 }
17662
17663 pub fn matvec_bf16_down_addscale_into(
17666 &self,
17667 w: &CudaSlice<u8>,
17668 x: &CudaSlice<f32>,
17669 scale: &CudaSlice<f32>,
17670 dst: &mut CudaSlice<f32>,
17671 in_f: usize,
17672 out_f: usize,
17673 ) -> Result<(), Box<dyn std::error::Error>> {
17674 if w.len() != in_f * out_f * 2
17675 || x.len() < in_f
17676 || in_f % 8 != 0
17677 || dst.len() < out_f
17678 || scale.is_empty()
17679 {
17680 return Err("matvec_bf16_down_addscale geometry".into());
17681 }
17682 let f = self.func("matvec_bf16_down_addscale");
17683 let cfg = LaunchConfig {
17684 grid_dim: (out_f as u32, 1, 1),
17685 block_dim: (mmv_block(), 1, 1),
17686 shared_mem_bytes: 0,
17687 };
17688 let ini = in_f as i32;
17689 let __s_b = self.gpu.stream();
17690 let mut b = __s_b.launch_builder(&f);
17691 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17692 unsafe {
17693 b.launch(cfg)?;
17694 }
17695 Ok(())
17696 }
17697
17698 pub fn matvec_bf16_dual_silu_into(
17701 &self,
17702 wg: &CudaSlice<u8>,
17703 wu: &CudaSlice<u8>,
17704 x: &CudaSlice<f32>,
17705 act: &mut CudaSlice<f32>,
17706 in_f: usize,
17707 out_f: usize,
17708 limit: Option<f32>,
17709 ) -> Result<(), Box<dyn std::error::Error>> {
17710 if wg.len() != in_f * out_f * 2
17711 || wu.len() != in_f * out_f * 2
17712 || x.len() < in_f
17713 || in_f % 8 != 0
17714 || act.len() < out_f
17715 {
17716 return Err("matvec_bf16_dual_silu geometry".into());
17717 }
17718 let f = self.func("matvec_bf16_dual_silu");
17719 let cfg = LaunchConfig {
17720 grid_dim: (out_f as u32, 1, 1),
17721 block_dim: (mmv_block(), 1, 1),
17722 shared_mem_bytes: 0,
17723 };
17724 let (ini, outi) = (in_f as i32, out_f as i32);
17725 let lim = limit.unwrap_or(0.0);
17726 let __s_b = self.gpu.stream();
17727 let mut b = __s_b.launch_builder(&f);
17728 b.arg(wg)
17729 .arg(wu)
17730 .arg(x)
17731 .arg(act)
17732 .arg(&ini)
17733 .arg(&outi)
17734 .arg(&lim);
17735 unsafe {
17736 b.launch(cfg)?;
17737 }
17738 Ok(())
17739 }
17740
17741 #[allow(clippy::too_many_arguments)]
17744 pub fn matvec_bf16_dual_view_into(
17745 &self,
17746 wg: &cudarc::driver::CudaView<'_, u8>,
17747 wu: &cudarc::driver::CudaView<'_, u8>,
17748 x: &CudaSlice<f32>,
17749 yg: &mut CudaSlice<f32>,
17750 yu: &mut CudaSlice<f32>,
17751 in_f: usize,
17752 out_f: usize,
17753 ) -> Result<(), Box<dyn std::error::Error>> {
17754 if wg.len() != in_f * out_f * 2
17755 || wu.len() != in_f * out_f * 2
17756 || x.len() < in_f
17757 || in_f % 8 != 0
17758 || yg.len() < out_f
17759 || yu.len() < out_f
17760 {
17761 return Err(format!(
17762 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17763 wg.len(),
17764 wu.len(),
17765 x.len()
17766 )
17767 .into());
17768 }
17769 let f = self.func("matvec_bf16_dual");
17770 let cfg = LaunchConfig {
17771 grid_dim: ((2 * out_f) as u32, 1, 1),
17772 block_dim: (mmv_block(), 1, 1),
17773 shared_mem_bytes: 0,
17774 };
17775 let (ini, outi) = (in_f as i32, out_f as i32);
17776 let __s_b = self.gpu.stream();
17777 let mut b = __s_b.launch_builder(&f);
17778 b.arg(wg)
17779 .arg(wu)
17780 .arg(x)
17781 .arg(yg)
17782 .arg(yu)
17783 .arg(&ini)
17784 .arg(&outi);
17785 unsafe {
17786 b.launch(cfg)?;
17787 }
17788 Ok(())
17789 }
17790
17791 #[allow(clippy::too_many_arguments)]
17793 pub fn matvec_bf16_dual_into(
17794 &self,
17795 wg: &CudaSlice<u8>,
17796 wu: &CudaSlice<u8>,
17797 x: &CudaSlice<f32>,
17798 yg: &mut CudaSlice<f32>,
17799 yu: &mut CudaSlice<f32>,
17800 in_f: usize,
17801 out_f: usize,
17802 ) -> Result<(), Box<dyn std::error::Error>> {
17803 if wg.len() != in_f * out_f * 2
17804 || wu.len() != in_f * out_f * 2
17805 || x.len() < in_f
17806 || in_f % 8 != 0
17807 || yg.len() < out_f
17808 || yu.len() < out_f
17809 {
17810 return Err(format!(
17811 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17812 wg.len(),
17813 wu.len(),
17814 x.len()
17815 )
17816 .into());
17817 }
17818 let f = self.func("matvec_bf16_dual");
17819 let cfg = LaunchConfig {
17820 grid_dim: ((2 * out_f) as u32, 1, 1),
17821 block_dim: (mmv_block(), 1, 1),
17822 shared_mem_bytes: 0,
17823 };
17824 let (ini, outi) = (in_f as i32, out_f as i32);
17825 let __s_b = self.gpu.stream();
17826 let mut b = __s_b.launch_builder(&f);
17827 b.arg(wg)
17828 .arg(wu)
17829 .arg(x)
17830 .arg(yg)
17831 .arg(yu)
17832 .arg(&ini)
17833 .arg(&outi);
17834 unsafe {
17835 b.launch(cfg)?;
17836 }
17837 Ok(())
17838 }
17839
17840 pub(crate) fn matvec_bf16_dual(
17843 &self,
17844 wg: &CudaSlice<u8>,
17845 wu: &CudaSlice<u8>,
17846 x: &CudaSlice<f32>,
17847 in_f: usize,
17848 out_f: usize,
17849 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17850 if wg.len() != in_f * out_f * 2
17851 || wu.len() != in_f * out_f * 2
17852 || x.len() < in_f
17853 || in_f % 8 != 0
17854 {
17855 return Err(format!(
17856 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
17857 wg.len(),
17858 wu.len(),
17859 x.len()
17860 )
17861 .into());
17862 }
17863 let mut yg = self.alloc_uninit::<f32>(out_f)?;
17864 let mut yu = self.alloc_uninit::<f32>(out_f)?;
17865 let f = self.func("matvec_bf16_dual");
17866 let cfg = LaunchConfig {
17867 grid_dim: ((2 * out_f) as u32, 1, 1),
17868 block_dim: (mmv_block(), 1, 1),
17869 shared_mem_bytes: 0,
17870 };
17871 let (ini, outi) = (in_f as i32, out_f as i32);
17872 let __s_b = self.gpu.stream();
17873 let mut b = __s_b.launch_builder(&f);
17874 b.arg(wg)
17875 .arg(wu)
17876 .arg(x)
17877 .arg(&mut yg)
17878 .arg(&mut yu)
17879 .arg(&ini)
17880 .arg(&outi);
17881 unsafe {
17882 b.launch(cfg)?;
17883 }
17884 Ok((yg, yu))
17885 }
17886
17887 #[allow(clippy::too_many_arguments)]
17888 fn linear_bf16_chunked_inner(
17889 &self,
17890 x: &CudaSlice<f32>,
17891 data: &CudaSlice<u8>,
17892 m: usize,
17893 in_f: usize,
17894 out_f: usize,
17895 exact: bool,
17896 canonical_chunk_rows: Option<usize>,
17897 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17898 const CHUNK_BYTES: usize = 256 << 20;
17899 if m == 1
17902 && !exact
17903 && canonical_chunk_rows.is_none()
17904 && in_f % 8 == 0
17905 && Self::bf16_mmv_on()
17906 {
17907 return self.matvec_bf16(data, x, in_f, out_f);
17908 }
17909 let row_bytes = in_f
17910 .checked_mul(std::mem::size_of::<f32>())
17911 .ok_or("BF16 chunk row byte count overflow")?;
17912 if row_bytes == 0 || out_f == 0 {
17913 return Err("BF16 chunk dimensions must be nonzero".into());
17914 }
17915 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
17916 let chunk_rows = match canonical_chunk_rows {
17917 Some(rows) if rows == 0 => {
17918 return Err("canonical BF16 chunk rows must be nonzero".into());
17919 }
17920 Some(rows) if rows > max_chunk_rows => {
17921 return Err(format!(
17922 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
17923 )
17924 .into());
17925 }
17926 Some(rows) if out_f % rows != 0 => {
17927 return Err(format!(
17928 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
17929 )
17930 .into());
17931 }
17932 Some(rows) => rows,
17933 None => max_chunk_rows,
17934 };
17935 if chunk_rows >= out_f {
17936 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
17937 return if exact {
17938 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
17939 } else {
17940 self.linear(x, &wf32, m, in_f, out_f)
17941 };
17942 }
17943 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17944 let mut r0 = 0usize;
17945 while r0 < out_f {
17946 let rows = chunk_rows.min(out_f - r0);
17947 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
17948 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
17949 let yc = if exact {
17950 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
17951 } else {
17952 self.linear(x, &wf32, m, in_f, rows)?
17953 };
17954 for mi in 0..m {
17956 let src = yc.slice(mi * rows..(mi + 1) * rows);
17957 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
17958 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
17959 }
17960 r0 += rows;
17961 }
17962 Ok(y)
17963 }
17964
17965 pub fn linear_bf16_resident(
17969 &self,
17970 x: &CudaSlice<f32>,
17971 data: &CudaSlice<u8>,
17972 m: usize,
17973 in_f: usize,
17974 out_f: usize,
17975 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17976 if data.len() != in_f * out_f * 2 {
17977 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17978 }
17979 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
17980 }
17981
17982 pub fn linear_bf16_resident_canonical_rows(
17988 &self,
17989 x: &CudaSlice<f32>,
17990 data: &CudaSlice<u8>,
17991 m: usize,
17992 in_f: usize,
17993 out_f: usize,
17994 canonical_chunk_rows: usize,
17995 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17996 if data.len() != in_f * out_f * 2 {
17997 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17998 }
17999 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18000 }
18001
18002 pub fn linear_f32_resident_canonical_rows(
18007 &self,
18008 x: &CudaSlice<f32>,
18009 data: &CudaSlice<f32>,
18010 m: usize,
18011 in_f: usize,
18012 out_f: usize,
18013 canonical_chunk_rows: usize,
18014 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18015 self.linear_f32_resident_canonical_rows_inner(
18016 x,
18017 data,
18018 m,
18019 in_f,
18020 out_f,
18021 canonical_chunk_rows,
18022 false,
18023 )
18024 }
18025
18026 pub fn linear_f32_resident_canonical_rows_strided(
18032 &self,
18033 x: &CudaSlice<f32>,
18034 data: &CudaSlice<f32>,
18035 m: usize,
18036 in_f: usize,
18037 out_f: usize,
18038 canonical_chunk_rows: usize,
18039 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18040 self.linear_f32_resident_canonical_rows_inner(
18041 x,
18042 data,
18043 m,
18044 in_f,
18045 out_f,
18046 canonical_chunk_rows,
18047 true,
18048 )
18049 }
18050
18051 fn linear_f32_resident_canonical_rows_inner(
18052 &self,
18053 x: &CudaSlice<f32>,
18054 data: &CudaSlice<f32>,
18055 m: usize,
18056 in_f: usize,
18057 out_f: usize,
18058 canonical_chunk_rows: usize,
18059 strided_output: bool,
18060 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18061 if data.len() != in_f * out_f {
18062 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18063 }
18064 if canonical_chunk_rows == 0
18065 || canonical_chunk_rows > out_f
18066 || out_f % canonical_chunk_rows != 0
18067 {
18068 return Err(format!(
18069 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18070 )
18071 .into());
18072 }
18073 if canonical_chunk_rows == out_f {
18074 return self.linear(x, data, m, in_f, out_f);
18075 }
18076
18077 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18078 let input = x.slice(0..x.len());
18079 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18080 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18081 if m == 1 {
18082 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18083 self.linear_device_into(
18084 &input,
18085 &weights,
18086 &mut destination,
18087 1,
18088 in_f,
18089 canonical_chunk_rows,
18090 )?;
18091 continue;
18092 }
18093 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18094 if strided_output {
18095 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18096 } else {
18097 for token in 0..m {
18098 let source = chunk
18099 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18100 let mut destination =
18101 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18102 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18103 }
18104 }
18105 }
18106 Ok(y)
18107 }
18108
18109 pub fn linear_f32_resident_canonical_rows_t1_into(
18114 &self,
18115 x: &CudaSlice<f32>,
18116 data: &CudaSlice<f32>,
18117 y: &mut CudaSlice<f32>,
18118 in_f: usize,
18119 out_f: usize,
18120 canonical_chunk_rows: usize,
18121 ) -> Result<(), Box<dyn std::error::Error>> {
18122 if data.len() != in_f * out_f {
18123 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18124 }
18125 if y.len() != out_f || x.len() != in_f {
18126 return Err(format!(
18127 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18128 x.len(),
18129 y.len()
18130 )
18131 .into());
18132 }
18133 if canonical_chunk_rows == 0
18134 || canonical_chunk_rows > out_f
18135 || out_f % canonical_chunk_rows != 0
18136 {
18137 return Err(format!(
18138 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18139 )
18140 .into());
18141 }
18142 let input = x.slice(0..x.len());
18143 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18144 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18145 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18146 self.linear_device_into(
18147 &input,
18148 &weights,
18149 &mut destination,
18150 1,
18151 in_f,
18152 canonical_chunk_rows,
18153 )?;
18154 }
18155 Ok(())
18156 }
18157
18158 pub fn linear_t1_into(
18161 &self,
18162 x: &cudarc::driver::CudaView<'_, f32>,
18163 w: &cudarc::driver::CudaView<'_, f32>,
18164 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18165 in_f: usize,
18166 out_f: usize,
18167 ) -> Result<(), Box<dyn std::error::Error>> {
18168 self.linear_device_into(x, w, y, 1, in_f, out_f)
18169 }
18170
18171 pub fn linear_decode_exact(
18178 &self,
18179 x: &CudaSlice<f32>,
18180 w: &CudaSlice<f32>,
18181 m_tokens: usize,
18182 in_f: usize,
18183 out_f: usize,
18184 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18185 if m_tokens == 1 {
18186 return self.linear(x, w, 1, in_f, out_f);
18187 }
18188 let xv = self.view(x, m_tokens * in_f);
18189 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18190 for t in 0..m_tokens {
18191 let row = xv.slice(t * in_f..(t + 1) * in_f);
18192 let mut xr = self.alloc_uninit::<f32>(in_f)?;
18193 self.copy_view_into(&mut xr, 0, &row, in_f)?;
18194 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18195 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18196 }
18197 Ok(y)
18198 }
18199
18200 pub fn linear(
18201 &self,
18202 x: &CudaSlice<f32>,
18203 w: &CudaSlice<f32>,
18204 m_tokens: usize,
18205 in_f: usize,
18206 out_f: usize,
18207 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18208 self.linear_device(x, w, m_tokens, in_f, out_f)
18209 }
18210
18211 fn linear_device<I>(
18212 &self,
18213 x: &I,
18214 w: &I,
18215 m_tokens: usize,
18216 in_f: usize,
18217 out_f: usize,
18218 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18219 where
18220 I: cudarc::driver::DevicePtr<f32>,
18221 {
18222 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)?;
18224 Ok(c)
18225 }
18226
18227 fn linear_device_into<I, O>(
18228 &self,
18229 x: &I,
18230 w: &I,
18231 c: &mut O,
18232 m_tokens: usize,
18233 in_f: usize,
18234 out_f: usize,
18235 ) -> Result<(), Box<dyn std::error::Error>>
18236 where
18237 I: cudarc::driver::DevicePtr<f32>,
18238 O: cudarc::driver::DevicePtrMut<f32>,
18239 {
18240 use cudarc::cublaslt::{Matmul, MatmulConfig};
18241 let cfg = MatmulConfig {
18242 transa: true,
18243 transb: false,
18244 transc: false,
18245 m: out_f as u64,
18246 n: m_tokens as u64,
18247 k: in_f as u64,
18248 alpha: 1.0,
18249 lda: in_f as i64,
18250 ldb: in_f as i64,
18251 beta: 0.0,
18252 ldc: out_f as i64,
18253 stride_a: None,
18254 stride_b: None,
18255 stride_c: None,
18256 stride_bias: None,
18257 batch_size: None,
18258 };
18259 let blas = self.gpu.blas();
18260 unsafe {
18261 blas.matmul(cfg, w, x, c, None, None)?;
18262 }
18263 Ok(())
18264 }
18265
18266 pub fn sdpa_naive(
18268 &self,
18269 q: &CudaSlice<f32>,
18270 k: &CudaSlice<f32>,
18271 v: &CudaSlice<f32>,
18272 o: &mut CudaSlice<f32>,
18273 head_dim: usize,
18274 n_head: usize,
18275 n_head_kv: usize,
18276 t: usize,
18277 t_kv: usize,
18278 scale: f32,
18279 causal: bool,
18280 ) -> Result<(), Box<dyn std::error::Error>> {
18281 let f = self.func("sdpa_naive_f32");
18282 let cfg = LaunchConfig {
18283 grid_dim: (n_head as u32, t as u32, 1),
18284 block_dim: (128, 1, 1),
18285 shared_mem_bytes: (t_kv * 4) as u32,
18286 };
18287 let (hd, nh, nhkv, ti, tkvi, cz) = (
18288 head_dim as i32,
18289 n_head as i32,
18290 n_head_kv as i32,
18291 t as i32,
18292 t_kv as i32,
18293 causal as i32,
18294 );
18295 let __s_b = self.gpu.stream();
18296 let mut b = __s_b.launch_builder(&f);
18297 b.arg(q)
18298 .arg(k)
18299 .arg(v)
18300 .arg(o)
18301 .arg(&hd)
18302 .arg(&nh)
18303 .arg(&nhkv)
18304 .arg(&ti)
18305 .arg(&tkvi)
18306 .arg(&scale)
18307 .arg(&cz);
18308 unsafe {
18309 b.launch(cfg)?;
18310 }
18311 Ok(())
18312 }
18313
18314 #[allow(clippy::too_many_arguments)]
18319 pub fn sdpa_naive_island(
18320 &self,
18321 q: &CudaSlice<f32>,
18322 k: &CudaSlice<f32>,
18323 v: &CudaSlice<f32>,
18324 o: &mut CudaSlice<f32>,
18325 span_id: &CudaSlice<i32>,
18326 head_dim: usize,
18327 n_head: usize,
18328 n_head_kv: usize,
18329 t: usize,
18330 t_kv: usize,
18331 scale: f32,
18332 window: usize,
18333 ) -> Result<(), Box<dyn std::error::Error>> {
18334 let f = self.func("sdpa_naive_island_f32");
18335 let cfg = LaunchConfig {
18336 grid_dim: (n_head as u32, t as u32, 1),
18337 block_dim: (128, 1, 1),
18338 shared_mem_bytes: (t_kv * 4) as u32,
18339 };
18340 let (hd, nh, nhkv, ti, tkvi, wi) = (
18341 head_dim as i32,
18342 n_head as i32,
18343 n_head_kv as i32,
18344 t as i32,
18345 t_kv as i32,
18346 window as i32,
18347 );
18348 let __s_b = self.gpu.stream();
18349 let mut b = __s_b.launch_builder(&f);
18350 b.arg(q)
18351 .arg(k)
18352 .arg(v)
18353 .arg(o)
18354 .arg(span_id)
18355 .arg(&hd)
18356 .arg(&nh)
18357 .arg(&nhkv)
18358 .arg(&ti)
18359 .arg(&tkvi)
18360 .arg(&scale)
18361 .arg(&wi);
18362 unsafe {
18363 b.launch(cfg)?;
18364 }
18365 Ok(())
18366 }
18367
18368 #[allow(clippy::too_many_arguments)]
18370 pub fn sdpa_naive_w(
18371 &self,
18372 q: &CudaSlice<f32>,
18373 k: &CudaSlice<f32>,
18374 v: &CudaSlice<f32>,
18375 o: &mut CudaSlice<f32>,
18376 head_dim: usize,
18377 n_head: usize,
18378 n_head_kv: usize,
18379 t: usize,
18380 t_kv: usize,
18381 scale: f32,
18382 causal: bool,
18383 window: usize,
18384 ) -> Result<(), Box<dyn std::error::Error>> {
18385 let f = self.func("sdpa_naive_w_f32");
18386 let cfg = LaunchConfig {
18387 grid_dim: (n_head as u32, t as u32, 1),
18388 block_dim: (128, 1, 1),
18389 shared_mem_bytes: (t_kv * 4) as u32,
18390 };
18391 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18392 head_dim as i32,
18393 n_head as i32,
18394 n_head_kv as i32,
18395 t as i32,
18396 t_kv as i32,
18397 causal as i32,
18398 window as i32,
18399 );
18400 let __s_b = self.gpu.stream();
18401 let mut b = __s_b.launch_builder(&f);
18402 b.arg(q)
18403 .arg(k)
18404 .arg(v)
18405 .arg(o)
18406 .arg(&hd)
18407 .arg(&nh)
18408 .arg(&nhkv)
18409 .arg(&ti)
18410 .arg(&tkvi)
18411 .arg(&scale)
18412 .arg(&cz)
18413 .arg(&wi);
18414 unsafe {
18415 b.launch(cfg)?;
18416 }
18417 Ok(())
18418 }
18419
18420 #[allow(clippy::too_many_arguments)]
18430 pub fn sdpa_naive_w_lo(
18431 &self,
18432 q: &CudaSlice<f32>,
18433 k: &CudaSlice<f32>,
18434 v: &CudaSlice<f32>,
18435 o: &mut CudaSlice<f32>,
18436 head_dim: usize,
18437 n_head: usize,
18438 n_head_kv: usize,
18439 t: usize,
18440 t_kv: usize,
18441 scale: f32,
18442 causal: bool,
18443 window: usize,
18444 ) -> Result<(), Box<dyn std::error::Error>> {
18445 let kv_lo = if window > 0 {
18446 (t_kv - t + 1).saturating_sub(window)
18447 } else {
18448 0
18449 };
18450 let smem = (t_kv - kv_lo) * 4;
18451 if smem > 48 * 1024 {
18452 return Err(format!(
18453 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18454 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18455 a window this wide needs the multi-pass long-ctx kernel"
18456 )
18457 .into());
18458 }
18459 let f = self.func("sdpa_naive_w_lo_f32");
18460 let cfg = LaunchConfig {
18461 grid_dim: (n_head as u32, t as u32, 1),
18462 block_dim: (128, 1, 1),
18463 shared_mem_bytes: smem as u32,
18464 };
18465 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18466 head_dim as i32,
18467 n_head as i32,
18468 n_head_kv as i32,
18469 t as i32,
18470 t_kv as i32,
18471 causal as i32,
18472 window as i32,
18473 kv_lo as i32,
18474 );
18475 let __s_b = self.gpu.stream();
18476 let mut b = __s_b.launch_builder(&f);
18477 b.arg(q)
18478 .arg(k)
18479 .arg(v)
18480 .arg(o)
18481 .arg(&hd)
18482 .arg(&nh)
18483 .arg(&nhkv)
18484 .arg(&ti)
18485 .arg(&tkvi)
18486 .arg(&scale)
18487 .arg(&cz)
18488 .arg(&wi)
18489 .arg(&lo);
18490 unsafe {
18491 b.launch(cfg)?;
18492 }
18493 Ok(())
18494 }
18495
18496 pub fn sdpa_naive_view(
18498 &self,
18499 q: &CudaSlice<f32>,
18500 k: &cudarc::driver::CudaView<f32>,
18501 v: &cudarc::driver::CudaView<f32>,
18502 o: &mut CudaSlice<f32>,
18503 head_dim: usize,
18504 n_head: usize,
18505 n_head_kv: usize,
18506 t: usize,
18507 t_kv: usize,
18508 scale: f32,
18509 causal: bool,
18510 ) -> Result<(), Box<dyn std::error::Error>> {
18511 let f = self.func("sdpa_naive_f32");
18512 let cfg = LaunchConfig {
18513 grid_dim: (n_head as u32, t as u32, 1),
18514 block_dim: (128, 1, 1),
18515 shared_mem_bytes: (t_kv * 4) as u32,
18516 };
18517 let (hd, nh, nhkv, ti, tkvi, cz) = (
18518 head_dim as i32,
18519 n_head as i32,
18520 n_head_kv as i32,
18521 t as i32,
18522 t_kv as i32,
18523 causal as i32,
18524 );
18525 let __s_b = self.gpu.stream();
18526 let mut b = __s_b.launch_builder(&f);
18527 b.arg(q)
18528 .arg(k)
18529 .arg(v)
18530 .arg(o)
18531 .arg(&hd)
18532 .arg(&nh)
18533 .arg(&nhkv)
18534 .arg(&ti)
18535 .arg(&tkvi)
18536 .arg(&scale)
18537 .arg(&cz);
18538 unsafe {
18539 b.launch(cfg)?;
18540 }
18541 Ok(())
18542 }
18543
18544 #[allow(clippy::too_many_arguments)]
18552 pub fn fa_dequant_kv_view_f32(
18553 &self,
18554 k: &cudarc::driver::CudaView<u8>,
18555 v: &cudarc::driver::CudaView<u8>,
18556 kf: &mut CudaSlice<f32>,
18557 vf: &mut CudaSlice<f32>,
18558 kv_dim_k: usize,
18559 kv_dim_v: usize,
18560 t_kv: usize,
18561 k_tok_bytes: usize,
18562 v_tok_bytes: usize,
18563 g: bool,
18564 ) -> Result<(), Box<dyn std::error::Error>> {
18565 let f = if g {
18566 self.func_g("fa_dequant_kv_ws_f32")
18567 } else {
18568 self.func("fa_dequant_kv_ws_f32")
18569 };
18570 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18571 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18572 let cfg = LaunchConfig {
18573 grid_dim: (nblk.max(1), 1, 1),
18574 block_dim: (256, 1, 1),
18575 shared_mem_bytes: 0,
18576 };
18577 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18578 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18579 let __s_b = self.gpu.stream();
18580 let mut b = __s_b.launch_builder(&f);
18581 b.arg(k)
18582 .arg(v)
18583 .arg(&mut *kf)
18584 .arg(&mut *vf)
18585 .arg(&kdk)
18586 .arg(&kdv)
18587 .arg(&tkvi)
18588 .arg(&ktb)
18589 .arg(&vtb);
18590 unsafe {
18591 b.launch(cfg)?;
18592 }
18593 Ok(())
18594 }
18595
18596 #[allow(clippy::too_many_arguments)]
18597 pub fn sdpa_naive_quantized_view(
18598 &self,
18599 q: &CudaSlice<f32>,
18600 k: &cudarc::driver::CudaView<u8>,
18601 v: &cudarc::driver::CudaView<u8>,
18602 o: &mut CudaSlice<f32>,
18603 head_dim: usize,
18604 n_head: usize,
18605 n_head_kv: usize,
18606 t: usize,
18607 t_kv: usize,
18608 scale: f32,
18609 causal: bool,
18610 k_tok_bytes: usize,
18611 v_tok_bytes: usize,
18612 ) -> Result<(), Box<dyn std::error::Error>> {
18613 let kv_dim = n_head_kv * head_dim;
18614 let mut kf = self.uninit(t_kv * kv_dim)?;
18615 let mut vf = self.uninit(t_kv * kv_dim)?;
18616 let f = self.func("fa_dequant_kv_ws_f32");
18617 let total = (2 * t_kv * kv_dim) as u64;
18618 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18619 let cfg = LaunchConfig {
18620 grid_dim: (nblk.max(1), 1, 1),
18621 block_dim: (256, 1, 1),
18622 shared_mem_bytes: 0,
18623 };
18624 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18625 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18626 let __s_b = self.gpu.stream();
18627 let mut b = __s_b.launch_builder(&f);
18628 b.arg(k)
18629 .arg(v)
18630 .arg(&mut kf)
18631 .arg(&mut vf)
18632 .arg(&kv_dim_i)
18633 .arg(&kv_dim_i)
18634 .arg(&t_kv_i)
18635 .arg(&k_tok_bytes_i)
18636 .arg(&v_tok_bytes_i);
18637 unsafe { b.launch(cfg)? };
18638 self.sdpa_naive(
18639 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18640 )
18641 }
18642
18643 #[allow(clippy::too_many_arguments)]
18655 pub fn sdpa_naive_w_quantized_view(
18656 &self,
18657 q: &CudaSlice<f32>,
18658 k: &cudarc::driver::CudaView<u8>,
18659 v: &cudarc::driver::CudaView<u8>,
18660 o: &mut CudaSlice<f32>,
18661 head_dim: usize,
18662 n_head: usize,
18663 n_head_kv: usize,
18664 t: usize,
18665 t_kv: usize,
18666 scale: f32,
18667 causal: bool,
18668 window: usize,
18669 k_tok_bytes: usize,
18670 v_tok_bytes: usize,
18671 ) -> Result<(), Box<dyn std::error::Error>> {
18672 let kv_dim = n_head_kv * head_dim;
18673 let mut kf = self.uninit(t_kv * kv_dim)?;
18674 let mut vf = self.uninit(t_kv * kv_dim)?;
18675 let f = self.func("fa_dequant_kv_ws_f32");
18676 let total = (2 * t_kv * kv_dim) as u64;
18677 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18678 let cfg = LaunchConfig {
18679 grid_dim: (nblk.max(1), 1, 1),
18680 block_dim: (256, 1, 1),
18681 shared_mem_bytes: 0,
18682 };
18683 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18684 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18685 let __s_b = self.gpu.stream();
18686 let mut b = __s_b.launch_builder(&f);
18687 b.arg(k)
18688 .arg(v)
18689 .arg(&mut kf)
18690 .arg(&mut vf)
18691 .arg(&kv_dim_i)
18692 .arg(&kv_dim_i)
18693 .arg(&t_kv_i)
18694 .arg(&k_tok_bytes_i)
18695 .arg(&v_tok_bytes_i);
18696 unsafe { b.launch(cfg)? };
18697 self.sdpa_naive_w(
18698 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18699 )
18700 }
18701
18702 pub fn fa_prefill(
18706 &self,
18707 q: &CudaSlice<f32>,
18708 k: &CudaSlice<f32>,
18709 v: &CudaSlice<f32>,
18710 o: &mut CudaSlice<f32>,
18711 head_dim: usize,
18712 n_head: usize,
18713 n_head_kv: usize,
18714 t: usize,
18715 t_kv: usize,
18716 scale: f32,
18717 causal: bool,
18718 ) -> Result<(), Box<dyn std::error::Error>> {
18719 if portable_mma_gated() {
18720 return self.sdpa_naive(
18721 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18722 );
18723 }
18724 let fa3_on = head_dim == 256
18732 && causal
18733 && t == t_kv
18734 && match std::env::var("MEMRA_FA3").as_deref() {
18735 Ok("0") => false,
18736 Ok("1") => {
18740 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
18741 true
18742 }
18743 _ => cfg!(memra_hopper_mma),
18744 };
18745 if fa3_on {
18746 let n = t * n_head * head_dim;
18747 let nkv = t * n_head_kv * head_dim;
18748 let mut q16 = self.alloc_u8_uninit(n * 2)?;
18749 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18750 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18751 self.f32_to_bf16_into(q, &mut q16, n)?;
18752 self.f32_to_bf16_into(k, &mut k16, nkv)?;
18753 self.f32_to_bf16_into(v, &mut v16, nkv)?;
18754 let rc = {
18755 use cudarc::driver::{DevicePtr, DevicePtrMut};
18756 let stream = self.gpu.stream();
18757 let (qp, _g1) = q16.device_ptr(&stream);
18758 let (kp, _g2) = k16.device_ptr(&stream);
18759 let (vp, _g3) = v16.device_ptr(&stream);
18760 let (op, _g4) = o.device_ptr_mut(&stream);
18761 unsafe {
18762 memra_fa3_prefill(
18763 qp as *const core::ffi::c_void,
18764 kp as *const core::ffi::c_void,
18765 vp as *const core::ffi::c_void,
18766 op as *mut f32,
18767 t as i32,
18768 n_head as i32,
18769 n_head_kv as i32,
18770 head_dim as i32,
18771 scale,
18772 stream.cu_stream() as *mut core::ffi::c_void,
18773 )
18774 }
18775 };
18776 if rc != 0 {
18777 return Err(format!("memra_fa3_prefill rc={rc}").into());
18778 }
18779 return Ok(());
18780 }
18781 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18786 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18787 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18788 const BLOCK_Q: usize = 64;
18789 const BKX: usize = 32;
18790 let f = self.func("fa_prefill_bf16_p1");
18791 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18792 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18793 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18794 f.set_attribute(
18795 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18796 shmem as i32,
18797 )?;
18798 let cfg = LaunchConfig {
18799 grid_dim: (
18800 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18801 n_head as u32,
18802 1,
18803 ),
18804 block_dim: (32, 4, 1),
18805 shared_mem_bytes: shmem,
18806 };
18807 let (hd, nh, nhkv, ti, tkvi, cz) = (
18808 head_dim as i32,
18809 n_head as i32,
18810 n_head_kv as i32,
18811 t as i32,
18812 t_kv as i32,
18813 causal as i32,
18814 );
18815 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18816 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18817 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18818 let __s_b = self.gpu.stream();
18819 let mut b = __s_b.launch_builder(&f);
18820 b.arg(&qb)
18821 .arg(&kb)
18822 .arg(&vb)
18823 .arg(o)
18824 .arg(&hd)
18825 .arg(&nh)
18826 .arg(&nhkv)
18827 .arg(&ti)
18828 .arg(&tkvi)
18829 .arg(&scale)
18830 .arg(&cz);
18831 unsafe {
18832 b.launch(cfg)?;
18833 }
18834 return Ok(());
18835 }
18836 const BK: usize = 32;
18842 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
18845 let (block_q, warps, w2_sfx): (usize, u32, &str) =
18846 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
18847 let hd_sfx = fa_hd_suffix(head_dim)?;
18851 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18852 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
18857 let (kb16, vb16) = if bf16kv {
18858 let n = t_kv * n_head_kv * head_dim;
18859 let mut kb = self.alloc_u8_uninit(n * 2)?;
18860 let mut vb = self.alloc_u8_uninit(n * 2)?;
18861 let fcv = self.func("f32_to_bf16_bulk");
18862 let ni = n as i64;
18863 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
18864 let __s_b = self.gpu.stream();
18865 let mut b = __s_b.launch_builder(&fcv);
18866 b.arg(k).arg(&mut kb).arg(&ni);
18867 unsafe {
18868 b.launch(cfgc)?;
18869 }
18870 let __s_b = self.gpu.stream();
18871 let mut b = __s_b.launch_builder(&fcv);
18872 b.arg(v).arg(&mut vb).arg(&ni);
18873 unsafe {
18874 b.launch(cfgc)?;
18875 }
18876 (Some(kb), Some(vb))
18877 } else {
18878 (None, None)
18879 };
18880 let f = self.func(&if bf16kv {
18881 format!("fa_prefill_bf16kv_pp{hd_sfx}")
18882 } else {
18883 format!(
18884 "fa_prefill_f32{}{}{hd_sfx}",
18885 if floor { "" } else { "_pp" },
18886 if floor { "" } else { w2_sfx }
18887 )
18888 });
18889 let kv_stages = if bf16kv { 2 } else { 1 };
18892 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
18893 + 4 * (block_q * BK + 2 * block_q)) as u32;
18894 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18895 f.set_attribute(
18896 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18897 shmem as i32,
18898 )?;
18899 let cfg = LaunchConfig {
18900 grid_dim: (
18901 (t as u32 + block_q as u32 - 1) / block_q as u32,
18902 n_head as u32,
18903 1,
18904 ),
18905 block_dim: (32, warps, 1),
18906 shared_mem_bytes: shmem,
18907 };
18908 let (hd, nh, nhkv, ti, tkvi, cz) = (
18909 head_dim as i32,
18910 n_head as i32,
18911 n_head_kv as i32,
18912 t as i32,
18913 t_kv as i32,
18914 causal as i32,
18915 );
18916 let __s_b = self.gpu.stream();
18917 let mut b = __s_b.launch_builder(&f);
18918 b.arg(q);
18919 match (&kb16, &vb16) {
18920 (Some(kb), Some(vb)) => {
18921 b.arg(kb).arg(vb);
18922 }
18923 _ => {
18924 b.arg(k).arg(v);
18925 }
18926 }
18927 b.arg(o)
18928 .arg(&hd)
18929 .arg(&nh)
18930 .arg(&nhkv)
18931 .arg(&ti)
18932 .arg(&tkvi)
18933 .arg(&scale)
18934 .arg(&cz);
18935 unsafe {
18936 b.launch(cfg)?;
18937 }
18938 Ok(())
18939 }
18940
18941 #[allow(clippy::too_many_arguments)]
18945 pub fn fa_prefill_w(
18946 &self,
18947 q: &CudaSlice<f32>,
18948 k: &CudaSlice<f32>,
18949 v: &CudaSlice<f32>,
18950 o: &mut CudaSlice<f32>,
18951 head_dim: usize,
18952 n_head: usize,
18953 n_head_kv: usize,
18954 t: usize,
18955 t_kv: usize,
18956 scale: f32,
18957 causal: bool,
18958 window: usize,
18959 ) -> Result<(), Box<dyn std::error::Error>> {
18960 if portable_mma_gated() {
18963 return self.sdpa_naive_w(
18964 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18965 );
18966 }
18967 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18971 let faw_f32 =
18972 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
18973 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18974 self.fa_prefill_w_arm(
18975 q,
18976 k,
18977 v,
18978 o,
18979 head_dim,
18980 n_head,
18981 n_head_kv,
18982 t,
18983 t_kv,
18984 scale,
18985 causal,
18986 window,
18987 floor || faw_f32,
18988 floor,
18989 )
18990 }
18991
18992 #[allow(clippy::too_many_arguments)]
18995 pub fn fa_prefill_w_pre(
18996 &self,
18997 qb: &CudaSlice<u8>,
18998 kb: &CudaSlice<u8>,
18999 vb: &CudaSlice<u8>,
19000 o: &mut CudaSlice<f32>,
19001 head_dim: usize,
19002 n_head: usize,
19003 n_head_kv: usize,
19004 t: usize,
19005 t_kv: usize,
19006 scale: f32,
19007 causal: bool,
19008 window: usize,
19009 v_f16: bool,
19010 ) -> Result<(), Box<dyn std::error::Error>> {
19011 const BLOCK_Q: usize = 64;
19012 const BK: usize = 32;
19013 debug_assert_eq!(head_dim, 256);
19014 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19015 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19016 if hp {
19017 const BLOCK_QH: usize = 32;
19018 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19021 let vh: &CudaSlice<u8> = if v_f16 {
19022 vb
19023 } else {
19024 let n = t_kv * n_head_kv * head_dim;
19025 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19026 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19027 }
19028 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19029 vguard.as_ref().unwrap()
19030 };
19031 let f = self.func("fa_prefill_w_bf16_p1h2");
19032 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19033 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19034 f.set_attribute(
19035 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19036 shmem as i32,
19037 )?;
19038 let cfg = LaunchConfig {
19039 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19040 block_dim: (32, 4, 1),
19041 shared_mem_bytes: shmem,
19042 };
19043 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19044 head_dim as i32,
19045 n_head as i32,
19046 n_head_kv as i32,
19047 t as i32,
19048 t_kv as i32,
19049 causal as i32,
19050 window as i32,
19051 );
19052 let __s_b = self.gpu.stream();
19053 let mut b = __s_b.launch_builder(&f);
19054 b.arg(qb)
19055 .arg(kb)
19056 .arg(vh)
19057 .arg(o)
19058 .arg(&hd)
19059 .arg(&nh)
19060 .arg(&nhkv)
19061 .arg(&ti)
19062 .arg(&tkvi)
19063 .arg(&scale)
19064 .arg(&cz)
19065 .arg(&wi);
19066 unsafe {
19067 b.launch(cfg)?;
19068 }
19069 return Ok(());
19070 }
19071 let f = self.func("fa_prefill_w_bf16_p1");
19072 let shmem =
19073 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19074 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19075 f.set_attribute(
19076 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19077 shmem as i32,
19078 )?;
19079 let cfg = LaunchConfig {
19080 grid_dim: (
19081 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19082 n_head as u32,
19083 1,
19084 ),
19085 block_dim: (32, 4, 1),
19086 shared_mem_bytes: shmem,
19087 };
19088 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19089 head_dim as i32,
19090 n_head as i32,
19091 n_head_kv as i32,
19092 t as i32,
19093 t_kv as i32,
19094 causal as i32,
19095 window as i32,
19096 );
19097 let __s_b = self.gpu.stream();
19098 let mut b = __s_b.launch_builder(&f);
19099 b.arg(qb)
19100 .arg(kb)
19101 .arg(vb)
19102 .arg(o)
19103 .arg(&hd)
19104 .arg(&nh)
19105 .arg(&nhkv)
19106 .arg(&ti)
19107 .arg(&tkvi)
19108 .arg(&scale)
19109 .arg(&cz)
19110 .arg(&wi);
19111 unsafe {
19112 b.launch(cfg)?;
19113 }
19114 Ok(())
19115 }
19116
19117 #[allow(clippy::too_many_arguments)]
19119 pub fn fa_prefill_w_arm(
19120 &self,
19121 q: &CudaSlice<f32>,
19122 k: &CudaSlice<f32>,
19123 v: &CudaSlice<f32>,
19124 o: &mut CudaSlice<f32>,
19125 head_dim: usize,
19126 n_head: usize,
19127 n_head_kv: usize,
19128 t: usize,
19129 t_kv: usize,
19130 scale: f32,
19131 causal: bool,
19132 window: usize,
19133 f32_stage: bool,
19134 floor: bool,
19135 ) -> Result<(), Box<dyn std::error::Error>> {
19136 const BLOCK_Q: usize = 64;
19137 const BK: usize = 32;
19138 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19139 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19143 let p1 = !floor
19144 && !f32_stage
19145 && *P1_ON.get_or_init(|| {
19146 std::env::var("MEMRA_FAW_P1")
19147 .map(|v| v != "0")
19148 .unwrap_or(true)
19149 });
19150 let hp =
19151 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19152 if hp {
19153 const BLOCK_QH: usize = 32;
19154 let f = self.func("fa_prefill_w_bf16_p1h2");
19155 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19156 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19157 f.set_attribute(
19158 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19159 shmem as i32,
19160 )?;
19161 let cfg = LaunchConfig {
19162 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19163 block_dim: (32, 4, 1),
19164 shared_mem_bytes: shmem,
19165 };
19166 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19167 head_dim as i32,
19168 n_head as i32,
19169 n_head_kv as i32,
19170 t as i32,
19171 t_kv as i32,
19172 causal as i32,
19173 window as i32,
19174 );
19175 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19176 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19177 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19178 let __s_b = self.gpu.stream();
19179 let mut b = __s_b.launch_builder(&f);
19180 b.arg(&qb)
19181 .arg(&kb)
19182 .arg(&vh)
19183 .arg(o)
19184 .arg(&hd)
19185 .arg(&nh)
19186 .arg(&nhkv)
19187 .arg(&ti)
19188 .arg(&tkvi)
19189 .arg(&scale)
19190 .arg(&cz)
19191 .arg(&wi);
19192 unsafe {
19193 b.launch(cfg)?;
19194 }
19195 return Ok(());
19196 }
19197 if p1 {
19198 let f = self.func("fa_prefill_w_bf16_p1");
19199 let shmem =
19200 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19201 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19202 f.set_attribute(
19203 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19204 shmem as i32,
19205 )?;
19206 let cfg = LaunchConfig {
19207 grid_dim: (
19208 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19209 n_head as u32,
19210 1,
19211 ),
19212 block_dim: (32, 4, 1),
19213 shared_mem_bytes: shmem,
19214 };
19215 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19216 head_dim as i32,
19217 n_head as i32,
19218 n_head_kv as i32,
19219 t as i32,
19220 t_kv as i32,
19221 causal as i32,
19222 window as i32,
19223 );
19224 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19225 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19226 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19227 let __s_b = self.gpu.stream();
19228 let mut b = __s_b.launch_builder(&f);
19229 b.arg(&qb)
19230 .arg(&kb)
19231 .arg(&vb)
19232 .arg(o)
19233 .arg(&hd)
19234 .arg(&nh)
19235 .arg(&nhkv)
19236 .arg(&ti)
19237 .arg(&tkvi)
19238 .arg(&scale)
19239 .arg(&cz)
19240 .arg(&wi);
19241 unsafe {
19242 b.launch(cfg)?;
19243 }
19244 return Ok(());
19245 }
19246 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19249 let g4 = !floor
19250 && !f32_stage
19251 && n_head_kv == 1
19252 && n_head % 4 == 0
19253 && *G4_ON.get_or_init(|| {
19254 std::env::var("MEMRA_FAW_G4")
19255 .map(|v| v != "0")
19256 .unwrap_or(true)
19257 });
19258 if g4 {
19259 const SP_M: usize = 16;
19260 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19263 let o2 = *O2_ON.get_or_init(|| {
19264 std::env::var("MEMRA_FAW_O2")
19265 .map(|v| v != "0")
19266 .unwrap_or(true)
19267 });
19268 let f = self.func(if o2 {
19269 "fa_prefill_w_bf16_g4o2"
19270 } else {
19271 "fa_prefill_w_bf16_g4"
19272 });
19273 let shmem = if o2 {
19274 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19275 } else {
19276 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19277 as u32
19278 };
19279 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19280 f.set_attribute(
19281 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19282 shmem as i32,
19283 )?;
19284 let cfg = LaunchConfig {
19285 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19286 block_dim: (32, 4, 1),
19287 shared_mem_bytes: shmem,
19288 };
19289 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19290 head_dim as i32,
19291 n_head as i32,
19292 n_head_kv as i32,
19293 t as i32,
19294 t_kv as i32,
19295 causal as i32,
19296 window as i32,
19297 );
19298 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19299 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19300 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19301 let __s_b = self.gpu.stream();
19302 let mut b = __s_b.launch_builder(&f);
19303 b.arg(&qb)
19304 .arg(&kb)
19305 .arg(&vb)
19306 .arg(o)
19307 .arg(&hd)
19308 .arg(&nh)
19309 .arg(&nhkv)
19310 .arg(&ti)
19311 .arg(&tkvi)
19312 .arg(&scale)
19313 .arg(&cz)
19314 .arg(&wi);
19315 unsafe {
19316 b.launch(cfg)?;
19317 }
19318 return Ok(());
19319 }
19320 let f = self.func(if floor {
19321 "fa_prefill_w_f32"
19322 } else if f32_stage {
19323 "fa_prefill_w_f32_pp"
19324 } else {
19325 "fa_prefill_w_bf16_pp"
19326 });
19327 let shmem =
19328 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19329 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19330 f.set_attribute(
19331 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19332 shmem as i32,
19333 )?;
19334 let cfg = LaunchConfig {
19335 grid_dim: (
19336 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19337 n_head as u32,
19338 1,
19339 ),
19340 block_dim: (32, 4, 1),
19341 shared_mem_bytes: shmem,
19342 };
19343 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19344 head_dim as i32,
19345 n_head as i32,
19346 n_head_kv as i32,
19347 t as i32,
19348 t_kv as i32,
19349 causal as i32,
19350 window as i32,
19351 );
19352 if f32_stage {
19353 let __s_b = self.gpu.stream();
19354 let mut b = __s_b.launch_builder(&f);
19355 b.arg(q)
19356 .arg(k)
19357 .arg(v)
19358 .arg(o)
19359 .arg(&hd)
19360 .arg(&nh)
19361 .arg(&nhkv)
19362 .arg(&ti)
19363 .arg(&tkvi)
19364 .arg(&scale)
19365 .arg(&cz)
19366 .arg(&wi);
19367 unsafe {
19368 b.launch(cfg)?;
19369 }
19370 } else {
19371 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19372 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19373 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19374 let __s_b = self.gpu.stream();
19375 let mut b = __s_b.launch_builder(&f);
19376 b.arg(&qb)
19377 .arg(&kb)
19378 .arg(&vb)
19379 .arg(o)
19380 .arg(&hd)
19381 .arg(&nh)
19382 .arg(&nhkv)
19383 .arg(&ti)
19384 .arg(&tkvi)
19385 .arg(&scale)
19386 .arg(&cz)
19387 .arg(&wi);
19388 unsafe {
19389 b.launch(cfg)?;
19390 }
19391 }
19392 Ok(())
19393 }
19394
19395 #[allow(clippy::too_many_arguments)]
19399 pub fn fa_prefill_hd512(
19400 &self,
19401 q: &CudaSlice<f32>,
19402 k: &CudaSlice<f32>,
19403 v: &CudaSlice<f32>,
19404 o: &mut CudaSlice<f32>,
19405 head_dim: usize,
19406 n_head: usize,
19407 n_head_kv: usize,
19408 t: usize,
19409 t_kv: usize,
19410 scale: f32,
19411 causal: bool,
19412 ) -> Result<(), Box<dyn std::error::Error>> {
19413 if portable_mma_gated() {
19415 return self.sdpa_naive(
19416 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19417 );
19418 }
19419 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19425 let f32_stage =
19426 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19427 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19431 let sp = !f32_stage
19432 && *SP_ON.get_or_init(|| {
19433 std::env::var("MEMRA_FA512_SP")
19434 .map(|v| v != "0")
19435 .unwrap_or(true)
19436 });
19437 self.fa_prefill_hd512_arm(
19438 q,
19439 k,
19440 v,
19441 o,
19442 head_dim,
19443 n_head,
19444 n_head_kv,
19445 t,
19446 t_kv,
19447 scale,
19448 causal,
19449 f32_stage,
19450 sp,
19451 sp && fa_f16pv_on(),
19452 )
19453 }
19454
19455 #[allow(clippy::too_many_arguments)]
19457 pub fn fa_prefill_hd512_pre(
19458 &self,
19459 qb: &CudaSlice<u8>,
19460 kb: &CudaSlice<u8>,
19461 vb: &CudaSlice<u8>,
19462 o: &mut CudaSlice<f32>,
19463 head_dim: usize,
19464 n_head: usize,
19465 n_head_kv: usize,
19466 t: usize,
19467 t_kv: usize,
19468 scale: f32,
19469 causal: bool,
19470 v_f16: bool,
19471 ) -> Result<(), Box<dyn std::error::Error>> {
19472 debug_assert_eq!(head_dim, 512);
19473 const SP_M: usize = 16;
19474 const BKS: usize = 32;
19475 let f16pv = fa_f16pv_on();
19479 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19480 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19481 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19482 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19483 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19484 let n = t_kv * n_head_kv * head_dim;
19486 let need = n * 2;
19487 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19488 *vguard = Some(self.alloc_uninit::<u8>(need)?);
19489 }
19490 let dst = vguard.as_mut().unwrap();
19491 self.bf16_to_f16_into(vb, n, dst)?;
19492 vguard.as_ref().unwrap()
19493 } else {
19494 vb
19495 };
19496 let f = self.func(if hp {
19497 "fa_prefill_bf16_hd512_sp16h2"
19498 } else {
19499 match (f16pv, nw) {
19500 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19501 (true, _) => "fa_prefill_bf16_hd512_sp16",
19502 _ => "fa_prefill_bf16_hd512_sp",
19503 }
19504 });
19505 let (nwarp, npart) = if hp {
19506 (4usize, 4usize)
19507 } else if nw > 2 {
19508 (nw, nw)
19509 } else {
19510 (2, 1)
19511 };
19512 let shmem = if hp {
19514 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19515 as u32
19516 } else {
19517 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19518 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19519 };
19520 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19521 f.set_attribute(
19522 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19523 shmem as i32,
19524 )?;
19525 let grid_y = if hp {
19526 (n_head / 2) as u32
19527 } else {
19528 n_head as u32
19529 };
19530 let cfg = LaunchConfig {
19531 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19532 block_dim: (32, nwarp as u32, 1),
19533 shared_mem_bytes: shmem,
19534 };
19535 let (hd, nh, nhkv, ti, tkvi, cz) = (
19536 head_dim as i32,
19537 n_head as i32,
19538 n_head_kv as i32,
19539 t as i32,
19540 t_kv as i32,
19541 causal as i32,
19542 );
19543 let __s_b = self.gpu.stream();
19544 let mut b = __s_b.launch_builder(&f);
19545 b.arg(qb)
19546 .arg(kb)
19547 .arg(vref)
19548 .arg(o)
19549 .arg(&hd)
19550 .arg(&nh)
19551 .arg(&nhkv)
19552 .arg(&ti)
19553 .arg(&tkvi)
19554 .arg(&scale)
19555 .arg(&cz);
19556 unsafe {
19557 b.launch(cfg)?;
19558 }
19559 Ok(())
19560 }
19561
19562 #[allow(clippy::too_many_arguments)]
19565 pub fn fa_prefill_hd512_arm(
19566 &self,
19567 q: &CudaSlice<f32>,
19568 k: &CudaSlice<f32>,
19569 v: &CudaSlice<f32>,
19570 o: &mut CudaSlice<f32>,
19571 head_dim: usize,
19572 n_head: usize,
19573 n_head_kv: usize,
19574 t: usize,
19575 t_kv: usize,
19576 scale: f32,
19577 causal: bool,
19578 f32_stage: bool,
19579 sp: bool,
19580 f16pv: bool,
19581 ) -> Result<(), Box<dyn std::error::Error>> {
19582 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19583 if sp && !f32_stage {
19584 const SP_M: usize = 16;
19588 const BKS: usize = 32;
19589 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19590 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19591 let f = self.func(if hp {
19592 "fa_prefill_bf16_hd512_sp16h2"
19593 } else {
19594 match (f16pv, nw) {
19595 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19596 (true, _) => "fa_prefill_bf16_hd512_sp16",
19597 _ => "fa_prefill_bf16_hd512_sp",
19598 }
19599 });
19600 let (nwarp, npart) = if hp {
19601 (4usize, 4usize)
19602 } else if nw > 2 {
19603 (nw, nw)
19604 } else {
19605 (2, 1)
19606 };
19607 let shmem = if hp {
19608 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19609 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19610 } else {
19611 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19612 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19613 };
19614 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19615 f.set_attribute(
19616 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19617 shmem as i32,
19618 )?;
19619 let grid_y = if hp {
19620 (n_head / 2) as u32
19621 } else {
19622 n_head as u32
19623 };
19624 let cfg = LaunchConfig {
19625 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19626 block_dim: (32, nwarp as u32, 1),
19627 shared_mem_bytes: shmem,
19628 };
19629 let (hd, nh, nhkv, ti, tkvi, cz) = (
19630 head_dim as i32,
19631 n_head as i32,
19632 n_head_kv as i32,
19633 t as i32,
19634 t_kv as i32,
19635 causal as i32,
19636 );
19637 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19638 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19639 let vb = if f16pv {
19640 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19641 } else {
19642 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19643 };
19644 let __s_b = self.gpu.stream();
19645 let mut b = __s_b.launch_builder(&f);
19646 b.arg(&qb)
19647 .arg(&kb)
19648 .arg(&vb)
19649 .arg(o)
19650 .arg(&hd)
19651 .arg(&nh)
19652 .arg(&nhkv)
19653 .arg(&ti)
19654 .arg(&tkvi)
19655 .arg(&scale)
19656 .arg(&cz);
19657 unsafe {
19658 b.launch(cfg)?;
19659 }
19660 return Ok(());
19661 }
19662 const BLOCK_Q: usize = 32;
19663 const BK: usize = 32;
19664 const HALF: usize = 256;
19665 let f = self.func(if f32_stage {
19666 "fa_prefill_f32_hd512"
19667 } else {
19668 "fa_prefill_bf16_hd512"
19669 });
19670 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19672 + 4 * BLOCK_Q) as u32;
19673 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19674 f.set_attribute(
19675 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19676 shmem as i32,
19677 )?;
19678 let cfg = LaunchConfig {
19679 grid_dim: (
19680 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19681 n_head as u32,
19682 2,
19683 ),
19684 block_dim: (32, 2, 1),
19685 shared_mem_bytes: shmem,
19686 };
19687 let (hd, nh, nhkv, ti, tkvi, cz) = (
19688 head_dim as i32,
19689 n_head as i32,
19690 n_head_kv as i32,
19691 t as i32,
19692 t_kv as i32,
19693 causal as i32,
19694 );
19695 if f32_stage {
19696 let __s_b = self.gpu.stream();
19697 let mut b = __s_b.launch_builder(&f);
19698 b.arg(q)
19699 .arg(k)
19700 .arg(v)
19701 .arg(o)
19702 .arg(&hd)
19703 .arg(&nh)
19704 .arg(&nhkv)
19705 .arg(&ti)
19706 .arg(&tkvi)
19707 .arg(&scale)
19708 .arg(&cz);
19709 unsafe {
19710 b.launch(cfg)?;
19711 }
19712 } else {
19713 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19714 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19715 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19716 let __s_b = self.gpu.stream();
19717 let mut b = __s_b.launch_builder(&f);
19718 b.arg(&qb)
19719 .arg(&kb)
19720 .arg(&vb)
19721 .arg(o)
19722 .arg(&hd)
19723 .arg(&nh)
19724 .arg(&nhkv)
19725 .arg(&ti)
19726 .arg(&tkvi)
19727 .arg(&scale)
19728 .arg(&cz);
19729 unsafe {
19730 b.launch(cfg)?;
19731 }
19732 }
19733 Ok(())
19734 }
19735
19736 #[allow(clippy::too_many_arguments)]
19740 pub fn rope_neox2_bf16e(
19741 &self,
19742 q: &mut CudaSlice<f32>,
19743 k: &mut CudaSlice<f32>,
19744 qb: &mut CudaSlice<u8>,
19745 kb: &mut CudaSlice<u8>,
19746 pos: &CudaSlice<i32>,
19747 head_dim: usize,
19748 n_dims: usize,
19749 nh_q: usize,
19750 nh_k: usize,
19751 n_tokens: usize,
19752 base: f32,
19753 freq_scale: f32,
19754 ff: Option<&CudaSlice<f32>>,
19755 ) -> Result<(), Box<dyn std::error::Error>> {
19756 let f = self.func("rope_neox2_bf16e_f32");
19757 let rows = ((nh_q + nh_k) * n_tokens) as u32;
19758 let cfg = LaunchConfig {
19759 grid_dim: (rows, 1, 1),
19760 block_dim: ((head_dim / 2) as u32, 1, 1),
19761 shared_mem_bytes: 0,
19762 };
19763 let theta_scale = base.powf(-2.0 / n_dims as f32);
19764 let (hd, nd, nhq, nhk, nt) = (
19765 head_dim as i32,
19766 n_dims as i32,
19767 nh_q as i32,
19768 nh_k as i32,
19769 n_tokens as i32,
19770 );
19771 let __s_b = self.gpu.stream();
19772 let mut b = __s_b.launch_builder(&f);
19773 match ff {
19774 Some(t) => {
19775 b.arg(&mut *q)
19776 .arg(&mut *k)
19777 .arg(&mut *qb)
19778 .arg(&mut *kb)
19779 .arg(pos)
19780 .arg(&hd)
19781 .arg(&nd)
19782 .arg(&nhq)
19783 .arg(&nhk)
19784 .arg(&nt)
19785 .arg(&theta_scale)
19786 .arg(&freq_scale)
19787 .arg(t);
19788 unsafe {
19789 b.launch(cfg)?;
19790 }
19791 }
19792 None => {
19793 let null: u64 = 0;
19794 b.arg(&mut *q)
19795 .arg(&mut *k)
19796 .arg(&mut *qb)
19797 .arg(&mut *kb)
19798 .arg(pos)
19799 .arg(&hd)
19800 .arg(&nd)
19801 .arg(&nhq)
19802 .arg(&nhk)
19803 .arg(&nt)
19804 .arg(&theta_scale)
19805 .arg(&freq_scale)
19806 .arg(&null);
19807 unsafe {
19808 b.launch(cfg)?;
19809 }
19810 }
19811 }
19812 Ok(())
19813 }
19814
19815 pub fn f32_to_bf16(
19818 &self,
19819 x: &CudaSlice<f32>,
19820 n: usize,
19821 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19822 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19823 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19824 let f = self.func("f32_to_bf16_flat");
19825 let n_i = n as i64;
19826 let cfg = LaunchConfig {
19827 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19828 block_dim: (256, 1, 1),
19829 shared_mem_bytes: 0,
19830 };
19831 let __s_b = self.gpu.stream();
19832 let mut b = __s_b.launch_builder(&f);
19833 b.arg(x).arg(&mut y).arg(&n_i);
19834 unsafe {
19835 b.launch(cfg)?;
19836 }
19837 Ok(y)
19838 }
19839
19840 pub fn f32_to_f16(
19841 &self,
19842 x: &CudaSlice<f32>,
19843 n: usize,
19844 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19845 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
19846 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19847 let f = self.func("f32_to_f16_flat");
19848 let n_i = n as i64;
19849 let cfg = LaunchConfig {
19850 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19851 block_dim: (256, 1, 1),
19852 shared_mem_bytes: 0,
19853 };
19854 let __s_b = self.gpu.stream();
19855 let mut b = __s_b.launch_builder(&f);
19856 b.arg(x).arg(&mut y).arg(&n_i);
19857 unsafe {
19858 b.launch(cfg)?;
19859 }
19860 Ok(y)
19861 }
19862
19863 pub fn bf16_to_f16(
19865 &self,
19866 xb: &CudaSlice<u8>,
19867 n: usize,
19868 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19869 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19870 self.bf16_to_f16_into(xb, n, &mut y)?;
19871 Ok(y)
19872 }
19873
19874 pub fn bf16_to_f16_into(
19876 &self,
19877 xb: &CudaSlice<u8>,
19878 n: usize,
19879 y: &mut CudaSlice<u8>,
19880 ) -> Result<(), Box<dyn std::error::Error>> {
19881 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
19882 assert!(y.len() >= n * 2);
19883 let f = self.func("bf16_to_f16_flat");
19884 let n2 = (n / 2) as i64;
19885 let cfg = LaunchConfig {
19886 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
19887 block_dim: (256, 1, 1),
19888 shared_mem_bytes: 0,
19889 };
19890 let __s_b = self.gpu.stream();
19891 let mut b = __s_b.launch_builder(&f);
19892 b.arg(xb).arg(y).arg(&n2);
19893 unsafe {
19894 b.launch(cfg)?;
19895 }
19896 Ok(())
19897 }
19898
19899 #[allow(clippy::too_many_arguments)]
19904 pub fn fa_prefill_vl8(
19905 &self,
19906 seqs: &[FaSeqVl],
19907 head_dim: usize,
19908 n_head: usize,
19909 n_head_kv: usize,
19910 scale: f32,
19911 ) -> Result<(), Box<dyn std::error::Error>> {
19912 const BK: usize = 32;
19913 let b = seqs.len();
19914 assert!(b >= 1 && b <= 8);
19915 let mut packed = [FaSeqVl::default(); 8];
19916 packed[..b].copy_from_slice(seqs);
19917 let v = FaVl8(packed);
19918 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19919 let ept = (n_head_kv * head_dim) as i32;
19920 {
19921 let f = self.func("fa_mirror_vl");
19922 let max_n = (max_t as i64) * ept as i64;
19923 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
19924 for which in 0..2i32 {
19925 let cfg = LaunchConfig {
19926 grid_dim: (blocks, 1, b as u32),
19927 block_dim: (256, 1, 1),
19928 shared_mem_bytes: 0,
19929 };
19930 let __s_lb = self.gpu.stream();
19931 let mut lb = __s_lb.launch_builder(&f);
19932 lb.arg(&v).arg(&ept).arg(&which);
19933 unsafe {
19934 lb.launch(cfg)?;
19935 }
19936 }
19937 }
19938 let hd_sfx = fa_hd_suffix(head_dim)?;
19939 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
19940 let block_q = 64usize;
19941 let kv_stages = 2usize;
19942 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19943 + 4 * (block_q * BK + 2 * block_q)) as u32;
19944 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19945 f.set_attribute(
19946 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19947 shmem as i32,
19948 )?;
19949 let cfg = LaunchConfig {
19950 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
19951 block_dim: (32, 4, 1),
19952 shared_mem_bytes: shmem,
19953 };
19954 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19955 let __s_lb = self.gpu.stream();
19956 let mut lb = __s_lb.launch_builder(&f);
19957 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
19958 unsafe {
19959 lb.launch(cfg)?;
19960 }
19961 Ok(())
19962 }
19963
19964 #[allow(clippy::too_many_arguments)]
19968 pub fn attn_pre_vl8(
19969 &self,
19970 seqs: &[AttnPreVl],
19971 wq: &CudaSlice<f32>,
19972 wk: &CudaSlice<f32>,
19973 head_dim: usize,
19974 rope_dims: usize,
19975 n_head: usize,
19976 n_head_kv: usize,
19977 eps: f32,
19978 freq_base: f32,
19979 freq_scale: f32,
19980 kv_dim_k: usize,
19981 kv_dim_v: usize,
19982 k_tok_bytes: usize,
19983 v_tok_bytes: usize,
19984 ) -> Result<(), Box<dyn std::error::Error>> {
19985 let b = seqs.len();
19986 assert!(b >= 1 && b <= 8);
19987 let mut packed = [AttnPreVl::default(); 8];
19988 packed[..b].copy_from_slice(seqs);
19989 let v = AttnPreVl8(packed);
19990 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19991 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19992 {
19993 let f = self.func("q_gate_split_vl");
19994 let n = max_t * (n_head * head_dim) as u32;
19995 let cfg = LaunchConfig {
19996 grid_dim: (n.div_ceil(256), 1, b as u32),
19997 block_dim: (256, 1, 1),
19998 shared_mem_bytes: 0,
19999 };
20000 let __s_lb = self.gpu.stream();
20001 let mut lb = __s_lb.launch_builder(&f);
20002 lb.arg(&v).arg(&hd).arg(&nh);
20003 unsafe {
20004 lb.launch(cfg)?;
20005 }
20006 }
20007 {
20008 let f = self.func("attn_rms_vl");
20009 let cfg = LaunchConfig {
20010 grid_dim: (max_t * n_head as u32, 2, b as u32),
20011 block_dim: (rms_block(), 1, 1),
20012 shared_mem_bytes: 0,
20013 };
20014 let __s_lb = self.gpu.stream();
20015 let mut lb = __s_lb.launch_builder(&f);
20016 lb.arg(&v)
20017 .arg(wq)
20018 .arg(wk)
20019 .arg(&hd)
20020 .arg(&nh)
20021 .arg(&nhkv)
20022 .arg(&eps);
20023 unsafe {
20024 lb.launch(cfg)?;
20025 }
20026 }
20027 {
20028 let f = self.func("attn_rope_vl");
20029 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20030 let nd = rope_dims as i32;
20031 let cfg = LaunchConfig {
20032 grid_dim: (max_t * n_head as u32, 2, b as u32),
20033 block_dim: ((head_dim / 2) as u32, 1, 1),
20034 shared_mem_bytes: 0,
20035 };
20036 let __s_lb = self.gpu.stream();
20037 let mut lb = __s_lb.launch_builder(&f);
20038 lb.arg(&v)
20039 .arg(&hd)
20040 .arg(&nd)
20041 .arg(&nh)
20042 .arg(&nhkv)
20043 .arg(&theta_scale)
20044 .arg(&freq_scale);
20045 unsafe {
20046 lb.launch(cfg)?;
20047 }
20048 }
20049 {
20050 let f = self.func("append_kv_vl");
20051 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20052 let cfg = LaunchConfig {
20053 grid_dim: (nblk, max_t, b as u32),
20054 block_dim: (32, 1, 1),
20055 shared_mem_bytes: 0,
20056 };
20057 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20058 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20059 let __s_lb = self.gpu.stream();
20060 let mut lb = __s_lb.launch_builder(&f);
20061 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20062 unsafe {
20063 lb.launch(cfg)?;
20064 }
20065 }
20066 Ok(())
20067 }
20068
20069 pub fn fa_prefill_view(
20074 &self,
20075 q: &CudaSlice<f32>,
20076 k: &cudarc::driver::CudaView<u8>,
20077 v: &cudarc::driver::CudaView<u8>,
20078 o: &mut CudaSlice<f32>,
20079 head_dim: usize,
20080 n_head: usize,
20081 n_head_kv: usize,
20082 t: usize,
20083 t_kv: usize,
20084 scale: f32,
20085 causal: bool,
20086 k_tok_bytes: usize,
20087 v_tok_bytes: usize,
20088 g: bool,
20089 ) -> Result<(), Box<dyn std::error::Error>> {
20090 if portable_mma_gated() {
20091 return self.sdpa_naive_quantized_view(
20092 q,
20093 k,
20094 v,
20095 o,
20096 head_dim,
20097 n_head,
20098 n_head_kv,
20099 t,
20100 t_kv,
20101 scale,
20102 causal,
20103 k_tok_bytes,
20104 v_tok_bytes,
20105 );
20106 }
20107 const BLOCK_Q: usize = 64;
20108 const BK: usize = 32;
20109 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20112 let f = if g {
20113 self.func_g(&name)
20114 } else {
20115 self.func(&name)
20116 };
20117 let shmem =
20118 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20119 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20120 f.set_attribute(
20121 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20122 shmem as i32,
20123 )?;
20124 let cfg = LaunchConfig {
20125 grid_dim: (
20126 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20127 n_head as u32,
20128 1,
20129 ),
20130 block_dim: (32, 4, 1),
20131 shared_mem_bytes: shmem,
20132 };
20133 let (hd, nh, nhkv, ti, tkvi, cz) = (
20134 head_dim as i32,
20135 n_head as i32,
20136 n_head_kv as i32,
20137 t as i32,
20138 t_kv as i32,
20139 causal as i32,
20140 );
20141 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20142 let __s_b = self.gpu.stream();
20143 let mut b = __s_b.launch_builder(&f);
20144 b.arg(q)
20145 .arg(k)
20146 .arg(v)
20147 .arg(o)
20148 .arg(&hd)
20149 .arg(&nh)
20150 .arg(&nhkv)
20151 .arg(&ti)
20152 .arg(&tkvi)
20153 .arg(&scale)
20154 .arg(&cz)
20155 .arg(&ktb)
20156 .arg(&vtb);
20157 unsafe {
20158 b.launch(cfg)?;
20159 }
20160 Ok(())
20161 }
20162
20163 #[allow(clippy::too_many_arguments)]
20173 pub fn fa_prefill_view_ws(
20174 &self,
20175 q: &CudaSlice<f32>,
20176 k: &cudarc::driver::CudaView<u8>,
20177 v: &cudarc::driver::CudaView<u8>,
20178 o: &mut CudaSlice<f32>,
20179 head_dim: usize,
20180 n_head: usize,
20181 n_head_kv: usize,
20182 t: usize,
20183 t_kv: usize,
20184 scale: f32,
20185 causal: bool,
20186 k_tok_bytes: usize,
20187 v_tok_bytes: usize,
20188 g: bool,
20189 ) -> Result<(), Box<dyn std::error::Error>> {
20190 if portable_mma_gated() {
20191 return self.sdpa_naive_quantized_view(
20192 q,
20193 k,
20194 v,
20195 o,
20196 head_dim,
20197 n_head,
20198 n_head_kv,
20199 t,
20200 t_kv,
20201 scale,
20202 causal,
20203 k_tok_bytes,
20204 v_tok_bytes,
20205 );
20206 }
20207 const BLOCK_Q: usize = 64;
20208 const BK: usize = 32;
20209 let kv_dim_k = n_head_kv * head_dim;
20210 let kv_dim_v = n_head_kv * head_dim;
20211 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20213 let mut guard = self.prime_deqw_ws.lock().unwrap();
20215 let need_grow = match guard.as_ref() {
20216 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20217 None => true,
20218 };
20219 if need_grow {
20220 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20221 let (ck, cv) = guard
20222 .as_ref()
20223 .map(|(a, b)| (a.len(), b.len()))
20224 .unwrap_or((0, 0));
20225 *guard = Some((
20226 self.alloc_u8(grow(ck, k_ws_bytes))?,
20227 self.alloc_u8(grow(cv, v_ws_bytes))?,
20228 ));
20229 }
20230 let (kw, vw) = guard.as_mut().unwrap();
20231 {
20233 let f = if g {
20235 self.func_g("fa_dequant_kv_ws_bf16")
20236 } else {
20237 self.func("fa_dequant_kv_ws_bf16")
20238 };
20239 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20240 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20241 let cfg = LaunchConfig {
20242 grid_dim: (nblk.max(1), 1, 1),
20243 block_dim: (256, 1, 1),
20244 shared_mem_bytes: 0,
20245 };
20246 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20247 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20248 let __s_b = self.gpu.stream();
20249 let mut b = __s_b.launch_builder(&f);
20250 b.arg(k)
20251 .arg(v)
20252 .arg(&mut *kw)
20253 .arg(&mut *vw)
20254 .arg(&kdk)
20255 .arg(&kdv)
20256 .arg(&tkvi)
20257 .arg(&ktb)
20258 .arg(&vtb);
20259 unsafe {
20260 b.launch(cfg)?;
20261 }
20262 }
20263 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20271 .map(|v| v != "0")
20272 .unwrap_or(true);
20273 {
20274 let hd_sfx = fa_hd_suffix(head_dim)?;
20275 let f = self.func(&format!(
20276 "fa_prefill_qw{}{hd_sfx}",
20277 if db { "_db" } else { "" }
20278 ));
20279 let shmem = if db {
20280 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20282 } else {
20283 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20284 };
20285 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20286 f.set_attribute(
20287 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20288 shmem as i32,
20289 )?;
20290 let cfg = LaunchConfig {
20291 grid_dim: (
20292 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20293 n_head as u32,
20294 1,
20295 ),
20296 block_dim: (32, 4, 1),
20297 shared_mem_bytes: shmem,
20298 };
20299 let (hd, nh, nhkv, ti, tkvi, cz) = (
20300 head_dim as i32,
20301 n_head as i32,
20302 n_head_kv as i32,
20303 t as i32,
20304 t_kv as i32,
20305 causal as i32,
20306 );
20307 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20308 let __s_b = self.gpu.stream();
20309 let mut b = __s_b.launch_builder(&f);
20310 b.arg(q)
20311 .arg(&*kw)
20312 .arg(&*vw)
20313 .arg(o)
20314 .arg(&hd)
20315 .arg(&nh)
20316 .arg(&nhkv)
20317 .arg(&ti)
20318 .arg(&tkvi)
20319 .arg(&scale)
20320 .arg(&cz)
20321 .arg(&kdk)
20322 .arg(&kdv);
20323 unsafe {
20324 b.launch(cfg)?;
20325 }
20326 }
20327 Ok(())
20328 }
20329
20330 #[allow(clippy::too_many_arguments)]
20346 pub fn fa_prefill_view_ws_w_hd128(
20347 &self,
20348 q: &CudaSlice<f32>,
20349 k: &cudarc::driver::CudaView<u8>,
20350 v: &cudarc::driver::CudaView<u8>,
20351 o: &mut CudaSlice<f32>,
20352 head_dim: usize,
20353 n_head: usize,
20354 n_head_kv: usize,
20355 t: usize,
20356 t_kv: usize,
20357 scale: f32,
20358 causal: bool,
20359 window: usize,
20360 k_tok_bytes: usize,
20361 v_tok_bytes: usize,
20362 ) -> Result<(), Box<dyn std::error::Error>> {
20363 assert_eq!(
20364 head_dim, 128,
20365 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20366 );
20367 if portable_mma_gated() {
20368 return self.sdpa_naive_w_quantized_view(
20369 q,
20370 k,
20371 v,
20372 o,
20373 head_dim,
20374 n_head,
20375 n_head_kv,
20376 t,
20377 t_kv,
20378 scale,
20379 causal,
20380 window,
20381 k_tok_bytes,
20382 v_tok_bytes,
20383 );
20384 }
20385 const BLOCK_Q: usize = 64;
20386 const BK: usize = 32;
20387 let kv_dim_k = n_head_kv * head_dim;
20388 let kv_dim_v = n_head_kv * head_dim;
20389 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20391 let mut guard = self.prime_deqw_ws.lock().unwrap();
20392 let need_grow = match guard.as_ref() {
20393 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20394 None => true,
20395 };
20396 if need_grow {
20397 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20398 let (ck, cv) = guard
20399 .as_ref()
20400 .map(|(a, b)| (a.len(), b.len()))
20401 .unwrap_or((0, 0));
20402 *guard = Some((
20403 self.alloc_u8(grow(ck, k_ws_bytes))?,
20404 self.alloc_u8(grow(cv, v_ws_bytes))?,
20405 ));
20406 }
20407 let (kw, vw) = guard.as_mut().unwrap();
20408 {
20411 let f = self.func("fa_dequant_kv_ws_bf16");
20412 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20413 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20414 let cfg = LaunchConfig {
20415 grid_dim: (nblk.max(1), 1, 1),
20416 block_dim: (256, 1, 1),
20417 shared_mem_bytes: 0,
20418 };
20419 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20420 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20421 let __s_b = self.gpu.stream();
20422 let mut b = __s_b.launch_builder(&f);
20423 b.arg(k)
20424 .arg(v)
20425 .arg(&mut *kw)
20426 .arg(&mut *vw)
20427 .arg(&kdk)
20428 .arg(&kdv)
20429 .arg(&tkvi)
20430 .arg(&ktb)
20431 .arg(&vtb);
20432 unsafe {
20433 b.launch(cfg)?;
20434 }
20435 }
20436 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20438 .map(|v| v != "0")
20439 .unwrap_or(true);
20440 {
20441 let f = self.func(if db {
20442 "fa_prefill_qw_db_w_hd128"
20443 } else {
20444 "fa_prefill_qw_w_hd128"
20445 });
20446 let shmem = if db {
20447 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20448 } else {
20449 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20450 };
20451 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20452 f.set_attribute(
20453 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20454 shmem as i32,
20455 )?;
20456 let cfg = LaunchConfig {
20457 grid_dim: (
20458 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20459 n_head as u32,
20460 1,
20461 ),
20462 block_dim: (32, 4, 1),
20463 shared_mem_bytes: shmem,
20464 };
20465 let (hd, nh, nhkv, ti, tkvi, cz) = (
20466 head_dim as i32,
20467 n_head as i32,
20468 n_head_kv as i32,
20469 t as i32,
20470 t_kv as i32,
20471 causal as i32,
20472 );
20473 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20474 let __s_b = self.gpu.stream();
20475 let mut b = __s_b.launch_builder(&f);
20476 b.arg(q)
20477 .arg(&*kw)
20478 .arg(&*vw)
20479 .arg(o)
20480 .arg(&hd)
20481 .arg(&nh)
20482 .arg(&nhkv)
20483 .arg(&ti)
20484 .arg(&tkvi)
20485 .arg(&scale)
20486 .arg(&cz)
20487 .arg(&kdk)
20488 .arg(&kdv)
20489 .arg(&wnd);
20490 unsafe {
20491 b.launch(cfg)?;
20492 }
20493 }
20494 Ok(())
20495 }
20496
20497 pub fn fa_decode(
20501 &self,
20502 q: &CudaSlice<f32>,
20503 k: &cudarc::driver::CudaView<u8>,
20504 v: &cudarc::driver::CudaView<u8>,
20505 o: &mut CudaSlice<f32>,
20506 head_dim: usize,
20507 n_head: usize,
20508 n_head_kv: usize,
20509 t_kv: usize,
20510 scale: f32,
20511 k_tok_bytes: usize,
20512 v_tok_bytes: usize,
20513 ) -> Result<(), Box<dyn std::error::Error>> {
20514 self.fa_decode_kvmod(
20515 q,
20516 k,
20517 v,
20518 o,
20519 head_dim,
20520 n_head,
20521 n_head_kv,
20522 t_kv,
20523 scale,
20524 k_tok_bytes,
20525 v_tok_bytes,
20526 false,
20527 )
20528 }
20529
20530 #[allow(clippy::too_many_arguments)]
20534 #[allow(clippy::too_many_arguments)]
20538 #[allow(clippy::too_many_arguments)]
20539 fn fa_decode_scalar_unified(
20540 &self,
20541 q: &cudarc::driver::CudaView<f32>,
20542 k: &cudarc::driver::CudaView<u8>,
20543 v: &cudarc::driver::CudaView<u8>,
20544 o: &mut cudarc::driver::CudaViewMut<f32>,
20545 head_dim: usize,
20546 n_head: usize,
20547 n_head_kv: usize,
20548 t_kv_host: usize,
20549 t_kv_dev: Option<&CudaSlice<i32>>,
20550 scale: f32,
20551 n_splits: usize,
20552 split_keys: usize,
20553 k_tok_bytes: usize,
20554 v_tok_bytes: usize,
20555 g: bool,
20556 part_o: &mut CudaSlice<f32>,
20557 part_m: &mut CudaSlice<f32>,
20558 part_l: &mut CudaSlice<f32>,
20559 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20560 ) -> Result<(), Box<dyn std::error::Error>> {
20561 let f = if g {
20562 self.func_g("fa_decode_f32")
20563 } else {
20564 self.fa_func("fa_decode_f32", head_dim)
20565 };
20566 let cfg = LaunchConfig {
20567 grid_dim: (n_head as u32, n_splits as u32, 1),
20568 block_dim: (head_dim as u32, 1, 1),
20569 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20570 };
20571 let (hd, nh, nhkv, nsp) = (
20572 head_dim as i32,
20573 n_head as i32,
20574 n_head_kv as i32,
20575 n_splits as i32,
20576 );
20577 let (ktb, vtb, tkvi, ski) = (
20578 k_tok_bytes as i64,
20579 v_tok_bytes as i64,
20580 t_kv_host as i32,
20581 split_keys as i32,
20582 );
20583 let __s_b = self.gpu.stream();
20584 let mut b = __s_b.launch_builder(&f);
20585 match t_kv_dev {
20586 Some(d) => {
20587 b.arg(q)
20588 .arg(k)
20589 .arg(v)
20590 .arg(&mut *part_o)
20591 .arg(&mut *part_m)
20592 .arg(&mut *part_l)
20593 .arg(&hd)
20594 .arg(&nh)
20595 .arg(&nhkv)
20596 .arg(&tkvi)
20597 .arg(d)
20598 .arg(&scale)
20599 .arg(&nsp)
20600 .arg(&ski)
20601 .arg(&ktb)
20602 .arg(&vtb);
20603 unsafe {
20604 b.launch(cfg)?;
20605 }
20606 }
20607 None => {
20608 let null: u64 = 0;
20609 b.arg(q)
20610 .arg(k)
20611 .arg(v)
20612 .arg(&mut *part_o)
20613 .arg(&mut *part_m)
20614 .arg(&mut *part_l)
20615 .arg(&hd)
20616 .arg(&nh)
20617 .arg(&nhkv)
20618 .arg(&tkvi)
20619 .arg(&null)
20620 .arg(&scale)
20621 .arg(&nsp)
20622 .arg(&ski)
20623 .arg(&ktb)
20624 .arg(&vtb);
20625 unsafe {
20626 b.launch(cfg)?;
20627 }
20628 }
20629 }
20630 let cfg2 = LaunchConfig {
20631 grid_dim: (n_head as u32, 1, 1),
20632 block_dim: (head_dim as u32, 1, 1),
20633 shared_mem_bytes: 0,
20634 };
20635 if let Some((oq, od)) = q8_out {
20636 let fc = if g {
20638 self.func_g("fa_decode_combine_q8_1")
20639 } else {
20640 self.fa_func("fa_decode_combine_q8_1", head_dim)
20641 };
20642 let __s_b2 = self.gpu.stream();
20643 let mut b2 = __s_b2.launch_builder(&fc);
20644 b2.arg(&*part_o)
20645 .arg(&*part_m)
20646 .arg(&*part_l)
20647 .arg(oq)
20648 .arg(od)
20649 .arg(&hd)
20650 .arg(&nh)
20651 .arg(&nsp);
20652 unsafe {
20653 b2.launch(cfg2)?;
20654 }
20655 return Ok(());
20656 }
20657 let fc = if g {
20658 self.func_g("fa_decode_combine_f32")
20659 } else {
20660 self.fa_func("fa_decode_combine_f32", head_dim)
20661 };
20662 let __s_b2 = self.gpu.stream();
20663 let mut b2 = __s_b2.launch_builder(&fc);
20664 b2.arg(&*part_o)
20665 .arg(&*part_m)
20666 .arg(&*part_l)
20667 .arg(o)
20668 .arg(&hd)
20669 .arg(&nh)
20670 .arg(&nsp);
20671 unsafe {
20672 b2.launch(cfg2)?;
20673 }
20674 Ok(())
20675 }
20676
20677 pub fn fa_decode_kvmod(
20678 &self,
20679 q: &CudaSlice<f32>,
20680 k: &cudarc::driver::CudaView<u8>,
20681 v: &cudarc::driver::CudaView<u8>,
20682 o: &mut CudaSlice<f32>,
20683 head_dim: usize,
20684 n_head: usize,
20685 n_head_kv: usize,
20686 t_kv: usize,
20687 scale: f32,
20688 k_tok_bytes: usize,
20689 v_tok_bytes: usize,
20690 g: bool,
20691 ) -> Result<(), Box<dyn std::error::Error>> {
20692 let q_view = q.as_view();
20693 let mut o_view = o.as_view_mut();
20694 self.fa_decode_kvmod_view(
20695 &q_view,
20696 k,
20697 v,
20698 &mut o_view,
20699 head_dim,
20700 n_head,
20701 n_head_kv,
20702 t_kv,
20703 scale,
20704 k_tok_bytes,
20705 v_tok_bytes,
20706 g,
20707 )
20708 }
20709
20710 #[allow(clippy::too_many_arguments)]
20715 pub fn fa_decode_kvmod_view(
20716 &self,
20717 q: &cudarc::driver::CudaView<f32>,
20718 k: &cudarc::driver::CudaView<u8>,
20719 v: &cudarc::driver::CudaView<u8>,
20720 o: &mut cudarc::driver::CudaViewMut<f32>,
20721 head_dim: usize,
20722 n_head: usize,
20723 n_head_kv: usize,
20724 t_kv: usize,
20725 scale: f32,
20726 k_tok_bytes: usize,
20727 v_tok_bytes: usize,
20728 g: bool,
20729 ) -> Result<(), Box<dyn std::error::Error>> {
20730 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20751 if g && head_dim == 256 && !fa_v4_at(t_kv) {
20755 fa_vec = false;
20756 }
20757 let sp = fa_split_keys(t_kv, n_head_kv);
20758 let n_splits = if fa_vec {
20759 ((t_kv + sp - 1) / sp).max(1)
20760 } else {
20761 ((t_kv + 255) / 256).max(1)
20762 };
20763 let o_len = n_head * n_splits * head_dim;
20764 let ml_len = n_head * n_splits;
20765 let mut part_guard = self.fa_part_pool.lock().unwrap();
20766 if part_guard
20767 .as_ref()
20768 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20769 .unwrap_or(true)
20770 {
20771 let old = part_guard.take();
20782 let (co, cm) = old
20783 .as_ref()
20784 .map(|pp| (pp.0.len(), pp.1.len()))
20785 .unwrap_or((0, 0));
20786 if let Some(old) = old {
20787 self.fa_part_retired.lock().unwrap().push(old);
20788 }
20789 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20790 eprintln!(
20791 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20792 co, o_len, cm, ml_len
20793 );
20794 }
20795 *part_guard = Some((
20796 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20797 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20798 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20799 ));
20800 }
20801 let pg = part_guard.as_mut().unwrap();
20802 self.gpu
20803 .stream()
20804 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20805 self.gpu
20806 .stream()
20807 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20808 self.gpu
20809 .stream()
20810 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20811 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20812 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20813 let (hd, nh, nhkv, tkvi, nsp) = (
20814 head_dim as i32,
20815 n_head as i32,
20816 n_head_kv as i32,
20817 t_kv as i32,
20818 n_splits as i32,
20819 );
20820 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20821 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20825 let fa512_min = fa512_min_tkv();
20830 let deep = fa_vec
20833 && head_dim == 256
20834 && fa_v4_at(t_kv)
20835 && !g
20836 && fa_deep_at(t_kv)
20837 && !matches!(fa_v4_mode(), "noB3" | "stage");
20838 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
20839 let gqa = (n_head / n_head_kv).max(1) as u32;
20842 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
20843 (
20844 fv,
20845 LaunchConfig {
20846 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20847 block_dim: (32, gqa, 1),
20848 shared_mem_bytes: 0,
20849 },
20850 )
20851 } else if fa_vec && head_dim <= 256 {
20852 let gqa = (n_head / n_head_kv).max(1) as u32;
20853 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20864 let smem_tkv = *SMEM_TKV.get_or_init(|| {
20865 std::env::var("MEMRA_FA_SMEM_TKV")
20866 .ok()
20867 .and_then(|v| v.parse().ok())
20868 .unwrap_or_else(|| {
20869 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
20870 })
20871 });
20872 if fa_v4_at(t_kv) && head_dim == 256 {
20873 let v4name = match fa_v4_mode() {
20877 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
20880 _ => "fa_decode_vec_q_v4",
20881 };
20882 let fv = if g {
20883 self.func_g(v4name)
20884 } else {
20885 self.func(v4name)
20886 };
20887 let shmem = (if deep { 12160 } else { 11520 }
20890 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20891 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20892 fv.set_attribute(
20893 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20894 shmem as i32,
20895 )?;
20896 (
20897 fv,
20898 LaunchConfig {
20899 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20900 block_dim: (32, gqa, 1),
20901 shared_mem_bytes: shmem,
20902 },
20903 )
20904 } else if fa_v3_active(head_dim) {
20905 let fv = if g {
20908 self.func_g("fa_decode_vec_q_v3")
20909 } else {
20910 self.func("fa_decode_vec_q_v3")
20911 };
20912 let shmem = (32 * head_dim * 2) as u32; (
20914 fv,
20915 LaunchConfig {
20916 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20917 block_dim: (32, gqa, 1),
20918 shared_mem_bytes: shmem,
20919 },
20920 )
20921 } else if fa_v2_on() {
20922 let fv = if g {
20926 self.func_g("fa_decode_vec_q_v2")
20927 } else {
20928 self.func("fa_decode_vec_q_v2")
20929 };
20930 let shmem = (2 * 32 * head_dim * 2) as u32; (
20932 fv,
20933 LaunchConfig {
20934 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20935 block_dim: (32, gqa, 1),
20936 shared_mem_bytes: shmem,
20937 },
20938 )
20939 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
20940 {
20941 let fv = if g {
20945 self.func_g("fa_decode_vec_q_smem")
20946 } else {
20947 self.func("fa_decode_vec_q_smem")
20948 };
20949 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
20951 fv.set_attribute(
20952 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20953 shmem as i32,
20954 )?;
20955 (
20956 fv,
20957 LaunchConfig {
20958 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20959 block_dim: (32, gqa, 1),
20960 shared_mem_bytes: shmem,
20961 },
20962 )
20963 } else {
20964 let fv = if g {
20967 self.func_g("fa_decode_vec_q")
20968 } else {
20969 self.func("fa_decode_vec_q")
20970 };
20971 (
20972 fv,
20973 LaunchConfig {
20974 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20975 block_dim: (32, gqa, 1),
20976 shared_mem_bytes: 0,
20977 },
20978 )
20979 }
20980 } else {
20981 return self.fa_decode_scalar_unified(
20984 q,
20985 k,
20986 v,
20987 o,
20988 head_dim,
20989 n_head,
20990 n_head_kv,
20991 t_kv,
20992 None,
20993 scale,
20994 n_splits,
20995 if fa_vec { sp } else { 256 },
20996 k_tok_bytes,
20997 v_tok_bytes,
20998 g,
20999 part_o,
21000 part_m,
21001 part_l,
21002 None,
21003 );
21004 };
21005 let __s_b = self.gpu.stream();
21006 let mut b = __s_b.launch_builder(&f);
21007 b.arg(q)
21008 .arg(k)
21009 .arg(v)
21010 .arg(&mut *part_o)
21011 .arg(&mut *part_m)
21012 .arg(&mut *part_l)
21013 .arg(&hd)
21014 .arg(&nh)
21015 .arg(&nhkv)
21016 .arg(&tkvi)
21017 .arg(&scale)
21018 .arg(&nsp)
21019 .arg(&ktb)
21020 .arg(&vtb);
21021 unsafe {
21022 b.launch(cfg)?;
21023 }
21024 let (fc, cfg2) = (
21027 if g {
21028 self.func_g("fa_decode_combine_f32")
21029 } else {
21030 self.fa_func("fa_decode_combine_f32", head_dim)
21031 },
21032 LaunchConfig {
21033 grid_dim: (n_head as u32, 1, 1),
21034 block_dim: (head_dim as u32, 1, 1),
21035 shared_mem_bytes: 0,
21036 },
21037 );
21038 let __s_b2 = self.gpu.stream();
21039 let mut b2 = __s_b2.launch_builder(&fc);
21040 b2.arg(&*part_o)
21041 .arg(&*part_m)
21042 .arg(&*part_l)
21043 .arg(o)
21044 .arg(&hd)
21045 .arg(&nh)
21046 .arg(&nsp);
21047 unsafe {
21048 b2.launch(cfg2)?;
21049 }
21050 Ok(())
21051 }
21052
21053 #[allow(clippy::too_many_arguments)]
21064 pub fn fa_decode_batch_seqs_v4(
21065 &self,
21066 q: &CudaSlice<f32>,
21067 kv_ptrs: &cudarc::driver::CudaView<u64>,
21068 pos_seq: &CudaSlice<i32>,
21069 o: &mut CudaSlice<f32>,
21070 head_dim: usize,
21071 n_head: usize,
21072 n_head_kv: usize,
21073 b_n: usize,
21074 t_kv_max: usize,
21075 scale: f32,
21076 split_keys: usize,
21077 k_tok_bytes: usize,
21078 v_tok_bytes: usize,
21079 ) -> Result<(), Box<dyn std::error::Error>> {
21080 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21081 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21082 let o_len = b_n * n_head * n_splits_max * head_dim;
21083 let ml_len = b_n * n_head * n_splits_max;
21084 let mut part_guard = self.fa_part_pool.lock().unwrap();
21085 if part_guard
21086 .as_ref()
21087 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21088 .unwrap_or(true)
21089 {
21090 let old = part_guard.take();
21101 let (co, cm) = old
21102 .as_ref()
21103 .map(|pp| (pp.0.len(), pp.1.len()))
21104 .unwrap_or((0, 0));
21105 if let Some(old) = old {
21106 self.fa_part_retired.lock().unwrap().push(old);
21107 }
21108 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21109 eprintln!(
21110 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21111 co, o_len, cm, ml_len
21112 );
21113 }
21114 *part_guard = Some((
21115 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21116 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21117 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21118 ));
21119 }
21120 let pg = part_guard.as_mut().unwrap();
21121 self.gpu
21122 .stream()
21123 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21124 self.gpu
21125 .stream()
21126 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21127 self.gpu
21128 .stream()
21129 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21130 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21131 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21132 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21133 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21134 let gqa = (n_head / n_head_kv).max(1) as u32;
21135 let f = self.func("fa_decode_vec_q_seqs_v4");
21136 let shmem = (11520 + 32 * head_dim * 2) as u32;
21138 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21139 f.set_attribute(
21140 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21141 shmem as i32,
21142 )?;
21143 let cfg = LaunchConfig {
21144 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21145 block_dim: (32, gqa, 1),
21146 shared_mem_bytes: shmem,
21147 };
21148 {
21149 let __s_b = self.gpu.stream();
21150 let mut b = __s_b.launch_builder(&f);
21151 b.arg(q)
21152 .arg(kv_ptrs)
21153 .arg(pos_seq)
21154 .arg(&mut *part_o)
21155 .arg(&mut *part_m)
21156 .arg(&mut *part_l)
21157 .arg(&hd)
21158 .arg(&nh)
21159 .arg(&nhkv)
21160 .arg(&scale)
21161 .arg(&nspm)
21162 .arg(&spk)
21163 .arg(&ktb)
21164 .arg(&vtb);
21165 unsafe {
21166 b.launch(cfg)?;
21167 }
21168 }
21169 let fc = self.func("fa_decode_combine_seqs");
21170 let cfg2 = LaunchConfig {
21171 grid_dim: (n_head as u32, b_n as u32, 1),
21172 block_dim: (head_dim as u32, 1, 1),
21173 shared_mem_bytes: 0,
21174 };
21175 let __s_b2 = self.gpu.stream();
21176 let mut b2 = __s_b2.launch_builder(&fc);
21177 b2.arg(&*part_o)
21178 .arg(&*part_m)
21179 .arg(&*part_l)
21180 .arg(o)
21181 .arg(&hd)
21182 .arg(&nh)
21183 .arg(pos_seq)
21184 .arg(&nspm)
21185 .arg(&spk);
21186 unsafe {
21187 b2.launch(cfg2)?;
21188 }
21189 Ok(())
21190 }
21191
21192 #[allow(clippy::too_many_arguments)]
21199 pub fn append_kv_quantized_seqs(
21200 &self,
21201 k_rows: &CudaSlice<f32>,
21202 v_rows: &CudaSlice<f32>,
21203 kv_ptrs: &cudarc::driver::CudaView<u64>,
21204 pos_seq: &CudaSlice<i32>,
21205 b_n: usize,
21206 kv_dim_k: usize,
21207 kv_dim_v: usize,
21208 k_tok_bytes: usize,
21209 v_tok_bytes: usize,
21210 ) -> Result<(), Box<dyn std::error::Error>> {
21211 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21212 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21213 let cfg = LaunchConfig {
21214 grid_dim: (nblk, b_n as u32, 1),
21215 block_dim: (32, 1, 1),
21216 shared_mem_bytes: 0,
21217 };
21218 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21219 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21220 let __s_b = self.gpu.stream();
21221 let mut b = __s_b.launch_builder(&f);
21222 b.arg(k_rows)
21223 .arg(v_rows)
21224 .arg(kv_ptrs)
21225 .arg(pos_seq)
21226 .arg(&kdk)
21227 .arg(&kdv)
21228 .arg(&ktb)
21229 .arg(&vtb);
21230 unsafe {
21231 b.launch(cfg)?;
21232 }
21233 Ok(())
21234 }
21235
21236 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21242 std::env::var("MEMRA_NO_FA_VEC").is_err()
21243 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21244 && base_len + 1 >= fa_vec_min_tkv()
21245 && head_dim <= 256
21246 && head_dim % 32 == 0
21247 }
21248
21249 #[allow(clippy::too_many_arguments)]
21258 pub fn fa_decode_rows(
21259 &self,
21260 q: &CudaSlice<f32>,
21261 k: &cudarc::driver::CudaView<u8>,
21262 v: &cudarc::driver::CudaView<u8>,
21263 o: &mut CudaSlice<f32>,
21264 head_dim: usize,
21265 n_head: usize,
21266 n_head_kv: usize,
21267 base_len: usize,
21268 t: usize,
21269 scale: f32,
21270 k_tok_bytes: usize,
21271 v_tok_bytes: usize,
21272 base_dev: Option<(&CudaSlice<i32>, i32)>,
21276 kv_shared: bool,
21279 g: bool,
21283 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21286 ) -> Result<(), Box<dyn std::error::Error>> {
21287 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21288 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
21295 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21296 let v = *SP512.get_or_init(|| {
21299 std::env::var("MEMRA_FA_SP512")
21300 .ok()
21301 .and_then(|x| x.parse().ok())
21302 .unwrap_or(0)
21303 });
21304 sp = if v >= 8 {
21305 v
21306 } else {
21307 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21308 };
21309 }
21310 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21311 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21312 let gqa = (n_head / n_head_kv).max(1) as u32;
21313 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21324 groups.push((0, t, sp));
21325 } else {
21326 let mut r0 = 0usize;
21327 while r0 < t {
21328 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21329 let mut r1 = r0 + 1;
21330 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21331 r1 += 1;
21332 }
21333 groups.push((r0, r1 - r0, sp_g));
21334 r0 = r1;
21335 }
21336 }
21337 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21341 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21342 std::env::var("MEMRA_FA_SMEM_TKV")
21343 .ok()
21344 .and_then(|v| v.parse().ok())
21345 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21346 });
21347 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21348 let v3 = fa_v3_active(head_dim);
21349 let smem_rows =
21350 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21351 let _ = kv_shared;
21356 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21359 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21373 let tb512 = head_dim == 512
21375 && sp <= 32
21376 && n_head / n_head_kv.max(1) <= 16
21377 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21378 let fname = if tb512 {
21379 "fa_decode_vec_q_rows_v4_512_tb"
21380 } else if i2 {
21381 "fa_decode_vec_q_rows_dpl16_i2"
21382 } else if head_dim == 512 {
21383 "fa_decode_vec_q_rows_dpl16"
21384 }
21385 else if v4 {
21387 "fa_decode_vec_q_rows_v4"
21388 } else if v3 {
21389 "fa_decode_vec_q_rows_v3"
21390 } else if fa_v2_on() {
21391 "fa_decode_vec_q_rows_v2"
21392 } else if smem_rows {
21393 "fa_decode_vec_q_rows_smem"
21394 } else {
21395 "fa_decode_vec_q_rows"
21396 };
21397 let f = if head_dim == 512 {
21398 self.fa_func(fname, head_dim)
21399 } else if g {
21400 self.func_g(if smem_rows {
21408 "fa_decode_vec_q_rows"
21409 } else {
21410 fname
21411 })
21412 } else {
21413 self.func(fname)
21414 };
21415 let shmem = if tb512 {
21416 let gk = Self::gkv_on();
21418 let sh =
21419 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21420 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21421 f.set_attribute(
21422 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21423 sh as i32,
21424 )?;
21425 sh
21426 } else if v4 || v3 || smem_rows || fa_v2_on() {
21427 let sh = (if v4 {
21429 11520 + 32 * head_dim * if g { 1 } else { 2 }
21430 } else if v3 {
21431 32 * head_dim * 2
21432 } else {
21433 2 * 32 * head_dim * 2
21434 }) as u32;
21435 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21436 f.set_attribute(
21437 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21438 sh as i32,
21439 )?;
21440 sh
21441 } else {
21442 0
21443 };
21444 for &(r0, t_g, sp_g) in &groups {
21448 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21449 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21450 let base_i = (base_len + r0) as i32;
21451 let o_len = t_g * n_head * n_splits_g * head_dim;
21452 let ml_len = t_g * n_head * n_splits_g;
21453 let mut part_guard = self.fa_part_pool.lock().unwrap();
21454 if part_guard
21455 .as_ref()
21456 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21457 .unwrap_or(true)
21458 {
21459 let old = part_guard.take();
21470 let (co, cm) = old
21471 .as_ref()
21472 .map(|pp| (pp.0.len(), pp.1.len()))
21473 .unwrap_or((0, 0));
21474 if let Some(old) = old {
21475 self.fa_part_retired.lock().unwrap().push(old);
21476 }
21477 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21478 eprintln!(
21479 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21480 co, o_len, cm, ml_len
21481 );
21482 }
21483 *part_guard = Some((
21484 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21485 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21486 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21487 ));
21488 }
21489 let pg = part_guard.as_mut().unwrap();
21490 self.gpu
21491 .stream()
21492 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21493 self.gpu
21494 .stream()
21495 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21496 self.gpu
21497 .stream()
21498 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21499 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21500 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21501 let qv = self.view(q, t * n_head * head_dim);
21502 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21503 let cfg = LaunchConfig {
21504 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21505 block_dim: (32, gqa, 1),
21506 shared_mem_bytes: shmem,
21507 };
21508 {
21509 let __s_b = self.gpu.stream();
21510 let mut b = __s_b.launch_builder(&f);
21511 if tb512 {
21512 let (bd, plus) =
21514 base_dev.expect("hd512 rows twin requires a device base counter");
21515 let plus_g = plus + r0 as i32;
21516 let nr = t_g as i32;
21517 if Self::pdl_on() && Self::pdl_wb_on() {
21518 use cudarc::driver::{DevicePtr, DevicePtrMut};
21520 let s = &self.gpu.stream();
21521 let (pq, _b0) = q_g.device_ptr(s);
21522 let (pk, _b1) = k.device_ptr(s);
21523 let (pv, _b2) = v.device_ptr(s);
21524 let (po, _b3) = part_o.device_ptr_mut(s);
21525 let (pm, _b4) = part_m.device_ptr_mut(s);
21526 let (pl, _b5) = part_l.device_ptr_mut(s);
21527 let (pb, _b6) = bd.device_ptr(s);
21528 let mut ps = [
21529 &pq as *const _ as *mut std::ffi::c_void,
21530 &pk as *const _ as *mut _,
21531 &pv as *const _ as *mut _,
21532 &po as *const _ as *mut _,
21533 &pm as *const _ as *mut _,
21534 &pl as *const _ as *mut _,
21535 &hd as *const _ as *mut _,
21536 &nh as *const _ as *mut _,
21537 &nhkv as *const _ as *mut _,
21538 &pb as *const _ as *mut _,
21539 &plus_g as *const _ as *mut _,
21540 &scale as *const _ as *mut _,
21541 &nspm as *const _ as *mut _,
21542 &spk as *const _ as *mut _,
21543 &ktb as *const _ as *mut _,
21544 &vtb as *const _ as *mut _,
21545 &nr as *const _ as *mut _,
21546 ];
21547 unsafe {
21548 self.launch_pdl_flash(
21549 Self::gkv_on(),
21550 "fa_decode_vec_q_rows_v4_512_tb",
21551 (n_head_kv as u32, n_splits_g as u32, 1),
21552 (32, gqa, 1),
21553 shmem,
21554 &mut ps,
21555 )?;
21556 }
21557 } else {
21558 let cfg_tb = LaunchConfig {
21559 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21560 block_dim: (32, gqa, 1),
21561 shared_mem_bytes: shmem,
21562 };
21563 b.arg(&q_g)
21564 .arg(k)
21565 .arg(v)
21566 .arg(&mut *part_o)
21567 .arg(&mut *part_m)
21568 .arg(&mut *part_l)
21569 .arg(&hd)
21570 .arg(&nh)
21571 .arg(&nhkv)
21572 .arg(bd)
21573 .arg(&plus_g)
21574 .arg(&scale)
21575 .arg(&nspm)
21576 .arg(&spk)
21577 .arg(&ktb)
21578 .arg(&vtb)
21579 .arg(&nr);
21580 unsafe {
21581 b.launch(cfg_tb)?;
21582 }
21583 }
21584 } else if head_dim == 512 {
21585 let (bd, plus) =
21586 base_dev.expect("hd512 rows twin requires a device base counter");
21587 let plus_g = plus + r0 as i32;
21588 b.arg(&q_g)
21589 .arg(k)
21590 .arg(v)
21591 .arg(&mut *part_o)
21592 .arg(&mut *part_m)
21593 .arg(&mut *part_l)
21594 .arg(&hd)
21595 .arg(&nh)
21596 .arg(&nhkv)
21597 .arg(bd)
21598 .arg(&plus_g)
21599 .arg(&scale)
21600 .arg(&nspm)
21601 .arg(&spk)
21602 .arg(&ktb)
21603 .arg(&vtb);
21604 unsafe {
21605 b.launch(cfg)?;
21606 }
21607 } else {
21608 b.arg(&q_g)
21609 .arg(k)
21610 .arg(v)
21611 .arg(&mut *part_o)
21612 .arg(&mut *part_m)
21613 .arg(&mut *part_l)
21614 .arg(&hd)
21615 .arg(&nh)
21616 .arg(&nhkv)
21617 .arg(&base_i)
21618 .arg(&scale)
21619 .arg(&nspm)
21620 .arg(&spk)
21621 .arg(&ktb)
21622 .arg(&vtb);
21623 unsafe {
21624 b.launch(cfg)?;
21625 }
21626 }
21627 }
21628 let cfg2 = LaunchConfig {
21629 grid_dim: (n_head as u32, t_g as u32, 1),
21630 block_dim: (head_dim as u32, 1, 1),
21631 shared_mem_bytes: 0,
21632 };
21633 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21634 if head_dim == 512 {
21635 let (bd, plus) = base_dev.unwrap();
21638 let plus_g = plus + r0 as i32;
21639 if let Some((oq, od)) = q8_out.as_mut() {
21640 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21642 if Self::pdl_on() && Self::pdl_wb_on() {
21643 use cudarc::driver::{DevicePtr, DevicePtrMut};
21645 let s = &self.gpu.stream();
21646 let (po, _g0) = part_o.device_ptr(s);
21647 let (pm, _g1) = part_m.device_ptr(s);
21648 let (pl, _g2) = part_l.device_ptr(s);
21649 let (pq, _g3) = oq.device_ptr_mut(s);
21650 let (pd, _g4) = od.device_ptr_mut(s);
21651 let (pb, _g5) = bd.device_ptr(s);
21652 let mut ps = [
21653 &po as *const _ as *mut std::ffi::c_void,
21654 &pm as *const _ as *mut _,
21655 &pl as *const _ as *mut _,
21656 &pq as *const _ as *mut _,
21657 &pd as *const _ as *mut _,
21658 &hd as *const _ as *mut _,
21659 &nh as *const _ as *mut _,
21660 &pb as *const _ as *mut _,
21661 &plus_g as *const _ as *mut _,
21662 &nspm as *const _ as *mut _,
21663 &spk as *const _ as *mut _,
21664 ];
21665 unsafe {
21666 self.launch_pdl_flash(
21667 Self::gkv_on(),
21668 "fa_decode_combine_rows_dc_q8_1",
21669 cfg2.grid_dim,
21670 cfg2.block_dim,
21671 0,
21672 &mut ps,
21673 )?;
21674 }
21675 continue;
21676 }
21677 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21678 let __s_b2 = self.gpu.stream();
21679 let mut b2 = __s_b2.launch_builder(&fc);
21680 b2.arg(&*part_o)
21681 .arg(&*part_m)
21682 .arg(&*part_l)
21683 .arg(&mut **oq)
21684 .arg(&mut **od)
21685 .arg(&hd)
21686 .arg(&nh)
21687 .arg(bd)
21688 .arg(&plus_g)
21689 .arg(&nspm)
21690 .arg(&spk);
21691 unsafe {
21692 b2.launch(cfg2)?;
21693 }
21694 continue;
21695 }
21696 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21697 let __s_b2 = self.gpu.stream();
21698 let mut b2 = __s_b2.launch_builder(&fc);
21699 b2.arg(&*part_o)
21700 .arg(&*part_m)
21701 .arg(&*part_l)
21702 .arg(&mut o_g)
21703 .arg(&hd)
21704 .arg(&nh)
21705 .arg(bd)
21706 .arg(&plus_g)
21707 .arg(&nspm)
21708 .arg(&spk);
21709 unsafe {
21710 b2.launch(cfg2)?;
21711 }
21712 } else {
21713 assert!(
21716 q8_out.is_none(),
21717 "rows q8 emit requires the hd512 dc combine"
21718 );
21719 let fc = self.func("fa_decode_combine_rows");
21720 let __s_b2 = self.gpu.stream();
21721 let mut b2 = __s_b2.launch_builder(&fc);
21722 b2.arg(&*part_o)
21723 .arg(&*part_m)
21724 .arg(&*part_l)
21725 .arg(&mut o_g)
21726 .arg(&hd)
21727 .arg(&nh)
21728 .arg(&base_i)
21729 .arg(&nspm)
21730 .arg(&spk);
21731 unsafe {
21732 b2.launch(cfg2)?;
21733 }
21734 }
21735 }
21736 Ok(())
21737 }
21738
21739 #[allow(clippy::too_many_arguments)]
21743 pub fn fa_decode_rows_w(
21744 &self,
21745 q: &CudaSlice<f32>,
21746 k: &cudarc::driver::CudaView<u8>,
21747 v: &cudarc::driver::CudaView<u8>,
21748 o: &mut CudaSlice<f32>,
21749 head_dim: usize,
21750 n_head: usize,
21751 n_head_kv: usize,
21752 base_dev: &CudaSlice<i32>,
21753 base_plus: i32,
21754 t: usize,
21755 scale: f32,
21756 window: usize,
21757 k_tok_bytes: usize,
21758 v_tok_bytes: usize,
21759 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21760 ) -> Result<(), Box<dyn std::error::Error>> {
21761 debug_assert!(head_dim == 256);
21766 let sp = {
21774 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21775 let v = *SPW.get_or_init(|| {
21776 std::env::var("MEMRA_FA_SPW")
21777 .ok()
21778 .and_then(|x| x.parse().ok())
21779 .unwrap_or(0)
21780 });
21781 if v >= 8 {
21782 v
21783 } else {
21784 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21785 }
21786 };
21787 let n_splits_max = (window + sp - 1) / sp;
21788 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21789 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21790 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21791 let gqa = (n_head / n_head_kv).max(1) as u32;
21792 let o_len = t * n_head * n_splits_max * head_dim;
21793 let ml_len = t * n_head * n_splits_max;
21794 let mut part_guard = self.fa_part_pool.lock().unwrap();
21795 if part_guard
21796 .as_ref()
21797 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21798 .unwrap_or(true)
21799 {
21800 let old = part_guard.take();
21811 let (co, cm) = old
21812 .as_ref()
21813 .map(|pp| (pp.0.len(), pp.1.len()))
21814 .unwrap_or((0, 0));
21815 if let Some(old) = old {
21816 self.fa_part_retired.lock().unwrap().push(old);
21817 }
21818 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21819 eprintln!(
21820 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21821 co, o_len, cm, ml_len
21822 );
21823 }
21824 *part_guard = Some((
21825 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21826 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21827 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21828 ));
21829 }
21830 let pg = part_guard.as_mut().unwrap();
21831 self.gpu
21832 .stream()
21833 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21834 self.gpu
21835 .stream()
21836 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21837 self.gpu
21838 .stream()
21839 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21840 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21841 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21847 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
21848 std::env::var("MEMRA_FA_SMEM_TKV")
21849 .ok()
21850 .and_then(|v| v.parse().ok())
21851 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21852 });
21853 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21859 let wg = Self::wkv_on();
21864 let sp2 =
21867 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
21868 if sp2 {
21869 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21870 if Self::pdl_on() && Self::pdl_wb_on() {
21871 use cudarc::driver::{DevicePtr, DevicePtrMut};
21873 let s = &self.gpu.stream();
21874 let (pq, _b0) = q.device_ptr(s);
21875 let (pk, _b1) = k.device_ptr(s);
21876 let (pv, _b2) = v.device_ptr(s);
21877 let (po, _b3) = part_o.device_ptr_mut(s);
21878 let (pm, _b4) = part_m.device_ptr_mut(s);
21879 let (pl, _b5) = part_l.device_ptr_mut(s);
21880 let (pb, _b6) = base_dev.device_ptr(s);
21881 let mut ps = [
21882 &pq as *const _ as *mut std::ffi::c_void,
21883 &pk as *const _ as *mut _,
21884 &pv as *const _ as *mut _,
21885 &po as *const _ as *mut _,
21886 &pm as *const _ as *mut _,
21887 &pl as *const _ as *mut _,
21888 &hd as *const _ as *mut _,
21889 &nh as *const _ as *mut _,
21890 &nhkv as *const _ as *mut _,
21891 &pb as *const _ as *mut _,
21892 &base_plus as *const _ as *mut _,
21893 &scale as *const _ as *mut _,
21894 &nspm as *const _ as *mut _,
21895 &spk as *const _ as *mut _,
21896 &ktb as *const _ as *mut _,
21897 &vtb as *const _ as *mut _,
21898 &wini as *const _ as *mut _,
21899 ];
21900 unsafe {
21901 self.launch_pdl_flash(
21902 wg,
21903 "fa_decode_vec_q_rows_v4_w_sp",
21904 (n_head_kv as u32, n_splits_max as u32, t as u32),
21905 (32, gqa + 1, 1),
21906 sh,
21907 &mut ps,
21908 )?;
21909 }
21910 } else {
21911 let f = if wg {
21912 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
21913 } else {
21914 self.func("fa_decode_vec_q_rows_v4_w_sp")
21915 };
21916 f.set_attribute(
21917 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21918 sh as i32,
21919 )?;
21920 let cfg = LaunchConfig {
21921 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21922 block_dim: (32, gqa + 1, 1),
21923 shared_mem_bytes: sh,
21924 };
21925 let __s_b = self.gpu.stream();
21926 let mut b = __s_b.launch_builder(&f);
21927 b.arg(q)
21928 .arg(k)
21929 .arg(v)
21930 .arg(&mut *part_o)
21931 .arg(&mut *part_m)
21932 .arg(&mut *part_l)
21933 .arg(&hd)
21934 .arg(&nh)
21935 .arg(&nhkv)
21936 .arg(base_dev)
21937 .arg(&base_plus)
21938 .arg(&scale)
21939 .arg(&nspm)
21940 .arg(&spk)
21941 .arg(&ktb)
21942 .arg(&vtb)
21943 .arg(&wini);
21944 unsafe {
21945 b.launch(cfg)?;
21946 }
21947 }
21948 } else {
21949 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
21950 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21952 use cudarc::driver::{DevicePtr, DevicePtrMut};
21953 let s = &self.gpu.stream();
21954 let (pq, _b0) = q.device_ptr(s);
21955 let (pk, _b1) = k.device_ptr(s);
21956 let (pv, _b2) = v.device_ptr(s);
21957 let (po, _b3) = part_o.device_ptr_mut(s);
21958 let (pm, _b4) = part_m.device_ptr_mut(s);
21959 let (pl, _b5) = part_l.device_ptr_mut(s);
21960 let (pb, _b6) = base_dev.device_ptr(s);
21961 let mut ps = [
21962 &pq as *const _ as *mut std::ffi::c_void,
21963 &pk as *const _ as *mut _,
21964 &pv as *const _ as *mut _,
21965 &po as *const _ as *mut _,
21966 &pm as *const _ as *mut _,
21967 &pl as *const _ as *mut _,
21968 &hd as *const _ as *mut _,
21969 &nh as *const _ as *mut _,
21970 &nhkv as *const _ as *mut _,
21971 &pb as *const _ as *mut _,
21972 &base_plus as *const _ as *mut _,
21973 &scale as *const _ as *mut _,
21974 &nspm as *const _ as *mut _,
21975 &spk as *const _ as *mut _,
21976 &ktb as *const _ as *mut _,
21977 &vtb as *const _ as *mut _,
21978 &wini as *const _ as *mut _,
21979 ];
21980 unsafe {
21981 self.launch_pdl_flash(
21982 wg,
21983 "fa_decode_vec_q_rows_v4_w",
21984 (n_head_kv as u32, n_splits_max as u32, t as u32),
21985 (32, gqa, 1),
21986 sh,
21987 &mut ps,
21988 )?;
21989 }
21990 } else {
21991 let pick = |name: &str| {
21992 if wg {
21993 self.func_g(name)
21994 } else {
21995 self.func(name)
21996 }
21997 };
21998 let (f, sh) = if fa_v4_at(window) {
21999 let f = pick("fa_decode_vec_q_rows_v4_w");
22000 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22001 } else if smem_tkv > 0 && window >= smem_tkv {
22002 (
22005 pick("fa_decode_vec_q_rows_smem_w"),
22006 (2 * 32 * head_dim * 2) as u32,
22007 )
22008 } else {
22009 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22010 };
22011 f.set_attribute(
22012 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22013 sh as i32,
22014 )?;
22015 let cfg = LaunchConfig {
22016 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22017 block_dim: (32, gqa, 1),
22018 shared_mem_bytes: sh,
22019 };
22020 let __s_b = self.gpu.stream();
22021 let mut b = __s_b.launch_builder(&f);
22022 b.arg(q)
22023 .arg(k)
22024 .arg(v)
22025 .arg(&mut *part_o)
22026 .arg(&mut *part_m)
22027 .arg(&mut *part_l)
22028 .arg(&hd)
22029 .arg(&nh)
22030 .arg(&nhkv)
22031 .arg(base_dev)
22032 .arg(&base_plus)
22033 .arg(&scale)
22034 .arg(&nspm)
22035 .arg(&spk)
22036 .arg(&ktb)
22037 .arg(&vtb)
22038 .arg(&wini);
22039 unsafe {
22040 b.launch(cfg)?;
22041 }
22042 }
22043 }
22044 let cfg2 = LaunchConfig {
22045 grid_dim: (n_head as u32, t as u32, 1),
22046 block_dim: (head_dim as u32, 1, 1),
22047 shared_mem_bytes: 0,
22048 };
22049 if let Some((oq, od)) = q8_out {
22050 if Self::pdl_on() && Self::pdl_wb_on() {
22053 use cudarc::driver::{DevicePtr, DevicePtrMut};
22055 let s = &self.gpu.stream();
22056 let (po, _g0) = part_o.device_ptr(s);
22057 let (pm, _g1) = part_m.device_ptr(s);
22058 let (pl, _g2) = part_l.device_ptr(s);
22059 let (pq, _g3) = oq.device_ptr_mut(s);
22060 let (pd, _g4) = od.device_ptr_mut(s);
22061 let mut ps = [
22062 &po as *const _ as *mut std::ffi::c_void,
22063 &pm as *const _ as *mut _,
22064 &pl as *const _ as *mut _,
22065 &pq as *const _ as *mut _,
22066 &pd as *const _ as *mut _,
22067 &hd as *const _ as *mut _,
22068 &nh as *const _ as *mut _,
22069 &nspm as *const _ as *mut _,
22070 &spk as *const _ as *mut _,
22071 &wini as *const _ as *mut _,
22072 ];
22073 unsafe {
22074 self.launch_pdl_flash(
22075 wg,
22076 "fa_decode_combine_rows_w_q8_1",
22077 cfg2.grid_dim,
22078 cfg2.block_dim,
22079 0,
22080 &mut ps,
22081 )?;
22082 }
22083 return Ok(());
22084 }
22085 let fc = if wg {
22086 self.func_g("fa_decode_combine_rows_w_q8_1")
22087 } else {
22088 self.func("fa_decode_combine_rows_w_q8_1")
22089 };
22090 let __s_b2 = self.gpu.stream();
22091 let mut b2 = __s_b2.launch_builder(&fc);
22092 b2.arg(&*part_o)
22093 .arg(&*part_m)
22094 .arg(&*part_l)
22095 .arg(oq)
22096 .arg(od)
22097 .arg(&hd)
22098 .arg(&nh)
22099 .arg(&nspm)
22100 .arg(&spk)
22101 .arg(&wini);
22102 unsafe {
22103 b2.launch(cfg2)?;
22104 }
22105 return Ok(());
22106 }
22107 let fc = if wg {
22108 self.func_g("fa_decode_combine_rows_w")
22109 } else {
22110 self.func("fa_decode_combine_rows_w")
22111 };
22112 let __s_b2 = self.gpu.stream();
22113 let mut b2 = __s_b2.launch_builder(&fc);
22114 b2.arg(&*part_o)
22115 .arg(&*part_m)
22116 .arg(&*part_l)
22117 .arg(o)
22118 .arg(&hd)
22119 .arg(&nh)
22120 .arg(&nspm)
22121 .arg(&spk)
22122 .arg(&wini);
22123 unsafe {
22124 b2.launch(cfg2)?;
22125 }
22126 Ok(())
22127 }
22128
22129 #[allow(clippy::too_many_arguments)]
22135 pub fn fa_decode_rows_dc(
22136 &self,
22137 q: &CudaSlice<f32>,
22138 k: &cudarc::driver::CudaView<u8>,
22139 v: &cudarc::driver::CudaView<u8>,
22140 o: &mut CudaSlice<f32>,
22141 head_dim: usize,
22142 n_head: usize,
22143 n_head_kv: usize,
22144 base_dev: &CudaSlice<i32>,
22145 t_kv_upper: usize,
22146 t: usize,
22147 scale: f32,
22148 k_tok_bytes: usize,
22149 v_tok_bytes: usize,
22150 base_plus: i32,
22151 g: bool,
22152 ) -> Result<(), Box<dyn std::error::Error>> {
22153 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22154 assert!(
22155 v4 || fa_v3_active(head_dim),
22156 "stream fa rows requires the v3 or v4 lane"
22157 );
22158 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22159 if v4 {
22160 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22161 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22162 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22163 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22164 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22165 let gqa = (n_head / n_head_kv).max(1) as u32;
22166 let o_len = t * n_head * n_splits_max * head_dim;
22167 let ml_len = t * n_head * n_splits_max;
22168 let mut part_guard = self.fa_part_pool.lock().unwrap();
22169 if part_guard
22170 .as_ref()
22171 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22172 .unwrap_or(true)
22173 {
22174 let old = part_guard.take();
22185 let (co, cm) = old
22186 .as_ref()
22187 .map(|pp| (pp.0.len(), pp.1.len()))
22188 .unwrap_or((0, 0));
22189 if let Some(old) = old {
22190 self.fa_part_retired.lock().unwrap().push(old);
22191 }
22192 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22193 eprintln!(
22194 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22195 co, o_len, cm, ml_len
22196 );
22197 }
22198 *part_guard = Some((
22199 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22200 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22201 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22202 ));
22203 }
22204 let pg = part_guard.as_mut().unwrap();
22205 self.gpu
22206 .stream()
22207 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22208 self.gpu
22209 .stream()
22210 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22211 self.gpu
22212 .stream()
22213 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22214 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22215 let f = if g {
22216 self.func_g("fa_decode_vec_q_rows_v4_dc")
22217 } else {
22218 self.func("fa_decode_vec_q_rows_v4_dc")
22219 };
22220 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22221 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22222 f.set_attribute(
22223 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22224 sh as i32,
22225 )?;
22226 let cfg = LaunchConfig {
22227 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22228 block_dim: (32, gqa, 1),
22229 shared_mem_bytes: sh,
22230 };
22231 let __s_b = self.gpu.stream();
22232 let mut b = __s_b.launch_builder(&f);
22233 b.arg(q)
22234 .arg(k)
22235 .arg(v)
22236 .arg(&mut *part_o)
22237 .arg(&mut *part_m)
22238 .arg(&mut *part_l)
22239 .arg(&hd)
22240 .arg(&nh)
22241 .arg(&nhkv)
22242 .arg(base_dev)
22243 .arg(&base_plus)
22244 .arg(&scale)
22245 .arg(&nspm)
22246 .arg(&spk)
22247 .arg(&ktb)
22248 .arg(&vtb);
22249 unsafe {
22250 b.launch(cfg)?;
22251 }
22252 let fc = self.func("fa_decode_combine_rows_dc");
22253 let cfg2 = LaunchConfig {
22254 grid_dim: (n_head as u32, t as u32, 1),
22255 block_dim: (head_dim as u32, 1, 1),
22256 shared_mem_bytes: 0,
22257 };
22258 let __s_b2 = self.gpu.stream();
22259 let mut b2 = __s_b2.launch_builder(&fc);
22260 b2.arg(&*part_o)
22261 .arg(&*part_m)
22262 .arg(&*part_l)
22263 .arg(o)
22264 .arg(&hd)
22265 .arg(&nh)
22266 .arg(base_dev)
22267 .arg(&base_plus)
22268 .arg(&nspm)
22269 .arg(&spk);
22270 unsafe {
22271 b2.launch(cfg2)?;
22272 }
22273 return Ok(());
22274 }
22275 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22276 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22277 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22278 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22279 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22280 let gqa = (n_head / n_head_kv).max(1) as u32;
22281 let o_len = t * n_head * n_splits_max * head_dim;
22282 let ml_len = t * n_head * n_splits_max;
22283 let mut part_guard = self.fa_part_pool.lock().unwrap();
22284 if part_guard
22285 .as_ref()
22286 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22287 .unwrap_or(true)
22288 {
22289 let old = part_guard.take();
22300 let (co, cm) = old
22301 .as_ref()
22302 .map(|pp| (pp.0.len(), pp.1.len()))
22303 .unwrap_or((0, 0));
22304 if let Some(old) = old {
22305 self.fa_part_retired.lock().unwrap().push(old);
22306 }
22307 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22308 eprintln!(
22309 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22310 co, o_len, cm, ml_len
22311 );
22312 }
22313 *part_guard = Some((
22314 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22315 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22316 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22317 ));
22318 }
22319 let pg = part_guard.as_mut().unwrap();
22320 self.gpu
22321 .stream()
22322 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22323 self.gpu
22324 .stream()
22325 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22326 self.gpu
22327 .stream()
22328 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22329 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22330 let f = self.func("fa_decode_vec_q_rows_v3_dc");
22331 let sh = (32 * head_dim * 2) as u32;
22332 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22333 f.set_attribute(
22334 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22335 sh as i32,
22336 )?;
22337 let cfg = LaunchConfig {
22338 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22339 block_dim: (32, gqa, 1),
22340 shared_mem_bytes: sh,
22341 };
22342 let __s_b = self.gpu.stream();
22343 let mut b = __s_b.launch_builder(&f);
22344 b.arg(q)
22345 .arg(k)
22346 .arg(v)
22347 .arg(&mut *part_o)
22348 .arg(&mut *part_m)
22349 .arg(&mut *part_l)
22350 .arg(&hd)
22351 .arg(&nh)
22352 .arg(&nhkv)
22353 .arg(base_dev)
22354 .arg(&scale)
22355 .arg(&nspm)
22356 .arg(&spk)
22357 .arg(&ktb)
22358 .arg(&vtb);
22359 unsafe {
22360 b.launch(cfg)?;
22361 }
22362 let fc = self.func("fa_decode_combine_rows_dc");
22363 let cfg2 = LaunchConfig {
22364 grid_dim: (n_head as u32, t as u32, 1),
22365 block_dim: (head_dim as u32, 1, 1),
22366 shared_mem_bytes: 0,
22367 };
22368 let plus0 = 0i32;
22369 let __s_b2 = self.gpu.stream();
22370 let mut b2 = __s_b2.launch_builder(&fc);
22371 b2.arg(&*part_o)
22372 .arg(&*part_m)
22373 .arg(&*part_l)
22374 .arg(o)
22375 .arg(&hd)
22376 .arg(&nh)
22377 .arg(base_dev)
22378 .arg(&plus0)
22379 .arg(&nspm)
22380 .arg(&spk);
22381 unsafe {
22382 b2.launch(cfg2)?;
22383 }
22384 Ok(())
22385 }
22386
22387 pub fn fa_decode_dc(
22398 &self,
22399 q: &CudaSlice<f32>,
22400 k: &cudarc::driver::CudaView<u8>,
22401 v: &cudarc::driver::CudaView<u8>,
22402 o: &mut CudaSlice<f32>,
22403 head_dim: usize,
22404 n_head: usize,
22405 n_head_kv: usize,
22406 t_kv_dev: &CudaSlice<i32>,
22407 bucket_max: usize,
22408 scale: f32,
22409 k_tok_bytes: usize,
22410 v_tok_bytes: usize,
22411 g: bool,
22412 ) -> Result<(), Box<dyn std::error::Error>> {
22413 self.fa_decode_dc_q8(
22414 q,
22415 k,
22416 v,
22417 o,
22418 head_dim,
22419 n_head,
22420 n_head_kv,
22421 t_kv_dev,
22422 bucket_max,
22423 scale,
22424 k_tok_bytes,
22425 v_tok_bytes,
22426 g,
22427 None,
22428 )
22429 }
22430
22431 #[allow(clippy::too_many_arguments)]
22434 pub fn fa_decode_dc_q8(
22435 &self,
22436 q: &CudaSlice<f32>,
22437 k: &cudarc::driver::CudaView<u8>,
22438 v: &cudarc::driver::CudaView<u8>,
22439 o: &mut CudaSlice<f32>,
22440 head_dim: usize,
22441 n_head: usize,
22442 n_head_kv: usize,
22443 t_kv_dev: &CudaSlice<i32>,
22444 bucket_max: usize,
22445 scale: f32,
22446 k_tok_bytes: usize,
22447 v_tok_bytes: usize,
22448 g: bool,
22449 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22450 ) -> Result<(), Box<dyn std::error::Error>> {
22451 let mut fa_vec =
22459 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22460 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22461 fa_vec = false;
22462 } let sp = fa_split_keys(bucket_max, n_head_kv);
22464 let n_splits = if fa_vec {
22465 ((bucket_max + sp - 1) / sp).max(1)
22466 } else {
22467 ((bucket_max + 255) / 256).max(1)
22468 };
22469 let o_len = n_head * n_splits * head_dim;
22470 let ml_len = n_head * n_splits;
22471 let mut part_guard = self.fa_part_pool.lock().unwrap();
22472 if part_guard
22473 .as_ref()
22474 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22475 .unwrap_or(true)
22476 {
22477 let old = part_guard.take();
22488 let (co, cm) = old
22489 .as_ref()
22490 .map(|pp| (pp.0.len(), pp.1.len()))
22491 .unwrap_or((0, 0));
22492 if let Some(old) = old {
22493 self.fa_part_retired.lock().unwrap().push(old);
22494 }
22495 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22496 eprintln!(
22497 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22498 co, o_len, cm, ml_len
22499 );
22500 }
22501 *part_guard = Some((
22502 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22503 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22504 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22505 ));
22506 }
22507 let pg = part_guard.as_mut().unwrap();
22508 self.gpu
22509 .stream()
22510 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22511 self.gpu
22512 .stream()
22513 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22514 self.gpu
22515 .stream()
22516 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22517 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22518 let (hd, nh, nhkv, nsp) = (
22519 head_dim as i32,
22520 n_head as i32,
22521 n_head_kv as i32,
22522 n_splits as i32,
22523 );
22524 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22525 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22526 let deep = fa_vec
22529 && head_dim == 256
22530 && fa_v4_at(bucket_max)
22531 && !g
22532 && fa_deep_at(bucket_max)
22533 && !matches!(fa_v4_mode(), "noB3" | "stage");
22534 let (f, cfg) = if fa_vec
22535 && head_dim == 512
22536 && bucket_max >= {
22537 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22538 *FA512_MIN_DC.get_or_init(|| {
22539 std::env::var("MEMRA_FA512_MIN")
22540 .ok()
22541 .and_then(|v| v.parse().ok())
22542 .unwrap_or(512)
22543 })
22544 } {
22545 let gqa = (n_head / n_head_kv).max(1) as u32;
22547 (
22548 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22549 LaunchConfig {
22550 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22551 block_dim: (32, gqa, 1),
22552 shared_mem_bytes: 0,
22553 },
22554 )
22555 } else if fa_vec && head_dim == 512 {
22556 let q_view = q.as_view();
22559 let mut o_view = o.as_view_mut();
22560 return self.fa_decode_scalar_unified(
22561 &q_view,
22562 k,
22563 v,
22564 &mut o_view,
22565 head_dim,
22566 n_head,
22567 n_head_kv,
22568 0,
22569 Some(t_kv_dev),
22570 scale,
22571 n_splits,
22572 sp,
22573 k_tok_bytes,
22574 v_tok_bytes,
22575 g,
22576 &mut *part_o,
22577 &mut *part_m,
22578 &mut *part_l,
22579 q8_out,
22580 );
22581 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22582 let gqa = (n_head / n_head_kv).max(1) as u32;
22585 let fv = if g {
22586 self.func_g("fa_decode_vec_q_v4_dc")
22587 } else if deep {
22588 self.func("fa_decode_vec_q_v4_deep_dc")
22589 } else {
22590 self.func("fa_decode_vec_q_v4_dc")
22591 };
22592 let shmem =
22593 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22594 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22595 fv.set_attribute(
22596 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22597 shmem as i32,
22598 )?;
22599 (
22600 fv,
22601 LaunchConfig {
22602 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22603 block_dim: (32, gqa, 1),
22604 shared_mem_bytes: shmem,
22605 },
22606 )
22607 } else if fa_vec && fa_v3_active(head_dim) {
22608 let gqa = (n_head / n_head_kv).max(1) as u32;
22611 let fv = if g {
22612 self.func_g("fa_decode_vec_q_v3_dc")
22613 } else {
22614 self.func("fa_decode_vec_q_v3_dc")
22615 };
22616 let shmem = (32 * head_dim * 2) as u32; (
22618 fv,
22619 LaunchConfig {
22620 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22621 block_dim: (32, gqa, 1),
22622 shared_mem_bytes: shmem,
22623 },
22624 )
22625 } else if fa_vec && fa_v2_on() {
22626 let gqa = (n_head / n_head_kv).max(1) as u32;
22630 let fv = if g {
22631 self.func_g("fa_decode_vec_q_v2_dc")
22632 } else {
22633 self.func("fa_decode_vec_q_v2_dc")
22634 };
22635 let shmem = (2 * 32 * head_dim * 2) as u32; (
22637 fv,
22638 LaunchConfig {
22639 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22640 block_dim: (32, gqa, 1),
22641 shared_mem_bytes: shmem,
22642 },
22643 )
22644 } else if fa_vec {
22645 let gqa = (n_head / n_head_kv).max(1) as u32;
22646 let fv = if g {
22648 self.func_g("fa_decode_vec_q_dc")
22649 } else {
22650 self.func("fa_decode_vec_q_dc")
22651 };
22652 (
22653 fv,
22654 LaunchConfig {
22655 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22656 block_dim: (32, gqa, 1),
22657 shared_mem_bytes: 0,
22658 },
22659 )
22660 } else {
22661 let q_view = q.as_view();
22662 let mut o_view = o.as_view_mut();
22663 return self.fa_decode_scalar_unified(
22664 &q_view,
22665 k,
22666 v,
22667 &mut o_view,
22668 head_dim,
22669 n_head,
22670 n_head_kv,
22671 0,
22672 Some(t_kv_dev),
22673 scale,
22674 n_splits,
22675 if fa_vec { sp } else { 256 },
22676 k_tok_bytes,
22677 v_tok_bytes,
22678 g,
22679 &mut *part_o,
22680 &mut *part_m,
22681 &mut *part_l,
22682 q8_out,
22683 );
22684 };
22685 let ski = sp as i32; let __s_b = self.gpu.stream();
22687 let mut b = __s_b.launch_builder(&f);
22688 b.arg(q)
22689 .arg(k)
22690 .arg(v)
22691 .arg(&mut *part_o)
22692 .arg(&mut *part_m)
22693 .arg(&mut *part_l)
22694 .arg(&hd)
22695 .arg(&nh)
22696 .arg(&nhkv)
22697 .arg(t_kv_dev)
22698 .arg(&scale)
22699 .arg(&nsp)
22700 .arg(&ski)
22701 .arg(&ktb)
22702 .arg(&vtb);
22703 unsafe {
22704 b.launch(cfg)?;
22705 }
22706 let cfg2 = LaunchConfig {
22707 grid_dim: (n_head as u32, 1, 1),
22708 block_dim: (head_dim as u32, 1, 1),
22709 shared_mem_bytes: 0,
22710 };
22711 if let Some((oq, od)) = q8_out {
22712 let fc = if g {
22713 self.func_g("fa_decode_combine_q8_1")
22714 } else {
22715 self.fa_func("fa_decode_combine_q8_1", head_dim)
22716 };
22717 let __s_b2 = self.gpu.stream();
22718 let mut b2 = __s_b2.launch_builder(&fc);
22719 b2.arg(&*part_o)
22720 .arg(&*part_m)
22721 .arg(&*part_l)
22722 .arg(oq)
22723 .arg(od)
22724 .arg(&hd)
22725 .arg(&nh)
22726 .arg(&nsp);
22727 unsafe {
22728 b2.launch(cfg2)?;
22729 }
22730 return Ok(());
22731 }
22732 let fc = if g {
22733 self.func_g("fa_decode_combine_f32")
22734 } else {
22735 self.fa_func("fa_decode_combine_f32", head_dim)
22736 };
22737 let __s_b2 = self.gpu.stream();
22738 let mut b2 = __s_b2.launch_builder(&fc);
22739 b2.arg(&*part_o)
22740 .arg(&*part_m)
22741 .arg(&*part_l)
22742 .arg(o)
22743 .arg(&hd)
22744 .arg(&nh)
22745 .arg(&nsp);
22746 unsafe {
22747 b2.launch(cfg2)?;
22748 }
22749 Ok(())
22750 }
22751
22752 #[allow(clippy::too_many_arguments)]
22756 pub fn append_kv_quantized_dcw(
22757 &self,
22758 k_row: &CudaSlice<f32>,
22759 v_row: &CudaSlice<f32>,
22760 kc: &mut CudaSlice<u8>,
22761 vc: &mut CudaSlice<u8>,
22762 len_dev: &CudaSlice<i32>,
22763 base_dev: Option<&CudaSlice<i32>>,
22764 kv_dim_k: usize,
22765 kv_dim_v: usize,
22766 k_tok_bytes: usize,
22767 v_tok_bytes: usize,
22768 ) -> Result<(), Box<dyn std::error::Error>> {
22769 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22770 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22771 let cfg = LaunchConfig {
22772 grid_dim: (nblk, 1, 1),
22773 block_dim: (32, 1, 1),
22774 shared_mem_bytes: 0,
22775 };
22776 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22777 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22778 let null: u64 = 0;
22779 let __s_b = self.gpu.stream();
22780 let mut b = __s_b.launch_builder(&f);
22781 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22782 match base_dev {
22783 Some(base) => {
22784 b.arg(base);
22785 }
22786 None => {
22787 b.arg(&null);
22788 }
22789 }
22790 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22791 unsafe {
22792 b.launch(cfg)?;
22793 }
22794 Ok(())
22795 }
22796
22797 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22799 let f = self.func("inc_i32");
22800 let cfg = LaunchConfig {
22801 grid_dim: (1, 1, 1),
22802 block_dim: (1, 1, 1),
22803 shared_mem_bytes: 0,
22804 };
22805 let __s_b = self.gpu.stream();
22806 let mut b = __s_b.launch_builder(&f);
22807 b.arg(counter);
22808 unsafe {
22809 b.launch(cfg)?;
22810 }
22811 Ok(())
22812 }
22813
22814 #[allow(clippy::too_many_arguments)]
22823 fn fa_part_pool_grow(
22827 &self,
22828 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22829 o_len: usize,
22830 ml_len: usize,
22831 ) -> Result<(), Box<dyn std::error::Error>> {
22832 if part_guard
22833 .as_ref()
22834 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22835 .unwrap_or(true)
22836 {
22837 let old = part_guard.take();
22838 let (co, cm) = old
22839 .as_ref()
22840 .map(|pp| (pp.0.len(), pp.1.len()))
22841 .unwrap_or((0, 0));
22842 if let Some(old) = old {
22843 self.fa_part_retired.lock().unwrap().push(old);
22844 }
22845 *part_guard = Some((
22846 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22847 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22848 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22849 ));
22850 }
22851 Ok(())
22852 }
22853
22854 pub fn fa_dcw_pool_ensure(
22857 &self,
22858 head_dim: usize,
22859 n_head: usize,
22860 n_head_kv: usize,
22861 bucket_max: usize,
22862 ) -> Result<(), Box<dyn std::error::Error>> {
22863 let sp = fa_split_keys(bucket_max, n_head_kv);
22864 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22865 let o_len = n_head * n_splits * head_dim;
22866 let ml_len = n_head * n_splits;
22867 let mut part_guard = self.fa_part_pool.lock().unwrap();
22868 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
22869 }
22870
22871 #[allow(clippy::too_many_arguments)]
22879 pub fn fa_decode_dcw2(
22880 &self,
22881 q2: &CudaSlice<f32>,
22882 k_ring: &cudarc::driver::CudaView<u8>,
22883 v_ring: &cudarc::driver::CudaView<u8>,
22884 o2: &mut CudaSlice<f32>,
22885 head_dim: usize,
22886 n_head: usize,
22887 n_head_kv: usize,
22888 len_dev: &CudaSlice<i32>,
22889 base_dev: Option<&CudaSlice<i32>>,
22890 window: usize,
22891 bucket_max: usize,
22892 scale: f32,
22893 k_tok_bytes: usize,
22894 v_tok_bytes: usize,
22895 gate2: &CudaSlice<f32>,
22896 ) -> Result<(), Box<dyn std::error::Error>> {
22897 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22898 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
22899 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
22900 }
22901 let sp = fa_split_keys(bucket_max, n_head_kv);
22902 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22903 let o_len = 2 * n_head * n_splits * head_dim;
22905 let ml_len = 2 * n_head * n_splits;
22906 let mut part_guard = self.fa_part_pool.lock().unwrap();
22907 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
22908 let pg = part_guard.as_mut().unwrap();
22909 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22910 let (hd, nh, nhkv, nsp) = (
22911 head_dim as i32,
22912 n_head as i32,
22913 n_head_kv as i32,
22914 n_splits as i32,
22915 );
22916 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22917 let (ski, win) = (sp as i32, window as i32);
22918 let gqa = (n_head / n_head_kv).max(1) as u32;
22919 let smem = (32 * head_dim * 2) as u32;
22920 let f = self.func("fa_decode_vec_q_v3_dcw2");
22921 let cfg = LaunchConfig {
22922 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22923 block_dim: (32, gqa, 1),
22924 shared_mem_bytes: smem,
22925 };
22926 let null: u64 = 0;
22927 {
22928 let __s_b = self.gpu.stream();
22929 let mut b = __s_b.launch_builder(&f);
22930 b.arg(q2)
22931 .arg(k_ring)
22932 .arg(v_ring)
22933 .arg(&mut *part_o)
22934 .arg(&mut *part_m)
22935 .arg(&mut *part_l)
22936 .arg(&hd)
22937 .arg(&nh)
22938 .arg(&nhkv)
22939 .arg(len_dev);
22940 match base_dev {
22941 Some(base) => {
22942 b.arg(base);
22943 }
22944 None => {
22945 b.arg(&null);
22946 }
22947 }
22948 b.arg(&win)
22949 .arg(&scale)
22950 .arg(&nsp)
22951 .arg(&ski)
22952 .arg(&ktb)
22953 .arg(&vtb);
22954 unsafe {
22955 b.launch(cfg)?;
22956 }
22957 }
22958 let fc = {
22962 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22963 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
22964 self.func("fa_decode_combine_gate_f32_s")
22965 } else {
22966 self.func("fa_decode_combine_gate_f32")
22967 }
22968 };
22969 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
22970 let nh2 = (2 * n_head) as i32;
22971 let cfg2 = LaunchConfig {
22972 grid_dim: ((2 * n_head) as u32, 1, 1),
22973 block_dim: (head_dim as u32, 1, 1),
22974 shared_mem_bytes: if combine_shared {
22975 (2 * n_splits * 4) as u32
22976 } else {
22977 0
22978 },
22979 };
22980 let __s_b2 = self.gpu.stream();
22981 let mut b2 = __s_b2.launch_builder(&fc);
22982 b2.arg(&*part_o)
22983 .arg(&*part_m)
22984 .arg(&*part_l)
22985 .arg(gate2)
22986 .arg(o2)
22987 .arg(&hd)
22988 .arg(&nh2)
22989 .arg(&nsp);
22990 unsafe {
22991 b2.launch(cfg2)?;
22992 }
22993 Ok(())
22994 }
22995
22996 pub fn fa_decode_dcw(
22997 &self,
22998 q: &CudaSlice<f32>,
22999 k_ring: &cudarc::driver::CudaView<u8>,
23000 v_ring: &cudarc::driver::CudaView<u8>,
23001 o: &mut CudaSlice<f32>,
23002 head_dim: usize,
23003 n_head: usize,
23004 n_head_kv: usize,
23005 len_dev: &CudaSlice<i32>,
23006 base_dev: Option<&CudaSlice<i32>>,
23007 window: usize,
23008 bucket_max: usize,
23009 scale: f32,
23010 k_tok_bytes: usize,
23011 v_tok_bytes: usize,
23012 fused_gate: Option<&CudaSlice<f32>>,
23016 ) -> Result<(), Box<dyn std::error::Error>> {
23017 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23018 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23019 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"
23020 .into());
23021 }
23022 let sp = fa_split_keys(bucket_max, n_head_kv);
23023 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23024 let o_len = n_head * n_splits * head_dim;
23025 let ml_len = n_head * n_splits;
23026 let mut part_guard = self.fa_part_pool.lock().unwrap();
23027 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23028 let pg = part_guard.as_mut().unwrap();
23029 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23034 let memset_on = *MEMSET_ON
23039 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23040 || crate::tp::token_graph_building();
23041 if memset_on {
23042 self.gpu
23043 .stream()
23044 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23045 self.gpu
23046 .stream()
23047 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23048 self.gpu
23049 .stream()
23050 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23051 }
23052 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23053 let (hd, nh, nhkv, nsp) = (
23054 head_dim as i32,
23055 n_head as i32,
23056 n_head_kv as i32,
23057 n_splits as i32,
23058 );
23059 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23060 let (ski, win) = (sp as i32, window as i32);
23061 let gqa = (n_head / n_head_kv).max(1) as u32;
23062 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23066 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23067 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23068 Ok("2") => 2,
23069 Ok("1") => 1,
23070 _ => 0,
23071 });
23072 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23077 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23078 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23079 std::sync::Mutex::new(None);
23080 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23084 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23085 && (n_head / n_head_kv) % 2 == 0
23086 && (n_head / n_head_kv) >= 2;
23087 let f = if fprof {
23088 self.func("fa_decode_vec_q_v3_dcw_prof")
23089 } else if hs2 {
23090 self.func("fa_decode_vec_q_v3_dcw_hs2")
23091 } else if hoist == 2 {
23092 self.func("fa_decode_vec_q_v3_dcw_hc")
23094 } else if hoist == 1 {
23095 self.func("fa_decode_vec_q_v3_dcw_h")
23097 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23098 self.func("fa_decode_vec_q_v3_dcw_u8")
23099 } else {
23100 self.func("fa_decode_vec_q_v3_dcw")
23101 };
23102 let cfg = LaunchConfig {
23103 grid_dim: if hs2 {
23104 ((2 * n_head_kv) as u32, n_splits as u32, 1)
23105 } else {
23106 (n_head_kv as u32, n_splits as u32, 1)
23107 },
23108 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23109 shared_mem_bytes: smem,
23110 };
23111 let null: u64 = 0;
23112 let __s_b = self.gpu.stream();
23113 let mut b = __s_b.launch_builder(&f);
23114 b.arg(q)
23115 .arg(k_ring)
23116 .arg(v_ring)
23117 .arg(&mut *part_o)
23118 .arg(&mut *part_m)
23119 .arg(&mut *part_l)
23120 .arg(&hd)
23121 .arg(&nh)
23122 .arg(&nhkv)
23123 .arg(len_dev);
23124 match base_dev {
23125 Some(base) => {
23126 b.arg(base);
23127 }
23128 None => {
23129 b.arg(&null);
23130 }
23131 }
23132 b.arg(&win)
23133 .arg(&scale)
23134 .arg(&nsp)
23135 .arg(&ski)
23136 .arg(&ktb)
23137 .arg(&vtb);
23138 if fprof {
23139 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
23140 if guard
23141 .as_ref()
23142 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
23143 {
23144 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
23145 }
23146 let (_, buf) = guard.as_mut().expect("armed above");
23147 b.arg(&*buf);
23148 unsafe {
23149 b.launch(cfg)?;
23150 }
23151 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23152 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
23153 if n % 430 == 0 {
23154 self.stream().synchronize()?;
23155 let h = self.dtoh_u64(buf)?;
23156 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
23157 let tot: u64 = h[..6].iter().sum();
23158 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
23159 for (i, name) in phases.iter().enumerate() {
23160 let pct = if tot > 0 {
23161 h[i] as f64 / tot as f64 * 100.0
23162 } else {
23163 0.0
23164 };
23165 line.push_str(&format!(" {name}={pct:.1}%"));
23166 }
23167 if h[6] > 0 {
23168 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
23169 }
23170 eprintln!("{line}");
23171 }
23172 } else {
23173 unsafe {
23174 b.launch(cfg)?;
23175 }
23176 }
23177 let mut combine_shared = false;
23178 let fc = if fused_gate.is_some() {
23179 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23182 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23183 combine_shared = true;
23184 self.func("fa_decode_combine_gate_f32_s")
23185 } else {
23186 self.func("fa_decode_combine_gate_f32")
23187 }
23188 } else {
23189 self.fa_func("fa_decode_combine_f32", head_dim)
23190 };
23191 let cfg2 = LaunchConfig {
23192 grid_dim: (n_head as u32, 1, 1),
23193 block_dim: (head_dim as u32, 1, 1),
23194 shared_mem_bytes: if combine_shared {
23195 (2 * n_splits * 4) as u32
23196 } else {
23197 0
23198 },
23199 };
23200 let __s_b2 = self.gpu.stream();
23201 let mut b2 = __s_b2.launch_builder(&fc);
23202 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
23203 if let Some(gate_row) = fused_gate {
23204 b2.arg(gate_row);
23205 }
23206 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
23207 unsafe {
23208 b2.launch(cfg2)?;
23209 }
23210 Ok(())
23211 }
23212
23213 pub fn fa_geom_eager(
23219 &self,
23220 t_kv: usize,
23221 head_dim: usize,
23222 n_head_kv: usize,
23223 g: bool,
23224 ) -> (bool, usize) {
23225 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
23229 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
23235 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
23236 if g && head_dim == 256 && !fa_v4_at(t_kv) {
23242 fa_vec = false;
23243 }
23244 let sp = fa_split_keys(t_kv, n_head_kv);
23245 let n_splits = if fa_vec {
23246 ((t_kv + sp - 1) / sp).max(1)
23247 } else {
23248 ((t_kv + 255) / 256).max(1)
23249 };
23250 (fa_vec, n_splits)
23251 }
23252
23253 pub fn fa_bucket_key(
23259 &self,
23260 t_kv: usize,
23261 head_dim: usize,
23262 n_head_kv: usize,
23263 g: bool,
23264 ) -> (bool, usize) {
23265 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
23266 }
23267
23268 pub fn capture_graph_retained<F>(
23280 &self,
23281 step: F,
23282 ) -> Result<
23283 (
23284 cudarc::driver::CudaGraph,
23285 Vec<Box<dyn std::any::Any + Send>>,
23286 ),
23287 Box<dyn std::error::Error>,
23288 >
23289 where
23290 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23291 {
23292 use cudarc::driver::sys::CUgraphInstantiate_flags;
23293 self.capture_graph_retained_flags(
23294 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23295 step,
23296 )
23297 }
23298
23299 pub fn capture_graph_retained_flags<F>(
23304 &self,
23305 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
23306 mut step: F,
23307 ) -> Result<
23308 (
23309 cudarc::driver::CudaGraph,
23310 Vec<Box<dyn std::any::Any + Send>>,
23311 ),
23312 Box<dyn std::error::Error>,
23313 >
23314 where
23315 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23316 {
23317 use cudarc::driver::sys::CUstreamCaptureMode;
23318 self.capture_keep.lock().unwrap().clear();
23326 let was_tracking = self.gpu.ctx.is_event_tracking();
23327 if was_tracking {
23328 unsafe {
23329 self.gpu.ctx.disable_event_tracking();
23330 }
23331 }
23332 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23333 self.capture_keep_on
23334 .store(true, std::sync::atomic::Ordering::Relaxed);
23335 let w = (|| {
23336 step(self)?;
23337 step(self)
23338 })();
23339 self.capture_keep_on
23340 .store(false, std::sync::atomic::Ordering::Relaxed);
23341 w?;
23342 self.gpu.stream().synchronize()?;
23343 self.gpu
23344 .stream()
23345 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23346 let r = step(self);
23347 let g = self.gpu.stream().end_capture(flags);
23348 r?;
23349 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23350 graph.upload()?;
23351 Ok(graph)
23352 };
23353 let result = run();
23354 self.capture_keep_on
23355 .store(false, std::sync::atomic::Ordering::Relaxed);
23356 if was_tracking {
23357 unsafe {
23358 self.gpu.ctx.enable_event_tracking();
23359 }
23360 }
23361 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
23362 Ok((result?, keeper))
23363 }
23364
23365 pub fn capture_graph_retained_nowarm<F>(
23371 &self,
23372 mut step: F,
23373 ) -> Result<
23374 (
23375 cudarc::driver::CudaGraph,
23376 Vec<Box<dyn std::any::Any + Send>>,
23377 ),
23378 Box<dyn std::error::Error>,
23379 >
23380 where
23381 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23382 {
23383 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23384 let was_tracking = self.gpu.ctx.is_event_tracking();
23385 if was_tracking {
23386 unsafe {
23387 self.gpu.ctx.disable_event_tracking();
23388 }
23389 }
23390 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23391 self.gpu.stream().synchronize()?;
23392 self.gpu
23393 .stream()
23394 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23395 let r = step(self);
23396 let g = self.gpu.stream().end_capture(
23397 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23398 );
23399 r?;
23400 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23401 graph.upload()?;
23402 Ok(graph)
23403 };
23404 let result = run();
23405 if was_tracking {
23406 unsafe {
23407 self.gpu.ctx.enable_event_tracking();
23408 }
23409 }
23410 Ok((result?, Vec::new()))
23411 }
23412
23413 pub fn capture_graph<F>(
23414 &self,
23415 mut step: F,
23416 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
23417 where
23418 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23419 {
23420 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23421 let was_tracking = self.gpu.ctx.is_event_tracking();
23429 if was_tracking {
23430 unsafe {
23431 self.gpu.ctx.disable_event_tracking();
23432 }
23433 }
23434 let iflag = {
23441 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23442 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23443 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23446 Ok("priority") => {
23447 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23448 }
23449 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23450 })
23451 };
23452 let ct = {
23459 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23460 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23461 };
23462 let warmups = {
23485 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23486 *W.get_or_init(|| {
23487 std::env::var("MEMRA_GRAPH_WARMUPS")
23488 .ok()
23489 .and_then(|v| v.parse().ok())
23490 .filter(|n| *n >= 1)
23491 .unwrap_or(1)
23492 })
23493 };
23494 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23495 let t_w = std::time::Instant::now();
23496 for _ in 0..warmups {
23498 step(self)?;
23499 }
23500 self.gpu.stream().synchronize()?;
23501 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23502 let t_c = std::time::Instant::now();
23504 self.gpu
23505 .stream()
23506 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23507 let r = step(self);
23510 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23511 let t_i = std::time::Instant::now();
23512 let g = self.gpu.stream().end_capture(iflag);
23513 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23514 r?;
23515 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23516 let t_u = std::time::Instant::now();
23517 graph.upload()?;
23518 if ct {
23519 println!(
23520 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
23521 instantiate {ms_inst:.2} ms upload {:.2} ms",
23522 t_u.elapsed().as_secs_f64() * 1e3
23523 );
23524 }
23525 Ok(graph)
23526 };
23527 let result = run();
23528 if was_tracking {
23529 unsafe {
23530 self.gpu.ctx.enable_event_tracking();
23531 }
23532 }
23533 result
23534 }
23535
23536 pub fn gdn_scan_s128_view(
23538 &self,
23539 q: &CudaSlice<f32>,
23540 k: &CudaSlice<f32>,
23541 v: &CudaSlice<f32>,
23542 g: &CudaSlice<f32>,
23543 beta: &CudaSlice<f32>,
23544 state_in: &cudarc::driver::CudaView<f32>,
23545 state_out: &mut cudarc::driver::CudaViewMut<f32>,
23546 o: &mut CudaSlice<f32>,
23547 n_head: usize,
23548 t: usize,
23549 scale: f32,
23550 ) -> Result<(), Box<dyn std::error::Error>> {
23551 let f = self.func("gdn_scan_s128");
23552 const S_V: u32 = 128;
23553 const WARP: u32 = 32;
23554 const COLS: u32 = 4;
23555 let cfg = LaunchConfig {
23556 grid_dim: (n_head as u32, 1, S_V / COLS),
23557 block_dim: (WARP, COLS, 1),
23558 shared_mem_bytes: 0,
23559 };
23560 let (h, ti) = (n_head as i32, t as i32);
23561 let __s_b = self.gpu.stream();
23562 let mut b = __s_b.launch_builder(&f);
23563 b.arg(q)
23564 .arg(k)
23565 .arg(v)
23566 .arg(g)
23567 .arg(beta)
23568 .arg(state_in)
23569 .arg(state_out)
23570 .arg(o)
23571 .arg(&h)
23572 .arg(&ti)
23573 .arg(&scale);
23574 unsafe {
23575 b.launch(cfg)?;
23576 }
23577 Ok(())
23578 }
23579
23580 pub fn ssm_conv1d_view(
23582 &self,
23583 x: &cudarc::driver::CudaView<f32>,
23584 w: &CudaSlice<f32>,
23585 y: &mut CudaSlice<f32>,
23586 conv_dim: usize,
23587 t: usize,
23588 d_conv: usize,
23589 silu: bool,
23590 ) -> Result<(), Box<dyn std::error::Error>> {
23591 let f = self.func("ssm_conv1d_silu_f32");
23592 let cfg = LaunchConfig {
23594 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23595 block_dim: (256, 1, 1),
23596 shared_mem_bytes: 0,
23597 };
23598 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23599 let __s_b = self.gpu.stream();
23600 let mut b = __s_b.launch_builder(&f);
23601 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23602 unsafe {
23603 b.launch(cfg)?;
23604 }
23605 Ok(())
23606 }
23607
23608 pub fn ssm_conv1d_tm(
23615 &self,
23616 qkv_tm: &CudaSlice<f32>,
23617 w: &CudaSlice<f32>,
23618 y: &mut CudaSlice<f32>,
23619 conv_dim: usize,
23620 t: usize,
23621 d_conv: usize,
23622 ) -> Result<(), Box<dyn std::error::Error>> {
23623 let f = self.func("ssm_conv1d_tm_f32");
23624 let cfg = LaunchConfig {
23625 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23626 block_dim: (256, 1, 1),
23627 shared_mem_bytes: 0,
23628 };
23629 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23630 let __s_b = self.gpu.stream();
23631 let mut b = __s_b.launch_builder(&f);
23632 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23633 unsafe {
23634 b.launch(cfg)?;
23635 }
23636 Ok(())
23637 }
23638
23639 pub fn ssm_conv1d_tm_state(
23647 &self,
23648 qkv_tm: &CudaSlice<f32>,
23649 conv_state: &mut CudaSlice<f32>,
23650 w: &CudaSlice<f32>,
23651 y: &mut CudaSlice<f32>,
23652 conv_dim: usize,
23653 t: usize,
23654 d_conv: usize,
23655 ) -> Result<(), Box<dyn std::error::Error>> {
23656 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23657 }
23658
23659 #[allow(clippy::too_many_arguments)]
23662 pub fn ssm_conv1d_tm_state_pad(
23663 &self,
23664 qkv_tm: &CudaSlice<f32>,
23665 conv_state: &mut CudaSlice<f32>,
23666 w: &CudaSlice<f32>,
23667 y: &mut CudaSlice<f32>,
23668 conv_dim: usize,
23669 t: usize,
23670 d_conv: usize,
23671 pad_len: Option<&CudaSlice<i32>>,
23672 ) -> Result<(), Box<dyn std::error::Error>> {
23673 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23674 let ring_old = if t < d_conv - 1 {
23678 Some(self.clone_dtod(conv_state)?)
23679 } else {
23680 None
23681 };
23682 {
23683 let f = self.func("ssm_conv1d_tm_state_f32");
23684 let cfg = LaunchConfig {
23685 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23686 block_dim: (256, 1, 1),
23687 shared_mem_bytes: 0,
23688 };
23689 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23690 let __s_b = self.gpu.stream();
23691 let mut b = __s_b.launch_builder(&f);
23692 b.arg(qkv_tm)
23693 .arg(&*conv_state)
23694 .arg(w)
23695 .arg(y)
23696 .arg(&cd)
23697 .arg(&ti)
23698 .arg(&dc);
23699 unsafe {
23700 b.launch(cfg)?;
23701 }
23702 }
23703 match (ring_old, pad_len) {
23704 (None, Some(len_d)) => {
23705 let f = self.func("ssm_conv_ring_update_dev_f32");
23706 let n = conv_dim * (d_conv - 1);
23707 let cfg = LaunchConfig::for_num_elems(n as u32);
23708 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23709 let __s_b = self.gpu.stream();
23710 let mut b = __s_b.launch_builder(&f);
23711 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23712 unsafe {
23713 b.launch(cfg)?;
23714 }
23715 }
23716 (None, None) => {
23717 let f = self.func("ssm_conv_ring_update_f32");
23718 let n = conv_dim * (d_conv - 1);
23719 let cfg = LaunchConfig::for_num_elems(n as u32);
23720 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23721 let __s_b = self.gpu.stream();
23722 let mut b = __s_b.launch_builder(&f);
23723 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23724 unsafe {
23725 b.launch(cfg)?;
23726 }
23727 }
23728 (Some(old), _) => {
23729 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23730 }
23731 }
23732 Ok(())
23733 }
23734
23735 pub fn ssm_conv1d_tm_state_pad_v(
23737 &self,
23738 qkv_tm: &cudarc::driver::CudaView<f32>,
23739 conv_state: &mut CudaSlice<f32>,
23740 w: &CudaSlice<f32>,
23741 y: &mut CudaSlice<f32>,
23742 conv_dim: usize,
23743 t: usize,
23744 d_conv: usize,
23745 pad_len: Option<&CudaSlice<i32>>,
23746 ) -> Result<(), Box<dyn std::error::Error>> {
23747 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23748 let ring_old = if t < d_conv - 1 {
23752 Some(self.clone_dtod(conv_state)?)
23753 } else {
23754 None
23755 };
23756 {
23757 let f = self.func("ssm_conv1d_tm_state_f32");
23758 let cfg = LaunchConfig {
23759 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23760 block_dim: (256, 1, 1),
23761 shared_mem_bytes: 0,
23762 };
23763 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23764 let __s_b = self.gpu.stream();
23765 let mut b = __s_b.launch_builder(&f);
23766 b.arg(qkv_tm)
23767 .arg(&*conv_state)
23768 .arg(w)
23769 .arg(y)
23770 .arg(&cd)
23771 .arg(&ti)
23772 .arg(&dc);
23773 unsafe {
23774 b.launch(cfg)?;
23775 }
23776 }
23777 match (ring_old, pad_len) {
23778 (None, Some(len_d)) => {
23779 let f = self.func("ssm_conv_ring_update_dev_f32");
23780 let n = conv_dim * (d_conv - 1);
23781 let cfg = LaunchConfig::for_num_elems(n as u32);
23782 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23783 let __s_b = self.gpu.stream();
23784 let mut b = __s_b.launch_builder(&f);
23785 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23786 unsafe {
23787 b.launch(cfg)?;
23788 }
23789 }
23790 (None, None) => {
23791 let f = self.func("ssm_conv_ring_update_f32");
23792 let n = conv_dim * (d_conv - 1);
23793 let cfg = LaunchConfig::for_num_elems(n as u32);
23794 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23795 let __s_b = self.gpu.stream();
23796 let mut b = __s_b.launch_builder(&f);
23797 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23798 unsafe {
23799 b.launch(cfg)?;
23800 }
23801 }
23802 (Some(_), _) => unreachable!(
23803 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23804 ),
23805 }
23806 Ok(())
23807 }
23808
23809 pub fn ssm_conv_ring_rebuild(
23814 &self,
23815 qkv_tm: &CudaSlice<f32>,
23816 ring_old: &CudaSlice<f32>,
23817 conv_state: &mut CudaSlice<f32>,
23818 conv_dim: usize,
23819 tc: usize,
23820 d_conv: usize,
23821 ) -> Result<(), Box<dyn std::error::Error>> {
23822 let f = self.func("ssm_conv_ring_rebuild_f32");
23823 let n = conv_dim * (d_conv - 1);
23824 let cfg = LaunchConfig::for_num_elems(n as u32);
23825 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23826 let __s_b = self.gpu.stream();
23827 let mut b = __s_b.launch_builder(&f);
23828 b.arg(qkv_tm)
23829 .arg(ring_old)
23830 .arg(conv_state)
23831 .arg(&cd)
23832 .arg(&ti)
23833 .arg(&dc);
23834 unsafe {
23835 b.launch(cfg)?;
23836 }
23837 Ok(())
23838 }
23839
23840 #[allow(clippy::too_many_arguments)]
23845 pub fn gdn_prep_decode(
23846 &self,
23847 conv_out: &CudaSlice<f32>,
23848 beta_raw: &CudaSlice<f32>,
23849 alpha: &CudaSlice<f32>,
23850 dt_bias: &CudaSlice<f32>,
23851 a: &CudaSlice<f32>,
23852 q_l2: &mut CudaSlice<f32>,
23853 k_l2: &mut CudaSlice<f32>,
23854 v_g: &mut CudaSlice<f32>,
23855 beta: &mut CudaSlice<f32>,
23856 g_log: &mut CudaSlice<f32>,
23857 d_state: usize,
23858 num_v: usize,
23859 num_k: usize,
23860 key_dim: usize,
23861 eps: f32,
23862 ) -> Result<(), Box<dyn std::error::Error>> {
23863 let f = self.func("gdn_prep_decode_f32");
23864 let cfg = LaunchConfig {
23865 grid_dim: (num_v as u32, 1, 1),
23866 block_dim: (32, 4, 1),
23867 shared_mem_bytes: 0,
23868 };
23869 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23870 let __s_b = self.gpu.stream();
23871 let mut b = __s_b.launch_builder(&f);
23872 b.arg(conv_out)
23873 .arg(beta_raw)
23874 .arg(alpha)
23875 .arg(dt_bias)
23876 .arg(a)
23877 .arg(q_l2)
23878 .arg(k_l2)
23879 .arg(v_g)
23880 .arg(beta)
23881 .arg(g_log)
23882 .arg(&ds)
23883 .arg(&nv)
23884 .arg(&nk)
23885 .arg(&kd)
23886 .arg(&eps);
23887 unsafe {
23888 b.launch(cfg)?;
23889 }
23890 Ok(())
23891 }
23892
23893 #[allow(clippy::too_many_arguments)]
23897 pub fn ssm_conv1d_gdn(
23898 &self,
23899 qkv_tm: &CudaSlice<f32>,
23900 w: &CudaSlice<f32>,
23901 q_g: &mut CudaSlice<f32>,
23902 k_g: &mut CudaSlice<f32>,
23903 v_g: &mut CudaSlice<f32>,
23904 conv_dim: usize,
23905 t: usize,
23906 d_conv: usize,
23907 d_state: usize,
23908 num_v: usize,
23909 num_k: usize,
23910 key_dim: usize,
23911 ) -> Result<(), Box<dyn std::error::Error>> {
23912 let f = self.func("ssm_conv1d_gdn_f32");
23913 let cfg = LaunchConfig {
23914 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23915 block_dim: (256, 1, 1),
23916 shared_mem_bytes: 0,
23917 };
23918 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23919 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23920 let __s_b = self.gpu.stream();
23921 let mut b = __s_b.launch_builder(&f);
23922 b.arg(qkv_tm)
23923 .arg(w)
23924 .arg(q_g)
23925 .arg(k_g)
23926 .arg(v_g)
23927 .arg(&cd)
23928 .arg(&ti)
23929 .arg(&dc)
23930 .arg(&ds)
23931 .arg(&nv)
23932 .arg(&nk)
23933 .arg(&kd);
23934 unsafe {
23935 b.launch(cfg)?;
23936 }
23937 Ok(())
23938 }
23939
23940 pub fn ssm_conv1d(
23941 &self,
23942 x: &CudaSlice<f32>,
23943 w: &CudaSlice<f32>,
23944 y: &mut CudaSlice<f32>,
23945 conv_dim: usize,
23946 t: usize,
23947 d_conv: usize,
23948 silu: bool,
23949 ) -> Result<(), Box<dyn std::error::Error>> {
23950 let f = self.func("ssm_conv1d_silu_f32");
23951 let cfg = LaunchConfig {
23952 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23953 block_dim: (256, 1, 1),
23954 shared_mem_bytes: 0,
23955 };
23956 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23957 let __s_b = self.gpu.stream();
23958 let mut b = __s_b.launch_builder(&f);
23959 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23960 unsafe {
23961 b.launch(cfg)?;
23962 }
23963 Ok(())
23964 }
23965
23966 pub fn gdn_scan_s128(
23969 &self,
23970 q: &CudaSlice<f32>,
23971 k: &CudaSlice<f32>,
23972 v: &CudaSlice<f32>,
23973 g: &CudaSlice<f32>,
23974 beta: &CudaSlice<f32>,
23975 state_in: &CudaSlice<f32>,
23976 state_out: &mut CudaSlice<f32>,
23977 o: &mut CudaSlice<f32>,
23978 n_head: usize,
23979 t: usize,
23980 scale: f32,
23981 ) -> Result<(), Box<dyn std::error::Error>> {
23982 let f = self.func("gdn_scan_s128");
23983 const S_V: u32 = 128;
23984 const WARP: u32 = 32;
23985 const COLS_PER_BLOCK: u32 = 4;
23986 let cfg = LaunchConfig {
23987 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
23988 block_dim: (WARP, COLS_PER_BLOCK, 1),
23989 shared_mem_bytes: 0,
23990 };
23991 let (h, ti) = (n_head as i32, t as i32);
23992 let __s_b = self.gpu.stream();
23993 let mut b = __s_b.launch_builder(&f);
23994 b.arg(q)
23995 .arg(k)
23996 .arg(v)
23997 .arg(g)
23998 .arg(beta)
23999 .arg(state_in)
24000 .arg(state_out)
24001 .arg(o)
24002 .arg(&h)
24003 .arg(&ti)
24004 .arg(&scale);
24005 unsafe {
24006 b.launch(cfg)?;
24007 }
24008 Ok(())
24009 }
24010
24011 #[allow(clippy::too_many_arguments)]
24016 pub fn ssm_conv1d_fused_decode_b(
24017 &self,
24018 qkv_cols: &CudaSlice<f32>,
24019 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24020 w: &CudaSlice<f32>,
24021 conv_outs: &mut CudaSlice<f32>,
24022 conv_dim: usize,
24023 d_conv: usize,
24024 b_n: usize,
24025 ) -> Result<(), Box<dyn std::error::Error>> {
24026 let f = self.func("ssm_conv1d_fused_decode_b_f32");
24027 let cfg = LaunchConfig {
24028 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24029 block_dim: (256, 1, 1),
24030 shared_mem_bytes: 0,
24031 };
24032 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24033 let __s_b = self.gpu.stream();
24034 let mut b = __s_b.launch_builder(&f);
24035 b.arg(qkv_cols)
24036 .arg(conv_state_ptrs)
24037 .arg(w)
24038 .arg(conv_outs)
24039 .arg(&cd)
24040 .arg(&dc);
24041 unsafe {
24042 b.launch(cfg)?;
24043 }
24044 Ok(())
24045 }
24046
24047 #[allow(clippy::too_many_arguments)]
24048 pub fn gdn_prep_decode_b(
24049 &self,
24050 conv_outs: &CudaSlice<f32>,
24051 beta_raws: &CudaSlice<f32>,
24052 alphas: &CudaSlice<f32>,
24053 dt_bias: &CudaSlice<f32>,
24054 a: &CudaSlice<f32>,
24055 q_l2: &mut CudaSlice<f32>,
24056 k_l2: &mut CudaSlice<f32>,
24057 v_g: &mut CudaSlice<f32>,
24058 beta: &mut CudaSlice<f32>,
24059 g_log: &mut CudaSlice<f32>,
24060 d_state: usize,
24061 num_v: usize,
24062 num_k: usize,
24063 key_dim: usize,
24064 eps: f32,
24065 conv_dim: usize,
24066 b_n: usize,
24067 ) -> Result<(), Box<dyn std::error::Error>> {
24068 let f = self.func("gdn_prep_decode_b_f32");
24069 let cfg = LaunchConfig {
24070 grid_dim: (num_v as u32, 1, b_n as u32),
24071 block_dim: (32, 4, 1),
24072 shared_mem_bytes: 0,
24073 };
24074 let (ds, nv, nk, kd, cd) = (
24075 d_state as i32,
24076 num_v as i32,
24077 num_k as i32,
24078 key_dim as i32,
24079 conv_dim as i32,
24080 );
24081 let __s_b = self.gpu.stream();
24082 let mut b = __s_b.launch_builder(&f);
24083 b.arg(conv_outs)
24084 .arg(beta_raws)
24085 .arg(alphas)
24086 .arg(dt_bias)
24087 .arg(a)
24088 .arg(q_l2)
24089 .arg(k_l2)
24090 .arg(v_g)
24091 .arg(beta)
24092 .arg(g_log)
24093 .arg(&ds)
24094 .arg(&nv)
24095 .arg(&nk)
24096 .arg(&kd)
24097 .arg(&eps)
24098 .arg(&cd);
24099 unsafe {
24100 b.launch(cfg)?;
24101 }
24102 Ok(())
24103 }
24104
24105 #[allow(clippy::too_many_arguments)]
24106 pub fn gdn_scan_s128_batched(
24107 &self,
24108 q: &CudaSlice<f32>,
24109 k: &CudaSlice<f32>,
24110 v: &CudaSlice<f32>,
24111 g: &CudaSlice<f32>,
24112 beta: &CudaSlice<f32>,
24113 state_in_ptrs: &cudarc::driver::CudaView<u64>,
24114 state_out_ptrs: &cudarc::driver::CudaView<u64>,
24115 o: &mut CudaSlice<f32>,
24116 n_head: usize,
24117 b_n: usize,
24118 scale: f32,
24119 ) -> Result<(), Box<dyn std::error::Error>> {
24120 let f = self.func("gdn_scan_s128_b");
24121 const S_V: u32 = 128;
24122 const WARP: u32 = 32;
24123 const COLS_PER_BLOCK: u32 = 4;
24124 let cfg = LaunchConfig {
24125 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24126 block_dim: (WARP, COLS_PER_BLOCK, 1),
24127 shared_mem_bytes: 0,
24128 };
24129 let h = n_head as i32;
24130 let __s_b = self.gpu.stream();
24131 let mut b = __s_b.launch_builder(&f);
24132 b.arg(q)
24133 .arg(k)
24134 .arg(v)
24135 .arg(g)
24136 .arg(beta)
24137 .arg(state_in_ptrs)
24138 .arg(state_out_ptrs)
24139 .arg(o)
24140 .arg(&h)
24141 .arg(&scale);
24142 unsafe {
24143 b.launch(cfg)?;
24144 }
24145 Ok(())
24146 }
24147
24148 #[allow(clippy::too_many_arguments)]
24154 pub fn ssm_conv1d_fused_decode_b_view(
24155 &self,
24156 qkv_cols: &cudarc::driver::CudaView<f32>,
24157 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24158 w: &CudaSlice<f32>,
24159 conv_outs: &mut CudaSlice<f32>,
24160 conv_dim: usize,
24161 d_conv: usize,
24162 b_n: usize,
24163 ) -> Result<(), Box<dyn std::error::Error>> {
24164 let f = self.func("ssm_conv1d_fused_decode_b_f32");
24165 let cfg = LaunchConfig {
24166 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24167 block_dim: (256, 1, 1),
24168 shared_mem_bytes: 0,
24169 };
24170 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24171 let __s_b = self.gpu.stream();
24172 let mut b = __s_b.launch_builder(&f);
24173 b.arg(qkv_cols)
24174 .arg(conv_state_ptrs)
24175 .arg(w)
24176 .arg(conv_outs)
24177 .arg(&cd)
24178 .arg(&dc);
24179 unsafe {
24180 b.launch(cfg)?;
24181 }
24182 Ok(())
24183 }
24184
24185 #[allow(clippy::too_many_arguments)]
24186 pub fn gdn_prep_decode_b_view(
24187 &self,
24188 conv_outs: &CudaSlice<f32>,
24189 beta_raws: &cudarc::driver::CudaView<f32>,
24190 alphas: &cudarc::driver::CudaView<f32>,
24191 dt_bias: &CudaSlice<f32>,
24192 a: &CudaSlice<f32>,
24193 q_l2: &mut CudaSlice<f32>,
24194 k_l2: &mut CudaSlice<f32>,
24195 v_g: &mut CudaSlice<f32>,
24196 beta: &mut CudaSlice<f32>,
24197 g_log: &mut CudaSlice<f32>,
24198 d_state: usize,
24199 num_v: usize,
24200 num_k: usize,
24201 key_dim: usize,
24202 eps: f32,
24203 conv_dim: usize,
24204 b_n: usize,
24205 ) -> Result<(), Box<dyn std::error::Error>> {
24206 let f = self.func("gdn_prep_decode_b_f32");
24207 let cfg = LaunchConfig {
24208 grid_dim: (num_v as u32, 1, b_n as u32),
24209 block_dim: (32, 4, 1),
24210 shared_mem_bytes: 0,
24211 };
24212 let (ds, nv, nk, kd, cd) = (
24213 d_state as i32,
24214 num_v as i32,
24215 num_k as i32,
24216 key_dim as i32,
24217 conv_dim as i32,
24218 );
24219 let __s_b = self.gpu.stream();
24220 let mut b = __s_b.launch_builder(&f);
24221 b.arg(conv_outs)
24222 .arg(beta_raws)
24223 .arg(alphas)
24224 .arg(dt_bias)
24225 .arg(a)
24226 .arg(q_l2)
24227 .arg(k_l2)
24228 .arg(v_g)
24229 .arg(beta)
24230 .arg(g_log)
24231 .arg(&ds)
24232 .arg(&nv)
24233 .arg(&nk)
24234 .arg(&kd)
24235 .arg(&eps)
24236 .arg(&cd);
24237 unsafe {
24238 b.launch(cfg)?;
24239 }
24240 Ok(())
24241 }
24242
24243 #[allow(clippy::too_many_arguments)]
24244 pub fn gdn_scan_s128_batched_view(
24245 &self,
24246 q: &CudaSlice<f32>,
24247 k: &CudaSlice<f32>,
24248 v: &CudaSlice<f32>,
24249 g: &CudaSlice<f32>,
24250 beta: &CudaSlice<f32>,
24251 state_in_ptrs: &cudarc::driver::CudaView<u64>,
24252 state_out_ptrs: &cudarc::driver::CudaView<u64>,
24253 o: &mut cudarc::driver::CudaViewMut<f32>,
24254 n_head: usize,
24255 b_n: usize,
24256 scale: f32,
24257 ) -> Result<(), Box<dyn std::error::Error>> {
24258 let f = self.func("gdn_scan_s128_b");
24259 const S_V: u32 = 128;
24260 const WARP: u32 = 32;
24261 const COLS_PER_BLOCK: u32 = 4;
24262 let cfg = LaunchConfig {
24263 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24264 block_dim: (WARP, COLS_PER_BLOCK, 1),
24265 shared_mem_bytes: 0,
24266 };
24267 let h = n_head as i32;
24268 let __s_b = self.gpu.stream();
24269 let mut b = __s_b.launch_builder(&f);
24270 b.arg(q)
24271 .arg(k)
24272 .arg(v)
24273 .arg(g)
24274 .arg(beta)
24275 .arg(state_in_ptrs)
24276 .arg(state_out_ptrs)
24277 .arg(o)
24278 .arg(&h)
24279 .arg(&scale);
24280 unsafe {
24281 b.launch(cfg)?;
24282 }
24283 Ok(())
24284 }
24285
24286 pub fn gdn_chunked_enabled() -> bool {
24295 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24296 *E.get_or_init(|| {
24297 std::env::var("MEMRA_GDN_CHUNKED")
24298 .map(|v| v != "0")
24299 .unwrap_or(true)
24300 })
24301 }
24302
24303 pub fn gdn_chunk_size() -> usize {
24308 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24309 *C.get_or_init(|| {
24310 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
24311 .ok()
24312 .and_then(|v| v.parse().ok())
24313 .unwrap_or(32);
24314 c.clamp(32, 128) / 32 * 32
24315 })
24316 }
24317
24318 #[allow(clippy::too_many_arguments)]
24323 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
24326 #[allow(clippy::too_many_arguments)]
24327 pub fn gdn_chunk_k123(
24328 &self,
24329 q: &CudaSlice<f32>,
24330 k: &CudaSlice<f32>,
24331 v: &CudaSlice<f32>,
24332 g: &CudaSlice<f32>,
24333 beta: &CudaSlice<f32>,
24334 wb16: Option<&mut CudaSlice<u8>>,
24335 n_head: usize,
24336 t: usize,
24337 c: usize,
24338 hk: usize,
24339 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
24340 ) -> Result<
24341 (
24342 CudaSlice<f32>,
24343 CudaSlice<f32>,
24344 CudaSlice<f32>,
24345 CudaSlice<f32>,
24346 ),
24347 Box<dyn std::error::Error>,
24348 > {
24349 const D: usize = 128;
24350 let h = n_head;
24351 let nc = (t + c - 1) / c;
24352 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24353 let mut gcum = self.uninit(t * h)?;
24354 let mut a = self.uninit(nc * h * c * c)?;
24355 let mut p = self.uninit(nc * h * c * c)?;
24356 let mut u = self.uninit(nc * h * c * D)?;
24357 let mut w = self.uninit(nc * h * c * D)?;
24358 {
24359 let f = self.func("gdn_chunk_cumgate_f32");
24361 let cfg = LaunchConfig {
24362 grid_dim: (nc as u32, h as u32, 1),
24363 block_dim: (32, 1, 1),
24364 shared_mem_bytes: 0,
24365 };
24366 let __s_b = self.gpu.stream();
24367 let mut b = __s_b.launch_builder(&f);
24368 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
24369 unsafe {
24370 b.launch(cfg)?;
24371 }
24372 }
24373 if let Some((qb, kb, pb)) = k2w {
24374 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
24377 let f = self.func("gdn_k2_wgmma");
24378 let cfg = LaunchConfig {
24379 grid_dim: (nc as u32, h as u32, 1),
24380 block_dim: (128, 1, 1),
24381 shared_mem_bytes: 0,
24382 };
24383 let hki = hk as i32;
24384 let __s_b = self.gpu.stream();
24385 let mut b = __s_b.launch_builder(&f);
24386 b.arg(qb)
24387 .arg(kb)
24388 .arg(&gcum)
24389 .arg(beta)
24390 .arg(&mut a)
24391 .arg(&mut *pb)
24392 .arg(&hi)
24393 .arg(&ti)
24394 .arg(&ci)
24395 .arg(&hki);
24396 unsafe {
24397 b.launch(cfg)?;
24398 }
24399 } else if c <= 64 && !portable_mma_gated() {
24400 let f = self.func("gdn_chunk_attn_f32");
24402 f.set_attribute(
24403 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24404 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24405 )?;
24406 let jt = ((c + 31) / 32) as u32;
24407 let cfg = LaunchConfig {
24408 grid_dim: (nc as u32, h as u32, jt),
24409 block_dim: (256, 1, 1),
24410 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24411 };
24412 let hki = hk as i32;
24413 let __s_b = self.gpu.stream();
24414 let mut b = __s_b.launch_builder(&f);
24415 b.arg(q)
24416 .arg(k)
24417 .arg(&gcum)
24418 .arg(beta)
24419 .arg(&mut a)
24420 .arg(&mut p)
24421 .arg(&hi)
24422 .arg(&ti)
24423 .arg(&ci)
24424 .arg(&hki);
24425 unsafe {
24426 b.launch(cfg)?;
24427 }
24428 } else {
24429 assert!(
24431 hk == h,
24432 "generic K2 is broadcast-only (de-broadcast rides C==32)"
24433 );
24434 let f = self.func("gdn_chunk_attn_g_f32");
24435 let cfg = LaunchConfig {
24436 grid_dim: (nc as u32, h as u32, 1),
24437 block_dim: (32, 8, 1),
24438 shared_mem_bytes: 0,
24439 };
24440 let __s_b = self.gpu.stream();
24441 let mut b = __s_b.launch_builder(&f);
24442 b.arg(q)
24443 .arg(k)
24444 .arg(&gcum)
24445 .arg(beta)
24446 .arg(&mut a)
24447 .arg(&mut p)
24448 .arg(&hi)
24449 .arg(&ti)
24450 .arg(&ci);
24451 unsafe {
24452 b.launch(cfg)?;
24453 }
24454 }
24455 {
24456 let cfg = LaunchConfig {
24458 grid_dim: (nc as u32, h as u32, 1),
24459 block_dim: (256, 1, 1),
24460 shared_mem_bytes: 0,
24461 };
24462 match c {
24463 32 | 64 => {
24464 let f = self.func(if c == 32 {
24465 "gdn_chunk_solve32_f32"
24466 } else {
24467 "gdn_chunk_solve64_f32"
24468 });
24469 let wb: u64 = match wb16 {
24471 Some(d) => self.addr_u8(d),
24472 None => 0,
24473 };
24474 let hki = hk as i32;
24475 let __s_b = self.gpu.stream();
24476 let mut b = __s_b.launch_builder(&f);
24477 b.arg(v)
24478 .arg(k)
24479 .arg(&a)
24480 .arg(&gcum)
24481 .arg(&mut u)
24482 .arg(&mut w)
24483 .arg(&wb)
24484 .arg(&hi)
24485 .arg(&ti)
24486 .arg(&hki);
24487 unsafe {
24488 b.launch(cfg)?;
24489 }
24490 }
24491 _ => {
24492 assert!(hk == h, "generic K3 is broadcast-only");
24493 let f = self.func("gdn_chunk_solve_f32");
24494 let __s_b = self.gpu.stream();
24495 let mut b = __s_b.launch_builder(&f);
24496 b.arg(v)
24497 .arg(k)
24498 .arg(&a)
24499 .arg(&gcum)
24500 .arg(&mut u)
24501 .arg(&mut w)
24502 .arg(&hi)
24503 .arg(&ti)
24504 .arg(&ci);
24505 unsafe {
24506 b.launch(cfg)?;
24507 }
24508 }
24509 }
24510 }
24511 Ok((gcum, p, u, w))
24512 }
24513
24514 pub fn gdn_db_on() -> bool {
24518 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24519 }
24520
24521 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24530 !portable_mma_gated()
24531 && c == 32
24532 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24533 Ok("1") => true,
24534 Ok("0") => false,
24535 _ => gdn_mma_default_on(),
24536 }
24537 }
24538
24539 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24546 cfg!(memra_hopper_mma)
24547 && self.gdn_mma_enabled(c)
24548 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24549 }
24550
24551 #[allow(clippy::too_many_arguments)]
24556 pub fn ssm_conv1d_gdn_state_pad(
24557 &self,
24558 qkv_tm: &cudarc::driver::CudaView<f32>,
24559 conv_state: &mut CudaSlice<f32>,
24560 w: &CudaSlice<f32>,
24561 q_g: &mut CudaSlice<f32>,
24562 k_g: &mut CudaSlice<f32>,
24563 v_g: &mut CudaSlice<f32>,
24564 conv_dim: usize,
24565 t: usize,
24566 d_conv: usize,
24567 d_state: usize,
24568 num_v: usize,
24569 num_k: usize,
24570 key_dim: usize,
24571 hk: usize,
24572 pad_len: Option<&CudaSlice<i32>>,
24573 ) -> Result<(), Box<dyn std::error::Error>> {
24574 assert!(
24575 t >= d_conv - 1,
24576 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24577 );
24578 {
24579 let f = self.func("ssm_conv1d_gdn_state_f32");
24580 let cfg = LaunchConfig {
24581 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24582 block_dim: (256, 1, 1),
24583 shared_mem_bytes: 0,
24584 };
24585 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24586 let (ds, nv, nk, kd, hki) = (
24587 d_state as i32,
24588 num_v as i32,
24589 num_k as i32,
24590 key_dim as i32,
24591 hk as i32,
24592 );
24593 let __s_b = self.gpu.stream();
24594 let mut b = __s_b.launch_builder(&f);
24595 b.arg(qkv_tm)
24596 .arg(&*conv_state)
24597 .arg(w)
24598 .arg(q_g)
24599 .arg(k_g)
24600 .arg(v_g)
24601 .arg(&cd)
24602 .arg(&ti)
24603 .arg(&dc)
24604 .arg(&ds)
24605 .arg(&nv)
24606 .arg(&nk)
24607 .arg(&kd)
24608 .arg(&hki);
24609 unsafe {
24610 b.launch(cfg)?;
24611 }
24612 }
24613 match pad_len {
24614 Some(len_d) => {
24615 let f = self.func("ssm_conv_ring_update_dev_f32");
24616 let n = conv_dim * (d_conv - 1);
24617 let cfg = LaunchConfig::for_num_elems(n as u32);
24618 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24619 let __s_b = self.gpu.stream();
24620 let mut b = __s_b.launch_builder(&f);
24621 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24622 unsafe {
24623 b.launch(cfg)?;
24624 }
24625 }
24626 None => {
24627 let f = self.func("ssm_conv_ring_update_f32");
24628 let n = conv_dim * (d_conv - 1);
24629 let cfg = LaunchConfig::for_num_elems(n as u32);
24630 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24631 let __s_b = self.gpu.stream();
24632 let mut b = __s_b.launch_builder(&f);
24633 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24634 unsafe {
24635 b.launch(cfg)?;
24636 }
24637 }
24638 }
24639 Ok(())
24640 }
24641
24642 pub fn gdn_chunk_alloc(
24646 &self,
24647 n_head: usize,
24648 t: usize,
24649 c: usize,
24650 hk: usize,
24651 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24652 const D: usize = 128;
24653 assert!(
24654 c == 32,
24655 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24656 );
24657 let h = n_head;
24658 let nc = (t + c - 1) / c;
24659 Ok(GdnChunkBufs {
24660 gcum: self.uninit(t * h)?,
24661 a: self.uninit(nc * h * c * c)?,
24662 p: self.uninit(nc * h * c * c)?,
24663 u: self.uninit(nc * h * c * D)?,
24664 w: self.uninit(nc * h * c * D)?,
24665 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24666 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24667 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24668 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24669 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24670 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24671 o: self.uninit(D * h * t)?,
24672 t,
24673 nc,
24674 })
24675 }
24676
24677 pub fn f32_to_bf16_v(
24679 &self,
24680 x: &cudarc::driver::CudaView<f32>,
24681 dst: &mut CudaSlice<u8>,
24682 n: usize,
24683 ) -> Result<(), Box<dyn std::error::Error>> {
24684 let f = self.func("f32_to_bf16_bulk");
24685 let ni = n as i64;
24686 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24687 let __s_b = self.gpu.stream();
24688 let mut b = __s_b.launch_builder(&f);
24689 b.arg(x).arg(dst).arg(&ni);
24690 unsafe {
24691 b.launch(cfg)?;
24692 }
24693 Ok(())
24694 }
24695
24696 pub fn f32_to_bf16_into(
24698 &self,
24699 x: &CudaSlice<f32>,
24700 dst: &mut CudaSlice<u8>,
24701 n: usize,
24702 ) -> Result<(), Box<dyn std::error::Error>> {
24703 let f = self.func("f32_to_bf16_bulk");
24704 let ni = n as i64;
24705 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24706 let __s_b = self.gpu.stream();
24707 let mut b = __s_b.launch_builder(&f);
24708 b.arg(x).arg(dst).arg(&ni);
24709 unsafe {
24710 b.launch(cfg)?;
24711 }
24712 Ok(())
24713 }
24714
24715 pub fn gdn_chunk_k123_vl8(
24718 &self,
24719 seqs: &[GdnSeqVl],
24720 n_head: usize,
24721 hk: usize,
24722 wq: Option<&GdnWVl8>,
24723 ) -> Result<(), Box<dyn std::error::Error>> {
24724 let b = seqs.len();
24725 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24726 let mut packed = [GdnSeqVl::default(); 8];
24727 packed[..b].copy_from_slice(seqs);
24728 let v = GdnVl8(packed);
24729 let (hi, ci) = (n_head as i32, 32i32);
24730 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24731 {
24732 let f = self.func("gdn_chunk_cumgate_vl");
24733 let cfg = LaunchConfig {
24734 grid_dim: (max_nc, n_head as u32, b as u32),
24735 block_dim: (32, 1, 1),
24736 shared_mem_bytes: 0,
24737 };
24738 let __s_lb = self.gpu.stream();
24739 let mut lb = __s_lb.launch_builder(&f);
24740 lb.arg(&v).arg(&hi).arg(&ci);
24741 unsafe {
24742 lb.launch(cfg)?;
24743 }
24744 }
24745 let hki = hk as i32;
24746 if let Some(w) = wq {
24747 let f = self.func("gdn_k2_wgmma_vl");
24749 let cfg = LaunchConfig {
24750 grid_dim: (max_nc, n_head as u32, b as u32),
24751 block_dim: (128, 1, 1),
24752 shared_mem_bytes: 0,
24753 };
24754 let __s_lb = self.gpu.stream();
24755 let mut lb = __s_lb.launch_builder(&f);
24756 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24757 unsafe {
24758 lb.launch(cfg)?;
24759 }
24760 } else {
24761 let f = self.func("gdn_chunk_attn_vl");
24762 f.set_attribute(
24763 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24764 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24765 )?;
24766 let cfg = LaunchConfig {
24767 grid_dim: (max_nc, n_head as u32, b as u32),
24768 block_dim: (256, 1, 1),
24769 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24770 };
24771 let __s_lb = self.gpu.stream();
24772 let mut lb = __s_lb.launch_builder(&f);
24773 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24774 unsafe {
24775 lb.launch(cfg)?;
24776 }
24777 }
24778 {
24779 let f = self.func("gdn_chunk_solve32_vl");
24780 let cfg = LaunchConfig {
24781 grid_dim: (max_nc, n_head as u32, b as u32),
24782 block_dim: (256, 1, 1),
24783 shared_mem_bytes: 0,
24784 };
24785 let __s_lb = self.gpu.stream();
24786 let mut lb = __s_lb.launch_builder(&f);
24787 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24788 unsafe {
24789 lb.launch(cfg)?;
24790 }
24791 }
24792 Ok(())
24793 }
24794
24795 #[allow(clippy::too_many_arguments)]
24799 pub fn gdn_prep_vl8(
24800 &self,
24801 seqs: &[GdnPrepVl],
24802 conv_w: &CudaSlice<f32>,
24803 dt_bias: &CudaSlice<f32>,
24804 a: &CudaSlice<f32>,
24805 conv_dim: usize,
24806 d_conv: usize,
24807 d_state: usize,
24808 num_v: usize,
24809 num_k: usize,
24810 key_dim: usize,
24811 hk: usize,
24812 eps: f32,
24813 ) -> Result<(), Box<dyn std::error::Error>> {
24814 let b = seqs.len();
24815 assert!(b >= 1 && b <= 8);
24816 let mut packed = [GdnPrepVl::default(); 8];
24817 packed[..b].copy_from_slice(seqs);
24818 let v = GdnPrepVl8(packed);
24819 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24820 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24821 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24822 assert!(
24823 conv_fuse || hk == num_v,
24824 "de-broadcast requires the fused conv"
24825 );
24826 if conv_fuse {
24827 let f = self.func("ssm_conv1d_gdn_state_vl");
24828 let cfg = LaunchConfig {
24829 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24830 block_dim: (256, 1, 1),
24831 shared_mem_bytes: 0,
24832 };
24833 let (dsi, nvi, nki, kdi, hki) = (
24834 d_state as i32,
24835 num_v as i32,
24836 num_k as i32,
24837 key_dim as i32,
24838 hk as i32,
24839 );
24840 let __s_lb = self.gpu.stream();
24841 let mut lb = __s_lb.launch_builder(&f);
24842 lb.arg(&v)
24843 .arg(conv_w)
24844 .arg(&cdi)
24845 .arg(&dci)
24846 .arg(&dsi)
24847 .arg(&nvi)
24848 .arg(&nki)
24849 .arg(&kdi)
24850 .arg(&hki);
24851 unsafe {
24852 lb.launch(cfg)?;
24853 }
24854 } else {
24855 let f = self.func("ssm_conv1d_tm_state_vl");
24856 let cfg = LaunchConfig {
24857 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24858 block_dim: (256, 1, 1),
24859 shared_mem_bytes: 0,
24860 };
24861 let __s_lb = self.gpu.stream();
24862 let mut lb = __s_lb.launch_builder(&f);
24863 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
24864 unsafe {
24865 lb.launch(cfg)?;
24866 }
24867 }
24868 {
24869 let f = self.func("ssm_conv_ring_update_vl");
24870 let n = (conv_dim * (d_conv - 1)) as u32;
24871 let cfg = LaunchConfig {
24872 grid_dim: (n.div_ceil(256), 1, b as u32),
24873 block_dim: (256, 1, 1),
24874 shared_mem_bytes: 0,
24875 };
24876 let __s_lb = self.gpu.stream();
24877 let mut lb = __s_lb.launch_builder(&f);
24878 lb.arg(&v).arg(&cdi).arg(&dci);
24879 unsafe {
24880 lb.launch(cfg)?;
24881 }
24882 }
24883 if !conv_fuse {
24884 let f = self.func("qkv_to_gdn_repack_vl");
24885 let n = max_t * (num_v * d_state) as u32;
24886 let cfg = LaunchConfig {
24887 grid_dim: (n.div_ceil(256), 1, b as u32),
24888 block_dim: (256, 1, 1),
24889 shared_mem_bytes: 0,
24890 };
24891 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24892 let __s_lb = self.gpu.stream();
24893 let mut lb = __s_lb.launch_builder(&f);
24894 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
24895 unsafe {
24896 lb.launch(cfg)?;
24897 }
24898 }
24899 if Self::l2_v2_on(d_state) {
24900 let f = self.func("gdn_l2_v2_vl");
24901 let cfg = LaunchConfig {
24902 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
24903 block_dim: (256, 1, 1),
24904 shared_mem_bytes: 0,
24905 };
24906 let (dsi, nvi) = (d_state as i32, hk as i32);
24907 let __s_lb = self.gpu.stream();
24908 let mut lb = __s_lb.launch_builder(&f);
24909 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24910 unsafe {
24911 lb.launch(cfg)?;
24912 }
24913 } else {
24914 let f = self.func("gdn_l2_vl");
24915 let cfg = LaunchConfig {
24916 grid_dim: (max_t * hk as u32, 2, b as u32),
24917 block_dim: (256, 1, 1),
24918 shared_mem_bytes: 0,
24919 };
24920 let (dsi, nvi) = (d_state as i32, hk as i32);
24921 let __s_lb = self.gpu.stream();
24922 let mut lb = __s_lb.launch_builder(&f);
24923 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24924 unsafe {
24925 lb.launch(cfg)?;
24926 }
24927 }
24928 {
24929 let f = self.func("gdn_gate_prep_vl");
24930 let n = max_t * num_v as u32;
24931 let cfg = LaunchConfig {
24932 grid_dim: (n.div_ceil(256), 1, b as u32),
24933 block_dim: (256, 1, 1),
24934 shared_mem_bytes: 0,
24935 };
24936 let nvi = num_v as i32;
24937 let __s_lb = self.gpu.stream();
24938 let mut lb = __s_lb.launch_builder(&f);
24939 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
24940 unsafe {
24941 lb.launch(cfg)?;
24942 }
24943 }
24944 Ok(())
24945 }
24946
24947 pub fn gdn_mirror_vl8(
24949 &self,
24950 seqs: &[GdnSeqVl],
24951 n_head: usize,
24952 which: i32,
24953 hk: usize,
24954 ) -> Result<(), Box<dyn std::error::Error>> {
24955 let b = seqs.len();
24956 assert!(b >= 1 && b <= 8);
24957 let mut packed = [GdnSeqVl::default(); 8];
24958 packed[..b].copy_from_slice(seqs);
24959 let v = GdnVl8(packed);
24960 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
24961 let max_n = seqs
24962 .iter()
24963 .map(|s| {
24964 if which == 0 {
24965 s.t as i64 * ept as i64
24966 } else {
24967 s.nc as i64 * ept as i64 * 32
24968 }
24969 })
24970 .max()
24971 .unwrap();
24972 let f = self.func("gdn_mirror_vl");
24973 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
24974 let cfg = LaunchConfig {
24975 grid_dim: (blocks, 1, b as u32),
24976 block_dim: (256, 1, 1),
24977 shared_mem_bytes: 0,
24978 };
24979 let __s_lb = self.gpu.stream();
24980 let mut lb = __s_lb.launch_builder(&f);
24981 lb.arg(&v).arg(&ept).arg(&which);
24982 unsafe {
24983 lb.launch(cfg)?;
24984 }
24985 Ok(())
24986 }
24987
24988 pub fn gdn_tail_vl8(
24990 &self,
24991 seqs: &[GdnPrepVl],
24992 norm_w: &CudaSlice<f32>,
24993 d_state: usize,
24994 num_v: usize,
24995 eps: f32,
24996 ) -> Result<(), Box<dyn std::error::Error>> {
24997 let b = seqs.len();
24998 assert!(b >= 1 && b <= 8);
24999 let mut packed = [GdnPrepVl::default(); 8];
25000 packed[..b].copy_from_slice(seqs);
25001 let v = GdnPrepVl8(packed);
25002 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25003 let f = self.func("gated_rmsnorm_f16out_vl");
25004 let cfg = LaunchConfig {
25006 grid_dim: (max_t * num_v as u32, 1, b as u32),
25007 block_dim: (128, 1, 1),
25008 shared_mem_bytes: 0,
25009 };
25010 let (dsi, nvi) = (d_state as i32, num_v as i32);
25011 let __s_lb = self.gpu.stream();
25012 let mut lb = __s_lb.launch_builder(&f);
25013 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25014 unsafe {
25015 lb.launch(cfg)?;
25016 }
25017 Ok(())
25018 }
25019
25020 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25023 use cudarc::driver::DevicePtr;
25024 let s = self.gpu.stream();
25025 let (p, _g) = x.device_ptr(&s);
25026 p as u64
25027 }
25028 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25029 use cudarc::driver::DevicePtrMut;
25030 let s = self.gpu.stream();
25031 let (p, _g) = x.device_ptr_mut(&s);
25032 p as u64
25033 }
25034 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25035 use cudarc::driver::DevicePtr;
25036 let s = self.gpu.stream();
25037 let (p, _g) = x.device_ptr(&s);
25038 p as u64
25039 }
25040 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25041 use cudarc::driver::DevicePtr;
25042 let s = self.gpu.stream();
25043 let (p, _g) = x.device_ptr(&s);
25044 p as u64
25045 }
25046
25047 pub fn gdn_chunk_vl8(
25051 &self,
25052 seqs: &[GdnSeqVl],
25053 n_head: usize,
25054 scale: f32,
25055 hk: usize,
25056 wq: Option<&GdnWVl8>,
25057 ) -> Result<(), Box<dyn std::error::Error>> {
25058 const NSPLIT: u32 = 4;
25059 let b = seqs.len();
25060 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25061 let mut packed = [GdnSeqVl::default(); 8];
25062 packed[..b].copy_from_slice(seqs);
25063 let v = GdnVl8(packed);
25064 let (hi, ci) = (n_head as i32, 32i32);
25065 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25066 let hki = hk as i32;
25067 if let Some(w) = wq {
25068 let f = self.func("gdn_k45_wgmma_vl");
25070 let cfg = LaunchConfig {
25071 grid_dim: (n_head as u32, NSPLIT, b as u32),
25072 block_dim: (256, 1, 1),
25073 shared_mem_bytes: 0,
25074 };
25075 let __s_lb = self.gpu.stream();
25076 let mut lb = __s_lb.launch_builder(&f);
25077 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25078 unsafe {
25079 lb.launch(cfg)?;
25080 }
25081 let _ = max_nc;
25082 return Ok(());
25083 }
25084 {
25085 let f = self.func("gdn_chunk_state_mma_vl");
25086 let cfg = LaunchConfig {
25087 grid_dim: (n_head as u32, NSPLIT, b as u32),
25088 block_dim: (256, 1, 1),
25089 shared_mem_bytes: 0,
25090 };
25091 let __s_lb = self.gpu.stream();
25092 let mut lb = __s_lb.launch_builder(&f);
25093 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25094 unsafe {
25095 lb.launch(cfg)?;
25096 }
25097 }
25098 {
25099 let f = self.func("gdn_chunk_output_mma_vl");
25100 let cfg = LaunchConfig {
25101 grid_dim: (max_nc, n_head as u32, b as u32),
25102 block_dim: (256, 1, 1),
25103 shared_mem_bytes: 0,
25104 };
25105 let __s_lb = self.gpu.stream();
25106 let mut lb = __s_lb.launch_builder(&f);
25107 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25108 unsafe {
25109 lb.launch(cfg)?;
25110 }
25111 }
25112 Ok(())
25113 }
25114 pub fn gdn_scan_chunked(
25115 &self,
25116 q: &CudaSlice<f32>,
25117 k: &CudaSlice<f32>,
25118 v: &CudaSlice<f32>,
25119 g: &CudaSlice<f32>,
25120 beta: &CudaSlice<f32>,
25121 kb16_pre: Option<&CudaSlice<u8>>,
25122 qb16_pre: Option<&CudaSlice<u8>>,
25123 state_in: &CudaSlice<f32>,
25124 state_out: &mut CudaSlice<f32>,
25125 o: &mut CudaSlice<f32>,
25126 n_head: usize,
25127 t: usize,
25128 scale: f32,
25129 c: usize,
25130 hk: usize,
25131 ) -> Result<(), Box<dyn std::error::Error>> {
25132 const D: usize = 128;
25133 const NSPLIT: u32 = 4;
25134 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25135 let h = n_head;
25136 let nc = (t + c - 1) / c;
25137 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25138 let gdn_mma_pre = !portable_mma_gated()
25143 && c == 32
25144 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25145 Ok("1") => true,
25146 Ok("0") => false,
25147 _ => gdn_mma_default_on(),
25148 };
25149 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
25150 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
25151 } else {
25152 None
25153 };
25154 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
25159 && gdn_mma_pre
25160 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
25161 let nk = t * hk * D;
25162 let mut kb16_local: Option<CudaSlice<u8>> = None;
25163 if gdn_mma_pre && kb16_pre.is_none() {
25164 let mut kb = self.alloc_u8_uninit(nk * 2)?;
25165 let f = self.func("f32_to_bf16_bulk");
25166 let n2 = nk as i64;
25167 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25168 let __s_b = self.gpu.stream();
25169 let mut b = __s_b.launch_builder(&f);
25170 b.arg(k).arg(&mut kb).arg(&n2);
25171 unsafe {
25172 b.launch(cfg2)?;
25173 }
25174 kb16_local = Some(kb);
25175 }
25176 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
25177 if let Some(kb) = kb16_pre {
25178 assert!(kb.len() >= nk * 2, "kb16_pre too small");
25179 }
25180 let mut qb16: Option<CudaSlice<u8>> = None;
25181 let mut pb16: Option<CudaSlice<u8>> = None;
25182 if gdn_wgmma_pre {
25183 if qb16_pre.is_none() {
25186 let mut qb = self.alloc_u8_uninit(nk * 2)?;
25187 let f = self.func("f32_to_bf16_bulk");
25188 let n2 = nk as i64;
25189 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25190 let __s_b = self.gpu.stream();
25191 let mut b = __s_b.launch_builder(&f);
25192 b.arg(q).arg(&mut qb).arg(&n2);
25193 unsafe {
25194 b.launch(cfg2)?;
25195 }
25196 qb16 = Some(qb);
25197 } else if let Some(qb) = qb16_pre {
25198 assert!(qb.len() >= nk * 2, "qb16_pre too small");
25199 }
25200 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
25201 }
25202 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
25203 let k2w = if gdn_wgmma_pre {
25204 Some((
25205 *qb16_ref0.as_ref().unwrap(),
25206 *kb16_ref0.as_ref().unwrap(),
25207 pb16.as_mut().unwrap(),
25208 ))
25209 } else {
25210 None
25211 };
25212 let (gcum, p, u, w) =
25213 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
25214 let _ = &w;
25215 let mut y = self.uninit(nc * h * c * D)?;
25216 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
25230 && c == 32
25231 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25232 Ok("1") => true,
25233 Ok("0") => false,
25234 _ => gdn_mma_default_on(),
25235 };
25236 if gdn_mma {
25237 let wb16 = wb16_pre
25238 .take()
25239 .expect("mma path pre-allocates wb16 (K3 store fold)");
25240 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
25241 if gdn_wgmma_pre {
25253 let qb16 = qb16_ref0.unwrap();
25255 let pb16 = pb16.as_ref().unwrap();
25256 {
25257 let f = self.func("gdn_k45_wgmma");
25258 let cfg = LaunchConfig {
25259 grid_dim: (h as u32, 4, 1),
25260 block_dim: (256, 1, 1),
25261 shared_mem_bytes: 0,
25262 };
25263 let hki = hk as i32;
25264 let __s_b = self.gpu.stream();
25265 let mut b = __s_b.launch_builder(&f);
25266 b.arg(kb16_ref)
25267 .arg(&gcum)
25268 .arg(beta)
25269 .arg(&u)
25270 .arg(&wb16)
25271 .arg(qb16)
25272 .arg(pb16)
25273 .arg(o)
25274 .arg(&scale)
25275 .arg(state_in)
25276 .arg(&mut *state_out)
25277 .arg(&hi)
25278 .arg(&ti)
25279 .arg(&ci)
25280 .arg(&hki);
25281 unsafe {
25282 b.launch(cfg)?;
25283 }
25284 }
25285 return Ok(());
25286 }
25287 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
25291 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
25292 {
25293 let f = self.func("gdn_chunk_state_mma");
25294 let cfg = LaunchConfig {
25295 grid_dim: (h as u32, NSPLIT, 1),
25296 block_dim: (256, 1, 1),
25297 shared_mem_bytes: 0,
25298 };
25299 let hki = hk as i32;
25300 let __s_b = self.gpu.stream();
25301 let mut b = __s_b.launch_builder(&f);
25302 b.arg(kb16_ref)
25303 .arg(&gcum)
25304 .arg(beta)
25305 .arg(&u)
25306 .arg(&wb16)
25307 .arg(&mut y16)
25308 .arg(&mut ssnap16)
25309 .arg(state_in)
25310 .arg(&mut *state_out)
25311 .arg(&hi)
25312 .arg(&ti)
25313 .arg(&ci)
25314 .arg(&hki);
25315 unsafe {
25316 b.launch(cfg)?;
25317 }
25318 }
25319 {
25320 let f = self.func("gdn_chunk_output_mma");
25322 let jt = ((c + 31) / 32) as u32;
25323 let cfg = LaunchConfig {
25324 grid_dim: (nc as u32, h as u32, jt),
25325 block_dim: (256, 1, 1),
25326 shared_mem_bytes: 0,
25327 };
25328 let hki = hk as i32;
25329 let __s_b = self.gpu.stream();
25330 let mut b = __s_b.launch_builder(&f);
25331 b.arg(q)
25332 .arg(&gcum)
25333 .arg(&p)
25334 .arg(&y16)
25335 .arg(&ssnap16)
25336 .arg(o)
25337 .arg(&hi)
25338 .arg(&ti)
25339 .arg(&ci)
25340 .arg(&scale)
25341 .arg(&hki);
25342 unsafe {
25343 b.launch(cfg)?;
25344 }
25345 }
25346 return Ok(());
25347 }
25348 {
25349 let f = self.func("gdn_chunk_state_f32");
25351 let cfg = LaunchConfig {
25352 grid_dim: (h as u32, NSPLIT, 1),
25353 block_dim: (256, 1, 1),
25354 shared_mem_bytes: 0,
25355 };
25356 let __s_b = self.gpu.stream();
25357 let mut b = __s_b.launch_builder(&f);
25358 b.arg(k)
25359 .arg(&gcum)
25360 .arg(beta)
25361 .arg(&u)
25362 .arg(&w)
25363 .arg(&mut y)
25364 .arg(&mut ssnap)
25365 .arg(state_in)
25366 .arg(&mut *state_out)
25367 .arg(&hi)
25368 .arg(&ti)
25369 .arg(&ci);
25370 unsafe {
25371 b.launch(cfg)?;
25372 }
25373 }
25374 {
25375 let f = self.func("gdn_chunk_output_f32");
25377 let jt = ((c + 31) / 32) as u32;
25378 let cfg = LaunchConfig {
25379 grid_dim: (nc as u32, h as u32, jt),
25380 block_dim: (256, 1, 1),
25381 shared_mem_bytes: 0,
25382 };
25383 let __s_b = self.gpu.stream();
25384 let mut b = __s_b.launch_builder(&f);
25385 b.arg(q)
25386 .arg(&gcum)
25387 .arg(&p)
25388 .arg(&y)
25389 .arg(&ssnap)
25390 .arg(o)
25391 .arg(&hi)
25392 .arg(&ti)
25393 .arg(&ci)
25394 .arg(&scale);
25395 unsafe {
25396 b.launch(cfg)?;
25397 }
25398 }
25399 Ok(())
25400 }
25401
25402 #[allow(clippy::too_many_arguments)]
25411 #[allow(clippy::too_many_arguments)]
25412 pub fn gdn_scan_prefill(
25413 &self,
25414 q: &CudaSlice<f32>,
25415 k: &CudaSlice<f32>,
25416 v: &CudaSlice<f32>,
25417 g: &CudaSlice<f32>,
25418 beta: &CudaSlice<f32>,
25419 kb16_pre: Option<&CudaSlice<u8>>,
25420 qb16_pre: Option<&CudaSlice<u8>>,
25421 state_in: &CudaSlice<f32>,
25422 state_out: &mut CudaSlice<f32>,
25423 o: &mut CudaSlice<f32>,
25424 n_head: usize,
25425 t: usize,
25426 scale: f32,
25427 hk: usize,
25428 ) -> Result<(), Box<dyn std::error::Error>> {
25429 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
25430 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
25431 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25432 }
25433 if Self::gdn_chunked_enabled() && t >= 16 {
25434 self.gdn_scan_chunked(
25435 q,
25436 k,
25437 v,
25438 g,
25439 beta,
25440 kb16_pre,
25441 qb16_pre,
25442 state_in,
25443 state_out,
25444 o,
25445 n_head,
25446 t,
25447 scale,
25448 Self::gdn_chunk_size(),
25449 hk,
25450 )
25451 } else {
25452 assert!(
25453 hk == n_head,
25454 "s128 scan is broadcast-only (prep guarantees by predicate)"
25455 );
25456 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25457 }
25458 }
25459
25460 #[allow(clippy::too_many_arguments)]
25462 fn gdn_scan_diff(
25463 &self,
25464 q: &CudaSlice<f32>,
25465 k: &CudaSlice<f32>,
25466 v: &CudaSlice<f32>,
25467 g: &CudaSlice<f32>,
25468 beta: &CudaSlice<f32>,
25469 state_in: &CudaSlice<f32>,
25470 state_out: &mut CudaSlice<f32>,
25471 o: &mut CudaSlice<f32>,
25472 n_head: usize,
25473 t: usize,
25474 scale: f32,
25475 ) -> Result<(), Box<dyn std::error::Error>> {
25476 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25477 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25478 let mut o_c = self.uninit(o.len())?;
25479 let mut st_c = self.uninit(state_out.len())?;
25480 self.gdn_scan_chunked(
25481 q,
25482 k,
25483 v,
25484 g,
25485 beta,
25486 None,
25487 None,
25488 state_in,
25489 &mut st_c,
25490 &mut o_c,
25491 n_head,
25492 t,
25493 scale,
25494 Self::gdn_chunk_size(),
25495 n_head,
25496 )?;
25497 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25498 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25499 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25500 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25501 let mut max_abs = 0f32;
25502 let mut max_rel = 0f32;
25503 let mut sum_rel = 0f64;
25504 for (x, y) in a.iter().zip(b) {
25505 let ad = (x - y).abs();
25506 let rel = ad / x.abs().max(y.abs()).max(1e-3);
25507 if ad > max_abs {
25508 max_abs = ad;
25509 }
25510 if rel > max_rel {
25511 max_rel = rel;
25512 }
25513 sum_rel += rel as f64;
25514 }
25515 (max_abs, max_rel, sum_rel / a.len() as f64)
25516 };
25517 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25518 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25519 println!(
25520 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25521 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25522 Self::gdn_chunk_size()
25523 );
25524 Ok(())
25525 }
25526
25527 pub fn gdn_glog(
25529 &self,
25530 alpha: &CudaSlice<f32>,
25531 dt_bias: &CudaSlice<f32>,
25532 a: &CudaSlice<f32>,
25533 g_log: &mut CudaSlice<f32>,
25534 n_head: usize,
25535 t: usize,
25536 ) -> Result<(), Box<dyn std::error::Error>> {
25537 let f = self.func("gdn_glog_f32");
25538 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25539 let (h, ti) = (n_head as i32, t as i32);
25540 let __s_b = self.gpu.stream();
25541 let mut b = __s_b.launch_builder(&f);
25542 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25543 unsafe {
25544 b.launch(cfg)?;
25545 }
25546 Ok(())
25547 }
25548
25549 pub fn sigmoid_v(
25552 &self,
25553 x: &cudarc::driver::CudaView<f32>,
25554 y: &mut CudaSlice<f32>,
25555 n: usize,
25556 ) -> Result<(), Box<dyn std::error::Error>> {
25557 let f = self.func("sigmoid_f32");
25558 let cfg = LaunchConfig::for_num_elems(n as u32);
25559 let ni = n as i32;
25560 let __s_b = self.gpu.stream();
25561 let mut b = __s_b.launch_builder(&f);
25562 b.arg(x).arg(y).arg(&ni);
25563 unsafe {
25564 b.launch(cfg)?;
25565 }
25566 Ok(())
25567 }
25568
25569 pub fn gdn_glog_v(
25570 &self,
25571 alpha: &cudarc::driver::CudaView<f32>,
25572 dt_bias: &CudaSlice<f32>,
25573 a: &CudaSlice<f32>,
25574 g_log: &mut CudaSlice<f32>,
25575 n_head: usize,
25576 t: usize,
25577 ) -> Result<(), Box<dyn std::error::Error>> {
25578 let f = self.func("gdn_glog_f32");
25579 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25580 let (h, ti) = (n_head as i32, t as i32);
25581 let __s_b = self.gpu.stream();
25582 let mut b = __s_b.launch_builder(&f);
25583 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25584 unsafe {
25585 b.launch(cfg)?;
25586 }
25587 Ok(())
25588 }
25589
25590 pub fn sigmoid(
25591 &self,
25592 x: &CudaSlice<f32>,
25593 y: &mut CudaSlice<f32>,
25594 n: usize,
25595 ) -> Result<(), Box<dyn std::error::Error>> {
25596 let f = self.func("sigmoid_f32");
25597 let cfg = LaunchConfig::for_num_elems(n as u32);
25598 let ni = n as i32;
25599 let __s_b = self.gpu.stream();
25600 let mut b = __s_b.launch_builder(&f);
25601 b.arg(x).arg(y).arg(&ni);
25602 unsafe {
25603 b.launch(cfg)?;
25604 }
25605 Ok(())
25606 }
25607
25608 pub fn sig_mul_f16out(
25611 &self,
25612 a: &CudaSlice<f32>,
25613 g: &CudaSlice<f32>,
25614 dst: &mut CudaSlice<f32>,
25615 dst16: &mut CudaSlice<u8>,
25616 n: usize,
25617 ) -> Result<(), Box<dyn std::error::Error>> {
25618 let f = self.func("sig_mul_f16out_f32");
25619 let cfg = LaunchConfig::for_num_elems(n as u32);
25620 let ni = n as i32;
25621 let __s_b = self.gpu.stream();
25622 let mut b = __s_b.launch_builder(&f);
25623 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25624 unsafe {
25625 b.launch(cfg)?;
25626 }
25627 Ok(())
25628 }
25629
25630 #[allow(clippy::too_many_arguments)]
25639 pub fn attn_head_gate(
25640 &self,
25641 a: &CudaSlice<f32>,
25642 g: &CudaSlice<f32>,
25643 dst: &mut CudaSlice<f32>,
25644 dst16: Option<&mut CudaSlice<u8>>,
25645 head_dim: usize,
25646 n_head: usize,
25647 t: usize,
25648 ) -> Result<(), Box<dyn std::error::Error>> {
25649 let f = self.func("attn_head_gate_f32");
25650 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25651 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25652 let d16: u64 = match dst16 {
25654 Some(d) => self.addr_u8(d),
25655 None => 0,
25656 };
25657 let __s_b = self.gpu.stream();
25658 let mut b = __s_b.launch_builder(&f);
25659 b.arg(a)
25660 .arg(g)
25661 .arg(dst)
25662 .arg(&d16)
25663 .arg(&hd)
25664 .arg(&nh)
25665 .arg(&ti);
25666 unsafe {
25667 b.launch(cfg)?;
25668 }
25669 Ok(())
25670 }
25671
25672 #[allow(clippy::too_many_arguments)]
25681 pub fn swiglu_clamped_mul_scaled(
25682 &self,
25683 gate: &CudaSlice<f32>,
25684 up: &CudaSlice<f32>,
25685 gs: f32,
25686 us: f32,
25687 limit: f32,
25688 dst: &mut CudaSlice<f32>,
25689 n: usize,
25690 ) -> Result<(), Box<dyn std::error::Error>> {
25691 debug_assert!(
25692 limit > 1e-6,
25693 "swiglu_clamped needs a live limit; use silu_mul_scaled"
25694 );
25695 let f = self.func("swiglu_clamped_mul_scaled_f32");
25696 let cfg = LaunchConfig::for_num_elems(n as u32);
25697 let ni = n as i32;
25698 let __s_b = self.gpu.stream();
25699 let mut b = __s_b.launch_builder(&f);
25700 b.arg(gate)
25701 .arg(up)
25702 .arg(&gs)
25703 .arg(&us)
25704 .arg(&limit)
25705 .arg(dst)
25706 .arg(&ni);
25707 unsafe {
25708 b.launch(cfg)?;
25709 }
25710 Ok(())
25711 }
25712
25713 pub fn gated_rmsnorm(
25715 &self,
25716 o: &CudaSlice<f32>,
25717 w: &CudaSlice<f32>,
25718 z: &CudaSlice<f32>,
25719 dst: &mut CudaSlice<f32>,
25720 ncols: usize,
25721 nrows: usize,
25722 eps: f32,
25723 ) -> Result<(), Box<dyn std::error::Error>> {
25724 let f = self.func("gated_rmsnorm_f32");
25725 let cfg = LaunchConfig {
25726 grid_dim: (nrows as u32, 1, 1),
25727 block_dim: (128, 1, 1),
25728 shared_mem_bytes: 0,
25729 };
25730 let (nc, e) = (ncols as i32, eps);
25731 let __s_b = self.gpu.stream();
25732 let mut b = __s_b.launch_builder(&f);
25733 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25734 unsafe {
25735 b.launch(cfg)?;
25736 }
25737 Ok(())
25738 }
25739
25740 pub fn gated_rmsnorm_f16out(
25743 &self,
25744 o: &CudaSlice<f32>,
25745 w: &CudaSlice<f32>,
25746 z: &CudaSlice<f32>,
25747 dst: &mut CudaSlice<f32>,
25748 dst16: &mut CudaSlice<u8>,
25749 ncols: usize,
25750 nrows: usize,
25751 eps: f32,
25752 ) -> Result<(), Box<dyn std::error::Error>> {
25753 let f = self.func("gated_rmsnorm_f16out_f32");
25754 let cfg = LaunchConfig {
25756 grid_dim: (nrows as u32, 1, 1),
25757 block_dim: (128, 1, 1),
25758 shared_mem_bytes: 0,
25759 };
25760 let (nc, e) = (ncols as i32, eps);
25761 let __s_b = self.gpu.stream();
25762 let mut b = __s_b.launch_builder(&f);
25763 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25764 unsafe {
25765 b.launch(cfg)?;
25766 }
25767 Ok(())
25768 }
25769
25770 #[allow(clippy::too_many_arguments)]
25774 pub fn add_rms_norm_zq8(
25775 &self,
25776 a: &CudaSlice<f32>,
25777 b_in: &CudaSlice<f32>,
25778 w: &CudaSlice<f32>,
25779 res: &mut CudaSlice<f32>,
25780 z: &mut CudaSlice<f32>,
25781 ncols: usize,
25782 nrows: usize,
25783 eps: f32,
25784 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25785 assert!(ncols % 32 == 0);
25786 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25787 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25788 let f = self.func("add_rms_norm_zq8");
25789 let cfg = LaunchConfig {
25790 grid_dim: (nrows as u32, 1, 1),
25791 block_dim: (1024, 1, 1),
25792 shared_mem_bytes: 0,
25793 };
25794 let (nc, ep) = (ncols as i32, eps);
25795 let __s_b = self.gpu.stream();
25796 let mut b = __s_b.launch_builder(&f);
25797 b.arg(a)
25798 .arg(b_in)
25799 .arg(w)
25800 .arg(res)
25801 .arg(z)
25802 .arg(&mut q)
25803 .arg(&mut d)
25804 .arg(&nc)
25805 .arg(&ep);
25806 unsafe {
25807 b.launch(cfg)?;
25808 }
25809 Ok((q, d))
25810 }
25811
25812 pub fn gated_rmsnorm_zv(
25817 &self,
25818 o: &CudaSlice<f32>,
25819 w: &CudaSlice<f32>,
25820 z: &cudarc::driver::CudaView<f32>,
25821 dst: &mut CudaSlice<f32>,
25822 ncols: usize,
25823 nrows: usize,
25824 eps: f32,
25825 ) -> Result<(), Box<dyn std::error::Error>> {
25826 let f = self.func("gated_rmsnorm_f32");
25827 let cfg = LaunchConfig {
25828 grid_dim: (nrows as u32, 1, 1),
25829 block_dim: (128, 1, 1),
25830 shared_mem_bytes: 0,
25831 };
25832 let (nc, e) = (ncols as i32, eps);
25833 let __s_b = self.gpu.stream();
25834 let mut b = __s_b.launch_builder(&f);
25835 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25836 unsafe {
25837 b.launch(cfg)?;
25838 }
25839 Ok(())
25840 }
25841
25842 pub fn gated_rmsnorm_f16out_zv(
25843 &self,
25844 o: &CudaSlice<f32>,
25845 w: &CudaSlice<f32>,
25846 z: &cudarc::driver::CudaView<f32>,
25847 dst: &mut CudaSlice<f32>,
25848 dst16: &mut CudaSlice<u8>,
25849 ncols: usize,
25850 nrows: usize,
25851 eps: f32,
25852 ) -> Result<(), Box<dyn std::error::Error>> {
25853 let f = self.func("gated_rmsnorm_f16out_f32");
25854 let cfg = LaunchConfig {
25856 grid_dim: (nrows as u32, 1, 1),
25857 block_dim: (128, 1, 1),
25858 shared_mem_bytes: 0,
25859 };
25860 let (nc, e) = (ncols as i32, eps);
25861 let __s_b = self.gpu.stream();
25862 let mut b = __s_b.launch_builder(&f);
25863 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25864 unsafe {
25865 b.launch(cfg)?;
25866 }
25867 Ok(())
25868 }
25869
25870 pub fn gated_rmsnorm_q8_1(
25871 &self,
25872 o: &CudaSlice<f32>,
25873 w: &CudaSlice<f32>,
25874 z: &CudaSlice<f32>,
25875 ncols: usize,
25876 nrows: usize,
25877 eps: f32,
25878 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25879 assert!(ncols % 32 == 0);
25880 let f = self.func("gated_rmsnorm_q8_1");
25881 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
25882 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25883 let cfg = LaunchConfig {
25884 grid_dim: (nrows as u32, 1, 1),
25885 block_dim: (128, 1, 1),
25886 shared_mem_bytes: 0,
25887 };
25888 let (nc, ep) = (ncols as i32, eps);
25889 let __s_b = self.gpu.stream();
25890 let mut b = __s_b.launch_builder(&f);
25891 b.arg(o)
25892 .arg(w)
25893 .arg(z)
25894 .arg(&mut out_q)
25895 .arg(&mut out_d)
25896 .arg(&nc)
25897 .arg(&ep);
25898 unsafe {
25899 b.launch(cfg)?;
25900 }
25901 Ok((out_q, out_d))
25902 }
25903
25904 pub fn transpose(
25906 &self,
25907 inp: &CudaSlice<f32>,
25908 rows: usize,
25909 cols: usize,
25910 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25911 let f = self.func("transpose_f32");
25912 let mut out = self.zeros(rows * cols)?;
25913 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
25914 let (r, c) = (rows as i32, cols as i32);
25915 let __s_b = self.gpu.stream();
25916 let mut b = __s_b.launch_builder(&f);
25917 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
25918 unsafe {
25919 b.launch(cfg)?;
25920 }
25921 Ok(out)
25922 }
25923
25924 pub fn repeat_heads(
25926 &self,
25927 inp: &CudaSlice<f32>,
25928 out: &mut CudaSlice<f32>,
25929 head_dim: usize,
25930 n_in: usize,
25931 n_out: usize,
25932 t: usize,
25933 ) -> Result<(), Box<dyn std::error::Error>> {
25934 let f = self.func("repeat_heads_f32");
25935 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
25936 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
25937 let __s_b = self.gpu.stream();
25938 let mut b = __s_b.launch_builder(&f);
25939 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
25940 unsafe {
25941 b.launch(cfg)?;
25942 }
25943 Ok(())
25944 }
25945
25946 pub fn q_gate_split(
25953 &self,
25954 qf: &CudaSlice<f32>,
25955 q_out: &mut CudaSlice<f32>,
25956 gate_out: &mut CudaSlice<f32>,
25957 head_dim: usize,
25958 n_head: usize,
25959 t: usize,
25960 ) -> Result<(), Box<dyn std::error::Error>> {
25961 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
25962 let out_need = head_dim * n_head * t;
25963 if q_out.len() < out_need || gate_out.len() < out_need {
25964 return Err(format!(
25965 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
25966 q_out.len(),
25967 gate_out.len()
25968 )
25969 .into());
25970 }
25971 let f = self.func("q_gate_split_f32");
25972 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25973 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25974 let __s_b = self.gpu.stream();
25975 let mut b = __s_b.launch_builder(&f);
25976 b.arg(qf)
25977 .arg(q_out)
25978 .arg(gate_out)
25979 .arg(&hd)
25980 .arg(&nh)
25981 .arg(&ti);
25982 unsafe {
25983 b.launch(cfg)?;
25984 }
25985 Ok(())
25986 }
25987
25988 pub fn qkv_to_gdn_repack(
25992 &self,
25993 conv_out: &CudaSlice<f32>,
25994 q_g: &mut CudaSlice<f32>,
25995 k_g: &mut CudaSlice<f32>,
25996 v_g: &mut CudaSlice<f32>,
25997 d_state: usize,
25998 num_v: usize,
25999 num_k: usize,
26000 key_dim: usize,
26001 t: usize,
26002 ) -> Result<(), Box<dyn std::error::Error>> {
26003 let f = self.func("qkv_to_gdn_repack_f32");
26004 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26005 let (ds, nv, nk, kd, ti) = (
26006 d_state as i32,
26007 num_v as i32,
26008 num_k as i32,
26009 key_dim as i32,
26010 t as i32,
26011 );
26012 let __s_b = self.gpu.stream();
26013 let mut b = __s_b.launch_builder(&f);
26014 b.arg(conv_out)
26015 .arg(q_g)
26016 .arg(k_g)
26017 .arg(v_g)
26018 .arg(&ds)
26019 .arg(&nv)
26020 .arg(&nk)
26021 .arg(&kd)
26022 .arg(&ti);
26023 unsafe {
26024 b.launch(cfg)?;
26025 }
26026 Ok(())
26027 }
26028
26029 pub fn conv_left_pad(
26032 &self,
26033 src: &CudaSlice<f32>,
26034 dst: &mut CudaSlice<f32>,
26035 conv_dim: usize,
26036 t: usize,
26037 pad: usize,
26038 ) -> Result<(), Box<dyn std::error::Error>> {
26039 let f = self.func("conv_left_pad_f32");
26040 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26041 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26042 let __s_b = self.gpu.stream();
26043 let mut b = __s_b.launch_builder(&f);
26044 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26045 unsafe {
26046 b.launch(cfg)?;
26047 }
26048 Ok(())
26049 }
26050
26051 pub fn conv_assemble_and_roll(
26055 &self,
26056 qkv_col: &CudaSlice<f32>,
26057 conv_state: &mut CudaSlice<f32>,
26058 conv_in: &mut CudaSlice<f32>,
26059 conv_dim: usize,
26060 pad: usize,
26061 ) -> Result<(), Box<dyn std::error::Error>> {
26062 let f = self.func("conv_assemble_and_roll_f32");
26063 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26064 let (cd, p) = (conv_dim as i32, pad as i32);
26065 let __s_b = self.gpu.stream();
26066 let mut b = __s_b.launch_builder(&f);
26067 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26068 unsafe {
26069 b.launch(cfg)?;
26070 }
26071 Ok(())
26072 }
26073
26074 pub fn ssm_conv1d_fused_decode(
26080 &self,
26081 qkv_col: &CudaSlice<f32>,
26082 conv_state: &mut CudaSlice<f32>,
26083 w: &CudaSlice<f32>,
26084 conv_out: &mut CudaSlice<f32>,
26085 conv_dim: usize,
26086 d_conv: usize,
26087 ) -> Result<(), Box<dyn std::error::Error>> {
26088 let f = self.func("ssm_conv1d_fused_decode_f32");
26089 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26090 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26091 let __s_b = self.gpu.stream();
26092 let mut b = __s_b.launch_builder(&f);
26093 b.arg(qkv_col)
26094 .arg(conv_state)
26095 .arg(w)
26096 .arg(conv_out)
26097 .arg(&cd)
26098 .arg(&dc);
26099 unsafe {
26100 b.launch(cfg)?;
26101 }
26102 Ok(())
26103 }
26104
26105 pub fn slice_range(
26108 &self,
26109 src: &CudaSlice<f32>,
26110 start: usize,
26111 len: usize,
26112 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26113 let host = self.gpu.stream().clone_dtoh(src)?;
26114 self.gpu.stream().synchronize()?;
26115 Ok(self.htod(&host[start..start + len])?)
26116 }
26117}
26118
26119#[cfg(test)]
26120mod target_dispatch_tests {
26121 use super::legacy_quant_gemm_allowed;
26122
26123 #[test]
26124 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26125 assert!(legacy_quant_gemm_allowed(false, false, false));
26127 assert!(!legacy_quant_gemm_allowed(false, false, true));
26128 assert!(!legacy_quant_gemm_allowed(true, false, false));
26130 assert!(!legacy_quant_gemm_allowed(true, false, true));
26131 assert!(legacy_quant_gemm_allowed(true, true, false));
26133 assert!(!legacy_quant_gemm_allowed(true, true, true));
26134 }
26135
26136 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26137 #[test]
26138 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26139 assert!(!legacy_quant_gemm_allowed(
26140 cfg!(memra_portable_cuda),
26141 cfg!(memra_hopper_mma),
26142 false
26143 ));
26144 }
26145
26146 #[cfg(memra_hopper_mma)]
26147 #[test]
26148 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
26149 assert!(legacy_quant_gemm_allowed(
26150 cfg!(memra_portable_cuda),
26151 cfg!(memra_hopper_mma),
26152 false
26153 ));
26154 assert!(super::portable_mma_gated() == false);
26155 }
26156}
26157
26158impl memra_kv::KvDev for Engine {
26161 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26162 Engine::zeros(self, n)
26163 }
26164 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26165 Engine::uninit(self, n)
26166 }
26167 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
26168 Engine::alloc_u8(self, n)
26169 }
26170 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
26171 Engine::htod_i32(self, v)
26172 }
26173 fn clone_dtod(
26174 &self,
26175 src: &CudaSlice<f32>,
26176 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26177 Engine::clone_dtod(self, src)
26178 }
26179 fn copy_into(
26180 &self,
26181 dst: &mut CudaSlice<f32>,
26182 off: usize,
26183 src: &CudaSlice<f32>,
26184 len: usize,
26185 ) -> Result<(), Box<dyn std::error::Error>> {
26186 Engine::copy_into(self, dst, off, src, len)
26187 }
26188 fn set_i32_one(
26189 &self,
26190 d: &mut CudaSlice<i32>,
26191 v: i32,
26192 ) -> Result<(), Box<dyn std::error::Error>> {
26193 Engine::set_i32_one(self, d, v)
26194 }
26195}
26196
26197#[cfg(test)]
26198mod fused_gate_bounds_tests {
26199 use super::*;
26200
26201 #[test]
26214 #[ignore = "requires a CUDA GPU"]
26215 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
26216 let e = Engine::new(0).unwrap();
26217 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
26218 let fused = 2 * head_dim * n_head * t;
26219 let out_n = head_dim * n_head * t;
26220
26221 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
26223 let mut q = e.uninit(out_n).unwrap();
26224 let mut gate = e.uninit(out_n).unwrap();
26225 let err = e
26226 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
26227 .expect_err("half-width wq must be refused, not read past")
26228 .to_string();
26229 assert!(err.contains("NO fused gate"), "{err}");
26230 assert!(err.contains(&format!("{fused}")), "{err}");
26231
26232 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
26235 let wide = e.htod(&host).unwrap();
26236 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
26237 .expect("full-width wq splits");
26238 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
26239 for tok in 0..t {
26240 for hh in 0..n_head {
26241 for d in 0..head_dim {
26242 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
26243 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
26244 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
26245 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
26246 }
26247 }
26248 }
26249
26250 let mut small = e.uninit(out_n - 1).unwrap();
26252 assert!(
26253 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
26254 .is_err()
26255 );
26256 }
26257}
26258
26259#[cfg(test)]
26263mod fused_rope_width_tests {
26264 use super::Engine;
26265
26266 #[test]
26269 fn full_width_is_accepted() {
26270 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
26271 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
26272 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
26273 }
26274
26275 #[test]
26288 fn gemma4_official_artifact_widths_pass() {
26289 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
26290 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
26291 }
26292
26293 #[test]
26296 fn partial_rotary_is_refused_with_the_geometry_named() {
26297 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
26299 .expect_err("partial rotary must refuse");
26300 let msg = err.to_string();
26301 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
26302 assert!(msg.contains("n_rot 64"), "{msg}");
26303 assert!(msg.contains("head_dim 256"), "{msg}");
26304 assert!(
26305 msg.contains("64..256"),
26306 "names the band it would corrupt: {msg}"
26307 );
26308 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
26310 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
26312 }
26313}