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
876pub(crate) fn 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 #[allow(clippy::too_many_arguments)]
7756 pub fn qmatvec_nvfp4_sel_gu_ep_into(
7757 &self,
7758 gate_bank: &CudaSlice<u8>,
7759 up_bank: &CudaSlice<u8>,
7760 sel: &CudaSlice<i32>,
7761 aq: &CudaSlice<i8>,
7762 ad: &CudaSlice<f32>,
7763 yg: &mut CudaSlice<f32>,
7764 yu: &mut CudaSlice<f32>,
7765 n_sel: usize,
7766 in_f: usize,
7767 out_f: usize,
7768 row_bytes: usize,
7769 expert_stride: usize,
7770 owner: usize,
7771 ) -> Result<(), Box<dyn std::error::Error>> {
7772 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7773 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7774 return Err("NVFP4 gu ep geometry".into());
7775 }
7776 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
7777 let cfg = LaunchConfig {
7778 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
7779 block_dim: (128, 1, 1),
7780 shared_mem_bytes: 0,
7781 };
7782 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
7783 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7784 let (ars, adrs) = (0i64, 0i64);
7785 let __s_b = self.gpu.stream();
7786 let mut b = __s_b.launch_builder(&f);
7787 b.arg(gate_bank)
7788 .arg(up_bank)
7789 .arg(sel)
7790 .arg(aq)
7791 .arg(ad)
7792 .arg(yg)
7793 .arg(yu)
7794 .arg(&inf)
7795 .arg(&outf)
7796 .arg(&ns)
7797 .arg(&rb)
7798 .arg(&es)
7799 .arg(&ars)
7800 .arg(&adrs)
7801 .arg(&own);
7802 unsafe {
7803 b.launch(cfg)?;
7804 }
7805 Ok(())
7806 }
7807
7808 #[allow(clippy::too_many_arguments)]
7810 pub fn silu_mul_scaled_q8_1_sel_ep_into(
7811 &self,
7812 gate: &CudaSlice<f32>,
7813 up: &CudaSlice<f32>,
7814 gmac: &CudaSlice<f32>,
7815 umac: &CudaSlice<f32>,
7816 sel: &CudaSlice<i32>,
7817 limit: Option<f32>,
7818 out_q: &mut CudaSlice<i8>,
7819 out_d: &mut CudaSlice<f32>,
7820 n_per: usize,
7821 n_sel: usize,
7822 owner: usize,
7823 ) -> Result<(), Box<dyn std::error::Error>> {
7824 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
7825 return Err("NVFP4 silu ep geometry".into());
7826 }
7827 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
7828 let warps = n_sel * n_per / 32;
7829 let cfg = LaunchConfig {
7830 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
7831 block_dim: (128, 1, 1),
7832 shared_mem_bytes: 0,
7833 };
7834 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
7835 let (lim, has) = match limit {
7836 Some(l) => (l, 1i32),
7837 None => (0.0f32, 0i32),
7838 };
7839 let __s_b = self.gpu.stream();
7840 let mut b = __s_b.launch_builder(&f);
7841 b.arg(gate)
7842 .arg(up)
7843 .arg(gmac)
7844 .arg(umac)
7845 .arg(sel)
7846 .arg(&lim)
7847 .arg(&has)
7848 .arg(out_q)
7849 .arg(out_d)
7850 .arg(&np)
7851 .arg(&ns)
7852 .arg(&own);
7853 unsafe {
7854 b.launch(cfg)?;
7855 }
7856 Ok(())
7857 }
7858
7859 #[allow(clippy::too_many_arguments)]
7861 pub fn qmatvec_nvfp4_sel_down8_ep_into(
7862 &self,
7863 bank: &CudaSlice<u8>,
7864 sel: &CudaSlice<i32>,
7865 aq: &CudaSlice<i8>,
7866 ad: &CudaSlice<f32>,
7867 route_w: &CudaSlice<f32>,
7868 md: &CudaSlice<f32>,
7869 dst: &mut CudaSlice<f32>,
7870 n_sel: usize,
7871 in_f: usize,
7872 out_f: usize,
7873 row_bytes: usize,
7874 expert_stride: usize,
7875 act_row_stride: usize,
7876 ad_row_stride: usize,
7877 owner: usize,
7878 ) -> Result<(), Box<dyn std::error::Error>> {
7879 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
7880 return Err("NVFP4 down8 ep geometry".into());
7881 }
7882 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
7883 let cfg = LaunchConfig {
7884 grid_dim: (out_f as u32, 1, 1),
7885 block_dim: (32, n_sel as u32, 1),
7886 shared_mem_bytes: 0,
7887 };
7888 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
7889 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7890 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
7891 let __s_b = self.gpu.stream();
7892 let mut b = __s_b.launch_builder(&f);
7893 b.arg(bank)
7894 .arg(sel)
7895 .arg(aq)
7896 .arg(ad)
7897 .arg(route_w)
7898 .arg(md)
7899 .arg(dst)
7900 .arg(&inf)
7901 .arg(&outf)
7902 .arg(&ns)
7903 .arg(&rb)
7904 .arg(&es)
7905 .arg(&ars)
7906 .arg(&adrs)
7907 .arg(&own);
7908 unsafe {
7909 b.launch(cfg)?;
7910 }
7911 Ok(())
7912 }
7913
7914 pub fn qmatvec_nvfp4_sel_into(
7915 &self,
7916 bank: &CudaSlice<u8>,
7917 sel: &CudaSlice<i32>,
7918 aq: &CudaSlice<i8>,
7919 ad: &CudaSlice<f32>,
7920 y: &mut CudaSlice<f32>,
7921 n_sel: usize,
7922 in_f: usize,
7923 out_f: usize,
7924 row_bytes: usize,
7925 expert_stride: usize,
7926 act_row_stride: usize,
7927 ad_row_stride: usize,
7928 ) -> Result<(), Box<dyn std::error::Error>> {
7929 assert!(
7930 in_f % 64 == 0,
7931 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7932 );
7933 if y.len() < n_sel * out_f || sel.len() < n_sel {
7934 return Err(format!(
7935 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7936 y.len(),
7937 sel.len()
7938 )
7939 .into());
7940 }
7941 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7948 let mode = *MR.get_or_init(|| {
7949 if crate::tp::nvfp4_bank_v2_on() {
7950 3
7951 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7952 2
7953 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7954 1
7955 } else {
7956 0
7957 }
7958 });
7959 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7960 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7964 let v2s = mode == 3
7965 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7966 && row_bytes % 16 == 0
7967 && in_f <= 4096;
7968 let f = match (mode, v2s) {
7969 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7970 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7971 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7972 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7973 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7974 };
7975 let nsb = in_f >> 5;
7980 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7981 32
7982 } else if mode == 1 {
7983 512
7984 } else {
7985 128
7986 };
7987 let cfg = LaunchConfig {
7988 grid_dim: (
7989 if v2s {
7990 (out_f as u32).div_ceil(8)
7991 } else {
7992 match mode {
7993 2 => (out_f as u32).div_ceil(16),
7994 1 => (out_f as u32).div_ceil(4),
7995 _ => out_f as u32,
7996 }
7997 },
7998 n_sel as u32,
7999 1,
8000 ),
8001 block_dim: (fit_block, 1, 1),
8002 shared_mem_bytes: 0,
8003 };
8004 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8005 let (rb, es, ars, adrs) = (
8006 row_bytes as i64,
8007 expert_stride as i64,
8008 act_row_stride as i64,
8009 ad_row_stride as i64,
8010 );
8011 let __s_b = self.gpu.stream();
8012 let mut b = __s_b.launch_builder(&f);
8013 b.arg(bank)
8014 .arg(sel)
8015 .arg(aq)
8016 .arg(ad)
8017 .arg(y)
8018 .arg(&inf)
8019 .arg(&outf)
8020 .arg(&ns)
8021 .arg(&rb)
8022 .arg(&es)
8023 .arg(&ars)
8024 .arg(&adrs);
8025 unsafe {
8026 b.launch(cfg)?;
8027 }
8028 Ok(())
8029 }
8030
8031 #[allow(clippy::too_many_arguments)]
8036 pub fn silu_mul_scaled_q8_1_sel_into(
8037 &self,
8038 gate: &CudaSlice<f32>,
8039 up: &CudaSlice<f32>,
8040 gmac: &CudaSlice<f32>,
8041 umac: &CudaSlice<f32>,
8042 sel: &CudaSlice<i32>,
8043 limit: Option<f32>,
8044 out_q: &mut CudaSlice<i8>,
8045 out_d: &mut CudaSlice<f32>,
8046 n_per: usize,
8047 n_sel: usize,
8048 ) -> Result<(), Box<dyn std::error::Error>> {
8049 let n = n_per * n_sel;
8050 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8051 return Err(format!(
8052 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8053 out_q.len(),
8054 out_d.len()
8055 )
8056 .into());
8057 }
8058 if let Some(limit) = limit {
8059 if limit <= 1e-6 {
8060 return Err(format!(
8061 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8062 )
8063 .into());
8064 }
8065 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8066 let cfg = LaunchConfig::for_num_elems(n as u32);
8067 let (np, ns) = (n_per as i32, n_sel as i32);
8068 let __s_b = self.gpu.stream();
8069 let mut b = __s_b.launch_builder(&f);
8070 b.arg(gate)
8071 .arg(up)
8072 .arg(gmac)
8073 .arg(umac)
8074 .arg(sel)
8075 .arg(&limit)
8076 .arg(out_q)
8077 .arg(out_d)
8078 .arg(&np)
8079 .arg(&ns);
8080 unsafe {
8081 b.launch(cfg)?;
8082 }
8083 return Ok(());
8084 }
8085 let f = self.func("silu_mul_scaled_q8_1_sel");
8086 let cfg = LaunchConfig::for_num_elems(n as u32);
8087 let (np, ns) = (n_per as i32, n_sel as i32);
8088 let __s_b = self.gpu.stream();
8089 let mut b = __s_b.launch_builder(&f);
8090 b.arg(gate)
8091 .arg(up)
8092 .arg(gmac)
8093 .arg(umac)
8094 .arg(sel)
8095 .arg(out_q)
8096 .arg(out_d)
8097 .arg(&np)
8098 .arg(&ns);
8099 unsafe {
8100 b.launch(cfg)?;
8101 }
8102 Ok(())
8103 }
8104
8105 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8106 Ok(self.gpu.stream().clone_htod(v)?)
8107 }
8108 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8109 Ok(self.gpu.stream().clone_htod(v)?)
8110 }
8111 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8113 Ok(self.gpu.stream().clone_htod(v)?)
8114 }
8115 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8116 Ok(self.gpu.stream().clone_htod(v)?)
8117 }
8118 pub fn dtoh_view(
8120 &self,
8121 d: &cudarc::driver::CudaView<f32>,
8122 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8123 let v = self.gpu.stream().clone_dtoh(d)?;
8124 self.gpu.stream().synchronize()?;
8125 Ok(v)
8126 }
8127 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8128 let v = self.gpu.stream().clone_dtoh(d)?;
8129 self.gpu.stream().synchronize()?;
8130 Ok(v)
8131 }
8132 pub fn dtoh_pair(
8136 &self,
8137 a: &CudaSlice<f32>,
8138 b: &CudaSlice<f32>,
8139 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8140 let av = self.gpu.stream().clone_dtoh(a)?;
8141 let bv = self.gpu.stream().clone_dtoh(b)?;
8142 self.gpu.stream().synchronize()?;
8143 Ok((av, bv))
8144 }
8145 pub fn dtoh_pair_views(
8148 &self,
8149 a: &cudarc::driver::CudaView<f32>,
8150 b: &cudarc::driver::CudaView<f32>,
8151 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8152 let av = self.gpu.stream().clone_dtoh(a)?;
8153 let bv = self.gpu.stream().clone_dtoh(b)?;
8154 self.gpu.stream().synchronize()?;
8155 Ok((av, bv))
8156 }
8157 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8159 let v = self.gpu.stream().clone_dtoh(d)?;
8160 self.gpu.stream().synchronize()?;
8161 Ok(v)
8162 }
8163 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8165 let v = self.gpu.stream().clone_dtoh(d)?;
8166 self.gpu.stream().synchronize()?;
8167 Ok(v)
8168 }
8169 pub fn dtoh_u8_view(
8170 &self,
8171 d: &cudarc::driver::CudaView<u8>,
8172 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8173 let v = self.gpu.stream().clone_dtoh(d)?;
8174 self.gpu.stream().synchronize()?;
8175 Ok(v)
8176 }
8177 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8178 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8179 self.keep_if_capturing(&s);
8180 Ok(s)
8181 }
8182
8183 pub fn prob_of_token_device(
8192 &self,
8193 logits: &CudaSlice<f32>,
8194 tok: &CudaSlice<u32>,
8195 n_vocab: usize,
8196 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8197 let nb = ARGMAX_NB;
8198 let mut part = self.alloc_uninit::<f32>(nb)?;
8199 let mut p = self.alloc_uninit::<f32>(1)?;
8200 let f1 = self.func("prob_of_token_partial_f32");
8201 let cfg1 = LaunchConfig {
8202 grid_dim: (nb as u32, 1, 1),
8203 block_dim: (256, 1, 1),
8204 shared_mem_bytes: 0,
8205 };
8206 let nv = n_vocab as i32;
8207 let __s_b1 = self.gpu.stream();
8208 let mut b1 = __s_b1.launch_builder(&f1);
8209 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8210 unsafe {
8211 b1.launch(cfg1)?;
8212 }
8213 let f2 = self.func("prob_of_token_final_f32");
8214 let cfg2 = LaunchConfig {
8215 grid_dim: (1, 1, 1),
8216 block_dim: (256, 1, 1),
8217 shared_mem_bytes: 0,
8218 };
8219 let nbi = nb as i32;
8220 let __s_b2 = self.gpu.stream();
8221 let mut b2 = __s_b2.launch_builder(&f2);
8222 b2.arg(&part).arg(&mut p).arg(&nbi);
8223 unsafe {
8224 b2.launch(cfg2)?;
8225 }
8226 Ok(p)
8227 }
8228
8229 pub fn prob_of_token_device_col(
8236 &self,
8237 logits: &CudaSlice<f32>,
8238 tok_all: &CudaSlice<u32>,
8239 tok_idx: usize,
8240 p_out: &mut CudaSlice<f32>,
8241 p_idx: usize,
8242 n_vocab: usize,
8243 ) -> Result<(), Box<dyn std::error::Error>> {
8244 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8245 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8246 let nb = ARGMAX_NB;
8247 let mut part = self.alloc_uninit::<f32>(nb)?;
8248 let f1 = self.func("prob_of_token_partial_f32");
8249 let cfg1 = LaunchConfig {
8250 grid_dim: (nb as u32, 1, 1),
8251 block_dim: (256, 1, 1),
8252 shared_mem_bytes: 0,
8253 };
8254 let nv = n_vocab as i32;
8255 let __s_b1 = self.gpu.stream();
8256 let mut b1 = __s_b1.launch_builder(&f1);
8257 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8258 unsafe {
8259 b1.launch(cfg1)?;
8260 }
8261 let f2 = self.func("prob_of_token_final_f32");
8262 let cfg2 = LaunchConfig {
8263 grid_dim: (1, 1, 1),
8264 block_dim: (256, 1, 1),
8265 shared_mem_bytes: 0,
8266 };
8267 let nbi = nb as i32;
8268 let __s_b2 = self.gpu.stream();
8269 let mut b2 = __s_b2.launch_builder(&f2);
8270 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8271 unsafe {
8272 b2.launch(cfg2)?;
8273 }
8274 Ok(())
8275 }
8276
8277 pub fn prob_of_token_device_into(
8278 &self,
8279 logits: &CudaSlice<f32>,
8280 tok: &CudaSlice<u32>,
8281 p_out: &mut CudaSlice<f32>,
8282 n_vocab: usize,
8283 ) -> Result<(), Box<dyn std::error::Error>> {
8284 let nb = ARGMAX_NB;
8285 let mut part = self.alloc_uninit::<f32>(nb)?;
8286 let f1 = self.func("prob_of_token_partial_f32");
8287 let cfg1 = LaunchConfig {
8288 grid_dim: (nb as u32, 1, 1),
8289 block_dim: (256, 1, 1),
8290 shared_mem_bytes: 0,
8291 };
8292 let nv = n_vocab as i32;
8293 let __s_b1 = self.gpu.stream();
8294 let mut b1 = __s_b1.launch_builder(&f1);
8295 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8296 unsafe {
8297 b1.launch(cfg1)?;
8298 }
8299 let f2 = self.func("prob_of_token_final_f32");
8300 let cfg2 = LaunchConfig {
8301 grid_dim: (1, 1, 1),
8302 block_dim: (256, 1, 1),
8303 shared_mem_bytes: 0,
8304 };
8305 let nbi = nb as i32;
8306 let __s_b2 = self.gpu.stream();
8307 let mut b2 = __s_b2.launch_builder(&f2);
8308 b2.arg(&part).arg(p_out).arg(&nbi);
8309 unsafe {
8310 b2.launch(cfg2)?;
8311 }
8312 Ok(())
8313 }
8314
8315 pub fn u32_hist_append(
8318 &self,
8319 tok: &CudaSlice<u32>,
8320 hist: &mut CudaSlice<u32>,
8321 idx: &mut CudaSlice<i32>,
8322 ) -> Result<(), Box<dyn std::error::Error>> {
8323 let f = self.func("u32_hist_append");
8324 let cfg = LaunchConfig {
8325 grid_dim: (1, 1, 1),
8326 block_dim: (32, 1, 1),
8327 shared_mem_bytes: 0,
8328 };
8329 let __s_b = self.gpu.stream();
8330 let mut b = __s_b.launch_builder(&f);
8331 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
8332 unsafe {
8333 b.launch(cfg)?;
8334 }
8335 Ok(())
8336 }
8337
8338 pub fn argmax_token_device(
8339 &self,
8340 logits: &CudaSlice<f32>,
8341 n_vocab: usize,
8342 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8343 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
8344 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
8345 Ok(tok)
8346 }
8347 pub fn argmax_token_device_into(
8354 &self,
8355 logits: &CudaSlice<f32>,
8356 tok: &mut CudaSlice<u32>,
8357 n_vocab: usize,
8358 ) -> Result<(), Box<dyn std::error::Error>> {
8359 let nb = ARGMAX_NB;
8360 let f1 = self.func("argmax_partial_f32");
8361 let f2 = self.func("argmax_final_f32");
8362 let mut guard = self.argmax_partials.lock().unwrap();
8363 if guard.is_none() {
8364 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8367 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8368 *guard = Some((pv, pi));
8369 }
8370 let (part_v, part_i) = guard.as_mut().unwrap();
8371 let nv = n_vocab as i32;
8372 let nbi = nb as i32;
8373 let cfg1 = LaunchConfig {
8375 grid_dim: (nb as u32, 1, 1),
8376 block_dim: (256, 1, 1),
8377 shared_mem_bytes: 0,
8378 };
8379 let __s_b1 = self.gpu.stream();
8380 let mut b1 = __s_b1.launch_builder(&f1);
8381 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8382 unsafe {
8383 b1.launch(cfg1)?;
8384 }
8385 let cfg2 = LaunchConfig {
8387 grid_dim: (1, 1, 1),
8388 block_dim: (256, 1, 1),
8389 shared_mem_bytes: 0,
8390 };
8391 let __s_b2 = self.gpu.stream();
8392 let mut b2 = __s_b2.launch_builder(&f2);
8393 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8394 unsafe {
8395 b2.launch(cfg2)?;
8396 }
8397 Ok(())
8398 }
8399 pub fn argmax_token_device_col(
8405 &self,
8406 logits: &CudaSlice<f32>,
8407 col: usize,
8408 n_vocab: usize,
8409 toks: &mut CudaSlice<u32>,
8410 out_idx: usize,
8411 ) -> Result<(), Box<dyn std::error::Error>> {
8412 let nb = ARGMAX_NB;
8413 let f1 = self.func("argmax_partial_f32");
8414 let f2 = self.func("argmax_final_f32");
8415 let mut guard = self.argmax_partials.lock().unwrap();
8416 if guard.is_none() {
8417 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8418 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8419 *guard = Some((pv, pi));
8420 }
8421 let (part_v, part_i) = guard.as_mut().unwrap();
8422 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8423 let nv = n_vocab as i32;
8424 let nbi = nb as i32;
8425 let cfg1 = LaunchConfig {
8426 grid_dim: (nb as u32, 1, 1),
8427 block_dim: (256, 1, 1),
8428 shared_mem_bytes: 0,
8429 };
8430 let __s_b1 = self.gpu.stream();
8431 let mut b1 = __s_b1.launch_builder(&f1);
8432 b1.arg(&col_view)
8433 .arg(&mut *part_v)
8434 .arg(&mut *part_i)
8435 .arg(&nv);
8436 unsafe {
8437 b1.launch(cfg1)?;
8438 }
8439 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8440 let cfg2 = LaunchConfig {
8441 grid_dim: (1, 1, 1),
8442 block_dim: (256, 1, 1),
8443 shared_mem_bytes: 0,
8444 };
8445 let __s_b2 = self.gpu.stream();
8446 let mut b2 = __s_b2.launch_builder(&f2);
8447 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8448 unsafe {
8449 b2.launch(cfg2)?;
8450 }
8451 Ok(())
8452 }
8453 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8455 Ok(self.gpu.stream().clone_htod(v)?)
8456 }
8457 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
8458 let v = self.gpu.stream().clone_dtoh(d)?;
8459 self.gpu.stream().synchronize()?;
8460 Ok(v)
8461 }
8462
8463 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8464 let v = self.gpu.stream().clone_dtoh(d)?;
8465 self.gpu.stream().synchronize()?;
8466 Ok(v)
8467 }
8468 pub fn htod_u32_into(
8472 &self,
8473 dst: &mut CudaSlice<u32>,
8474 src: &[u32],
8475 ) -> Result<(), Box<dyn std::error::Error>> {
8476 let mut view = dst.slice_mut(0..src.len());
8477 self.gpu.stream().memcpy_htod(src, &mut view)?;
8478 Ok(())
8479 }
8480
8481 pub fn htod_i32_into(
8484 &self,
8485 dst: &mut CudaSlice<i32>,
8486 src: &[i32],
8487 ) -> Result<(), Box<dyn std::error::Error>> {
8488 let mut view = dst.slice_mut(0..src.len());
8489 self.gpu.stream().memcpy_htod(src, &mut view)?;
8490 Ok(())
8491 }
8492
8493 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8494 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8495 self.keep_if_capturing(&s);
8496 Ok(s)
8497 }
8498 pub fn embed_gather_device_into(
8501 &self,
8502 embd: &CudaSlice<u8>,
8503 token_d: &CudaSlice<u32>,
8504 x_out: &mut CudaSlice<f32>,
8505 n_embd: usize,
8506 qtype: i32,
8507 row_bytes: usize,
8508 ) -> Result<(), Box<dyn std::error::Error>> {
8509 let f = self.func("embed_gather_u32");
8510 let cfg = LaunchConfig {
8511 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8512 block_dim: (256, 1, 1),
8513 shared_mem_bytes: 0,
8514 };
8515 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8516 let __s_b = self.gpu.stream();
8517 let mut b = __s_b.launch_builder(&f);
8518 b.arg(embd)
8519 .arg(token_d)
8520 .arg(x_out)
8521 .arg(&ne)
8522 .arg(&qt)
8523 .arg(&rb);
8524 unsafe {
8525 b.launch(cfg)?;
8526 }
8527 Ok(())
8528 }
8529 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8531 let v = self.gpu.stream().clone_dtoh(d)?;
8532 self.gpu.stream().synchronize()?;
8533 Ok(v[0])
8534 }
8535 pub fn i32_set_k(
8542 &self,
8543 dst: &mut CudaSlice<i32>,
8544 v: i32,
8545 ) -> Result<(), Box<dyn std::error::Error>> {
8546 let f = self.func("i32_set_k");
8547 let cfg = LaunchConfig {
8548 grid_dim: (1, 1, 1),
8549 block_dim: (1, 1, 1),
8550 shared_mem_bytes: 0,
8551 };
8552 let idx = 0i32;
8553 let __s_b = self.gpu.stream();
8554 let mut b = __s_b.launch_builder(&f);
8555 b.arg(dst).arg(&v).arg(&idx);
8556 unsafe {
8557 b.launch(cfg)?;
8558 }
8559 Ok(())
8560 }
8561
8562 pub fn set_i32_one(
8563 &self,
8564 d: &mut CudaSlice<i32>,
8565 v: i32,
8566 ) -> Result<(), Box<dyn std::error::Error>> {
8567 self.gpu.stream().memcpy_htod(&[v], d)?;
8568 Ok(())
8569 }
8570 pub fn set_u32_one(
8573 &self,
8574 d: &mut CudaSlice<u32>,
8575 v: u32,
8576 ) -> Result<(), Box<dyn std::error::Error>> {
8577 self.gpu.stream().memcpy_htod(&[v], d)?;
8578 Ok(())
8579 }
8580 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8582 let v = self.gpu.stream().clone_dtoh(d)?;
8583 self.gpu.stream().synchronize()?;
8584 Ok(v[0])
8585 }
8586 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8588 Ok(self.gpu.stream().clone_htod(bytes)?)
8589 }
8590 pub fn embed_gather_device(
8594 &self,
8595 embd: &CudaSlice<u8>,
8596 token_d: &CudaSlice<u32>,
8597 n_embd: usize,
8598 qtype: i32,
8599 row_bytes: usize,
8600 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8601 let f = self.func("embed_gather_u32");
8602 let mut x = self.alloc_uninit::<f32>(n_embd)?;
8603 let cfg = LaunchConfig {
8604 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8605 block_dim: (256, 1, 1),
8606 shared_mem_bytes: 0,
8607 };
8608 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8609 let __s_b = self.gpu.stream();
8610 let mut b = __s_b.launch_builder(&f);
8611 b.arg(embd)
8612 .arg(token_d)
8613 .arg(&mut x)
8614 .arg(&ne)
8615 .arg(&qt)
8616 .arg(&rb);
8617 unsafe {
8618 b.launch(cfg)?;
8619 }
8620 Ok(x)
8621 }
8622
8623 pub fn embed_gather_device_t(
8627 &self,
8628 embd: &CudaSlice<u8>,
8629 tokens: &[u32],
8630 n_embd: usize,
8631 qtype: i32,
8632 row_bytes: usize,
8633 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8634 let t = tokens.len();
8635 let tok_d = self.gpu.stream().clone_htod(tokens)?;
8636 let f = self.func("embed_gather_u32_t");
8637 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8638 let cfg = LaunchConfig {
8639 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8640 block_dim: (256, 1, 1),
8641 shared_mem_bytes: 0,
8642 };
8643 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8644 let __s_b = self.gpu.stream();
8645 let mut b = __s_b.launch_builder(&f);
8646 b.arg(embd)
8647 .arg(&tok_d)
8648 .arg(&mut x)
8649 .arg(&ne)
8650 .arg(&qt)
8651 .arg(&rb)
8652 .arg(&ti);
8653 unsafe {
8654 b.launch(cfg)?;
8655 }
8656 Ok(x)
8657 }
8658
8659 pub fn embed_gather_device_tv(
8664 &self,
8665 embd: &CudaSlice<u8>,
8666 tok_v: &cudarc::driver::CudaView<u32>,
8667 t: usize,
8668 n_embd: usize,
8669 qtype: i32,
8670 row_bytes: usize,
8671 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8672 let f = self.func("embed_gather_u32_t");
8673 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8674 let cfg = LaunchConfig {
8675 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8676 block_dim: (256, 1, 1),
8677 shared_mem_bytes: 0,
8678 };
8679 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8680 let __s_b = self.gpu.stream();
8681 let mut b = __s_b.launch_builder(&f);
8682 b.arg(embd)
8683 .arg(tok_v)
8684 .arg(&mut x)
8685 .arg(&ne)
8686 .arg(&qt)
8687 .arg(&rb)
8688 .arg(&ti);
8689 unsafe {
8690 b.launch(cfg)?;
8691 }
8692 Ok(x)
8693 }
8694
8695 pub fn embed_gather_device_td(
8696 &self,
8697 embd: &CudaSlice<u8>,
8698 tok_d: &CudaSlice<u32>,
8699 t: usize,
8700 n_embd: usize,
8701 qtype: i32,
8702 row_bytes: usize,
8703 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8704 let f = self.func("embed_gather_u32_t");
8705 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8706 let cfg = LaunchConfig {
8707 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8708 block_dim: (256, 1, 1),
8709 shared_mem_bytes: 0,
8710 };
8711 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8712 let __s_b = self.gpu.stream();
8713 let mut b = __s_b.launch_builder(&f);
8714 b.arg(embd)
8715 .arg(tok_d)
8716 .arg(&mut x)
8717 .arg(&ne)
8718 .arg(&qt)
8719 .arg(&rb)
8720 .arg(&ti);
8721 unsafe {
8722 b.launch(cfg)?;
8723 }
8724 Ok(x)
8725 }
8726
8727 #[inline]
8733 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8735 if self
8736 .capture_keep_on
8737 .load(std::sync::atomic::Ordering::Relaxed)
8738 {
8739 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8740 }
8741 }
8742
8743 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8744 &self,
8745 n: usize,
8746 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8747 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8748 {
8752 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8753 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8754 use cudarc::driver::DevicePtrMut;
8756 let n_bytes = s.len() * std::mem::size_of::<T>();
8757 let stream = self.gpu.stream();
8758 let (p_, _g) = s.device_ptr_mut(&stream);
8759 unsafe {
8760 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8761 .result()?;
8762 }
8763 }
8764 }
8765 self.keep_if_capturing(&s);
8766 Ok(s)
8767 }
8768
8769 pub fn uninit_q8_pair(
8774 &self,
8775 n: usize,
8776 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8777 Ok((
8778 self.alloc_uninit::<i8>(n)?,
8779 self.alloc_uninit::<f32>(n / 32)?,
8780 ))
8781 }
8782
8783 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8784 self.alloc_uninit::<f32>(n)
8785 }
8786
8787 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8789 self.alloc_uninit::<i8>(n)
8790 }
8791
8792 #[allow(clippy::too_many_arguments)]
8796 pub fn rms_norm3(
8797 &self,
8798 x: &CudaSlice<f32>,
8799 w0: &CudaSlice<f32>,
8800 w1: &CudaSlice<f32>,
8801 w2: &CudaSlice<f32>,
8802 d0: &mut CudaSlice<f32>,
8803 d1: &mut CudaSlice<f32>,
8804 d2: &mut CudaSlice<f32>,
8805 ncols: usize,
8806 nrows: usize,
8807 eps: f32,
8808 ) -> Result<(), Box<dyn std::error::Error>> {
8809 let f = self.func("rms_norm3_f32");
8810 let cfg = LaunchConfig {
8811 grid_dim: (nrows as u32, 1, 1),
8812 block_dim: (rms_block(), 1, 1),
8813 shared_mem_bytes: 0,
8814 };
8815 let (nc, e) = (ncols as i32, eps);
8816 let __s_b = self.gpu.stream();
8817 let mut b = __s_b.launch_builder(&f);
8818 b.arg(x)
8819 .arg(w0)
8820 .arg(w1)
8821 .arg(w2)
8822 .arg(d0)
8823 .arg(d1)
8824 .arg(d2)
8825 .arg(&nc)
8826 .arg(&e);
8827 unsafe {
8828 b.launch(cfg)?;
8829 }
8830 Ok(())
8831 }
8832
8833 #[allow(clippy::too_many_arguments)]
8835 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8838 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8839 *WARP_ON.get_or_init(|| {
8840 std::env::var("MEMRA_QKVNORM_W")
8841 .map(|v| v != "0")
8842 .unwrap_or(true)
8843 }) && ncols % 4 == 0
8844 && rows >= 64
8845 }
8846
8847 #[allow(clippy::too_many_arguments)]
8850 pub fn rms_norm_qkv_w4b(
8851 &self,
8852 q: &CudaSlice<f32>,
8853 k: &CudaSlice<f32>,
8854 v: &CudaSlice<f32>,
8855 wq: &CudaSlice<f32>,
8856 wk: &CudaSlice<f32>,
8857 wv: &CudaSlice<f32>,
8858 dq: &mut CudaSlice<f32>,
8859 dk: &mut CudaSlice<f32>,
8860 dv: &mut CudaSlice<f32>,
8861 dvb: &mut CudaSlice<u8>,
8862 ncols: usize,
8863 rq: usize,
8864 rk: usize,
8865 eps: f32,
8866 vf16: bool,
8867 ) -> Result<(), Box<dyn std::error::Error>> {
8868 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8869 let f = self.func("rms_norm_qkv_w4b_f32");
8870 let rows = (rq + 2 * rk) as u32;
8871 let cfg = LaunchConfig {
8872 grid_dim: (rows.div_ceil(8), 1, 1),
8873 block_dim: (256, 1, 1),
8874 shared_mem_bytes: 0,
8875 };
8876 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8877 let vf = vf16 as i32;
8878 let __s_b = self.gpu.stream();
8879 let mut b = __s_b.launch_builder(&f);
8880 b.arg(q)
8881 .arg(k)
8882 .arg(v)
8883 .arg(wq)
8884 .arg(wk)
8885 .arg(wv)
8886 .arg(dq)
8887 .arg(dk)
8888 .arg(dv)
8889 .arg(&mut *dvb)
8890 .arg(&nc)
8891 .arg(&rqi)
8892 .arg(&rki)
8893 .arg(&rvi)
8894 .arg(&e)
8895 .arg(&vf);
8896 unsafe {
8897 b.launch(cfg)?;
8898 }
8899 Ok(())
8900 }
8901
8902 pub fn rms_norm_qkv(
8903 &self,
8904 q: &CudaSlice<f32>,
8905 k: &CudaSlice<f32>,
8906 v: &CudaSlice<f32>,
8907 wq: &CudaSlice<f32>,
8908 wk: &CudaSlice<f32>,
8909 wv: &CudaSlice<f32>,
8910 dq: &mut CudaSlice<f32>,
8911 dk: &mut CudaSlice<f32>,
8912 dv: &mut CudaSlice<f32>,
8913 ncols: usize,
8914 rq: usize,
8915 rk: usize,
8916 eps: f32,
8917 ) -> Result<(), Box<dyn std::error::Error>> {
8918 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8922 let warp_on = *WARP_ON.get_or_init(|| {
8923 std::env::var("MEMRA_QKVNORM_W")
8924 .map(|v| v != "0")
8925 .unwrap_or(true)
8926 });
8927 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8930 let f = self.func("rms_norm_qkv_w4_f32");
8931 let rows = (rq + 2 * rk) as u32;
8932 let cfg = LaunchConfig {
8933 grid_dim: (rows.div_ceil(8), 1, 1),
8934 block_dim: (256, 1, 1),
8935 shared_mem_bytes: 0,
8936 };
8937 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8938 let __s_b = self.gpu.stream();
8939 let mut b = __s_b.launch_builder(&f);
8940 b.arg(q)
8941 .arg(k)
8942 .arg(v)
8943 .arg(wq)
8944 .arg(wk)
8945 .arg(wv)
8946 .arg(dq)
8947 .arg(dk)
8948 .arg(dv)
8949 .arg(&nc)
8950 .arg(&rqi)
8951 .arg(&rki)
8952 .arg(&rvi)
8953 .arg(&e);
8954 unsafe {
8955 b.launch(cfg)?;
8956 }
8957 return Ok(());
8958 }
8959 let f = self.func("rms_norm_qkv_f32");
8960 let grid = (rq + 2 * rk) as u32;
8961 let cfg = LaunchConfig {
8962 grid_dim: (grid, 1, 1),
8963 block_dim: (rms_block(), 1, 1),
8964 shared_mem_bytes: 0,
8965 };
8966 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8967 let __s_b = self.gpu.stream();
8968 let mut b = __s_b.launch_builder(&f);
8969 b.arg(q)
8970 .arg(k)
8971 .arg(v)
8972 .arg(wq)
8973 .arg(wk)
8974 .arg(wv)
8975 .arg(dq)
8976 .arg(dk)
8977 .arg(dv)
8978 .arg(&nc)
8979 .arg(&rqi)
8980 .arg(&rki)
8981 .arg(&e);
8982 unsafe {
8983 b.launch(cfg)?;
8984 }
8985 Ok(())
8986 }
8987
8988 #[allow(clippy::too_many_arguments)]
8990 pub fn rms_norm2x(
8991 &self,
8992 a: &CudaSlice<f32>,
8993 bb: &CudaSlice<f32>,
8994 wa: &CudaSlice<f32>,
8995 wb: &CudaSlice<f32>,
8996 da: &mut CudaSlice<f32>,
8997 db: &mut CudaSlice<f32>,
8998 ncols: usize,
8999 nrows: usize,
9000 eps: f32,
9001 ) -> Result<(), Box<dyn std::error::Error>> {
9002 let f = self.func("rms_norm2x_f32");
9003 let cfg = LaunchConfig {
9004 grid_dim: (2 * nrows as u32, 1, 1),
9005 block_dim: (rms_block(), 1, 1),
9006 shared_mem_bytes: 0,
9007 };
9008 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9009 let __s_b = self.gpu.stream();
9010 let mut b = __s_b.launch_builder(&f);
9011 b.arg(a)
9012 .arg(bb)
9013 .arg(wa)
9014 .arg(wb)
9015 .arg(da)
9016 .arg(db)
9017 .arg(&nc)
9018 .arg(&nr)
9019 .arg(&e);
9020 unsafe {
9021 b.launch(cfg)?;
9022 }
9023 Ok(())
9024 }
9025
9026 pub fn softcap(
9028 &self,
9029 y: &mut CudaSlice<f32>,
9030 cap: f32,
9031 n: usize,
9032 ) -> Result<(), Box<dyn std::error::Error>> {
9033 let f = self.func("softcap_f32");
9034 let cfg = LaunchConfig::for_num_elems(n as u32);
9035 let ni = n as i32;
9036 let __s_b = self.gpu.stream();
9037 let mut b = __s_b.launch_builder(&f);
9038 b.arg(y).arg(&cap).arg(&ni);
9039 unsafe {
9040 b.launch(cfg)?;
9041 }
9042 Ok(())
9043 }
9044
9045 pub fn mask_ids_rows(
9048 &self,
9049 y: &mut CudaSlice<f32>,
9050 ids: &CudaSlice<i32>,
9051 n_ids: usize,
9052 n_vocab: usize,
9053 t: usize,
9054 ) -> Result<(), Box<dyn std::error::Error>> {
9055 let f = self.func("mask_ids_rows_f32");
9056 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9057 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9058 let __s_b = self.gpu.stream();
9059 let mut b = __s_b.launch_builder(&f);
9060 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9061 unsafe {
9062 b.launch(cfg)?;
9063 }
9064 Ok(())
9065 }
9066
9067 #[allow(clippy::too_many_arguments)]
9069 pub fn add_scale_rms_norm(
9070 &self,
9071 a: &CudaSlice<f32>,
9072 b_in: &CudaSlice<f32>,
9073 c: f32,
9074 w: &CudaSlice<f32>,
9075 res: &mut CudaSlice<f32>,
9076 dst: &mut CudaSlice<f32>,
9077 ncols: usize,
9078 nrows: usize,
9079 eps: f32,
9080 ) -> Result<(), Box<dyn std::error::Error>> {
9081 let f = self.func("add_scale_rms_norm_f32");
9082 let cfg = LaunchConfig {
9083 grid_dim: (nrows as u32, 1, 1),
9084 block_dim: (rms_block(), 1, 1),
9085 shared_mem_bytes: 0,
9086 };
9087 let (nc, e2) = (ncols as i32, eps);
9088 let __s_b = self.gpu.stream();
9089 let mut b = __s_b.launch_builder(&f);
9090 b.arg(a)
9091 .arg(b_in)
9092 .arg(&c)
9093 .arg(w)
9094 .arg(res)
9095 .arg(dst)
9096 .arg(&nc)
9097 .arg(&e2);
9098 unsafe {
9099 b.launch(cfg)?;
9100 }
9101 Ok(())
9102 }
9103
9104 #[allow(clippy::too_many_arguments)]
9107 pub fn add_scale_rms_norm_q8_1(
9108 &self,
9109 a: &CudaSlice<f32>,
9110 b_in: &CudaSlice<f32>,
9111 c: f32,
9112 w: &CudaSlice<f32>,
9113 res: &mut CudaSlice<f32>,
9114 ncols: usize,
9115 nrows: usize,
9116 eps: f32,
9117 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9118 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9119 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9120 let (nc, e2) = (ncols as i32, eps);
9121 if Self::pdl_on() && Self::pdl_wb_on() {
9122 {
9123 use cudarc::driver::{DevicePtr, DevicePtrMut};
9124 let s = &self.gpu.stream();
9125 let (pa, _g0) = a.device_ptr(s);
9126 let (pb, _g1) = b_in.device_ptr(s);
9127 let (pw, _g2) = w.device_ptr(s);
9128 let (pr, _g3) = res.device_ptr_mut(s);
9129 let (pq, _g4) = out_q.device_ptr_mut(s);
9130 let (pd, _g5) = out_d.device_ptr_mut(s);
9131 let mut ps = [
9132 &pa as *const _ as *mut std::ffi::c_void,
9133 &pb as *const _ as *mut _,
9134 &c as *const _ as *mut _,
9135 &pw as *const _ as *mut _,
9136 &pr as *const _ as *mut _,
9137 &pq as *const _ as *mut _,
9138 &pd as *const _ as *mut _,
9139 &nc as *const _ as *mut _,
9140 &e2 as *const _ as *mut _,
9141 ];
9142 unsafe {
9143 self.launch_pdl(
9144 "add_scale_rms_norm_q8_1",
9145 (nrows as u32, 1, 1),
9146 (rms_block(), 1, 1),
9147 &mut ps,
9148 )?;
9149 }
9150 }
9151 return Ok((out_q, out_d));
9152 }
9153 let f = self.func("add_scale_rms_norm_q8_1");
9154 let cfg = LaunchConfig {
9155 grid_dim: (nrows as u32, 1, 1),
9156 block_dim: (rms_block(), 1, 1),
9157 shared_mem_bytes: 0,
9158 };
9159 let __s_b = self.gpu.stream();
9160 let mut b = __s_b.launch_builder(&f);
9161 b.arg(a)
9162 .arg(b_in)
9163 .arg(&c)
9164 .arg(w)
9165 .arg(res)
9166 .arg(&mut out_q)
9167 .arg(&mut out_d)
9168 .arg(&nc)
9169 .arg(&e2);
9170 unsafe {
9171 b.launch(cfg)?;
9172 }
9173 Ok((out_q, out_d))
9174 }
9175
9176 #[allow(clippy::too_many_arguments)]
9178 pub fn add_scale_rms_norm_q8_1_into(
9179 &self,
9180 a: &CudaSlice<f32>,
9181 b_in: &CudaSlice<f32>,
9182 c: f32,
9183 w: &CudaSlice<f32>,
9184 res: &mut CudaSlice<f32>,
9185 ncols: usize,
9186 nrows: usize,
9187 eps: f32,
9188 out_q: &mut CudaSlice<i8>,
9189 out_d: &mut CudaSlice<f32>,
9190 ) -> Result<(), Box<dyn std::error::Error>> {
9191 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9192 let (nc, e2) = (ncols as i32, eps);
9193 if Self::pdl_on() && Self::pdl_wb_on() {
9194 use cudarc::driver::{DevicePtr, DevicePtrMut};
9195 let s = &self.gpu.stream();
9196 let (pa, _g0) = a.device_ptr(s);
9197 let (pb, _g1) = b_in.device_ptr(s);
9198 let (pw, _g2) = w.device_ptr(s);
9199 let (pr, _g3) = res.device_ptr_mut(s);
9200 let (pq, _g4) = out_q.device_ptr_mut(s);
9201 let (pd, _g5) = out_d.device_ptr_mut(s);
9202 let mut ps = [
9203 &pa as *const _ as *mut std::ffi::c_void,
9204 &pb as *const _ as *mut _,
9205 &c as *const _ as *mut _,
9206 &pw as *const _ as *mut _,
9207 &pr as *const _ as *mut _,
9208 &pq as *const _ as *mut _,
9209 &pd as *const _ as *mut _,
9210 &nc as *const _ as *mut _,
9211 &e2 as *const _ as *mut _,
9212 ];
9213 unsafe {
9214 self.launch_pdl(
9215 "add_scale_rms_norm_q8_1",
9216 (nrows as u32, 1, 1),
9217 (rms_block(), 1, 1),
9218 &mut ps,
9219 )?;
9220 }
9221 return Ok(());
9222 }
9223 let f = self.func("add_scale_rms_norm_q8_1");
9224 let cfg = LaunchConfig {
9225 grid_dim: (nrows as u32, 1, 1),
9226 block_dim: (rms_block(), 1, 1),
9227 shared_mem_bytes: 0,
9228 };
9229 let __s_b = self.gpu.stream();
9230 let mut b = __s_b.launch_builder(&f);
9231 b.arg(a)
9232 .arg(b_in)
9233 .arg(&c)
9234 .arg(w)
9235 .arg(res)
9236 .arg(&mut *out_q)
9237 .arg(&mut *out_d)
9238 .arg(&nc)
9239 .arg(&e2);
9240 unsafe {
9241 b.launch(cfg)?;
9242 }
9243 Ok(())
9244 }
9245
9246 #[allow(clippy::too_many_arguments)]
9249 pub fn rms_pre_add_scale_rms_norm_q8_1(
9250 &self,
9251 a: &CudaSlice<f32>,
9252 wa: &CudaSlice<f32>,
9253 b_in: &CudaSlice<f32>,
9254 c: f32,
9255 w: &CudaSlice<f32>,
9256 res: &mut CudaSlice<f32>,
9257 ncols: usize,
9258 nrows: usize,
9259 eps: f32,
9260 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9261 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9262 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9263 let (nc, e2) = (ncols as i32, eps);
9264 if Self::pdl_on() {
9265 {
9266 use cudarc::driver::{DevicePtr, DevicePtrMut};
9267 let s = &self.gpu.stream();
9268 let (pa, _g0) = a.device_ptr(s);
9269 let (pwa, _g1) = wa.device_ptr(s);
9270 let (pb, _g2) = b_in.device_ptr(s);
9271 let (pw, _g3) = w.device_ptr(s);
9272 let (pr, _g4) = res.device_ptr_mut(s);
9273 let (pq, _g5) = out_q.device_ptr_mut(s);
9274 let (pd, _g6) = out_d.device_ptr_mut(s);
9275 let mut ps = [
9276 &pa as *const _ as *mut std::ffi::c_void,
9277 &pwa as *const _ as *mut _,
9278 &pb as *const _ as *mut _,
9279 &c as *const _ as *mut _,
9280 &pw as *const _ as *mut _,
9281 &pr as *const _ as *mut _,
9282 &pq as *const _ as *mut _,
9283 &pd as *const _ as *mut _,
9284 &nc as *const _ as *mut _,
9285 &e2 as *const _ as *mut _,
9286 ];
9287 unsafe {
9288 self.launch_pdl(
9289 "rms_pre_add_scale_rms_norm_q8_1",
9290 (nrows as u32, 1, 1),
9291 (rms_block(), 1, 1),
9292 &mut ps,
9293 )?;
9294 }
9295 }
9296 return Ok((out_q, out_d));
9297 }
9298 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9299 let cfg = LaunchConfig {
9300 grid_dim: (nrows as u32, 1, 1),
9301 block_dim: (rms_block(), 1, 1),
9302 shared_mem_bytes: 0,
9303 };
9304 let __s_b = self.gpu.stream();
9305 let mut b = __s_b.launch_builder(&f);
9306 b.arg(a)
9307 .arg(wa)
9308 .arg(b_in)
9309 .arg(&c)
9310 .arg(w)
9311 .arg(res)
9312 .arg(&mut out_q)
9313 .arg(&mut out_d)
9314 .arg(&nc)
9315 .arg(&e2);
9316 unsafe {
9317 b.launch(cfg)?;
9318 }
9319 Ok((out_q, out_d))
9320 }
9321
9322 pub fn gelu_tanh_mul_q8_1(
9325 &self,
9326 gate: &CudaSlice<f32>,
9327 up: &cudarc::driver::CudaView<f32>,
9328 act: &mut CudaSlice<f32>,
9329 ncols: usize,
9330 nrows: usize,
9331 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9332 debug_assert!(ncols % 128 == 0);
9333 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9334 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9335 let nc = ncols as i32;
9336 if Self::pdl_on() {
9337 {
9338 use cudarc::driver::{DevicePtr, DevicePtrMut};
9339 let s = &self.gpu.stream();
9340 let (pg, _g0) = gate.device_ptr(s);
9341 let (pu, _g1) = up.device_ptr(s);
9342 let (pact, _g2) = act.device_ptr_mut(s);
9343 let (pq, _g3) = out_q.device_ptr_mut(s);
9344 let (pd, _g4) = out_d.device_ptr_mut(s);
9345 let mut ps = [
9346 &pg as *const _ as *mut std::ffi::c_void,
9347 &pu as *const _ as *mut _,
9348 &pact as *const _ as *mut _,
9349 &pq as *const _ as *mut _,
9350 &pd as *const _ as *mut _,
9351 &nc as *const _ as *mut _,
9352 ];
9353 unsafe {
9354 self.launch_pdl(
9355 "gelu_tanh_mul_q8_1",
9356 (nrows as u32, 1, 1),
9357 (rms_block(), 1, 1),
9358 &mut ps,
9359 )?;
9360 }
9361 }
9362 return Ok((out_q, out_d));
9363 }
9364 let f = self.func("gelu_tanh_mul_q8_1");
9365 let cfg = LaunchConfig {
9366 grid_dim: (nrows as u32, 1, 1),
9367 block_dim: (rms_block(), 1, 1),
9368 shared_mem_bytes: 0,
9369 };
9370 let __s_b = self.gpu.stream();
9371 let mut b = __s_b.launch_builder(&f);
9372 b.arg(gate)
9373 .arg(up)
9374 .arg(act)
9375 .arg(&mut out_q)
9376 .arg(&mut out_d)
9377 .arg(&nc);
9378 unsafe {
9379 b.launch(cfg)?;
9380 }
9381 Ok((out_q, out_d))
9382 }
9383
9384 #[allow(clippy::too_many_arguments)]
9386 pub fn gelu_tanh_mul_q8_1_into(
9387 &self,
9388 gate: &CudaSlice<f32>,
9389 up: &cudarc::driver::CudaView<f32>,
9390 act: &mut CudaSlice<f32>,
9391 ncols: usize,
9392 nrows: usize,
9393 out_q: &mut CudaSlice<i8>,
9394 out_d: &mut CudaSlice<f32>,
9395 ) -> Result<(), Box<dyn std::error::Error>> {
9396 debug_assert!(ncols % 128 == 0);
9397 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9398 let nc = ncols as i32;
9399 if Self::pdl_on() {
9400 use cudarc::driver::{DevicePtr, DevicePtrMut};
9401 let s = &self.gpu.stream();
9402 let (pg, _g0) = gate.device_ptr(s);
9403 let (pu, _g1) = up.device_ptr(s);
9404 let (pact, _g2) = act.device_ptr_mut(s);
9405 let (pq, _g3) = out_q.device_ptr_mut(s);
9406 let (pd, _g4) = out_d.device_ptr_mut(s);
9407 let mut ps = [
9408 &pg as *const _ as *mut std::ffi::c_void,
9409 &pu as *const _ as *mut _,
9410 &pact as *const _ as *mut _,
9411 &pq as *const _ as *mut _,
9412 &pd as *const _ as *mut _,
9413 &nc as *const _ as *mut _,
9414 ];
9415 unsafe {
9416 self.launch_pdl(
9417 "gelu_tanh_mul_q8_1",
9418 (nrows as u32, 1, 1),
9419 (rms_block(), 1, 1),
9420 &mut ps,
9421 )?;
9422 }
9423 return Ok(());
9424 }
9425 let f = self.func("gelu_tanh_mul_q8_1");
9426 let cfg = LaunchConfig {
9427 grid_dim: (nrows as u32, 1, 1),
9428 block_dim: (rms_block(), 1, 1),
9429 shared_mem_bytes: 0,
9430 };
9431 let __s_b = self.gpu.stream();
9432 let mut b = __s_b.launch_builder(&f);
9433 b.arg(gate)
9434 .arg(up)
9435 .arg(&mut *act)
9436 .arg(&mut *out_q)
9437 .arg(&mut *out_d)
9438 .arg(&nc);
9439 unsafe {
9440 b.launch(cfg)?;
9441 }
9442 Ok(())
9443 }
9444
9445 #[allow(clippy::too_many_arguments)]
9447 pub fn add_rms_norm3_q8z(
9448 &self,
9449 a: &CudaSlice<f32>,
9450 b_in: &CudaSlice<f32>,
9451 w0: &CudaSlice<f32>,
9452 w1: &CudaSlice<f32>,
9453 w2: &CudaSlice<f32>,
9454 res: &mut CudaSlice<f32>,
9455 out1: &mut CudaSlice<f32>,
9456 ncols: usize,
9457 nrows: usize,
9458 eps: f32,
9459 ) -> Result<
9460 (
9461 (CudaSlice<i8>, CudaSlice<f32>),
9462 (CudaSlice<i8>, CudaSlice<f32>),
9463 ),
9464 Box<dyn std::error::Error>,
9465 > {
9466 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9467 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9468 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9469 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9470 let f = self.func("add_rms_norm3_q8z_f32");
9471 let cfg = LaunchConfig {
9472 grid_dim: (nrows as u32, 1, 1),
9473 block_dim: (rms_block(), 1, 1),
9474 shared_mem_bytes: 0,
9475 };
9476 let (nc, e2) = (ncols as i32, eps);
9477 let __s_b = self.gpu.stream();
9478 let mut b = __s_b.launch_builder(&f);
9479 b.arg(a)
9480 .arg(b_in)
9481 .arg(w0)
9482 .arg(w1)
9483 .arg(w2)
9484 .arg(res)
9485 .arg(&mut q0)
9486 .arg(&mut d0)
9487 .arg(out1)
9488 .arg(&mut q2)
9489 .arg(&mut d2)
9490 .arg(&nc)
9491 .arg(&e2);
9492 unsafe {
9493 b.launch(cfg)?;
9494 }
9495 Ok(((q0, d0), (q2, d2)))
9496 }
9497
9498 #[allow(clippy::too_many_arguments)]
9500 pub fn add_rms_norm3(
9501 &self,
9502 a: &CudaSlice<f32>,
9503 b_in: &CudaSlice<f32>,
9504 w0: &CudaSlice<f32>,
9505 w1: &CudaSlice<f32>,
9506 w2: &CudaSlice<f32>,
9507 res: &mut CudaSlice<f32>,
9508 d0: &mut CudaSlice<f32>,
9509 d1: &mut CudaSlice<f32>,
9510 d2: &mut CudaSlice<f32>,
9511 ncols: usize,
9512 nrows: usize,
9513 eps: f32,
9514 ) -> Result<(), Box<dyn std::error::Error>> {
9515 let f = self.func("add_rms_norm3_f32");
9516 let cfg = LaunchConfig {
9517 grid_dim: (nrows as u32, 1, 1),
9518 block_dim: (rms_block(), 1, 1),
9519 shared_mem_bytes: 0,
9520 };
9521 let (nc, e2) = (ncols as i32, eps);
9522 let __s_b = self.gpu.stream();
9523 let mut b = __s_b.launch_builder(&f);
9524 b.arg(a)
9525 .arg(b_in)
9526 .arg(w0)
9527 .arg(w1)
9528 .arg(w2)
9529 .arg(res)
9530 .arg(d0)
9531 .arg(d1)
9532 .arg(d2)
9533 .arg(&nc)
9534 .arg(&e2);
9535 unsafe {
9536 b.launch(cfg)?;
9537 }
9538 Ok(())
9539 }
9540
9541 pub fn add_scale(
9543 &self,
9544 a: &CudaSlice<f32>,
9545 b_in: &CudaSlice<f32>,
9546 c: f32,
9547 dst: &mut CudaSlice<f32>,
9548 n: usize,
9549 ) -> Result<(), Box<dyn std::error::Error>> {
9550 let f = self.func("add_scale_f32");
9551 let cfg = LaunchConfig::for_num_elems(n as u32);
9552 let ni = n as i32;
9553 let __s_b = self.gpu.stream();
9554 let mut b = __s_b.launch_builder(&f);
9555 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9556 unsafe {
9557 b.launch(cfg)?;
9558 }
9559 Ok(())
9560 }
9561
9562 pub fn layer_norm_bias(
9564 &self,
9565 x: &CudaSlice<f32>,
9566 w: &CudaSlice<f32>,
9567 b: &CudaSlice<f32>,
9568 dst: &mut CudaSlice<f32>,
9569 ncols: usize,
9570 nrows: usize,
9571 eps: f32,
9572 ) -> Result<(), Box<dyn std::error::Error>> {
9573 let f = self.func("layer_norm_bias_f32");
9574 let (nc, e) = (ncols as i32, eps);
9575 let cfg = LaunchConfig {
9576 grid_dim: (nrows as u32, 1, 1),
9577 block_dim: (256, 1, 1),
9578 shared_mem_bytes: 0,
9579 };
9580 let __s_b = self.gpu.stream();
9581 let mut lb = __s_b.launch_builder(&f);
9582 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9583 unsafe {
9584 lb.launch(cfg)?;
9585 }
9586 Ok(())
9587 }
9588
9589 pub fn gelu_tanh(
9591 &self,
9592 x: &CudaSlice<f32>,
9593 dst: &mut CudaSlice<f32>,
9594 n: usize,
9595 ) -> Result<(), Box<dyn std::error::Error>> {
9596 let f = self.func("gelu_tanh_f32");
9597 let ni = n as i64;
9598 let cfg = LaunchConfig {
9599 grid_dim: (n.div_ceil(256) as u32, 1, 1),
9600 block_dim: (256, 1, 1),
9601 shared_mem_bytes: 0,
9602 };
9603 let __s_b = self.gpu.stream();
9604 let mut lb = __s_b.launch_builder(&f);
9605 lb.arg(x).arg(&mut *dst).arg(&ni);
9606 unsafe {
9607 lb.launch(cfg)?;
9608 }
9609 Ok(())
9610 }
9611
9612 pub fn row_softmax(
9614 &self,
9615 x: &mut CudaSlice<f32>,
9616 ncols: usize,
9617 nrows: usize,
9618 ) -> Result<(), Box<dyn std::error::Error>> {
9619 let f = self.func("row_softmax_f32");
9620 let nc = ncols as i32;
9621 let cfg = LaunchConfig {
9622 grid_dim: (nrows as u32, 1, 1),
9623 block_dim: (256, 1, 1),
9624 shared_mem_bytes: 0,
9625 };
9626 let __s_b = self.gpu.stream();
9627 let mut lb = __s_b.launch_builder(&f);
9628 lb.arg(&mut *x).arg(&nc);
9629 unsafe {
9630 lb.launch(cfg)?;
9631 }
9632 Ok(())
9633 }
9634
9635 pub fn rms_norm(
9636 &self,
9637 x: &CudaSlice<f32>,
9638 w: &CudaSlice<f32>,
9639 dst: &mut CudaSlice<f32>,
9640 ncols: usize,
9641 nrows: usize,
9642 eps: f32,
9643 ) -> Result<(), Box<dyn std::error::Error>> {
9644 let (nc, e) = (ncols as i32, eps);
9645 let kname = if Self::norm_ilp_on() {
9646 "rms_norm_f32_v2"
9647 } else {
9648 "rms_norm_f32"
9649 };
9650 if Self::pdl_on() && Self::pdl_wb_on() {
9651 use cudarc::driver::{DevicePtr, DevicePtrMut};
9652 let s = &self.gpu.stream();
9653 let (px, _g0) = x.device_ptr(s);
9654 let (pw, _g1) = w.device_ptr(s);
9655 let (pd, _g2) = dst.device_ptr_mut(s);
9656 let mut ps = [
9657 &px as *const _ as *mut std::ffi::c_void,
9658 &pw as *const _ as *mut _,
9659 &pd as *const _ as *mut _,
9660 &nc as *const _ as *mut _,
9661 &e as *const _ as *mut _,
9662 ];
9663 unsafe {
9664 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9665 }
9666 return Ok(());
9667 }
9668 let f = self.func(kname);
9669 let cfg = LaunchConfig {
9670 grid_dim: (nrows as u32, 1, 1),
9671 block_dim: (rms_block(), 1, 1),
9672 shared_mem_bytes: 0,
9673 };
9674 let __s_b = self.gpu.stream();
9675 let mut b = __s_b.launch_builder(&f);
9676 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9677 unsafe {
9678 b.launch(cfg)?;
9679 }
9680 Ok(())
9681 }
9682
9683 pub fn rms_norm_decode(
9691 &self,
9692 x: &CudaSlice<f32>,
9693 w: &CudaSlice<f32>,
9694 dst: &mut CudaSlice<f32>,
9695 ncols: usize,
9696 nrows: usize,
9697 eps: f32,
9698 ) -> Result<(), Box<dyn std::error::Error>> {
9699 let f = self.func(if Self::norm_ilp_on() {
9700 "rms_norm_f32_v2"
9701 } else {
9702 "rms_norm_f32"
9703 });
9704 let cfg = LaunchConfig {
9705 grid_dim: (nrows as u32, 1, 1),
9706 block_dim: (1024, 1, 1),
9707 shared_mem_bytes: 0,
9708 };
9709 let (nc, e) = (ncols as i32, eps);
9710 let __s_b = self.gpu.stream();
9711 let mut b = __s_b.launch_builder(&f);
9712 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9713 unsafe {
9714 b.launch(cfg)?;
9715 }
9716 Ok(())
9717 }
9718
9719 pub fn rms_norm_q8_1(
9723 &self,
9724 x: &CudaSlice<f32>,
9725 w: &CudaSlice<f32>,
9726 ncols: usize,
9727 nrows: usize,
9728 eps: f32,
9729 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9730 let nblk = ncols / 32;
9731 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9732 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9733 let (nc, e) = (ncols as i32, eps);
9734 if Self::pdl_on() {
9735 {
9736 use cudarc::driver::{DevicePtr, DevicePtrMut};
9737 let s = &self.gpu.stream();
9738 let (px, _g0) = x.device_ptr(s);
9739 let (pw, _g1) = w.device_ptr(s);
9740 let (pq, _g2) = q.device_ptr_mut(s);
9741 let (pd, _g3) = d.device_ptr_mut(s);
9742 let mut ps = [
9743 &px as *const _ as *mut std::ffi::c_void,
9744 &pw as *const _ as *mut _,
9745 &pq as *const _ as *mut _,
9746 &pd as *const _ as *mut _,
9747 &nc as *const _ as *mut _,
9748 &e as *const _ as *mut _,
9749 ];
9750 unsafe {
9751 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9752 }
9753 }
9754 return Ok((q, d));
9755 }
9756 let f = self.func("rms_norm_q8_1");
9757 let cfg = LaunchConfig {
9760 grid_dim: (nrows as u32, 1, 1),
9761 block_dim: (1024, 1, 1),
9762 shared_mem_bytes: 0,
9763 };
9764 let __s_b = self.gpu.stream();
9765 let mut b = __s_b.launch_builder(&f);
9766 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9767 unsafe {
9768 b.launch(cfg)?;
9769 }
9770 Ok((q, d))
9771 }
9772
9773 pub fn rms_norm_q8_1_into(
9776 &self,
9777 x: &CudaSlice<f32>,
9778 w: &CudaSlice<f32>,
9779 ncols: usize,
9780 nrows: usize,
9781 eps: f32,
9782 q: &mut CudaSlice<i8>,
9783 d: &mut CudaSlice<f32>,
9784 ) -> Result<(), Box<dyn std::error::Error>> {
9785 let nblk = ncols / 32;
9786 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9787 let (nc, e) = (ncols as i32, eps);
9788 if Self::pdl_on() {
9789 use cudarc::driver::{DevicePtr, DevicePtrMut};
9790 let s = &self.gpu.stream();
9791 let (px, _g0) = x.device_ptr(s);
9792 let (pw, _g1) = w.device_ptr(s);
9793 let (pq, _g2) = q.device_ptr_mut(s);
9794 let (pd, _g3) = d.device_ptr_mut(s);
9795 let mut ps = [
9796 &px as *const _ as *mut std::ffi::c_void,
9797 &pw as *const _ as *mut _,
9798 &pq as *const _ as *mut _,
9799 &pd as *const _ as *mut _,
9800 &nc as *const _ as *mut _,
9801 &e as *const _ as *mut _,
9802 ];
9803 unsafe {
9804 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9805 }
9806 return Ok(());
9807 }
9808 let f = self.func("rms_norm_q8_1");
9809 let cfg = LaunchConfig {
9810 grid_dim: (nrows as u32, 1, 1),
9811 block_dim: (1024, 1, 1),
9812 shared_mem_bytes: 0,
9813 };
9814 let __s_b = self.gpu.stream();
9815 let mut b = __s_b.launch_builder(&f);
9816 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9817 unsafe {
9818 b.launch(cfg)?;
9819 }
9820 Ok(())
9821 }
9822
9823 pub fn quantize_q8_1_into(
9825 &self,
9826 x: &CudaSlice<f32>,
9827 m: usize,
9828 in_f: usize,
9829 q: &mut CudaSlice<i8>,
9830 d: &mut CudaSlice<f32>,
9831 ) -> Result<(), Box<dyn std::error::Error>> {
9832 let nblk = in_f / 32;
9833 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9834 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9835 let (inf, mi) = (in_f as i32, m as i32);
9836 if Self::pdl_on() && Self::pdl_wb_on() {
9837 use cudarc::driver::{DevicePtr, DevicePtrMut};
9838 let s = &self.gpu.stream();
9839 let (px, _g0) = x.device_ptr(s);
9840 let (pq, _g1) = q.device_ptr_mut(s);
9841 let (pd, _g2) = d.device_ptr_mut(s);
9842 let mut ps = [
9843 &px as *const _ as *mut std::ffi::c_void,
9844 &pq as *const _ as *mut _,
9845 &pd as *const _ as *mut _,
9846 &inf as *const _ as *mut _,
9847 &mi as *const _ as *mut _,
9848 ];
9849 unsafe {
9850 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9851 }
9852 return Ok(());
9853 }
9854 let f = self.func("quantize_q8_1");
9855 let __s_b = self.gpu.stream();
9856 let mut b = __s_b.launch_builder(&f);
9857 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9858 unsafe {
9859 b.launch(cfg)?;
9860 }
9861 Ok(())
9862 }
9863
9864 pub fn add_rms_norm_q8_1(
9868 &self,
9869 a: &CudaSlice<f32>,
9870 b_in: &CudaSlice<f32>,
9871 w: &CudaSlice<f32>,
9872 res: &mut CudaSlice<f32>,
9873 ncols: usize,
9874 nrows: usize,
9875 eps: f32,
9876 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9877 let nblk = ncols / 32;
9878 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9879 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9880 let f = self.func("add_rms_norm_q8_1");
9881 let cfg = LaunchConfig {
9883 grid_dim: (nrows as u32, 1, 1),
9884 block_dim: (1024, 1, 1),
9885 shared_mem_bytes: 0,
9886 };
9887 let (nc, e) = (ncols as i32, eps);
9888 let __s_bld = self.gpu.stream();
9889 let mut bld = __s_bld.launch_builder(&f);
9890 bld.arg(a)
9891 .arg(b_in)
9892 .arg(w)
9893 .arg(res)
9894 .arg(&mut q)
9895 .arg(&mut d)
9896 .arg(&nc)
9897 .arg(&e);
9898 unsafe {
9899 bld.launch(cfg)?;
9900 }
9901 Ok((q, d))
9902 }
9903
9904 #[allow(clippy::too_many_arguments)]
9910 pub fn join_add_rms_norm_raw(
9911 &self,
9912 a0_raw: u64,
9913 a1_raw: u64,
9914 x: &CudaSlice<f32>,
9915 w: &CudaSlice<f32>,
9916 res: &mut CudaSlice<f32>,
9917 dst: &mut CudaSlice<f32>,
9918 ncols: usize,
9919 eps: f32,
9920 ) -> Result<(), Box<dyn std::error::Error>> {
9921 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9922 return Err("join_add_rms_norm geometry".into());
9923 }
9924 let f = self.func("join_add_rms_norm_f32");
9925 let cfg = LaunchConfig {
9926 grid_dim: (1, 1, 1),
9927 block_dim: (rms_block(), 1, 1),
9928 shared_mem_bytes: 0,
9929 };
9930 let (nc, e) = (ncols as i32, eps);
9931 let __s_b = self.gpu.stream();
9932 let mut b = __s_b.launch_builder(&f);
9933 b.arg(&a0_raw)
9934 .arg(&a1_raw)
9935 .arg(x)
9936 .arg(w)
9937 .arg(&mut *res)
9938 .arg(&mut *dst)
9939 .arg(&nc)
9940 .arg(&e);
9941 unsafe {
9942 b.launch(cfg)?;
9943 }
9944 Ok(())
9945 }
9946
9947 pub fn add_rms_norm(
9948 &self,
9949 a: &CudaSlice<f32>,
9950 b: &CudaSlice<f32>,
9951 w: &CudaSlice<f32>,
9952 res: &mut CudaSlice<f32>,
9953 dst: &mut CudaSlice<f32>,
9954 ncols: usize,
9955 nrows: usize,
9956 eps: f32,
9957 ) -> Result<(), Box<dyn std::error::Error>> {
9958 let (nc, e) = (ncols as i32, eps);
9959 let kname = if Self::norm_ilp_on() {
9960 "add_rms_norm_f32_v2"
9961 } else {
9962 "add_rms_norm_f32"
9963 };
9964 if Self::pdl_on() && Self::pdl_wb_on() {
9965 use cudarc::driver::{DevicePtr, DevicePtrMut};
9966 let s = &self.gpu.stream();
9967 let (pa, _g0) = a.device_ptr(s);
9968 let (pb, _g1) = b.device_ptr(s);
9969 let (pw, _g2) = w.device_ptr(s);
9970 let (pr, _g3) = res.device_ptr_mut(s);
9971 let (pd, _g4) = dst.device_ptr_mut(s);
9972 let mut ps = [
9973 &pa as *const _ as *mut std::ffi::c_void,
9974 &pb as *const _ as *mut _,
9975 &pw as *const _ as *mut _,
9976 &pr as *const _ as *mut _,
9977 &pd as *const _ as *mut _,
9978 &nc as *const _ as *mut _,
9979 &e as *const _ as *mut _,
9980 ];
9981 unsafe {
9982 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9983 }
9984 return Ok(());
9985 }
9986 let f = self.func(kname);
9987 let cfg = LaunchConfig {
9988 grid_dim: (nrows as u32, 1, 1),
9989 block_dim: (rms_block(), 1, 1),
9990 shared_mem_bytes: 0,
9991 };
9992 let __s_b2 = self.gpu.stream();
9993 let mut b2 = __s_b2.launch_builder(&f);
9994 b2.arg(a)
9995 .arg(b)
9996 .arg(w)
9997 .arg(&mut *res)
9998 .arg(&mut *dst)
9999 .arg(&nc)
10000 .arg(&e);
10001 unsafe {
10002 b2.launch(cfg)?;
10003 }
10004 Ok(())
10005 }
10006
10007 #[allow(clippy::too_many_arguments)]
10010 pub fn rms_pre_add_rms_norm(
10011 &self,
10012 a: &CudaSlice<f32>,
10013 wa: &CudaSlice<f32>,
10014 b: &CudaSlice<f32>,
10015 w: &CudaSlice<f32>,
10016 res: &mut CudaSlice<f32>,
10017 dst: &mut CudaSlice<f32>,
10018 ncols: usize,
10019 nrows: usize,
10020 eps: f32,
10021 ) -> Result<(), Box<dyn std::error::Error>> {
10022 let f = self.func("rms_pre_add_rms_norm_f32");
10023 let cfg = LaunchConfig {
10024 grid_dim: (nrows as u32, 1, 1),
10025 block_dim: (rms_block(), 1, 1),
10026 shared_mem_bytes: 0,
10027 };
10028 let (nc, e) = (ncols as i32, eps);
10029 let __s_b2 = self.gpu.stream();
10030 let mut b2 = __s_b2.launch_builder(&f);
10031 b2.arg(a)
10032 .arg(wa)
10033 .arg(b)
10034 .arg(w)
10035 .arg(&mut *res)
10036 .arg(&mut *dst)
10037 .arg(&nc)
10038 .arg(&e);
10039 unsafe {
10040 b2.launch(cfg)?;
10041 }
10042 Ok(())
10043 }
10044
10045 #[allow(clippy::too_many_arguments)]
10047 pub fn rms_pre_add_rms_norm_q8z(
10048 &self,
10049 a: &CudaSlice<f32>,
10050 wa: &CudaSlice<f32>,
10051 b: &CudaSlice<f32>,
10052 w: &CudaSlice<f32>,
10053 res: &mut CudaSlice<f32>,
10054 dst: &mut CudaSlice<f32>,
10055 ncols: usize,
10056 nrows: usize,
10057 eps: f32,
10058 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10059 debug_assert!(ncols % 128 == 0);
10060 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10061 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10062 let (nc, e) = (ncols as i32, eps);
10063 if Self::pdl_on() {
10064 {
10065 use cudarc::driver::{DevicePtr, DevicePtrMut};
10066 let s = &self.gpu.stream();
10067 let (pa, _g0) = a.device_ptr(s);
10068 let (pwa, _g1) = wa.device_ptr(s);
10069 let (pb, _g2) = b.device_ptr(s);
10070 let (pw, _g3) = w.device_ptr(s);
10071 let (pr, _g4) = res.device_ptr_mut(s);
10072 let (pdst, _g5) = dst.device_ptr_mut(s);
10073 let (pq, _g6) = out_q.device_ptr_mut(s);
10074 let (pd, _g7) = out_d.device_ptr_mut(s);
10075 let mut ps = [
10076 &pa as *const _ as *mut std::ffi::c_void,
10077 &pwa as *const _ as *mut _,
10078 &pb as *const _ as *mut _,
10079 &pw as *const _ as *mut _,
10080 &pr as *const _ as *mut _,
10081 &pdst as *const _ as *mut _,
10082 &pq as *const _ as *mut _,
10083 &pd as *const _ as *mut _,
10084 &nc as *const _ as *mut _,
10085 &e as *const _ as *mut _,
10086 ];
10087 unsafe {
10088 self.launch_pdl(
10089 "rms_pre_add_rms_norm_q8z_f32",
10090 (nrows as u32, 1, 1),
10091 (rms_block(), 1, 1),
10092 &mut ps,
10093 )?;
10094 }
10095 }
10096 return Ok((out_q, out_d));
10097 }
10098 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10099 let cfg = LaunchConfig {
10100 grid_dim: (nrows as u32, 1, 1),
10101 block_dim: (rms_block(), 1, 1),
10102 shared_mem_bytes: 0,
10103 };
10104 let __s_b2 = self.gpu.stream();
10105 let mut b2 = __s_b2.launch_builder(&f);
10106 b2.arg(a)
10107 .arg(wa)
10108 .arg(b)
10109 .arg(w)
10110 .arg(&mut *res)
10111 .arg(&mut *dst)
10112 .arg(&mut out_q)
10113 .arg(&mut out_d)
10114 .arg(&nc)
10115 .arg(&e);
10116 unsafe {
10117 b2.launch(cfg)?;
10118 }
10119 Ok((out_q, out_d))
10120 }
10121
10122 #[allow(clippy::too_many_arguments)]
10126 pub fn rms_pre_add_rms_norm_q8z_into(
10127 &self,
10128 a: &CudaSlice<f32>,
10129 wa: &CudaSlice<f32>,
10130 b: &CudaSlice<f32>,
10131 w: &CudaSlice<f32>,
10132 res: &mut CudaSlice<f32>,
10133 dst: &mut CudaSlice<f32>,
10134 ncols: usize,
10135 nrows: usize,
10136 eps: f32,
10137 out_q: &mut CudaSlice<i8>,
10138 out_d: &mut CudaSlice<f32>,
10139 ) -> Result<(), Box<dyn std::error::Error>> {
10140 debug_assert!(ncols % 128 == 0);
10141 let (nc, e) = (ncols as i32, eps);
10142 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10143 let cfg = LaunchConfig {
10144 grid_dim: (nrows as u32, 1, 1),
10145 block_dim: (rms_block(), 1, 1),
10146 shared_mem_bytes: 0,
10147 };
10148 let __s_b = self.gpu.stream();
10149 let mut b2 = __s_b.launch_builder(&f);
10150 b2.arg(a)
10151 .arg(wa)
10152 .arg(b)
10153 .arg(w)
10154 .arg(&mut *res)
10155 .arg(&mut *dst)
10156 .arg(&mut *out_q)
10157 .arg(&mut *out_d)
10158 .arg(&nc)
10159 .arg(&e);
10160 unsafe {
10161 b2.launch(cfg)?;
10162 }
10163 Ok(())
10164 }
10165
10166 #[allow(clippy::too_many_arguments)]
10169 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10170 &self,
10171 a: &CudaSlice<f32>,
10172 wa: &CudaSlice<f32>,
10173 b_in: &CudaSlice<f32>,
10174 c: f32,
10175 w: &CudaSlice<f32>,
10176 res: &mut CudaSlice<f32>,
10177 ncols: usize,
10178 nrows: usize,
10179 eps: f32,
10180 out_q: &mut CudaSlice<i8>,
10181 out_d: &mut CudaSlice<f32>,
10182 ) -> Result<(), Box<dyn std::error::Error>> {
10183 debug_assert!(ncols % 128 == 0);
10184 let (nc, e2) = (ncols as i32, eps);
10185 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10186 let cfg = LaunchConfig {
10187 grid_dim: (nrows as u32, 1, 1),
10188 block_dim: (rms_block(), 1, 1),
10189 shared_mem_bytes: 0,
10190 };
10191 let __s_b = self.gpu.stream();
10192 let mut b2 = __s_b.launch_builder(&f);
10193 b2.arg(a)
10194 .arg(wa)
10195 .arg(b_in)
10196 .arg(&c)
10197 .arg(w)
10198 .arg(&mut *res)
10199 .arg(&mut *out_q)
10200 .arg(&mut *out_d)
10201 .arg(&nc)
10202 .arg(&e2);
10203 unsafe {
10204 b2.launch(cfg)?;
10205 }
10206 Ok(())
10207 }
10208
10209 pub fn g4_pnfold_on() -> bool {
10217 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10218 *ON.get_or_init(|| {
10219 std::env::var("MEMRA_G4_PNFOLD")
10220 .map(|v| v != "0")
10221 .unwrap_or(true)
10222 })
10223 }
10224
10225 pub fn build_q4_out_concat3(
10229 &self,
10230 w0: &crate::model::GpuTensor,
10231 w1: &crate::model::GpuTensor,
10232 w2: &crate::model::GpuTensor,
10233 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10234 use crate::model::GpuTensor;
10235 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10236 match w {
10237 GpuTensor::Quant {
10238 qtype,
10239 row_bytes,
10240 rp,
10241 ..
10242 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10243 _ => None,
10244 }
10245 };
10246 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10247 else {
10248 return Ok(None);
10249 };
10250 if rb0 != rb1
10251 || rb0 != rb2
10252 || w0.in_features() != w1.in_features()
10253 || w0.in_features() != w2.in_features()
10254 {
10255 return Ok(None);
10256 }
10257 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10258 match w {
10259 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10260 _ => unreachable!(),
10261 }
10262 }
10263 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10264 let total = rb0 * (o0 + o1 + o2);
10265 let mut cat = self.alloc_u8(total)?;
10266 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10267 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10268 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10269 Ok(Some(GpuTensor::Quant {
10270 bytes: cat,
10271 qtype: QT_Q4_0,
10272 row_bytes: rb0,
10273 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10274 scale: 1.0,
10275 rp: false,
10276 #[cfg(memra_cutlass)]
10277 cutlass: None,
10278 fp8: None,
10279 blk: None,
10280 rp4: None,
10281 f16: None,
10282 }))
10283 }
10284
10285 fn full_width_rope_only(
10303 kernel: &str,
10304 n_rot: usize,
10305 head_dim: usize,
10306 ) -> Result<(), Box<dyn std::error::Error>> {
10307 if n_rot == head_dim {
10308 return Ok(());
10309 }
10310 Err(format!(
10311 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
10312 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
10313 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
10314 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
10315 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
10316 )
10317 .into())
10318 }
10319
10320 #[allow(clippy::too_many_arguments)]
10324 pub fn rms_norm_qkv_rope_cat(
10325 &self,
10326 qkv: &CudaSlice<f32>,
10327 wq: &CudaSlice<f32>,
10328 wk: &CudaSlice<f32>,
10329 wv: &CudaSlice<f32>,
10330 q: &mut CudaSlice<f32>,
10331 k: &mut CudaSlice<f32>,
10332 v: &mut CudaSlice<f32>,
10333 head_dim: usize,
10334 n_rot: usize,
10335 rq: usize,
10336 rk: usize,
10337 pos: &CudaSlice<i32>,
10338 nh_q: usize,
10339 nh_k: usize,
10340 base: f32,
10341 freq_scale: f32,
10342 ff: Option<&CudaSlice<f32>>,
10343 eps: f32,
10344 ) -> Result<(), Box<dyn std::error::Error>> {
10345 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
10346 let rows = rq + rk + rk;
10347 let theta_scale = base.powf(-2.0 / head_dim as f32);
10348 let (nc, rqi, rki, nhq, nhk) = (
10349 head_dim as i32,
10350 rq as i32,
10351 rk as i32,
10352 nh_q as i32,
10353 nh_k as i32,
10354 );
10355 if Self::pdl_on() {
10356 use cudarc::driver::{DevicePtr, DevicePtrMut};
10357 let s = &self.gpu.stream();
10358 let (pqkv, _g0) = qkv.device_ptr(s);
10359 let (pwq, _g1) = wq.device_ptr(s);
10360 let (pwk, _g2) = wk.device_ptr(s);
10361 let (pwv, _g3) = wv.device_ptr(s);
10362 let (pq, _g4) = q.device_ptr_mut(s);
10363 let (pk, _g5) = k.device_ptr_mut(s);
10364 let (pv, _g6) = v.device_ptr_mut(s);
10365 let (ppos, _g7) = pos.device_ptr(s);
10366 let (pff, _g8) = match ff {
10367 Some(t) => {
10368 let (p, g) = t.device_ptr(s);
10369 (p, Some(g))
10370 }
10371 None => (0, None),
10372 };
10373 let mut ps = [
10374 &pqkv as *const _ as *mut std::ffi::c_void,
10375 &pwq as *const _ as *mut _,
10376 &pwk as *const _ as *mut _,
10377 &pwv as *const _ as *mut _,
10378 &pq as *const _ as *mut _,
10379 &pk as *const _ as *mut _,
10380 &pv as *const _ as *mut _,
10381 &nc as *const _ as *mut _,
10382 &rqi as *const _ as *mut _,
10383 &rki as *const _ as *mut _,
10384 &ppos as *const _ as *mut _,
10385 &nhq as *const _ as *mut _,
10386 &nhk as *const _ as *mut _,
10387 &theta_scale as *const _ as *mut _,
10388 &freq_scale as *const _ as *mut _,
10389 &pff as *const _ as *mut _,
10390 &eps as *const _ as *mut _,
10391 ];
10392 unsafe {
10393 self.launch_pdl(
10394 "rms_norm_qkv_rope_cat_f32",
10395 (rows as u32, 1, 1),
10396 (rms_block(), 1, 1),
10397 &mut ps,
10398 )?;
10399 }
10400 return Ok(());
10401 }
10402 let f = self.func("rms_norm_qkv_rope_cat_f32");
10403 let cfg = LaunchConfig {
10404 grid_dim: (rows as u32, 1, 1),
10405 block_dim: (rms_block(), 1, 1),
10406 shared_mem_bytes: 0,
10407 };
10408 let __s_b = self.gpu.stream();
10409 let mut b = __s_b.launch_builder(&f);
10410 match ff {
10411 Some(t) => {
10412 b.arg(qkv)
10413 .arg(wq)
10414 .arg(wk)
10415 .arg(wv)
10416 .arg(&mut *q)
10417 .arg(&mut *k)
10418 .arg(&mut *v)
10419 .arg(&nc)
10420 .arg(&rqi)
10421 .arg(&rki)
10422 .arg(pos)
10423 .arg(&nhq)
10424 .arg(&nhk)
10425 .arg(&theta_scale)
10426 .arg(&freq_scale)
10427 .arg(t)
10428 .arg(&eps);
10429 unsafe {
10430 b.launch(cfg)?;
10431 }
10432 }
10433 None => {
10434 let null: u64 = 0;
10435 b.arg(qkv)
10436 .arg(wq)
10437 .arg(wk)
10438 .arg(wv)
10439 .arg(&mut *q)
10440 .arg(&mut *k)
10441 .arg(&mut *v)
10442 .arg(&nc)
10443 .arg(&rqi)
10444 .arg(&rki)
10445 .arg(pos)
10446 .arg(&nhq)
10447 .arg(&nhk)
10448 .arg(&theta_scale)
10449 .arg(&freq_scale)
10450 .arg(&null)
10451 .arg(&eps);
10452 unsafe {
10453 b.launch(cfg)?;
10454 }
10455 }
10456 }
10457 Ok(())
10458 }
10459
10460 #[allow(clippy::too_many_arguments)]
10464 pub fn rms_norm_qkv_rope(
10465 &self,
10466 q0: &CudaSlice<f32>,
10467 k0: &CudaSlice<f32>,
10468 v0: &CudaSlice<f32>,
10469 wq: &CudaSlice<f32>,
10470 wk: &CudaSlice<f32>,
10471 wv: &CudaSlice<f32>,
10472 q: &mut CudaSlice<f32>,
10473 k: &mut CudaSlice<f32>,
10474 v: &mut CudaSlice<f32>,
10475 head_dim: usize,
10476 n_rot: usize,
10477 rq: usize,
10478 rk: usize,
10479 pos: &CudaSlice<i32>,
10480 nh_q: usize,
10481 nh_k: usize,
10482 base: f32,
10483 freq_scale: f32,
10484 ff: Option<&CudaSlice<f32>>,
10485 eps: f32,
10486 ) -> Result<(), Box<dyn std::error::Error>> {
10487 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10488 let f = self.func("rms_norm_qkv_rope_f32");
10489 let rows = rq + rk + rk; let cfg = LaunchConfig {
10491 grid_dim: (rows as u32, 1, 1),
10492 block_dim: (rms_block(), 1, 1),
10493 shared_mem_bytes: 0,
10494 };
10495 let theta_scale = base.powf(-2.0 / head_dim as f32);
10496 let (nc, rqi, rki, nhq, nhk) = (
10497 head_dim as i32,
10498 rq as i32,
10499 rk as i32,
10500 nh_q as i32,
10501 nh_k as i32,
10502 );
10503 let __s_b = self.gpu.stream();
10504 let mut b = __s_b.launch_builder(&f);
10505 match ff {
10506 Some(t) => {
10507 b.arg(q0)
10508 .arg(k0)
10509 .arg(v0)
10510 .arg(wq)
10511 .arg(wk)
10512 .arg(wv)
10513 .arg(&mut *q)
10514 .arg(&mut *k)
10515 .arg(&mut *v)
10516 .arg(&nc)
10517 .arg(&rqi)
10518 .arg(&rki)
10519 .arg(pos)
10520 .arg(&nhq)
10521 .arg(&nhk)
10522 .arg(&theta_scale)
10523 .arg(&freq_scale)
10524 .arg(t)
10525 .arg(&eps);
10526 unsafe {
10527 b.launch(cfg)?;
10528 }
10529 }
10530 None => {
10531 let null: u64 = 0;
10532 b.arg(q0)
10533 .arg(k0)
10534 .arg(v0)
10535 .arg(wq)
10536 .arg(wk)
10537 .arg(wv)
10538 .arg(&mut *q)
10539 .arg(&mut *k)
10540 .arg(&mut *v)
10541 .arg(&nc)
10542 .arg(&rqi)
10543 .arg(&rki)
10544 .arg(pos)
10545 .arg(&nhq)
10546 .arg(&nhk)
10547 .arg(&theta_scale)
10548 .arg(&freq_scale)
10549 .arg(&null)
10550 .arg(&eps);
10551 unsafe {
10552 b.launch(cfg)?;
10553 }
10554 }
10555 }
10556 Ok(())
10557 }
10558
10559 #[allow(clippy::too_many_arguments)]
10565 pub fn rms_norm_qkv_rope_append_dc(
10566 &self,
10567 q0: &CudaSlice<f32>,
10568 k0: &CudaSlice<f32>,
10569 v0: &CudaSlice<f32>,
10570 wq: &CudaSlice<f32>,
10571 wk: &CudaSlice<f32>,
10572 wv: &CudaSlice<f32>,
10573 q: &mut CudaSlice<f32>,
10574 k: &mut CudaSlice<f32>,
10575 v: &mut CudaSlice<f32>,
10576 head_dim: usize,
10577 n_rot: usize,
10578 rq: usize,
10579 rk: usize,
10580 pos: &CudaSlice<i32>,
10581 nh_q: usize,
10582 nh_k: usize,
10583 base: f32,
10584 freq_scale: f32,
10585 ff: Option<&CudaSlice<f32>>,
10586 eps: f32,
10587 kc: &mut CudaSlice<u8>,
10588 vc: &mut CudaSlice<u8>,
10589 t_dev: &CudaSlice<i32>,
10590 k_tok_bytes: usize,
10591 v_tok_bytes: usize,
10592 g: bool,
10593 ) -> Result<(), Box<dyn std::error::Error>> {
10594 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10595 let rows = rq + rk + rk;
10596 let theta_scale = base.powf(-2.0 / head_dim as f32);
10597 let (nc, rqi, rki, nhq, nhk) = (
10598 head_dim as i32,
10599 rq as i32,
10600 rk as i32,
10601 nh_q as i32,
10602 nh_k as i32,
10603 );
10604 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10605 if Self::pdl_on() && Self::pdl_wb_on() {
10606 use cudarc::driver::{DevicePtr, DevicePtrMut};
10607 let s = &self.gpu.stream();
10608 let (p0, _a0) = q0.device_ptr(s);
10609 let (p1, _a1) = k0.device_ptr(s);
10610 let (p2, _a2) = v0.device_ptr(s);
10611 let (pwq, _a3) = wq.device_ptr(s);
10612 let (pwk, _a4) = wk.device_ptr(s);
10613 let (pwv, _a5) = wv.device_ptr(s);
10614 let (pq, _a6) = q.device_ptr_mut(s);
10615 let (pk, _a7) = k.device_ptr_mut(s);
10616 let (pv, _a8) = v.device_ptr_mut(s);
10617 let (pp, _a9) = pos.device_ptr(s);
10618 let pff: u64 = match ff {
10619 Some(t) => {
10620 let (p, _gg) = t.device_ptr(s);
10621 p as u64
10622 }
10623 None => 0,
10624 };
10625 let (pkc, _a10) = kc.device_ptr_mut(s);
10626 let (pvc, _a11) = vc.device_ptr_mut(s);
10627 let (pt, _a12) = t_dev.device_ptr(s);
10628 let mut ps = [
10629 &p0 as *const _ as *mut std::ffi::c_void,
10630 &p1 as *const _ as *mut _,
10631 &p2 as *const _ as *mut _,
10632 &pwq as *const _ as *mut _,
10633 &pwk as *const _ as *mut _,
10634 &pwv as *const _ as *mut _,
10635 &pq as *const _ as *mut _,
10636 &pk as *const _ as *mut _,
10637 &pv as *const _ as *mut _,
10638 &nc as *const _ as *mut _,
10639 &rqi as *const _ as *mut _,
10640 &rki as *const _ as *mut _,
10641 &pp as *const _ as *mut _,
10642 &nhq as *const _ as *mut _,
10643 &nhk as *const _ as *mut _,
10644 &theta_scale as *const _ as *mut _,
10645 &freq_scale as *const _ as *mut _,
10646 &pff as *const _ as *mut _,
10647 &eps as *const _ as *mut _,
10648 &pkc as *const _ as *mut _,
10649 &pvc as *const _ as *mut _,
10650 &pt as *const _ as *mut _,
10651 &ktb as *const _ as *mut _,
10652 &vtb as *const _ as *mut _,
10653 ];
10654 unsafe {
10655 self.launch_pdl_flash(
10656 g,
10657 "rms_norm_qkv_rope_append_dc_f32",
10658 (rows as u32, 1, 1),
10659 (rms_block(), 1, 1),
10660 0,
10661 &mut ps,
10662 )?;
10663 }
10664 return Ok(());
10665 }
10666 let f = if g {
10667 self.func_g("rms_norm_qkv_rope_append_dc_f32")
10668 } else {
10669 self.func("rms_norm_qkv_rope_append_dc_f32")
10670 };
10671 let cfg = LaunchConfig {
10672 grid_dim: (rows as u32, 1, 1),
10673 block_dim: (rms_block(), 1, 1),
10674 shared_mem_bytes: 0,
10675 };
10676 let __s_b = self.gpu.stream();
10677 let mut b = __s_b.launch_builder(&f);
10678 match ff {
10679 Some(t) => {
10680 b.arg(q0)
10681 .arg(k0)
10682 .arg(v0)
10683 .arg(wq)
10684 .arg(wk)
10685 .arg(wv)
10686 .arg(&mut *q)
10687 .arg(&mut *k)
10688 .arg(&mut *v)
10689 .arg(&nc)
10690 .arg(&rqi)
10691 .arg(&rki)
10692 .arg(pos)
10693 .arg(&nhq)
10694 .arg(&nhk)
10695 .arg(&theta_scale)
10696 .arg(&freq_scale)
10697 .arg(t)
10698 .arg(&eps)
10699 .arg(&mut *kc)
10700 .arg(&mut *vc)
10701 .arg(t_dev)
10702 .arg(&ktb)
10703 .arg(&vtb);
10704 unsafe {
10705 b.launch(cfg)?;
10706 }
10707 }
10708 None => {
10709 let null: u64 = 0;
10710 b.arg(q0)
10711 .arg(k0)
10712 .arg(v0)
10713 .arg(wq)
10714 .arg(wk)
10715 .arg(wv)
10716 .arg(&mut *q)
10717 .arg(&mut *k)
10718 .arg(&mut *v)
10719 .arg(&nc)
10720 .arg(&rqi)
10721 .arg(&rki)
10722 .arg(pos)
10723 .arg(&nhq)
10724 .arg(&nhk)
10725 .arg(&theta_scale)
10726 .arg(&freq_scale)
10727 .arg(&null)
10728 .arg(&eps)
10729 .arg(&mut *kc)
10730 .arg(&mut *vc)
10731 .arg(t_dev)
10732 .arg(&ktb)
10733 .arg(&vtb);
10734 unsafe {
10735 b.launch(cfg)?;
10736 }
10737 }
10738 }
10739 Ok(())
10740 }
10741
10742 #[allow(clippy::too_many_arguments)]
10750 pub fn rms_norm_qkv_rope_append(
10751 &self,
10752 q0: &CudaSlice<f32>,
10753 k0: &CudaSlice<f32>,
10754 v0: &CudaSlice<f32>,
10755 wq: &CudaSlice<f32>,
10756 wk: &CudaSlice<f32>,
10757 wv: &CudaSlice<f32>,
10758 q: &mut CudaSlice<f32>,
10759 k: &mut CudaSlice<f32>,
10760 v: &mut CudaSlice<f32>,
10761 head_dim: usize,
10762 n_rot: usize,
10763 rq: usize,
10764 rk: usize,
10765 pos: &CudaSlice<i32>,
10766 nh_q: usize,
10767 nh_k: usize,
10768 base: f32,
10769 freq_scale: f32,
10770 ff: Option<&CudaSlice<f32>>,
10771 eps: f32,
10772 kc: &mut CudaSlice<u8>,
10773 vc: &mut CudaSlice<u8>,
10774 t: usize,
10775 k_tok_bytes: usize,
10776 v_tok_bytes: usize,
10777 g: bool,
10778 ) -> Result<(), Box<dyn std::error::Error>> {
10779 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10780 let rows = rq + rk + rk;
10781 let theta_scale = base.powf(-2.0 / head_dim as f32);
10782 let (nc, rqi, rki, nhq, nhk) = (
10783 head_dim as i32,
10784 rq as i32,
10785 rk as i32,
10786 nh_q as i32,
10787 nh_k as i32,
10788 );
10789 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10790 let ti = t as i32;
10791 if Self::pdl_on() && Self::pdl_wb_on() {
10792 use cudarc::driver::{DevicePtr, DevicePtrMut};
10793 let s = &self.gpu.stream();
10794 let (p0, _a0) = q0.device_ptr(s);
10795 let (p1, _a1) = k0.device_ptr(s);
10796 let (p2, _a2) = v0.device_ptr(s);
10797 let (pwq, _a3) = wq.device_ptr(s);
10798 let (pwk, _a4) = wk.device_ptr(s);
10799 let (pwv, _a5) = wv.device_ptr(s);
10800 let (pq, _a6) = q.device_ptr_mut(s);
10801 let (pk, _a7) = k.device_ptr_mut(s);
10802 let (pv, _a8) = v.device_ptr_mut(s);
10803 let (pp, _a9) = pos.device_ptr(s);
10804 let pff: u64 = match ff {
10805 Some(t) => {
10806 let (p, _gg) = t.device_ptr(s);
10807 p as u64
10808 }
10809 None => 0,
10810 };
10811 let (pkc, _a10) = kc.device_ptr_mut(s);
10812 let (pvc, _a11) = vc.device_ptr_mut(s);
10813 let mut ps = [
10814 &p0 as *const _ as *mut std::ffi::c_void,
10815 &p1 as *const _ as *mut _,
10816 &p2 as *const _ as *mut _,
10817 &pwq as *const _ as *mut _,
10818 &pwk as *const _ as *mut _,
10819 &pwv as *const _ as *mut _,
10820 &pq as *const _ as *mut _,
10821 &pk as *const _ as *mut _,
10822 &pv as *const _ as *mut _,
10823 &nc as *const _ as *mut _,
10824 &rqi as *const _ as *mut _,
10825 &rki as *const _ as *mut _,
10826 &pp as *const _ as *mut _,
10827 &nhq as *const _ as *mut _,
10828 &nhk as *const _ as *mut _,
10829 &theta_scale as *const _ as *mut _,
10830 &freq_scale as *const _ as *mut _,
10831 &pff as *const _ as *mut _,
10832 &eps as *const _ as *mut _,
10833 &pkc as *const _ as *mut _,
10834 &pvc as *const _ as *mut _,
10835 &ti as *const _ as *mut _,
10836 &ktb as *const _ as *mut _,
10837 &vtb as *const _ as *mut _,
10838 ];
10839 unsafe {
10840 self.launch_pdl_flash(
10841 g,
10842 "rms_norm_qkv_rope_append_f32",
10843 (rows as u32, 1, 1),
10844 (rms_block(), 1, 1),
10845 0,
10846 &mut ps,
10847 )?;
10848 }
10849 return Ok(());
10850 }
10851 let f = if g {
10852 self.func_g("rms_norm_qkv_rope_append_f32")
10853 } else {
10854 self.func("rms_norm_qkv_rope_append_f32")
10855 };
10856 let cfg = LaunchConfig {
10857 grid_dim: (rows as u32, 1, 1),
10858 block_dim: (rms_block(), 1, 1),
10859 shared_mem_bytes: 0,
10860 };
10861 let __s_b = self.gpu.stream();
10862 let mut b = __s_b.launch_builder(&f);
10863 let null: u64 = 0;
10864 b.arg(q0)
10865 .arg(k0)
10866 .arg(v0)
10867 .arg(wq)
10868 .arg(wk)
10869 .arg(wv)
10870 .arg(&mut *q)
10871 .arg(&mut *k)
10872 .arg(&mut *v)
10873 .arg(&nc)
10874 .arg(&rqi)
10875 .arg(&rki)
10876 .arg(pos)
10877 .arg(&nhq)
10878 .arg(&nhk)
10879 .arg(&theta_scale)
10880 .arg(&freq_scale);
10881 match ff {
10882 Some(t) => {
10883 b.arg(t);
10884 }
10885 None => {
10886 b.arg(&null);
10887 }
10888 }
10889 b.arg(&eps)
10890 .arg(&mut *kc)
10891 .arg(&mut *vc)
10892 .arg(&ti)
10893 .arg(&ktb)
10894 .arg(&vtb);
10895 unsafe {
10896 b.launch(cfg)?;
10897 }
10898 Ok(())
10899 }
10900
10901 pub fn add_q8_1(
10902 &self,
10903 a: &CudaSlice<f32>,
10904 b: &CudaSlice<f32>,
10905 res: &mut CudaSlice<f32>,
10906 ncols: usize,
10907 nrows: usize,
10908 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10909 debug_assert!(ncols % 128 == 0);
10910 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10911 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10912 let f = self.func("add_q8_1_f32");
10913 let cfg = LaunchConfig {
10914 grid_dim: (nrows as u32, 1, 1),
10915 block_dim: (rms_block(), 1, 1),
10916 shared_mem_bytes: 0,
10917 };
10918 let nc = ncols as i32;
10919 let __s_b2 = self.gpu.stream();
10920 let mut b2 = __s_b2.launch_builder(&f);
10921 b2.arg(a)
10922 .arg(b)
10923 .arg(&mut *res)
10924 .arg(&mut out_q)
10925 .arg(&mut out_d)
10926 .arg(&nc);
10927 unsafe {
10928 b2.launch(cfg)?;
10929 }
10930 Ok((out_q, out_d))
10931 }
10932
10933 pub fn rms_pre_add_q8_1(
10937 &self,
10938 a: &CudaSlice<f32>,
10939 wa: &CudaSlice<f32>,
10940 b: &CudaSlice<f32>,
10941 res: &mut CudaSlice<f32>,
10942 ncols: usize,
10943 nrows: usize,
10944 eps: f32,
10945 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10946 debug_assert!(ncols % 128 == 0);
10947 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10948 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10949 let f = self.func("rms_pre_add_q8_1_f32");
10950 let cfg = LaunchConfig {
10951 grid_dim: (nrows as u32, 1, 1),
10952 block_dim: (rms_block(), 1, 1),
10953 shared_mem_bytes: 0,
10954 };
10955 let (nc, ep) = (ncols as i32, eps);
10956 let __s_b2 = self.gpu.stream();
10957 let mut b2 = __s_b2.launch_builder(&f);
10958 b2.arg(a)
10959 .arg(wa)
10960 .arg(b)
10961 .arg(&mut *res)
10962 .arg(&mut out_q)
10963 .arg(&mut out_d)
10964 .arg(&nc)
10965 .arg(&ep);
10966 unsafe {
10967 b2.launch(cfg)?;
10968 }
10969 Ok((out_q, out_d))
10970 }
10971
10972 pub fn l2_v2_on(ncols: usize) -> bool {
10976 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10977 }
10978
10979 pub fn l2_norm_pp(
10980 &self,
10981 x: &CudaSlice<f32>,
10982 dst: &mut CudaSlice<f32>,
10983 dst16: Option<&mut CudaSlice<u8>>,
10984 ncols: usize,
10985 nrows: usize,
10986 eps: f32,
10987 ) -> Result<(), Box<dyn std::error::Error>> {
10988 if Self::l2_v2_on(ncols) {
10989 let f = self.func("l2_norm_pp_v2_f32");
10990 let rows_per_block = 8u32; let cfg = LaunchConfig {
10992 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10993 block_dim: (256, 1, 1),
10994 shared_mem_bytes: 0,
10995 };
10996 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10997 let d16: u64 = match dst16 {
10999 Some(d) => self.addr_u8(d),
11000 None => 0,
11001 };
11002 let __s_b = self.gpu.stream();
11003 let mut b = __s_b.launch_builder(&f);
11004 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11005 unsafe {
11006 b.launch(cfg)?;
11007 }
11008 return Ok(());
11009 }
11010 self.l2_norm(x, dst, ncols, nrows, eps)
11011 }
11012
11013 pub fn l2_norm(
11014 &self,
11015 x: &CudaSlice<f32>,
11016 dst: &mut CudaSlice<f32>,
11017 ncols: usize,
11018 nrows: usize,
11019 eps: f32,
11020 ) -> Result<(), Box<dyn std::error::Error>> {
11021 let f = self.func("l2_norm_f32");
11022 let cfg = LaunchConfig {
11023 grid_dim: (nrows as u32, 1, 1),
11024 block_dim: (256, 1, 1),
11025 shared_mem_bytes: 0,
11026 };
11027 let (nc, e) = (ncols as i32, eps);
11028 let __s_b = self.gpu.stream();
11029 let mut b = __s_b.launch_builder(&f);
11030 b.arg(x).arg(dst).arg(&nc).arg(&e);
11031 unsafe {
11032 b.launch(cfg)?;
11033 }
11034 Ok(())
11035 }
11036
11037 pub fn l2_norm_decode(
11043 &self,
11044 x: &CudaSlice<f32>,
11045 dst: &mut CudaSlice<f32>,
11046 ncols: usize,
11047 nrows: usize,
11048 eps: f32,
11049 ) -> Result<(), Box<dyn std::error::Error>> {
11050 let f = self.func("l2_norm_f32");
11051 let cfg = LaunchConfig {
11052 grid_dim: (nrows as u32, 1, 1),
11053 block_dim: (32, 1, 1),
11054 shared_mem_bytes: 0,
11055 };
11056 let (nc, e) = (ncols as i32, eps);
11057 let __s_b = self.gpu.stream();
11058 let mut b = __s_b.launch_builder(&f);
11059 b.arg(x).arg(dst).arg(&nc).arg(&e);
11060 unsafe {
11061 b.launch(cfg)?;
11062 }
11063 Ok(())
11064 }
11065
11066 pub fn rope_neox(
11068 &self,
11069 x: &mut CudaSlice<f32>,
11070 pos: &CudaSlice<i32>,
11071 head_dim: usize,
11072 n_dims: usize,
11073 n_heads: usize,
11074 n_tokens: usize,
11075 freq_base: f32,
11076 freq_scale: f32,
11077 ) -> Result<(), Box<dyn std::error::Error>> {
11078 let f = self.func("rope_neox_f32");
11079 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11080 let grid = (n_heads * n_tokens) as u32;
11081 let cfg = LaunchConfig {
11082 grid_dim: (grid, 1, 1),
11083 block_dim: ((head_dim / 2) as u32, 1, 1),
11084 shared_mem_bytes: 0,
11085 };
11086 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11087 let __s_b = self.gpu.stream();
11088 let mut b = __s_b.launch_builder(&f);
11089 b.arg(x)
11090 .arg(pos)
11091 .arg(&hd)
11092 .arg(&nd)
11093 .arg(&nh)
11094 .arg(&theta_scale)
11095 .arg(&freq_scale);
11096 unsafe {
11097 b.launch(cfg)?;
11098 }
11099 Ok(())
11100 }
11101
11102 pub fn rope_neox_ff(
11104 &self,
11105 x: &mut CudaSlice<f32>,
11106 pos: &CudaSlice<i32>,
11107 head_dim: usize,
11108 n_dims: usize,
11109 n_heads: usize,
11110 n_tokens: usize,
11111 freq_base: f32,
11112 freq_scale: f32,
11113 ff: &CudaSlice<f32>,
11114 ) -> Result<(), Box<dyn std::error::Error>> {
11115 let f = self.func("rope_neox_ff_f32");
11116 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11117 let grid = (n_heads * n_tokens) as u32;
11118 let cfg = LaunchConfig {
11119 grid_dim: (grid, 1, 1),
11120 block_dim: ((head_dim / 2) as u32, 1, 1),
11121 shared_mem_bytes: 0,
11122 };
11123 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11124 let __s_b = self.gpu.stream();
11125 let mut b = __s_b.launch_builder(&f);
11126 b.arg(x)
11127 .arg(pos)
11128 .arg(&hd)
11129 .arg(&nd)
11130 .arg(&nh)
11131 .arg(&theta_scale)
11132 .arg(&freq_scale)
11133 .arg(ff);
11134 unsafe {
11135 b.launch(cfg)?;
11136 }
11137 Ok(())
11138 }
11139
11140 #[allow(clippy::too_many_arguments)]
11142 pub fn rope_neox2(
11143 &self,
11144 q: &mut CudaSlice<f32>,
11145 k: &mut CudaSlice<f32>,
11146 pos: &CudaSlice<i32>,
11147 head_dim: usize,
11148 n_dims: usize,
11149 nh_q: usize,
11150 nh_k: usize,
11151 n_tokens: usize,
11152 freq_base: f32,
11153 freq_scale: f32,
11154 ff: Option<&CudaSlice<f32>>,
11155 ) -> Result<(), Box<dyn std::error::Error>> {
11156 let f = self.func("rope_neox2_f32");
11157 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11158 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11159 let cfg = LaunchConfig {
11160 grid_dim: (grid, 1, 1),
11161 block_dim: ((head_dim / 2) as u32, 1, 1),
11162 shared_mem_bytes: 0,
11163 };
11164 let (hd, nd, nq, nk, nt) = (
11165 head_dim as i32,
11166 n_dims as i32,
11167 nh_q as i32,
11168 nh_k as i32,
11169 n_tokens as i32,
11170 );
11171 let __s_b = self.gpu.stream();
11172 let mut b = __s_b.launch_builder(&f);
11173 b.arg(q)
11174 .arg(k)
11175 .arg(pos)
11176 .arg(&hd)
11177 .arg(&nd)
11178 .arg(&nq)
11179 .arg(&nk)
11180 .arg(&nt)
11181 .arg(&theta_scale)
11182 .arg(&freq_scale);
11183 match ff {
11184 Some(ffv) => {
11185 b.arg(ffv);
11186 unsafe {
11187 b.launch(cfg)?;
11188 }
11189 }
11190 None => {
11191 let null: u64 = 0;
11192 b.arg(&null);
11193 unsafe {
11194 b.launch(cfg)?;
11195 }
11196 }
11197 }
11198 Ok(())
11199 }
11200
11201 pub fn gelu_tanh_mul(
11203 &self,
11204 gate: &CudaSlice<f32>,
11205 up: &CudaSlice<f32>,
11206 dst: &mut CudaSlice<f32>,
11207 n: usize,
11208 ) -> Result<(), Box<dyn std::error::Error>> {
11209 let f = self.func("gelu_tanh_mul_f32");
11210 let cfg = LaunchConfig::for_num_elems(n as u32);
11211 let ni = n as i32;
11212 let __s_b = self.gpu.stream();
11213 let mut b = __s_b.launch_builder(&f);
11214 b.arg(gate).arg(up).arg(dst).arg(&ni);
11215 unsafe {
11216 b.launch(cfg)?;
11217 }
11218 Ok(())
11219 }
11220
11221 pub fn silu_mul(
11222 &self,
11223 gate: &CudaSlice<f32>,
11224 up: &CudaSlice<f32>,
11225 dst: &mut CudaSlice<f32>,
11226 n: usize,
11227 ) -> Result<(), Box<dyn std::error::Error>> {
11228 let f = self.func("silu_mul_f32");
11229 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11231 let ni = n as i32;
11232 let __s_b = self.gpu.stream();
11233 let mut b = __s_b.launch_builder(&f);
11234 b.arg(gate).arg(up).arg(dst).arg(&ni);
11235 unsafe {
11236 b.launch(cfg)?;
11237 }
11238 Ok(())
11239 }
11240
11241 pub fn silu_mul_host_expf(
11243 &self,
11244 gate: &CudaSlice<f32>,
11245 up: &CudaSlice<f32>,
11246 dst: &mut CudaSlice<f32>,
11247 n: usize,
11248 ) -> Result<(), Box<dyn std::error::Error>> {
11249 let f = self.func("silu_mul_host_expf_f32");
11250 let cfg = LaunchConfig::for_num_elems(n as u32);
11251 let ni = n as i32;
11252 let __s_b = self.gpu.stream();
11253 let mut b = __s_b.launch_builder(&f);
11254 b.arg(gate).arg(up).arg(dst).arg(&ni);
11255 unsafe {
11256 b.launch(cfg)?;
11257 }
11258 Ok(())
11259 }
11260
11261 pub fn silu_clamped_mul_host_expf(
11263 &self,
11264 gate: &CudaSlice<f32>,
11265 up: &CudaSlice<f32>,
11266 limit: f32,
11267 dst: &mut CudaSlice<f32>,
11268 n: usize,
11269 ) -> Result<(), Box<dyn std::error::Error>> {
11270 if !limit.is_finite() || limit <= 0.0 {
11271 return Err(
11272 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11273 );
11274 }
11275 let f = self.func("silu_clamped_mul_host_expf_f32");
11276 let cfg = LaunchConfig::for_num_elems(n as u32);
11277 let ni = n as i32;
11278 let __s_b = self.gpu.stream();
11279 let mut b = __s_b.launch_builder(&f);
11280 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
11281 unsafe {
11282 b.launch(cfg)?;
11283 }
11284 Ok(())
11285 }
11286
11287 pub fn silu_mul_f16out(
11290 &self,
11291 gate: &CudaSlice<f32>,
11292 up: &CudaSlice<f32>,
11293 dst: &mut CudaSlice<f32>,
11294 dst16: &mut CudaSlice<u8>,
11295 n: usize,
11296 ) -> Result<(), Box<dyn std::error::Error>> {
11297 let f = self.func("silu_mul_f16out_f32");
11298 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11299 let ni = n as i32;
11300 let __s_b = self.gpu.stream();
11301 let mut b = __s_b.launch_builder(&f);
11302 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
11303 unsafe {
11304 b.launch(cfg)?;
11305 }
11306 Ok(())
11307 }
11308
11309 pub fn silu_mul_scaled(
11316 &self,
11317 gate: &CudaSlice<f32>,
11318 up: &CudaSlice<f32>,
11319 gs: f32,
11320 us: f32,
11321 dst: &mut CudaSlice<f32>,
11322 n: usize,
11323 ) -> Result<(), Box<dyn std::error::Error>> {
11324 let f = self.func("silu_mul_scaled_f32");
11325 let cfg = LaunchConfig::for_num_elems(n as u32);
11326 let ni = n as i32;
11327 let (gsf, usf) = (gs, us);
11328 let __s_b = self.gpu.stream();
11329 let mut b = __s_b.launch_builder(&f);
11330 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
11331 unsafe {
11332 b.launch(cfg)?;
11333 }
11334 Ok(())
11335 }
11336
11337 #[allow(clippy::too_many_arguments)]
11341 pub fn swigluoai_mul_scaled(
11342 &self,
11343 gate: &CudaSlice<f32>,
11344 up: &CudaSlice<f32>,
11345 gs: f32,
11346 us: f32,
11347 alpha: f32,
11348 limit: f32,
11349 dst: &mut CudaSlice<f32>,
11350 n: usize,
11351 ) -> Result<(), Box<dyn std::error::Error>> {
11352 let f = self.func("swigluoai_mul_scaled_f32");
11353 let cfg = LaunchConfig::for_num_elems(n as u32);
11354 let ni = n as i32;
11355 let __s_b = self.gpu.stream();
11356 let mut b = __s_b.launch_builder(&f);
11357 b.arg(gate)
11358 .arg(up)
11359 .arg(&gs)
11360 .arg(&us)
11361 .arg(&alpha)
11362 .arg(&limit)
11363 .arg(dst)
11364 .arg(&ni);
11365 unsafe {
11366 b.launch(cfg)?;
11367 }
11368 Ok(())
11369 }
11370
11371 pub fn silu_mul_scaled_q8_1(
11379 &self,
11380 gate: &CudaSlice<f32>,
11381 up: &CudaSlice<f32>,
11382 gs: f32,
11383 us: f32,
11384 n: usize,
11385 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11386 let f = self.func("silu_mul_scaled_q8_1");
11387 let nblk = n / 32;
11388 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);
11392 let (gsf, usf, ni) = (gs, us, n as i32);
11393 let __s_b = self.gpu.stream();
11394 let mut b = __s_b.launch_builder(&f);
11395 b.arg(gate)
11396 .arg(up)
11397 .arg(&gsf)
11398 .arg(&usf)
11399 .arg(&mut aq)
11400 .arg(&mut ad)
11401 .arg(&ni);
11402 unsafe {
11403 b.launch(cfg)?;
11404 }
11405 Ok((aq, ad))
11406 }
11407
11408 pub fn add(
11409 &self,
11410 a: &CudaSlice<f32>,
11411 b_in: &CudaSlice<f32>,
11412 dst: &mut CudaSlice<f32>,
11413 n: usize,
11414 ) -> Result<(), Box<dyn std::error::Error>> {
11415 let f = self.func("add_f32");
11416 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11418 let ni = n as i32;
11419 let __s_bld = self.gpu.stream();
11420 let mut bld = __s_bld.launch_builder(&f);
11421 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11422 unsafe {
11423 bld.launch(cfg)?;
11424 }
11425 Ok(())
11426 }
11427
11428 pub fn mul(
11429 &self,
11430 a: &CudaSlice<f32>,
11431 b_in: &CudaSlice<f32>,
11432 dst: &mut CudaSlice<f32>,
11433 n: usize,
11434 ) -> Result<(), Box<dyn std::error::Error>> {
11435 let f = self.func("mul_f32");
11436 let cfg = LaunchConfig::for_num_elems(n as u32);
11437 let ni = n as i32;
11438 let __s_bld = self.gpu.stream();
11439 let mut bld = __s_bld.launch_builder(&f);
11440 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11441 unsafe {
11442 bld.launch(cfg)?;
11443 }
11444 Ok(())
11445 }
11446
11447 pub fn matmul(
11450 &self,
11451 w: &crate::model::GpuTensor,
11452 x: &CudaSlice<f32>,
11453 m: usize,
11454 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11455 use crate::model::GpuTensor;
11456 let in_f = w.in_features();
11457 let out_f = w.out_features();
11458 #[allow(non_snake_case)]
11466 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11469 usize::MAX
11470 } else {
11471 16usize
11472 };
11473
11474 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11499 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11500 return Ok(y);
11501 }
11502 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11509 return Ok(y);
11510 }
11511 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11514 return Ok(y);
11515 }
11516 }
11517 if let GpuTensor::Quant { qtype, .. } = w {
11532 if *qtype == QT_F8_E4M3_BLK {
11533 if m >= GEMM_M_THRESHOLD {
11534 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11535 return Ok(y);
11536 }
11537 }
11538 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11539 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11540 return Ok(y);
11541 }
11542 }
11543 }
11544 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11545 return self.qmatvec_mmq(w, x, m);
11546 }
11547 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11548 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11549 return self.qmatvec_gemm(w, &aq, &ad, m);
11550 }
11551 if m >= GEMM_M_THRESHOLD {
11554 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11555 return Ok(y);
11556 }
11557 }
11558 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11562 if m == 1 && fast {
11567 if let GpuTensor::Quant {
11568 bytes,
11569 qtype,
11570 row_bytes,
11571 rp,
11572 rp4,
11573 scale,
11574 ..
11575 } = w
11576 {
11577 if self.mmvq_supports(*qtype) {
11578 let (bytes, rp) = match rp4 {
11582 Some(m4) => (m4, true),
11583 None => (bytes, *rp),
11584 };
11585 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11586 return self.qmatvec_mmvq(
11587 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11588 );
11589 }
11590 }
11591 }
11592 if (2..=16).contains(&m)
11608 && fast
11609 && std::env::var("MEMRA_NO_BATCHED").is_err()
11610 && (m <= 4 || Self::b8_enabled())
11611 {
11612 let m_ok = m <= 8
11622 || matches!(w, GpuTensor::Quant { qtype, .. }
11623 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11624 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11625 if m_ok {
11626 if let GpuTensor::Quant {
11627 bytes,
11628 qtype,
11629 row_bytes,
11630 rp,
11631 rp4,
11632 ..
11633 } = w
11634 {
11635 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11636 let (bytes, rp) = match rp4 {
11637 Some(m4) => (m4, true),
11638 None => (bytes, *rp),
11639 };
11640 let mcols = Self::batched_mcols(m);
11641 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11642 let mut y = self.qmatvec_mmvq_batched(
11643 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11644 )?;
11645 if let GpuTensor::Quant { scale, .. } = w {
11646 if *scale != 1.0 {
11647 self.scale_inplace(&mut y, *scale, m * out_f)?;
11648 }
11649 }
11650 return Ok(y);
11651 }
11652 }
11653 }
11654 }
11655 if fast {
11661 if let GpuTensor::Quant {
11662 bytes,
11663 qtype,
11664 row_bytes,
11665 scale,
11666 ..
11667 } = w
11668 {
11669 if *qtype == QT_F8_E4M3 {
11670 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11671 return self.qmatvec_mmvq(
11672 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11673 );
11674 }
11675 }
11676 }
11677 let mut y = match w {
11678 GpuTensor::Quant {
11679 bytes,
11680 qtype,
11681 row_bytes,
11682 ..
11683 } if fast && *qtype == QT_Q8_0 => {
11684 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11685 }
11686 GpuTensor::Quant {
11687 bytes,
11688 qtype,
11689 row_bytes,
11690 ..
11691 } if fast && *qtype == QT_Q4_K => {
11692 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11693 }
11694 GpuTensor::Quant {
11695 bytes,
11696 qtype,
11697 row_bytes,
11698 ..
11699 } if fast && *qtype == QT_Q6_K => {
11700 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11701 }
11702 GpuTensor::Quant {
11703 bytes,
11704 qtype,
11705 row_bytes,
11706 ..
11707 } if fast && *qtype == QT_Q5_K => {
11708 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11709 }
11710 GpuTensor::Quant {
11711 bytes,
11712 qtype,
11713 row_bytes,
11714 ..
11715 } if fast && *qtype == QT_Q3_K => {
11716 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11717 }
11718 GpuTensor::Quant {
11719 bytes,
11720 qtype,
11721 row_bytes,
11722 rp,
11723 ..
11724 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11725 if *rp {
11726 "qmatvec_nvfp4_dp4a_rp"
11727 } else {
11728 "qmatvec_nvfp4_dp4a"
11729 },
11730 &bytes.slice(0..bytes.len()),
11731 x,
11732 m,
11733 in_f,
11734 out_f,
11735 *row_bytes,
11736 )?,
11737 GpuTensor::Quant {
11741 bytes,
11742 qtype,
11743 row_bytes,
11744 ..
11745 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11746 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11747 }
11748 GpuTensor::Quant {
11753 bytes,
11754 qtype,
11755 row_bytes,
11756 rp,
11757 ..
11758 } =>
11759 {
11762 self.qmatvec(
11763 bytes,
11764 x,
11765 m,
11766 in_f,
11767 out_f,
11768 if *rp && *qtype == QT_NVFP4 {
11769 QT_NVFP4_RP
11770 } else {
11771 *qtype
11772 },
11773 *row_bytes,
11774 )?
11775 }
11776 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11777 GpuTensor::FloatBf16 { data, .. } => {
11780 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11781 }
11782 };
11783 if let GpuTensor::Quant { scale, .. } = w {
11785 if *scale != 1.0 {
11786 self.scale_inplace(&mut y, *scale, m * out_f)?;
11787 }
11788 }
11789 Ok(y)
11790 }
11791
11792 pub fn stage_a_raw_needed() -> bool {
11802 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11803 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11804 }
11805
11806 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11809 use crate::model::GpuTensor;
11810 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11811 return false;
11812 }
11813 match w {
11814 GpuTensor::Quant { qtype, .. } => {
11821 matches!(
11822 *qtype,
11823 QT_Q8_0
11824 | QT_Q4_K
11825 | QT_Q6_K
11826 | QT_Q5_K
11827 | QT_Q3_K
11828 | QT_NVFP4
11829 | QT_F8_E4M3
11830 | QT_F8_E4M3_BLK
11831 | QT_Q4_0
11832 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11833 }
11834 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11835 }
11836 }
11837
11838 pub fn matmul_pre(
11843 &self,
11844 w: &crate::model::GpuTensor,
11845 aq: &CudaSlice<i8>,
11846 ad: &CudaSlice<f32>,
11847 x_fallback: &CudaSlice<f32>,
11848 m: usize,
11849 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11850 use crate::model::GpuTensor;
11851 let x_raw_ok = x_fallback.len() >= m * w.in_features();
11857 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11860 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11861 return Ok(y);
11862 }
11863 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11866 return Ok(y);
11867 }
11868 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11870 return Ok(y);
11871 }
11872 }
11873 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11879 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11880 return Ok(y);
11881 }
11882 }
11883 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11884 return Ok(y);
11885 }
11886 if m >= 16
11891 && w.out_features() >= 128
11892 && self.mmq_supports(w)
11893 && !self.verify_exact_on()
11894 && x_raw_ok
11895 {
11896 return self.qmatvec_mmq(w, x_fallback, m);
11897 }
11898 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11901 if let Some(y) =
11902 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11903 {
11904 return Ok(y);
11905 }
11906 }
11907 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11910 return self.qmatvec_gemm(w, aq, ad, m);
11911 }
11912 if !self.uses_q8_1_fast(w) {
11931 if !x_raw_ok {
11932 return Err(format!(
11933 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11934 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11935 dropped the f32, so there is nothing to fall back to — pass the real f32 \
11936 activation (see Engine::rms_norm_decode, which is bit-identical to \
11937 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11938 x_fallback.len(),
11939 m,
11940 w.in_features(),
11941 m * w.in_features()
11942 )
11943 .into());
11944 }
11945 return self.matmul(w, x_fallback, m);
11946 }
11947 let in_f = w.in_features();
11948 let out_f = w.out_features();
11949 let (bytes, qtype, row_bytes, scale, rp) = match w {
11950 GpuTensor::Quant {
11951 bytes,
11952 qtype,
11953 row_bytes,
11954 scale,
11955 rp,
11956 ..
11957 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11958 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11959 };
11960 let (mbytes, mrp) = match w {
11963 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11964 _ => (bytes, rp),
11965 };
11966 if m == 1 && self.mmvq_supports(qtype) {
11970 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11971 }
11972 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11985 && std::env::var("MEMRA_NO_BATCHED").is_err()
11986 && (m <= 4 || Self::b8_enabled())
11987 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11991 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11992 {
11993 let mcols = Self::batched_mcols(m);
11994 return self.qmatvec_mmvq_batched(
11995 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11996 );
11997 }
11998 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12004 let (b2, r2) = if qtype == QT_Q4_0 {
12005 (mbytes, mrp)
12006 } else {
12007 (bytes, rp)
12008 };
12009 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12010 }
12011 let name = match qtype {
12012 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12013 QT_Q4_K => "qmatvec_q4_K_dp4a",
12014 QT_Q6_K => "qmatvec_q6_K_dp4a",
12015 QT_Q5_K => "qmatvec_q5_K_dp4a",
12016 QT_Q3_K => "qmatvec_q3_K_dp4a",
12017 QT_NVFP4 => {
12018 if rp {
12019 "qmatvec_nvfp4_dp4a_rp"
12020 } else {
12021 "qmatvec_nvfp4_dp4a"
12022 }
12023 }
12024 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12025 _ => unreachable!(),
12026 };
12027 let f = self.func(name);
12028 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12030 grid_dim: (out_f as u32, m as u32, 1),
12031 block_dim: (128, 1, 1),
12032 shared_mem_bytes: 0,
12033 };
12034 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12035 let __s_b = self.gpu.stream();
12036 let mut b = __s_b.launch_builder(&f);
12037 b.arg(bytes)
12038 .arg(aq)
12039 .arg(ad)
12040 .arg(&mut y)
12041 .arg(&inf)
12042 .arg(&outf)
12043 .arg(&mi)
12044 .arg(&rb);
12045 unsafe {
12046 b.launch(cfg)?;
12047 }
12048 if scale != 1.0 {
12049 self.scale_inplace(&mut y, scale, m * out_f)?;
12050 }
12051 Ok(y)
12052 }
12053
12054 pub fn matmul_decode_exact(
12062 &self,
12063 w: &crate::model::GpuTensor,
12064 x: &CudaSlice<f32>,
12065 m: usize,
12066 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12067 use crate::model::GpuTensor;
12068 if let GpuTensor::Float { data, .. } = w {
12076 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12077 }
12078 if let GpuTensor::FloatBf16 { data, .. } = w {
12081 let (in_f, out_f) = (w.in_features(), w.out_features());
12082 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12083 }
12084 if !self.uses_q8_1_fast(w) {
12085 return self.matmul(w, x, m);
12086 }
12087 let in_f = w.in_features();
12088 let out_f = w.out_features();
12089 let (bytes, qtype, row_bytes, scale, rp) = match w {
12090 GpuTensor::Quant {
12091 bytes,
12092 qtype,
12093 row_bytes,
12094 scale,
12095 rp,
12096 ..
12097 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12098 _ => return self.matmul(w, x, m),
12099 };
12100 let (bytes, rp) = match w {
12103 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12104 _ => (bytes, rp),
12105 };
12106 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12107 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12111 return Ok(y);
12112 }
12113 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12122 && std::env::var("MEMRA_NO_BATCHED").is_err()
12123 && (m <= 4 || Self::b8_enabled())
12124 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12127 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12128 {
12129 let mcols = Self::batched_mcols(m);
12130 return self.qmatvec_mmvq_batched(
12131 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12132 );
12133 }
12134 if self.mmvq_supports(qtype) {
12135 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12138 }
12139 self.matmul_pre(w, &aq, &ad, x, m)
12142 }
12143
12144 pub fn matmul_decode_exact_pre(
12154 &self,
12155 w: &crate::model::GpuTensor,
12156 aq: &CudaSlice<i8>,
12157 ad: &CudaSlice<f32>,
12158 m: usize,
12159 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12160 use crate::model::GpuTensor;
12161 debug_assert!(
12162 self.uses_q8_1_fast(w),
12163 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12164 );
12165 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12167 return Ok(y);
12168 }
12169 let in_f = w.in_features();
12170 let out_f = w.out_features();
12171 let (bytes, qtype, row_bytes, scale, rp) = match w {
12172 GpuTensor::Quant {
12173 bytes,
12174 qtype,
12175 row_bytes,
12176 scale,
12177 rp,
12178 ..
12179 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12180 _ => {
12181 return Err(
12182 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12183 );
12184 }
12185 };
12186 let (bytes, rp) = match w {
12188 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12189 _ => (bytes, rp),
12190 };
12191 if (2..=16).contains(&m)
12193 && self.batched_supports(qtype)
12194 && self.mmvq_supports(qtype)
12195 && std::env::var("MEMRA_NO_BATCHED").is_err()
12196 && (m <= 4 || Self::b8_enabled())
12197 && (m <= 8
12198 || qtype == QT_Q4_0
12199 || qtype == QT_Q6_K
12200 || qtype == QT_F8_E4M3
12201 || qtype == QT_NVFP4
12202 || qtype == QT_Q4_K
12203 || qtype == QT_Q5_K
12204 || qtype == QT_Q8_0)
12205 {
12206 let mcols = Self::batched_mcols(m);
12207 return self.qmatvec_mmvq_batched(
12208 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12209 );
12210 }
12211 if self.mmvq_supports(qtype) {
12212 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12213 }
12214 let x0 = self.zeros(0)?;
12217 self.matmul_pre(w, aq, ad, &x0, m)
12218 }
12219
12220 pub fn matmul_decode_exact_dual_pre(
12229 &self,
12230 w0: &crate::model::GpuTensor,
12231 w1: &crate::model::GpuTensor,
12232 aq: &CudaSlice<i8>,
12233 ad: &CudaSlice<f32>,
12234 m: usize,
12235 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12236 {
12237 use crate::model::GpuTensor;
12238 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12239 let on = *ON.get_or_init(|| {
12240 std::env::var("MEMRA_SPEC_DUAL_T")
12241 .map(|v| v != "0")
12242 .unwrap_or(true)
12243 });
12244 if !on
12245 || !(2..=7).contains(&m)
12246 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12247 || !self.uses_q8_1_fast(w0)
12248 || !self.uses_q8_1_fast(w1)
12249 {
12250 return Ok(None);
12251 }
12252 if !self.mmvq_supports(QT_NVFP4) {
12257 return Ok(None);
12258 }
12259 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12260 if w1.in_features() != in_f || w1.out_features() != out_f {
12261 return Ok(None);
12262 }
12263 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12264 (
12265 GpuTensor::Quant {
12266 bytes: b0,
12267 qtype: q0,
12268 row_bytes: rb0,
12269 scale: s0,
12270 rp: rp0,
12271 rp4: None,
12272 ..
12273 },
12274 GpuTensor::Quant {
12275 bytes: b1,
12276 qtype: q1,
12277 row_bytes: rb1,
12278 scale: s1,
12279 rp: rp1,
12280 rp4: None,
12281 ..
12282 },
12283 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12284 (b0, b1, *rb0, *s0, *s1, *rp0)
12285 }
12286 _ => return Ok(None),
12287 };
12288 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
12291 {
12292 return Ok(None);
12293 }
12294 let (y0, y1) =
12295 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
12296 Ok(Some(((y0, s0), (y1, s1))))
12297 }
12298
12299 pub fn matmul_decode_exact_group4_pre(
12311 &self,
12312 ws: [&crate::model::GpuTensor; 4],
12313 aq: &CudaSlice<i8>,
12314 ad: &CudaSlice<f32>,
12315 m: usize,
12316 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12317 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12318 let on = *ON.get_or_init(|| {
12319 std::env::var("MEMRA_TK_GDN_GROUP")
12320 .map(|v| v != "0")
12321 .unwrap_or(true)
12322 });
12323 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
12324 }
12325
12326 pub fn matmul_decode_exact_group3_pre(
12331 &self,
12332 ws: [&crate::model::GpuTensor; 3],
12333 aq: &CudaSlice<i8>,
12334 ad: &CudaSlice<f32>,
12335 m: usize,
12336 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12337 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12338 let on = *ON.get_or_init(|| {
12339 std::env::var("MEMRA_TK_FA_GROUP")
12340 .map(|v| v != "0")
12341 .unwrap_or(true)
12342 });
12343 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
12344 }
12345
12346 fn matmul_decode_exact_group_pre(
12350 &self,
12351 ws: &[&crate::model::GpuTensor],
12352 aq: &CudaSlice<i8>,
12353 ad: &CudaSlice<f32>,
12354 m: usize,
12355 on: bool,
12356 tag: &'static str,
12357 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
12358 use crate::model::GpuTensor;
12359 if !on
12360 || !(2..=16).contains(&m)
12361 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12362 || (m > 4 && !Self::b8_enabled())
12363 || !self.mmvq_supports(QT_NVFP4)
12364 || !self.batched_supports(QT_NVFP4)
12365 {
12366 return Ok(None);
12367 }
12368 let in_f = ws[0].in_features();
12369 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12370 for w in ws {
12371 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12372 return Ok(None);
12373 }
12374 match w {
12375 GpuTensor::Quant {
12376 bytes,
12377 qtype,
12378 scale,
12379 rp: true,
12380 rp4: None,
12381 ..
12382 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12383 parts.push((bytes, w.out_features(), *scale));
12384 }
12385 _ => return Ok(None),
12386 }
12387 }
12388 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12390 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12391 let mcols = if (5..=7).contains(&m) && b567 {
12392 m
12393 } else {
12394 Self::batched_mcols(m)
12395 };
12396 let kname: &'static str = match mcols {
12397 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12398 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12399 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12400 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12401 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12402 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12403 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12404 _ => return Ok(None),
12405 };
12406 if std::env::var("MEMRA_DEBUG").is_ok() {
12409 use std::sync::Mutex;
12410 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12411 let mut seen = SEEN.lock().unwrap();
12412 if !seen.contains(&tag) {
12413 seen.push(tag);
12414 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12415 }
12416 }
12417 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12420 let three = parts.len() == 3;
12421 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12422 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12423 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12424 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12427 let cfg = LaunchConfig {
12428 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12429 block_dim: (32, ROWS_PER_BLOCK, 1),
12430 shared_mem_bytes: 0,
12431 };
12432 let (inf, mi) = (in_f as i32, m as i32);
12433 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12434 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12435 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12436 let s3 = if three { 1.0f32 } else { parts[3].2 };
12437 let w3 = if three { parts[0].0 } else { parts[3].0 };
12438 let f = self.func(kname);
12439 let __s_b = self.gpu.stream();
12440 let mut b = __s_b.launch_builder(&f);
12441 b.arg(parts[0].0)
12442 .arg(parts[1].0)
12443 .arg(parts[2].0)
12444 .arg(w3)
12445 .arg(aq)
12446 .arg(ad)
12447 .arg(&mut y0)
12448 .arg(&mut y1)
12449 .arg(&mut y2)
12450 .arg(&mut y3)
12451 .arg(&inf)
12452 .arg(&n0)
12453 .arg(&n1)
12454 .arg(&n2)
12455 .arg(&n3)
12456 .arg(&mi)
12457 .arg(&s0)
12458 .arg(&s1)
12459 .arg(&s2)
12460 .arg(&s3);
12461 unsafe {
12462 b.launch(cfg)?;
12463 }
12464 Ok(Some(if three {
12465 vec![y0, y1, y2]
12466 } else {
12467 vec![y0, y1, y2, y3]
12468 }))
12469 }
12470
12471 pub fn matmul_decode_exact_dual(
12487 &self,
12488 w0: &crate::model::GpuTensor,
12489 w1: &crate::model::GpuTensor,
12490 x: &CudaSlice<f32>,
12491 m: usize,
12492 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12493 use crate::model::GpuTensor;
12494 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12495 let on = *ON.get_or_init(|| {
12496 std::env::var("MEMRA_SPEC_DUAL_T")
12497 .map(|v| v != "0")
12498 .unwrap_or(true)
12499 });
12500 if !on
12501 || !(2..=4).contains(&m)
12502 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12503 || !self.uses_q8_1_fast(w0)
12504 || !self.uses_q8_1_fast(w1)
12505 {
12506 return Ok(None);
12507 }
12508 if !self.mmvq_supports(QT_NVFP4) {
12513 return Ok(None);
12514 }
12515 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12516 if w1.in_features() != in_f || w1.out_features() != out_f {
12517 return Ok(None);
12518 }
12519 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12520 (
12521 GpuTensor::Quant {
12522 bytes: b0,
12523 qtype: q0,
12524 row_bytes: rb0,
12525 scale: s0,
12526 rp: rp0,
12527 rp4: None,
12528 ..
12529 },
12530 GpuTensor::Quant {
12531 bytes: b1,
12532 qtype: q1,
12533 row_bytes: rb1,
12534 scale: s1,
12535 rp: rp1,
12536 rp4: None,
12537 ..
12538 },
12539 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12540 (b0, b1, *rb0, *s0, *s1, *rp0)
12541 }
12542 _ => return Ok(None),
12543 };
12544 if std::env::var("MEMRA_DEBUG").is_ok() {
12547 static ONCE: std::sync::Once = std::sync::Once::new();
12548 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12549 }
12550 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12551 let (y0, y1) =
12552 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12553 let mut y0 = y0;
12554 let mut y1 = y1;
12555 if s0 != 1.0 {
12556 self.scale_inplace(&mut y0, s0, m * out_f)?;
12557 }
12558 if s1 != 1.0 {
12559 self.scale_inplace(&mut y1, s1, m * out_f)?;
12560 }
12561 Ok(Some((y0, y1)))
12562 }
12563
12564 #[allow(clippy::too_many_arguments)]
12569 pub fn qmatvec_batched_dual_raw(
12570 &self,
12571 b0: &CudaSlice<u8>,
12572 b1: &CudaSlice<u8>,
12573 aq: &CudaSlice<i8>,
12574 ad: &CudaSlice<f32>,
12575 m: usize,
12576 in_f: usize,
12577 out_f: usize,
12578 row_bytes: usize,
12579 rp: bool,
12580 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12581 const ROWS_PER_BLOCK: u32 = 4;
12582 let mcols = Self::batched_mcols(m);
12583 let tiny_rp1 = rp
12586 && mcols == 4
12587 && out_f <= 128
12588 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12589 let (name, rows_per_block) = if tiny_rp1 {
12590 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12591 } else {
12592 match (mcols, rp, m) {
12593 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12594 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12595 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12596 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12597 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12598 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12599 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12600 _ => {
12601 return Err(
12602 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12603 );
12604 }
12605 }
12606 };
12607 let f = self.func(name);
12608 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12609 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12610 let cfg = LaunchConfig {
12611 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12612 block_dim: (32, ROWS_PER_BLOCK, 1),
12613 shared_mem_bytes: 0,
12614 };
12615 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12616 let __s_b = self.gpu.stream();
12617 let mut b = __s_b.launch_builder(&f);
12618 b.arg(b0)
12619 .arg(b1)
12620 .arg(aq)
12621 .arg(ad)
12622 .arg(&mut y0)
12623 .arg(&mut y1)
12624 .arg(&inf)
12625 .arg(&outf)
12626 .arg(&mi)
12627 .arg(&rb);
12628 unsafe {
12629 b.launch(cfg)?;
12630 }
12631 Ok((y0, y1))
12632 }
12633
12634 pub fn matmul_pre_dual_noscale(
12646 &self,
12647 w0: &crate::model::GpuTensor,
12648 w1: &crate::model::GpuTensor,
12649 aq: &CudaSlice<i8>,
12650 ad: &CudaSlice<f32>,
12651 m: usize,
12652 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12653 {
12654 use crate::model::GpuTensor;
12655 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12656 return Ok(None);
12657 }
12658 if !self.mmvq_supports(QT_NVFP4) {
12668 return Ok(None);
12669 }
12670 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12671 if w1.in_features() != in_f || w1.out_features() != out_f {
12672 return Ok(None);
12673 }
12674 let no_mirror =
12687 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12688 if self.q8_ffn_fuse2_on()
12689 && no_mirror(w0)
12690 && no_mirror(w1)
12691 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12692 {
12693 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12694 return Ok(Some(((y0, 1.0), (y1, 1.0))));
12695 }
12696 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12706 let (y0, y1) =
12707 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12708 return Ok(Some(((y0, p0.3), (y1, p1.3))));
12709 }
12710 let (b0, q0, rb0, s0, rp0) = match w0 {
12711 GpuTensor::Quant {
12712 bytes,
12713 qtype,
12714 row_bytes,
12715 scale,
12716 rp,
12717 ..
12718 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12719 _ => return Ok(None),
12720 };
12721 let (b1, q1, rb1, s1, rp1) = match w1 {
12722 GpuTensor::Quant {
12723 bytes,
12724 qtype,
12725 row_bytes,
12726 scale,
12727 rp,
12728 ..
12729 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12730 _ => return Ok(None),
12731 };
12732 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12733 return Ok(None);
12734 }
12735 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12737 let rows_per_block = ROWS_PER_BLOCK * RPW;
12738 let f = self.func(if rp0 {
12739 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12740 } else {
12741 "qmatvec_nvfp4_mmvq_dual_mr2"
12742 });
12743 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12744 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12745 let cfg = LaunchConfig {
12746 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12747 block_dim: (32, ROWS_PER_BLOCK, 1),
12748 shared_mem_bytes: 0,
12749 };
12750 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12751 let one = 1.0f32;
12754 let __s_b = self.gpu.stream();
12755 let mut b = __s_b.launch_builder(&f);
12756 b.arg(b0)
12757 .arg(b1)
12758 .arg(aq)
12759 .arg(ad)
12760 .arg(&mut y0)
12761 .arg(&mut y1)
12762 .arg(&inf)
12763 .arg(&outf)
12764 .arg(&mi)
12765 .arg(&rb)
12766 .arg(&one)
12767 .arg(&one);
12768 unsafe {
12769 b.launch(cfg)?;
12770 }
12771 Ok(Some(((y0, s0), (y1, s1))))
12772 }
12773
12774 #[allow(clippy::too_many_arguments)]
12782 pub fn matmul_nvfp4_fused3(
12783 &self,
12784 w0: &crate::model::GpuTensor,
12785 w1: &crate::model::GpuTensor,
12786 w2: &crate::model::GpuTensor,
12787 aq: &CudaSlice<i8>,
12788 ad: &CudaSlice<f32>,
12789 m: usize,
12790 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12791 {
12792 use crate::model::GpuTensor;
12793 if !self.mmvq_supports(QT_NVFP4)
12800 || !self.uses_q8_1_fast(w0)
12801 || !self.uses_q8_1_fast(w1)
12802 || !self.uses_q8_1_fast(w2)
12803 {
12804 return Ok(None);
12805 }
12806 if (9..=16).contains(&m) {
12809 return Ok(
12810 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12811 Some(mut ys) => {
12812 let y2 = ys.pop().unwrap();
12813 let y1 = ys.pop().unwrap();
12814 let y0 = ys.pop().unwrap();
12815 Some((y0, y1, y2))
12816 }
12817 None => None,
12818 },
12819 );
12820 }
12821 if !(1..=8).contains(&m) {
12822 return Ok(None);
12823 }
12824 if m > 1 {
12825 let in_f = w0.in_features();
12826 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12827 || !self.batched_supports(QT_NVFP4)
12828 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12829 || (m > 4 && !Self::b8_enabled())
12830 || in_f % 512 != 0
12831 || in_f / 64 > 272
12832 {
12833 return Ok(None);
12834 }
12835 }
12836 let unpack = |w: &crate::model::GpuTensor| match w {
12837 GpuTensor::Quant {
12838 bytes,
12839 qtype,
12840 scale,
12841 rp,
12842 ..
12843 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12844 _ => None,
12845 };
12846 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12847 return Ok(None);
12848 };
12849 let in_f = w0.in_features();
12850 if w1.in_features() != in_f || w2.in_features() != in_f {
12851 return Ok(None);
12852 }
12853 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12854 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12856 let rows_pb = ROWS_PER_BLOCK * RPW;
12857 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12858 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12859 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12860 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12861 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12862 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12865 if m > 1 {
12866 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12868 return Ok(None);
12869 }
12870 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12871 let cfg = LaunchConfig {
12872 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12873 block_dim: (32, ROWS_PER_BLOCK, 1),
12874 shared_mem_bytes: 0,
12875 };
12876 let __s_b = self.gpu.stream();
12877 let mut b = __s_b.launch_builder(&f);
12878 b.arg(b0)
12879 .arg(b1)
12880 .arg(b2)
12881 .arg(aq)
12882 .arg(ad)
12883 .arg(&mut y0)
12884 .arg(&mut y1)
12885 .arg(&mut y2)
12886 .arg(&inf)
12887 .arg(&oi0)
12888 .arg(&oi1)
12889 .arg(&oi2)
12890 .arg(&mi);
12891 unsafe {
12892 b.launch(cfg)?;
12893 }
12894 return Ok(Some((y0, y1, y2)));
12895 }
12896 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12897 let cfg = LaunchConfig {
12898 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12899 block_dim: (32, ROWS_PER_BLOCK, 1),
12900 shared_mem_bytes: 0,
12901 };
12902 let __s_b = self.gpu.stream();
12903 let mut b = __s_b.launch_builder(&f);
12904 b.arg(b0)
12905 .arg(b1)
12906 .arg(b2)
12907 .arg(aq)
12908 .arg(ad)
12909 .arg(&mut y0)
12910 .arg(&mut y1)
12911 .arg(&mut y2)
12912 .arg(&inf)
12913 .arg(&oi0)
12914 .arg(&oi1)
12915 .arg(&oi2)
12916 .arg(&mi)
12917 .arg(&p0.1)
12918 .arg(&p1.1)
12919 .arg(&p2.1);
12920 unsafe {
12921 b.launch(cfg)?;
12922 }
12923 Ok(Some((y0, y1, y2)))
12924 }
12925
12926 pub fn matmul_nvfp4_fused2(
12935 &self,
12936 w0: &crate::model::GpuTensor,
12937 w1: &crate::model::GpuTensor,
12938 aq: &CudaSlice<i8>,
12939 ad: &CudaSlice<f32>,
12940 m: usize,
12941 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12942 use crate::model::GpuTensor;
12943 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12944 let off =
12945 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12946 if off
12949 || m != 1
12950 || !self.mmvq_supports(QT_NVFP4)
12951 || !self.uses_q8_1_fast(w0)
12952 || !self.uses_q8_1_fast(w1)
12953 {
12954 return Ok(None);
12955 }
12956 let unpack = |w: &crate::model::GpuTensor| match w {
12957 GpuTensor::Quant {
12958 bytes,
12959 qtype,
12960 scale,
12961 rp,
12962 ..
12963 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12964 _ => None,
12965 };
12966 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12967 return Ok(None);
12968 };
12969 let in_f = w0.in_features();
12970 if w1.in_features() != in_f {
12971 return Ok(None);
12972 }
12973 let (o0, o1) = (w0.out_features(), w1.out_features());
12974 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12976 let rows_pb = ROWS_PER_BLOCK * RPW;
12977 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12978 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12979 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12980 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12981 let cfg = LaunchConfig {
12982 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12983 block_dim: (32, ROWS_PER_BLOCK, 1),
12984 shared_mem_bytes: 0,
12985 };
12986 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12987 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12990 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12993 {
12994 use cudarc::driver::{DevicePtr, DevicePtrMut};
12995 let s = &self.gpu.stream();
12996 let (pw0, _g0) = b0.device_ptr(s);
12997 let (pw1, _g1) = b1.device_ptr(s);
12998 let (paq, _g2) = aq.device_ptr(s);
12999 let (pad, _g3) = ad.device_ptr(s);
13000 let (py0, _g4) = y0.device_ptr_mut(s);
13001 let (py1, _g5) = y1.device_ptr_mut(s);
13002 let (s0, s1) = (p0.1, p1.1);
13003 let mut ps = [
13004 &pw0 as *const _ as *mut std::ffi::c_void,
13005 &pw1 as *const _ as *mut _,
13006 &paq as *const _ as *mut _,
13007 &pad as *const _ as *mut _,
13008 &py0 as *const _ as *mut _,
13009 &py1 as *const _ as *mut _,
13010 &inf as *const _ as *mut _,
13011 &oi0 as *const _ as *mut _,
13012 &oi1 as *const _ as *mut _,
13013 &mi as *const _ as *mut _,
13014 &s0 as *const _ as *mut _,
13015 &s1 as *const _ as *mut _,
13016 ];
13017 unsafe {
13018 self.launch_pdl(
13019 "qmatvec_nvfp4_mmvq_fused2_rp",
13020 cfg.grid_dim,
13021 cfg.block_dim,
13022 &mut ps,
13023 )?;
13024 }
13025 }
13026 return Ok(Some((y0, y1)));
13027 }
13028 let __s_b = self.gpu.stream();
13029 let mut b = __s_b.launch_builder(&f);
13030 b.arg(b0)
13031 .arg(b1)
13032 .arg(aq)
13033 .arg(ad)
13034 .arg(&mut y0)
13035 .arg(&mut y1)
13036 .arg(&inf)
13037 .arg(&oi0)
13038 .arg(&oi1)
13039 .arg(&mi)
13040 .arg(&p0.1)
13041 .arg(&p1.1);
13042 unsafe {
13043 b.launch(cfg)?;
13044 }
13045 Ok(Some((y0, y1)))
13046 }
13047
13048 pub fn matmul_nvfp4_fused2_into(
13053 &self,
13054 w0: &crate::model::GpuTensor,
13055 w1: &crate::model::GpuTensor,
13056 aq: &CudaSlice<i8>,
13057 ad: &CudaSlice<f32>,
13058 y0: &mut CudaSlice<f32>,
13059 y1: &mut CudaSlice<f32>,
13060 ) -> Result<bool, Box<dyn std::error::Error>> {
13061 use crate::model::GpuTensor;
13062 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13063 let off =
13064 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13065 if off
13066 || !self.mmvq_supports(QT_NVFP4)
13067 || !self.uses_q8_1_fast(w0)
13068 || !self.uses_q8_1_fast(w1)
13069 {
13070 return Ok(false);
13071 }
13072 let unpack = |w: &crate::model::GpuTensor| match w {
13073 GpuTensor::Quant {
13074 bytes,
13075 qtype,
13076 scale,
13077 rp,
13078 ..
13079 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13080 _ => None,
13081 };
13082 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13083 return Ok(false);
13084 };
13085 let in_f = w0.in_features();
13086 if w1.in_features() != in_f {
13087 return Ok(false);
13088 }
13089 let (o0, o1) = (w0.out_features(), w1.out_features());
13090 if y0.len() < o0 || y1.len() < o1 {
13091 return Ok(false);
13092 }
13093 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13095 let rows_pb = ROWS_PER_BLOCK * RPW;
13096 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13097 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13098 let cfg = LaunchConfig {
13099 grid_dim: (nb(o0) + nb(o1), 1, 1),
13100 block_dim: (32, ROWS_PER_BLOCK, 1),
13101 shared_mem_bytes: 0,
13102 };
13103 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13104 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13107 let __s_b = self.gpu.stream();
13108 let mut b = __s_b.launch_builder(&f);
13109 b.arg(b0)
13110 .arg(b1)
13111 .arg(aq)
13112 .arg(ad)
13113 .arg(&mut *y0)
13114 .arg(&mut *y1)
13115 .arg(&inf)
13116 .arg(&oi0)
13117 .arg(&oi1)
13118 .arg(&mi)
13119 .arg(&p0.1)
13120 .arg(&p1.1);
13121 unsafe {
13122 b.launch(cfg)?;
13123 }
13124 Ok(true)
13125 }
13126
13127 #[allow(clippy::type_complexity)]
13132 pub fn matmul_nvfp4_fused4(
13133 &self,
13134 w0: &crate::model::GpuTensor,
13135 w1: &crate::model::GpuTensor,
13136 w2: &crate::model::GpuTensor,
13137 w3: &crate::model::GpuTensor,
13138 aq: &CudaSlice<i8>,
13139 ad: &CudaSlice<f32>,
13140 m: usize,
13141 ) -> Result<
13142 Option<(
13143 CudaSlice<f32>,
13144 CudaSlice<f32>,
13145 CudaSlice<f32>,
13146 CudaSlice<f32>,
13147 )>,
13148 Box<dyn std::error::Error>,
13149 > {
13150 use crate::model::GpuTensor;
13151 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13158 || !self.mmvq_supports(QT_NVFP4)
13159 || !self.uses_q8_1_fast(w0)
13160 || !self.uses_q8_1_fast(w1)
13161 || !self.uses_q8_1_fast(w2)
13162 || !self.uses_q8_1_fast(w3)
13163 {
13164 return Ok(None);
13165 }
13166 if (9..=16).contains(&m) {
13171 return Ok(
13172 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13173 Some(mut ys) => {
13174 let y3 = ys.pop().unwrap();
13175 let y2 = ys.pop().unwrap();
13176 let y1 = ys.pop().unwrap();
13177 let y0 = ys.pop().unwrap();
13178 Some((y0, y1, y2, y3))
13179 }
13180 None => None,
13181 },
13182 );
13183 }
13184 if !(1..=8).contains(&m) {
13185 return Ok(None);
13186 }
13187 if m > 1 {
13188 let in_f = w0.in_features();
13191 if !self.batched_supports(QT_NVFP4)
13192 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13193 || (m > 4 && !Self::b8_enabled())
13194 || in_f % 512 != 0
13195 || in_f / 64 > 272
13196 {
13197 return Ok(None);
13198 }
13199 }
13200 let unpack = |w: &crate::model::GpuTensor| match w {
13201 GpuTensor::Quant {
13202 bytes,
13203 qtype,
13204 scale,
13205 rp,
13206 ..
13207 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13208 _ => None,
13209 };
13210 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13211 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13212 else {
13213 return Ok(None);
13214 };
13215 let in_f = w0.in_features();
13216 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13217 return Ok(None);
13218 }
13219 let (o0, o1, o2, o3) = (
13220 w0.out_features(),
13221 w1.out_features(),
13222 w2.out_features(),
13223 w3.out_features(),
13224 );
13225 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13227 let rows_pb = ROWS_PER_BLOCK * RPW;
13228 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13229 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13230 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13231 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13232 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13233 let (inf, oi0, oi1, oi2, oi3, mi) = (
13234 in_f as i32,
13235 o0 as i32,
13236 o1 as i32,
13237 o2 as i32,
13238 o3 as i32,
13239 m as i32,
13240 );
13241 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13244 if m > 1 {
13245 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13248 return Ok(None);
13249 }
13250 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13251 let cfg = LaunchConfig {
13252 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13253 block_dim: (32, ROWS_PER_BLOCK, 1),
13254 shared_mem_bytes: 0,
13255 };
13256 let __s_b = self.gpu.stream();
13257 let mut b = __s_b.launch_builder(&f);
13258 b.arg(b0)
13259 .arg(b1)
13260 .arg(b2)
13261 .arg(b3)
13262 .arg(aq)
13263 .arg(ad)
13264 .arg(&mut y0)
13265 .arg(&mut y1)
13266 .arg(&mut y2)
13267 .arg(&mut y3)
13268 .arg(&inf)
13269 .arg(&oi0)
13270 .arg(&oi1)
13271 .arg(&oi2)
13272 .arg(&oi3)
13273 .arg(&mi);
13274 unsafe {
13275 b.launch(cfg)?;
13276 }
13277 return Ok(Some((y0, y1, y2, y3)));
13278 }
13279 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
13280 let cfg = LaunchConfig {
13281 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
13282 block_dim: (32, ROWS_PER_BLOCK, 1),
13283 shared_mem_bytes: 0,
13284 };
13285 let __s_b = self.gpu.stream();
13286 let mut b = __s_b.launch_builder(&f);
13287 b.arg(b0)
13288 .arg(b1)
13289 .arg(b2)
13290 .arg(b3)
13291 .arg(aq)
13292 .arg(ad)
13293 .arg(&mut y0)
13294 .arg(&mut y1)
13295 .arg(&mut y2)
13296 .arg(&mut y3)
13297 .arg(&inf)
13298 .arg(&oi0)
13299 .arg(&oi1)
13300 .arg(&oi2)
13301 .arg(&oi3)
13302 .arg(&mi)
13303 .arg(&p0.1)
13304 .arg(&p1.1)
13305 .arg(&p2.1)
13306 .arg(&p3.1);
13307 unsafe {
13308 b.launch(cfg)?;
13309 }
13310 Ok(Some((y0, y1, y2, y3)))
13311 }
13312
13313 pub fn matmul_q8_fused2(
13321 &self,
13322 w0: &crate::model::GpuTensor,
13323 w1: &crate::model::GpuTensor,
13324 aq: &CudaSlice<i8>,
13325 ad: &CudaSlice<f32>,
13326 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13327 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13333 return Ok(Some(self.e4m3_fused2_core(
13334 p0.0,
13335 p1.0,
13336 aq,
13337 ad,
13338 w0.in_features(),
13339 p0.1,
13340 p1.1,
13341 p0.2,
13342 p0.3,
13343 p1.3,
13344 )?));
13345 }
13346 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13347 return Ok(None);
13348 };
13349 Ok(Some(self.q8_fused2_core(
13350 p0.0,
13351 p1.0,
13352 aq,
13353 ad,
13354 w0.in_features(),
13355 p0.1,
13356 p1.1,
13357 p0.2,
13358 )?))
13359 }
13360
13361 #[allow(clippy::too_many_arguments)]
13362 fn q8_fused2_core(
13363 &self,
13364 b0: &CudaSlice<u8>,
13365 b1: &CudaSlice<u8>,
13366 aq: &CudaSlice<i8>,
13367 ad: &CudaSlice<f32>,
13368 in_f: usize,
13369 out0: usize,
13370 out1: usize,
13371 row_bytes: usize,
13372 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13373 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13375 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13376 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13377 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13378 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13379 let cfg = LaunchConfig {
13380 grid_dim: (nb0 + nb1, 1, 1),
13381 block_dim: (32, ROWS_PER_BLOCK, 1),
13382 shared_mem_bytes: 0,
13383 };
13384 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13385 let __s_b = self.gpu.stream();
13386 let mut b = __s_b.launch_builder(&f);
13387 b.arg(b0)
13388 .arg(b1)
13389 .arg(aq)
13390 .arg(ad)
13391 .arg(&mut y0)
13392 .arg(&mut y1)
13393 .arg(&inf)
13394 .arg(&o0)
13395 .arg(&o1)
13396 .arg(&rbl);
13397 unsafe {
13398 b.launch(cfg)?;
13399 }
13400 Ok((y0, y1))
13401 }
13402
13403 pub fn matmul_q8_fused2_x(
13409 &self,
13410 w0: &crate::model::GpuTensor,
13411 w1: &crate::model::GpuTensor,
13412 x: &CudaSlice<f32>,
13413 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13414 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13415 return Ok(None);
13416 }
13417 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13418 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13419 return Ok(Some(self.e4m3_fused2_core(
13420 p0.0,
13421 p1.0,
13422 &aq,
13423 &ad,
13424 w0.in_features(),
13425 p0.1,
13426 p1.1,
13427 p0.2,
13428 p0.3,
13429 p1.3,
13430 )?));
13431 }
13432 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13433 return Ok(None);
13434 };
13435 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13436 Ok(Some(self.q8_fused2_core(
13437 p0.0,
13438 p1.0,
13439 &aq,
13440 &ad,
13441 w0.in_features(),
13442 p0.1,
13443 p1.1,
13444 p0.2,
13445 )?))
13446 }
13447
13448 #[allow(clippy::too_many_arguments)]
13451 pub fn qmatvec_q8_fused2_raw(
13452 &self,
13453 b0: &CudaSlice<u8>,
13454 b1: &CudaSlice<u8>,
13455 x: &CudaSlice<f32>,
13456 in_f: usize,
13457 out0: usize,
13458 out1: usize,
13459 row_bytes: usize,
13460 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13461 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13462 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13463 }
13464
13465 pub fn matmul_q4_fused3(
13471 &self,
13472 w0: &crate::model::GpuTensor,
13473 w1: &crate::model::GpuTensor,
13474 w2: &crate::model::GpuTensor,
13475 aq: &CudaSlice<i8>,
13476 ad: &CudaSlice<f32>,
13477 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13478 {
13479 use crate::model::GpuTensor;
13480 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13481 match w {
13482 GpuTensor::Quant {
13483 qtype, row_bytes, ..
13484 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13485 _ => None,
13486 }
13487 };
13488 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13489 return Ok(None);
13490 };
13491 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13492 return Ok(None);
13493 }
13494 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13498 match w {
13499 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13500 Some(m) => (m, true),
13501 None => (bytes, *rp),
13502 },
13503 _ => unreachable!(),
13504 }
13505 }
13506 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13507 if rp0 != rp1 || rp1 != rp2 {
13508 return Ok(None);
13509 }
13510 let rp = rp0;
13511 let rpb: u32 = 4;
13512 let mr1 = rp && Self::q40_mr1_on();
13516 let nb = |o: usize| {
13517 if mr1 {
13518 (o as u32).div_ceil(rpb)
13519 } else {
13520 (o as u32).div_ceil(2).div_ceil(rpb)
13521 }
13522 };
13523 let grid = nb(o0) + nb(o1) + nb(o2);
13524 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13525 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13526 let mut y2 = self.alloc_uninit::<f32>(o2)?;
13527 let f = self.func(if mr1 {
13528 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13529 } else if rp {
13530 "qmatvec_q4_0_mmvq_fused3_rp"
13531 } else {
13532 "qmatvec_q4_0_mmvq_fused3"
13533 });
13534 let cfg = LaunchConfig {
13535 grid_dim: (grid, 1, 1),
13536 block_dim: (32, rpb, 1),
13537 shared_mem_bytes: 0,
13538 };
13539 let inf = w0.in_features() as i32;
13540 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13541 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13542 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13545 {
13546 use cudarc::driver::{DevicePtr, DevicePtrMut};
13547 let s = &self.gpu.stream();
13548 let (p0, _g0) = b0.device_ptr(s);
13549 let (p1, _g1) = b1.device_ptr(s);
13550 let (p2, _g2) = b2.device_ptr(s);
13551 let (paq, _g3) = aq.device_ptr(s);
13552 let (pad, _g4) = ad.device_ptr(s);
13553 let (py0, _g5) = y0.device_ptr_mut(s);
13554 let (py1, _g6) = y1.device_ptr_mut(s);
13555 let (py2, _g7) = y2.device_ptr_mut(s);
13556 let mut ps = [
13557 &p0 as *const _ as *mut std::ffi::c_void,
13558 &p1 as *const _ as *mut _,
13559 &p2 as *const _ as *mut _,
13560 &paq as *const _ as *mut _,
13561 &pad as *const _ as *mut _,
13562 &py0 as *const _ as *mut _,
13563 &py1 as *const _ as *mut _,
13564 &py2 as *const _ as *mut _,
13565 &inf as *const _ as *mut _,
13566 &oo0 as *const _ as *mut _,
13567 &oo1 as *const _ as *mut _,
13568 &oo2 as *const _ as *mut _,
13569 &r0 as *const _ as *mut _,
13570 &r1 as *const _ as *mut _,
13571 &r2 as *const _ as *mut _,
13572 ];
13573 unsafe {
13574 self.launch_pdl(
13575 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13576 (grid, 1, 1),
13577 (32, rpb, 1),
13578 &mut ps,
13579 )?;
13580 }
13581 }
13582 return Ok(Some((y0, y1, y2)));
13583 }
13584 let __s_b = self.gpu.stream();
13585 let mut b = __s_b.launch_builder(&f);
13586 b.arg(b0)
13587 .arg(b1)
13588 .arg(b2)
13589 .arg(aq)
13590 .arg(ad)
13591 .arg(&mut y0)
13592 .arg(&mut y1)
13593 .arg(&mut y2)
13594 .arg(&inf)
13595 .arg(&oo0)
13596 .arg(&oo1)
13597 .arg(&oo2)
13598 .arg(&r0)
13599 .arg(&r1)
13600 .arg(&r2);
13601 unsafe {
13602 b.launch(cfg)?;
13603 }
13604 Ok(Some((y0, y1, y2)))
13605 }
13606
13607 #[allow(clippy::too_many_arguments)]
13610 pub fn matmul_q4_fused3_into(
13611 &self,
13612 w0: &crate::model::GpuTensor,
13613 w1: &crate::model::GpuTensor,
13614 w2: &crate::model::GpuTensor,
13615 aq: &CudaSlice<i8>,
13616 ad: &CudaSlice<f32>,
13617 y0: &mut CudaSlice<f32>,
13618 y1: &mut CudaSlice<f32>,
13619 y2: &mut CudaSlice<f32>,
13620 ) -> Result<bool, Box<dyn std::error::Error>> {
13621 use crate::model::GpuTensor;
13622 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13623 match w {
13624 GpuTensor::Quant {
13625 qtype, row_bytes, ..
13626 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13627 _ => None,
13628 }
13629 };
13630 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13631 return Ok(false);
13632 };
13633 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13634 return Ok(false);
13635 }
13636 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13637 match w {
13638 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13639 Some(m) => (m, true),
13640 None => (bytes, *rp),
13641 },
13642 _ => unreachable!(),
13643 }
13644 }
13645 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13646 if rp0 != rp1 || rp1 != rp2 {
13647 return Ok(false);
13648 }
13649 let rp = rp0;
13650 let rpb: u32 = 4;
13651 let mr1 = rp && Self::q40_mr1_on();
13652 let nb = |o: usize| {
13653 if mr1 {
13654 (o as u32).div_ceil(rpb)
13655 } else {
13656 (o as u32).div_ceil(2).div_ceil(rpb)
13657 }
13658 };
13659 let grid = nb(o0) + nb(o1) + nb(o2);
13660 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13661 let f = self.func(if mr1 {
13662 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13663 } else if rp {
13664 "qmatvec_q4_0_mmvq_fused3_rp"
13665 } else {
13666 "qmatvec_q4_0_mmvq_fused3"
13667 });
13668 let cfg = LaunchConfig {
13669 grid_dim: (grid, 1, 1),
13670 block_dim: (32, rpb, 1),
13671 shared_mem_bytes: 0,
13672 };
13673 let inf = w0.in_features() as i32;
13674 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13675 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13676 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13678 use cudarc::driver::{DevicePtr, DevicePtrMut};
13679 let s = &self.gpu.stream();
13680 let (p0, _g0) = b0.device_ptr(s);
13681 let (p1, _g1) = b1.device_ptr(s);
13682 let (p2, _g2) = b2.device_ptr(s);
13683 let (paq, _g3) = aq.device_ptr(s);
13684 let (pad, _g4) = ad.device_ptr(s);
13685 let (py0, _g5) = y0.device_ptr_mut(s);
13686 let (py1, _g6) = y1.device_ptr_mut(s);
13687 let (py2, _g7) = y2.device_ptr_mut(s);
13688 let mut ps = [
13689 &p0 as *const _ as *mut std::ffi::c_void,
13690 &p1 as *const _ as *mut _,
13691 &p2 as *const _ as *mut _,
13692 &paq as *const _ as *mut _,
13693 &pad as *const _ as *mut _,
13694 &py0 as *const _ as *mut _,
13695 &py1 as *const _ as *mut _,
13696 &py2 as *const _ as *mut _,
13697 &inf as *const _ as *mut _,
13698 &oo0 as *const _ as *mut _,
13699 &oo1 as *const _ as *mut _,
13700 &oo2 as *const _ as *mut _,
13701 &r0 as *const _ as *mut _,
13702 &r1 as *const _ as *mut _,
13703 &r2 as *const _ as *mut _,
13704 ];
13705 unsafe {
13706 self.launch_pdl(
13707 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13708 (grid, 1, 1),
13709 (32, rpb, 1),
13710 &mut ps,
13711 )?;
13712 }
13713 return Ok(true);
13714 }
13715 let __s_b = self.gpu.stream();
13716 let mut b = __s_b.launch_builder(&f);
13717 b.arg(b0)
13718 .arg(b1)
13719 .arg(b2)
13720 .arg(aq)
13721 .arg(ad)
13722 .arg(&mut *y0)
13723 .arg(&mut *y1)
13724 .arg(&mut *y2)
13725 .arg(&inf)
13726 .arg(&oo0)
13727 .arg(&oo1)
13728 .arg(&oo2)
13729 .arg(&r0)
13730 .arg(&r1)
13731 .arg(&r2);
13732 unsafe {
13733 b.launch(cfg)?;
13734 }
13735 Ok(true)
13736 }
13737
13738 pub fn matmul_q4_fused2(
13740 &self,
13741 w0: &crate::model::GpuTensor,
13742 w1: &crate::model::GpuTensor,
13743 aq: &CudaSlice<i8>,
13744 ad: &CudaSlice<f32>,
13745 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13746 use crate::model::GpuTensor;
13747 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13748 match w {
13749 GpuTensor::Quant {
13750 qtype, row_bytes, ..
13751 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13752 _ => None,
13753 }
13754 };
13755 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13756 return Ok(None);
13757 };
13758 if w0.in_features() != w1.in_features() {
13759 return Ok(None);
13760 }
13761 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13763 match w {
13764 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13765 Some(m) => (m, true),
13766 None => (bytes, *rp),
13767 },
13768 _ => unreachable!(),
13769 }
13770 }
13771 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13772 if rp0 != rp1 {
13773 return Ok(None);
13774 }
13775 let rp = rp0;
13776 let rpb: u32 = 4;
13777 let mr1 = rp && Self::q40_mr1_on();
13779 let nb = |o: usize| {
13780 if mr1 {
13781 (o as u32).div_ceil(rpb)
13782 } else {
13783 (o as u32).div_ceil(2).div_ceil(rpb)
13784 }
13785 };
13786 let grid = nb(o0) + nb(o1);
13787 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13788 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13789 let f = self.func(if mr1 {
13790 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13791 } else if rp {
13792 "qmatvec_q4_0_mmvq_fused2_rp"
13793 } else {
13794 "qmatvec_q4_0_mmvq_fused2"
13795 });
13796 let cfg = LaunchConfig {
13797 grid_dim: (grid, 1, 1),
13798 block_dim: (32, rpb, 1),
13799 shared_mem_bytes: 0,
13800 };
13801 let inf = w0.in_features() as i32;
13802 let (oo0, oo1) = (o0 as i32, o1 as i32);
13803 let (r0, r1) = (rb0 as i64, rb1 as i64);
13804 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13806 {
13807 use cudarc::driver::{DevicePtr, DevicePtrMut};
13808 let s = &self.gpu.stream();
13809 let (p0, _g0) = b0.device_ptr(s);
13810 let (p1, _g1) = b1.device_ptr(s);
13811 let (paq, _g2) = aq.device_ptr(s);
13812 let (pad, _g3) = ad.device_ptr(s);
13813 let (py0, _g4) = y0.device_ptr_mut(s);
13814 let (py1, _g5) = y1.device_ptr_mut(s);
13815 let mut ps = [
13816 &p0 as *const _ as *mut std::ffi::c_void,
13817 &p1 as *const _ as *mut _,
13818 &paq as *const _ as *mut _,
13819 &pad as *const _ as *mut _,
13820 &py0 as *const _ as *mut _,
13821 &py1 as *const _ as *mut _,
13822 &inf as *const _ as *mut _,
13823 &oo0 as *const _ as *mut _,
13824 &oo1 as *const _ as *mut _,
13825 &r0 as *const _ as *mut _,
13826 &r1 as *const _ as *mut _,
13827 ];
13828 unsafe {
13829 self.launch_pdl(
13830 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13831 (grid, 1, 1),
13832 (32, rpb, 1),
13833 &mut ps,
13834 )?;
13835 }
13836 }
13837 return Ok(Some((y0, y1)));
13838 }
13839 let __s_b = self.gpu.stream();
13840 let mut b = __s_b.launch_builder(&f);
13841 b.arg(b0)
13842 .arg(b1)
13843 .arg(aq)
13844 .arg(ad)
13845 .arg(&mut y0)
13846 .arg(&mut y1)
13847 .arg(&inf)
13848 .arg(&oo0)
13849 .arg(&oo1)
13850 .arg(&r0)
13851 .arg(&r1);
13852 unsafe {
13853 b.launch(cfg)?;
13854 }
13855 Ok(Some((y0, y1)))
13856 }
13857
13858 pub fn matmul_q4_fused2_into(
13860 &self,
13861 w0: &crate::model::GpuTensor,
13862 w1: &crate::model::GpuTensor,
13863 aq: &CudaSlice<i8>,
13864 ad: &CudaSlice<f32>,
13865 y0: &mut CudaSlice<f32>,
13866 y1: &mut CudaSlice<f32>,
13867 ) -> Result<bool, Box<dyn std::error::Error>> {
13868 use crate::model::GpuTensor;
13869 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13870 match w {
13871 GpuTensor::Quant {
13872 qtype, row_bytes, ..
13873 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13874 _ => None,
13875 }
13876 };
13877 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13878 return Ok(false);
13879 };
13880 if w0.in_features() != w1.in_features() {
13881 return Ok(false);
13882 }
13883 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13884 match w {
13885 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13886 Some(m) => (m, true),
13887 None => (bytes, *rp),
13888 },
13889 _ => unreachable!(),
13890 }
13891 }
13892 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13893 if rp0 != rp1 {
13894 return Ok(false);
13895 }
13896 let rp = rp0;
13897 let rpb: u32 = 4;
13898 let mr1 = rp && Self::q40_mr1_on();
13899 let nb = |o: usize| {
13900 if mr1 {
13901 (o as u32).div_ceil(rpb)
13902 } else {
13903 (o as u32).div_ceil(2).div_ceil(rpb)
13904 }
13905 };
13906 let grid = nb(o0) + nb(o1);
13907 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13908 let f = self.func(if mr1 {
13909 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13910 } else if rp {
13911 "qmatvec_q4_0_mmvq_fused2_rp"
13912 } else {
13913 "qmatvec_q4_0_mmvq_fused2"
13914 });
13915 let cfg = LaunchConfig {
13916 grid_dim: (grid, 1, 1),
13917 block_dim: (32, rpb, 1),
13918 shared_mem_bytes: 0,
13919 };
13920 let inf = w0.in_features() as i32;
13921 let (oo0, oo1) = (o0 as i32, o1 as i32);
13922 let (r0, r1) = (rb0 as i64, rb1 as i64);
13923 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13925 use cudarc::driver::{DevicePtr, DevicePtrMut};
13926 let s = &self.gpu.stream();
13927 let (p0, _g0) = b0.device_ptr(s);
13928 let (p1, _g1) = b1.device_ptr(s);
13929 let (paq, _g2) = aq.device_ptr(s);
13930 let (pad, _g3) = ad.device_ptr(s);
13931 let (py0, _g4) = y0.device_ptr_mut(s);
13932 let (py1, _g5) = y1.device_ptr_mut(s);
13933 let mut ps = [
13934 &p0 as *const _ as *mut std::ffi::c_void,
13935 &p1 as *const _ as *mut _,
13936 &paq as *const _ as *mut _,
13937 &pad as *const _ as *mut _,
13938 &py0 as *const _ as *mut _,
13939 &py1 as *const _ as *mut _,
13940 &inf as *const _ as *mut _,
13941 &oo0 as *const _ as *mut _,
13942 &oo1 as *const _ as *mut _,
13943 &r0 as *const _ as *mut _,
13944 &r1 as *const _ as *mut _,
13945 ];
13946 unsafe {
13947 self.launch_pdl(
13948 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13949 (grid, 1, 1),
13950 (32, rpb, 1),
13951 &mut ps,
13952 )?;
13953 }
13954 return Ok(true);
13955 }
13956 let __s_b = self.gpu.stream();
13957 let mut b = __s_b.launch_builder(&f);
13958 b.arg(b0)
13959 .arg(b1)
13960 .arg(aq)
13961 .arg(ad)
13962 .arg(&mut *y0)
13963 .arg(&mut *y1)
13964 .arg(&inf)
13965 .arg(&oo0)
13966 .arg(&oo1)
13967 .arg(&r0)
13968 .arg(&r1);
13969 unsafe {
13970 b.launch(cfg)?;
13971 }
13972 Ok(true)
13973 }
13974
13975 pub fn matmul_q4_fused2_batched(
13980 &self,
13981 w0: &crate::model::GpuTensor,
13982 w1: &crate::model::GpuTensor,
13983 aq: &CudaSlice<i8>,
13984 ad: &CudaSlice<f32>,
13985 m: usize,
13986 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13987 use crate::model::GpuTensor;
13988 if m < 2 || m > 8 {
13989 return Ok(None);
13990 }
13991 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13992 match w {
13993 GpuTensor::Quant {
13994 qtype, row_bytes, ..
13995 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13996 _ => None,
13997 }
13998 };
13999 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14000 return Ok(None);
14001 };
14002 if w0.in_features() != w1.in_features() {
14003 return Ok(None);
14004 }
14005 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14006 match w {
14007 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14008 Some(mr) => (mr, true),
14009 None => (bytes, *rp),
14010 },
14011 _ => unreachable!(),
14012 }
14013 }
14014 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14015 if !rp0 || !rp1 {
14016 return Ok(None);
14017 }
14018 let mcols = Self::batched_mcols(m);
14019 let rpb: u32 = 4;
14020 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14021 let grid = nb(o0) + nb(o1);
14022 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14023 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14024 let f = self.func(match mcols {
14025 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14026 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14027 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14028 });
14029 let cfg = LaunchConfig {
14030 grid_dim: (grid, 1, 1),
14031 block_dim: (32, rpb, 1),
14032 shared_mem_bytes: 0,
14033 };
14034 let inf = w0.in_features() as i32;
14035 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14036 let rb = rb0 as i64;
14037 let __s_b = self.gpu.stream();
14038 let mut b = __s_b.launch_builder(&f);
14039 b.arg(b0)
14040 .arg(b1)
14041 .arg(aq)
14042 .arg(ad)
14043 .arg(&mut y0)
14044 .arg(&mut y1)
14045 .arg(&inf)
14046 .arg(&oo0)
14047 .arg(&oo1)
14048 .arg(&mi)
14049 .arg(&rb);
14050 unsafe {
14051 b.launch(cfg)?;
14052 }
14053 Ok(Some((y0, y1)))
14054 }
14055
14056 #[allow(clippy::too_many_arguments)]
14059 pub fn matmul_q4_fused3_batched(
14060 &self,
14061 w0: &crate::model::GpuTensor,
14062 w1: &crate::model::GpuTensor,
14063 w2: &crate::model::GpuTensor,
14064 aq: &CudaSlice<i8>,
14065 ad: &CudaSlice<f32>,
14066 m: usize,
14067 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14068 {
14069 use crate::model::GpuTensor;
14070 if m < 2 || m > 8 {
14071 return Ok(None);
14072 }
14073 let q4 = |w: &GpuTensor| -> Option<usize> {
14074 match w {
14075 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14076 _ => None,
14077 }
14078 };
14079 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14080 return Ok(None);
14081 };
14082 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14083 return Ok(None);
14084 }
14085 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14086 match w {
14087 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14088 Some(mr) => (mr, true),
14089 None => (bytes, *rp),
14090 },
14091 _ => unreachable!(),
14092 }
14093 }
14094 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14095 if !rp0 || !rp1 || !rp2 {
14096 return Ok(None);
14097 }
14098 let mcols = Self::batched_mcols(m);
14099 let rpb: u32 = 4;
14100 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14101 let grid = nb(o0) + nb(o1) + nb(o2);
14102 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14103 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14104 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14105 let f = self.func(match mcols {
14106 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14107 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14108 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14109 });
14110 let cfg = LaunchConfig {
14111 grid_dim: (grid, 1, 1),
14112 block_dim: (32, rpb, 1),
14113 shared_mem_bytes: 0,
14114 };
14115 let inf = w0.in_features() as i32;
14116 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14117 let rb = 0i64;
14118 let __s_b = self.gpu.stream();
14119 let mut b = __s_b.launch_builder(&f);
14120 b.arg(b0)
14121 .arg(b1)
14122 .arg(b2)
14123 .arg(aq)
14124 .arg(ad)
14125 .arg(&mut y0)
14126 .arg(&mut y1)
14127 .arg(&mut y2)
14128 .arg(&inf)
14129 .arg(&oo0)
14130 .arg(&oo1)
14131 .arg(&oo2)
14132 .arg(&mi)
14133 .arg(&rb);
14134 unsafe {
14135 b.launch(cfg)?;
14136 }
14137 Ok(Some((y0, y1, y2)))
14138 }
14139
14140 pub fn matmul_q8_fused3(
14141 &self,
14142 w0: &crate::model::GpuTensor,
14143 w1: &crate::model::GpuTensor,
14144 w2: &crate::model::GpuTensor,
14145 aq: &CudaSlice<i8>,
14146 ad: &CudaSlice<f32>,
14147 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14148 {
14149 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14152 return Ok(Some(self.e4m3_fused3_core(
14153 p0.0,
14154 p1.0,
14155 p2.0,
14156 aq,
14157 ad,
14158 w0.in_features(),
14159 p0.1,
14160 p1.1,
14161 p2.1,
14162 p0.2,
14163 p0.3,
14164 p1.3,
14165 p2.3,
14166 )?));
14167 }
14168 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14169 return Ok(None);
14170 };
14171 Ok(Some(self.q8_fused3_core(
14172 p0.0,
14173 p1.0,
14174 p2.0,
14175 aq,
14176 ad,
14177 w0.in_features(),
14178 p0.1,
14179 p1.1,
14180 p2.1,
14181 p0.2,
14182 )?))
14183 }
14184
14185 #[allow(clippy::too_many_arguments)]
14186 fn q8_fused3_core(
14187 &self,
14188 b0: &CudaSlice<u8>,
14189 b1: &CudaSlice<u8>,
14190 b2: &CudaSlice<u8>,
14191 aq: &CudaSlice<i8>,
14192 ad: &CudaSlice<f32>,
14193 in_f: usize,
14194 out0: usize,
14195 out1: usize,
14196 out2: usize,
14197 row_bytes: usize,
14198 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14199 const ROWS_PER_BLOCK: u32 = 4;
14200 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14201 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14202 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14203 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14204 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14205 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14206 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14207 let cfg = LaunchConfig {
14208 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14209 block_dim: (32, ROWS_PER_BLOCK, 1),
14210 shared_mem_bytes: 0,
14211 };
14212 let (inf, o0, o1, o2, rbl) = (
14213 in_f as i32,
14214 out0 as i32,
14215 out1 as i32,
14216 out2 as i32,
14217 row_bytes as i64,
14218 );
14219 let __s_b = self.gpu.stream();
14220 let mut b = __s_b.launch_builder(&f);
14221 b.arg(b0)
14222 .arg(b1)
14223 .arg(b2)
14224 .arg(aq)
14225 .arg(ad)
14226 .arg(&mut y0)
14227 .arg(&mut y1)
14228 .arg(&mut y2)
14229 .arg(&inf)
14230 .arg(&o0)
14231 .arg(&o1)
14232 .arg(&o2)
14233 .arg(&rbl);
14234 unsafe {
14235 b.launch(cfg)?;
14236 }
14237 Ok((y0, y1, y2))
14238 }
14239
14240 #[allow(clippy::too_many_arguments)]
14242 pub fn qmatvec_q8_fused3_raw(
14243 &self,
14244 b0: &CudaSlice<u8>,
14245 b1: &CudaSlice<u8>,
14246 b2: &CudaSlice<u8>,
14247 x: &CudaSlice<f32>,
14248 in_f: usize,
14249 out0: usize,
14250 out1: usize,
14251 out2: usize,
14252 row_bytes: usize,
14253 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14254 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14255 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14256 }
14257
14258 pub fn matmul_q8_fused2_t(
14269 &self,
14270 w0: &crate::model::GpuTensor,
14271 w1: &crate::model::GpuTensor,
14272 aq: &CudaSlice<i8>,
14273 ad: &CudaSlice<f32>,
14274 m: usize,
14275 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14276 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14280 return Ok(None);
14281 }
14282 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14285 if m > 4 && !Self::b8_enabled() {
14286 return Ok(None);
14287 }
14288 return Ok(Some(self.e4m3_fused2_t_core(
14289 p0.0,
14290 p1.0,
14291 aq,
14292 ad,
14293 m,
14294 w0.in_features(),
14295 p0.1,
14296 p1.1,
14297 p0.2,
14298 p0.3,
14299 p1.3,
14300 )?));
14301 }
14302 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14303 return Ok(None);
14304 };
14305 Ok(Some(self.q8_fused2_t_core(
14306 p0.0,
14307 p1.0,
14308 aq,
14309 ad,
14310 m,
14311 w0.in_features(),
14312 p0.1,
14313 p1.1,
14314 p0.2,
14315 )?))
14316 }
14317
14318 #[allow(clippy::too_many_arguments)]
14319 fn q8_fused2_t_core(
14320 &self,
14321 b0: &CudaSlice<u8>,
14322 b1: &CudaSlice<u8>,
14323 aq: &CudaSlice<i8>,
14324 ad: &CudaSlice<f32>,
14325 m: usize,
14326 in_f: usize,
14327 out0: usize,
14328 out1: usize,
14329 row_bytes: usize,
14330 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14331 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14333 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14334 let f = self.func(match Self::batched_mcols(m) {
14335 2 => "qmatvec_q8_0_mmvq_fused2_b2",
14336 4 => "qmatvec_q8_0_mmvq_fused2_b4",
14337 _ => "qmatvec_q8_0_mmvq_fused2_b8",
14339 });
14340 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14341 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14342 let cfg = LaunchConfig {
14343 grid_dim: (nb0 + nb1, 1, 1),
14344 block_dim: (32, ROWS_PER_BLOCK, 1),
14345 shared_mem_bytes: 0,
14346 };
14347 let (inf, o0, o1, mi, rbl) = (
14348 in_f as i32,
14349 out0 as i32,
14350 out1 as i32,
14351 m as i32,
14352 row_bytes as i64,
14353 );
14354 let __s_b = self.gpu.stream();
14355 let mut b = __s_b.launch_builder(&f);
14356 b.arg(b0)
14357 .arg(b1)
14358 .arg(aq)
14359 .arg(ad)
14360 .arg(&mut y0)
14361 .arg(&mut y1)
14362 .arg(&inf)
14363 .arg(&o0)
14364 .arg(&o1)
14365 .arg(&mi)
14366 .arg(&rbl);
14367 unsafe {
14368 b.launch(cfg)?;
14369 }
14370 Ok((y0, y1))
14371 }
14372
14373 #[allow(clippy::too_many_arguments)]
14376 pub fn qmatvec_q8_fused2_t_raw(
14377 &self,
14378 b0: &CudaSlice<u8>,
14379 b1: &CudaSlice<u8>,
14380 x: &CudaSlice<f32>,
14381 m: usize,
14382 in_f: usize,
14383 out0: usize,
14384 out1: usize,
14385 row_bytes: usize,
14386 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14387 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14388 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14389 }
14390
14391 #[allow(clippy::too_many_arguments)]
14394 pub fn matmul_q8_fused3_t(
14395 &self,
14396 w0: &crate::model::GpuTensor,
14397 w1: &crate::model::GpuTensor,
14398 w2: &crate::model::GpuTensor,
14399 aq: &CudaSlice<i8>,
14400 ad: &CudaSlice<f32>,
14401 m: usize,
14402 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14403 {
14404 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14405 return Ok(None);
14406 }
14407 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14408 return Ok(Some(self.e4m3_fused3_t_core(
14409 p0.0,
14410 p1.0,
14411 p2.0,
14412 aq,
14413 ad,
14414 m,
14415 w0.in_features(),
14416 p0.1,
14417 p1.1,
14418 p2.1,
14419 p0.2,
14420 p0.3,
14421 p1.3,
14422 p2.3,
14423 )?));
14424 }
14425 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14426 return Ok(None);
14427 };
14428 Ok(Some(self.q8_fused3_t_core(
14429 p0.0,
14430 p1.0,
14431 p2.0,
14432 aq,
14433 ad,
14434 m,
14435 w0.in_features(),
14436 p0.1,
14437 p1.1,
14438 p2.1,
14439 p0.2,
14440 )?))
14441 }
14442
14443 #[allow(clippy::too_many_arguments)]
14444 fn q8_fused3_t_core(
14445 &self,
14446 b0: &CudaSlice<u8>,
14447 b1: &CudaSlice<u8>,
14448 b2: &CudaSlice<u8>,
14449 aq: &CudaSlice<i8>,
14450 ad: &CudaSlice<f32>,
14451 m: usize,
14452 in_f: usize,
14453 out0: usize,
14454 out1: usize,
14455 out2: usize,
14456 row_bytes: usize,
14457 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14458 const ROWS_PER_BLOCK: u32 = 4;
14459 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14460 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14461 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14462 let f = self.func(if Self::batched_mcols(m) == 2 {
14463 "qmatvec_q8_0_mmvq_fused3_b2"
14464 } else {
14465 "qmatvec_q8_0_mmvq_fused3_b4"
14466 });
14467 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14468 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14469 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14470 let cfg = LaunchConfig {
14471 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14472 block_dim: (32, ROWS_PER_BLOCK, 1),
14473 shared_mem_bytes: 0,
14474 };
14475 let (inf, o0, o1, o2, mi, rbl) = (
14476 in_f as i32,
14477 out0 as i32,
14478 out1 as i32,
14479 out2 as i32,
14480 m as i32,
14481 row_bytes as i64,
14482 );
14483 let __s_b = self.gpu.stream();
14484 let mut b = __s_b.launch_builder(&f);
14485 b.arg(b0)
14486 .arg(b1)
14487 .arg(b2)
14488 .arg(aq)
14489 .arg(ad)
14490 .arg(&mut y0)
14491 .arg(&mut y1)
14492 .arg(&mut y2)
14493 .arg(&inf)
14494 .arg(&o0)
14495 .arg(&o1)
14496 .arg(&o2)
14497 .arg(&mi)
14498 .arg(&rbl);
14499 unsafe {
14500 b.launch(cfg)?;
14501 }
14502 Ok((y0, y1, y2))
14503 }
14504
14505 #[allow(clippy::too_many_arguments)]
14507 pub fn qmatvec_q8_fused3_t_raw(
14508 &self,
14509 b0: &CudaSlice<u8>,
14510 b1: &CudaSlice<u8>,
14511 b2: &CudaSlice<u8>,
14512 x: &CudaSlice<f32>,
14513 m: usize,
14514 in_f: usize,
14515 out0: usize,
14516 out1: usize,
14517 out2: usize,
14518 row_bytes: usize,
14519 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14520 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14521 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14522 }
14523
14524 pub fn q8_ffn_fuse2_on(&self) -> bool {
14528 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14529 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14530 }
14531
14532 #[allow(clippy::type_complexity)]
14538 fn q8_fused_params<'w, const N: usize>(
14539 &self,
14540 ws: &[&'w crate::model::GpuTensor; N],
14541 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14542 use crate::model::GpuTensor;
14543 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14544 return None;
14545 }
14546 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14547 return None;
14548 }
14549 let in_f = ws[0].in_features();
14550 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14551 for (i, w) in ws.iter().enumerate() {
14552 match w {
14553 GpuTensor::Quant {
14554 bytes,
14555 qtype,
14556 row_bytes,
14557 scale,
14558 ..
14559 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14560 out[i] = Some((bytes, w.out_features(), *row_bytes))
14561 }
14562 _ => return None,
14563 }
14564 }
14565 Some(out.map(|o| o.unwrap()))
14566 }
14567
14568 pub fn e4m3_dual_on(&self) -> bool {
14571 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14572 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14573 }
14574
14575 #[allow(clippy::type_complexity)]
14587 fn e4m3_fused_params<'w, const N: usize>(
14588 &self,
14589 ws: &[&'w crate::model::GpuTensor; N],
14590 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14591 use crate::model::GpuTensor;
14592 if !self.e4m3_dual_on() {
14593 return None;
14594 }
14595 let in_f = ws[0].in_features();
14596 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14597 for (i, w) in ws.iter().enumerate() {
14598 match w {
14599 GpuTensor::Quant {
14600 bytes,
14601 qtype,
14602 row_bytes,
14603 scale,
14604 rp,
14605 rp4,
14606 ..
14607 } if *qtype == QT_F8_E4M3
14608 && w.in_features() == in_f
14609 && *row_bytes == in_f
14610 && !*rp
14611 && rp4.is_none() =>
14612 {
14613 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14614 }
14615 _ => return None,
14616 }
14617 }
14618 Some(out.map(|o| o.unwrap()))
14619 }
14620
14621 #[allow(clippy::too_many_arguments)]
14625 fn e4m3_fused2_core(
14626 &self,
14627 b0: &CudaSlice<u8>,
14628 b1: &CudaSlice<u8>,
14629 aq: &CudaSlice<i8>,
14630 ad: &CudaSlice<f32>,
14631 in_f: usize,
14632 out0: usize,
14633 out1: usize,
14634 row_bytes: usize,
14635 ws0: f32,
14636 ws1: f32,
14637 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14638 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14640 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14641 let f = self.func("qmatvec_e4m3_mmvq_fused2");
14642 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14643 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14644 let cfg = LaunchConfig {
14645 grid_dim: (nb0 + nb1, 1, 1),
14646 block_dim: (32, ROWS_PER_BLOCK, 1),
14647 shared_mem_bytes: 0,
14648 };
14649 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14650 let __s_b = self.gpu.stream();
14651 let mut b = __s_b.launch_builder(&f);
14652 b.arg(b0)
14653 .arg(b1)
14654 .arg(aq)
14655 .arg(ad)
14656 .arg(&mut y0)
14657 .arg(&mut y1)
14658 .arg(&inf)
14659 .arg(&o0)
14660 .arg(&o1)
14661 .arg(&rbl)
14662 .arg(&ws0)
14663 .arg(&ws1);
14664 unsafe {
14665 b.launch(cfg)?;
14666 }
14667 Ok((y0, y1))
14668 }
14669
14670 #[allow(clippy::too_many_arguments)]
14672 fn e4m3_fused3_core(
14673 &self,
14674 b0: &CudaSlice<u8>,
14675 b1: &CudaSlice<u8>,
14676 b2: &CudaSlice<u8>,
14677 aq: &CudaSlice<i8>,
14678 ad: &CudaSlice<f32>,
14679 in_f: usize,
14680 out0: usize,
14681 out1: usize,
14682 out2: usize,
14683 row_bytes: usize,
14684 ws0: f32,
14685 ws1: f32,
14686 ws2: f32,
14687 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14688 const ROWS_PER_BLOCK: u32 = 4;
14689 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14690 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14691 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14692 let f = self.func("qmatvec_e4m3_mmvq_fused3");
14693 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14694 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14695 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14696 let cfg = LaunchConfig {
14697 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14698 block_dim: (32, ROWS_PER_BLOCK, 1),
14699 shared_mem_bytes: 0,
14700 };
14701 let (inf, o0, o1, o2, rbl) = (
14702 in_f as i32,
14703 out0 as i32,
14704 out1 as i32,
14705 out2 as i32,
14706 row_bytes as i64,
14707 );
14708 let __s_b = self.gpu.stream();
14709 let mut b = __s_b.launch_builder(&f);
14710 b.arg(b0)
14711 .arg(b1)
14712 .arg(b2)
14713 .arg(aq)
14714 .arg(ad)
14715 .arg(&mut y0)
14716 .arg(&mut y1)
14717 .arg(&mut y2)
14718 .arg(&inf)
14719 .arg(&o0)
14720 .arg(&o1)
14721 .arg(&o2)
14722 .arg(&rbl)
14723 .arg(&ws0)
14724 .arg(&ws1)
14725 .arg(&ws2);
14726 unsafe {
14727 b.launch(cfg)?;
14728 }
14729 Ok((y0, y1, y2))
14730 }
14731
14732 #[allow(clippy::too_many_arguments)]
14736 fn e4m3_fused2_t_core(
14737 &self,
14738 b0: &CudaSlice<u8>,
14739 b1: &CudaSlice<u8>,
14740 aq: &CudaSlice<i8>,
14741 ad: &CudaSlice<f32>,
14742 m: usize,
14743 in_f: usize,
14744 out0: usize,
14745 out1: usize,
14746 row_bytes: usize,
14747 ws0: f32,
14748 ws1: f32,
14749 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14750 const ROWS_PER_BLOCK: u32 = 4;
14751 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14752 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14753 let f = self.func(match Self::batched_mcols(m) {
14754 2 => "qmatvec_e4m3_mmvq_fused2_b2",
14755 4 => "qmatvec_e4m3_mmvq_fused2_b4",
14756 _ => "qmatvec_e4m3_mmvq_fused2_b8",
14757 });
14758 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14759 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14760 let cfg = LaunchConfig {
14761 grid_dim: (nb0 + nb1, 1, 1),
14762 block_dim: (32, ROWS_PER_BLOCK, 1),
14763 shared_mem_bytes: 0,
14764 };
14765 let (inf, o0, o1, mi, rbl) = (
14766 in_f as i32,
14767 out0 as i32,
14768 out1 as i32,
14769 m as i32,
14770 row_bytes as i64,
14771 );
14772 let __s_b = self.gpu.stream();
14773 let mut b = __s_b.launch_builder(&f);
14774 b.arg(b0)
14775 .arg(b1)
14776 .arg(aq)
14777 .arg(ad)
14778 .arg(&mut y0)
14779 .arg(&mut y1)
14780 .arg(&inf)
14781 .arg(&o0)
14782 .arg(&o1)
14783 .arg(&mi)
14784 .arg(&rbl);
14785 unsafe {
14786 b.launch(cfg)?;
14787 }
14788 if ws0 != 1.0 {
14789 self.scale_inplace(&mut y0, ws0, m * out0)?;
14790 }
14791 if ws1 != 1.0 {
14792 self.scale_inplace(&mut y1, ws1, m * out1)?;
14793 }
14794 Ok((y0, y1))
14795 }
14796
14797 #[allow(clippy::too_many_arguments)]
14799 fn e4m3_fused3_t_core(
14800 &self,
14801 b0: &CudaSlice<u8>,
14802 b1: &CudaSlice<u8>,
14803 b2: &CudaSlice<u8>,
14804 aq: &CudaSlice<i8>,
14805 ad: &CudaSlice<f32>,
14806 m: usize,
14807 in_f: usize,
14808 out0: usize,
14809 out1: usize,
14810 out2: usize,
14811 row_bytes: usize,
14812 ws0: f32,
14813 ws1: f32,
14814 ws2: f32,
14815 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14816 const ROWS_PER_BLOCK: u32 = 4;
14817 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14818 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14819 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14820 let f = self.func(if Self::batched_mcols(m) == 2 {
14821 "qmatvec_e4m3_mmvq_fused3_b2"
14822 } else {
14823 "qmatvec_e4m3_mmvq_fused3_b4"
14824 });
14825 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14826 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14827 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14828 let cfg = LaunchConfig {
14829 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14830 block_dim: (32, ROWS_PER_BLOCK, 1),
14831 shared_mem_bytes: 0,
14832 };
14833 let (inf, o0, o1, o2, mi, rbl) = (
14834 in_f as i32,
14835 out0 as i32,
14836 out1 as i32,
14837 out2 as i32,
14838 m as i32,
14839 row_bytes as i64,
14840 );
14841 let __s_b = self.gpu.stream();
14842 let mut b = __s_b.launch_builder(&f);
14843 b.arg(b0)
14844 .arg(b1)
14845 .arg(b2)
14846 .arg(aq)
14847 .arg(ad)
14848 .arg(&mut y0)
14849 .arg(&mut y1)
14850 .arg(&mut y2)
14851 .arg(&inf)
14852 .arg(&o0)
14853 .arg(&o1)
14854 .arg(&o2)
14855 .arg(&mi)
14856 .arg(&rbl);
14857 unsafe {
14858 b.launch(cfg)?;
14859 }
14860 if ws0 != 1.0 {
14861 self.scale_inplace(&mut y0, ws0, m * out0)?;
14862 }
14863 if ws1 != 1.0 {
14864 self.scale_inplace(&mut y1, ws1, m * out1)?;
14865 }
14866 if ws2 != 1.0 {
14867 self.scale_inplace(&mut y2, ws2, m * out2)?;
14868 }
14869 Ok((y0, y1, y2))
14870 }
14871
14872 pub fn qmatvec_e4m3_blk_mmvq(
14882 &self,
14883 bytes: &CudaSlice<u8>,
14884 aq: &CudaSlice<i8>,
14885 ad: &CudaSlice<f32>,
14886 scales: &CudaSlice<f32>,
14887 m: usize,
14888 in_f: usize,
14889 out_f: usize,
14890 row_bytes: usize,
14891 scale_cols: usize,
14892 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14893 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
14895 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14896 )?;
14897 Ok(y)
14898 }
14899
14900 #[allow(clippy::too_many_arguments)]
14902 pub fn qmatvec_e4m3_blk_mmvq_into(
14903 &self,
14904 bytes: &CudaSlice<u8>,
14905 aq: &CudaSlice<i8>,
14906 ad: &CudaSlice<f32>,
14907 scales: &CudaSlice<f32>,
14908 m: usize,
14909 in_f: usize,
14910 out_f: usize,
14911 row_bytes: usize,
14912 scale_cols: usize,
14913 y: &mut CudaSlice<f32>,
14914 ) -> Result<(), Box<dyn std::error::Error>> {
14915 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
14917 let cfg = LaunchConfig {
14918 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14919 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
14922 let (inf, outf, mi, rb, sc) = (
14923 in_f as i32,
14924 out_f as i32,
14925 m as i32,
14926 row_bytes as i64,
14927 scale_cols as i32,
14928 );
14929 let __s_b = self.gpu.stream();
14930 let mut b = __s_b.launch_builder(&f);
14931 b.arg(bytes)
14932 .arg(aq)
14933 .arg(ad)
14934 .arg(scales)
14935 .arg(&mut *y)
14936 .arg(&inf)
14937 .arg(&outf)
14938 .arg(&mi)
14939 .arg(&rb)
14940 .arg(&sc);
14941 unsafe {
14942 b.launch(cfg)?;
14943 }
14944 Ok(())
14945 }
14946
14947 #[allow(clippy::too_many_arguments)]
14953 pub fn qmatvec_e4m3_blk_mmvq_batched(
14954 &self,
14955 bytes: &CudaSlice<u8>,
14956 aq: &CudaSlice<i8>,
14957 ad: &CudaSlice<f32>,
14958 scales: &CudaSlice<f32>,
14959 m: usize,
14960 in_f: usize,
14961 out_f: usize,
14962 row_bytes: usize,
14963 scale_cols: usize,
14964 mcols: usize,
14965 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14966 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14968 let name = match mcols {
14969 2 => "qmatvec_e4m3_blk_mmvq_b2",
14970 4 => "qmatvec_e4m3_blk_mmvq_b4",
14971 8 => "qmatvec_e4m3_blk_mmvq_b8",
14972 16 => "qmatvec_e4m3_blk_mmvq_b16",
14973 _ => {
14974 return Err(
14975 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14976 );
14977 }
14978 };
14979 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14980 let f = self.func(name);
14981 let cfg = LaunchConfig {
14982 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14983 block_dim: (32, ROWS_PER_BLOCK, 1),
14984 shared_mem_bytes: 0,
14985 };
14986 let (inf, outf, mi, rb, sc) = (
14987 in_f as i32,
14988 out_f as i32,
14989 m as i32,
14990 row_bytes as i64,
14991 scale_cols as i32,
14992 );
14993 let __s_b = self.gpu.stream();
14994 let mut b = __s_b.launch_builder(&f);
14995 b.arg(bytes)
14996 .arg(aq)
14997 .arg(ad)
14998 .arg(scales)
14999 .arg(&mut y)
15000 .arg(&inf)
15001 .arg(&outf)
15002 .arg(&mi)
15003 .arg(&rb)
15004 .arg(&sc);
15005 unsafe {
15006 b.launch(cfg)?;
15007 }
15008 Ok(y)
15009 }
15010
15011 #[allow(clippy::too_many_arguments)]
15014 pub fn qmatvec_e4m3_blk_batched_raw(
15015 &self,
15016 bytes: &CudaSlice<u8>,
15017 x: &CudaSlice<f32>,
15018 scales: &CudaSlice<f32>,
15019 m: usize,
15020 in_f: usize,
15021 out_f: usize,
15022 row_bytes: usize,
15023 scale_cols: usize,
15024 mcols: usize,
15025 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15026 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15027 self.qmatvec_e4m3_blk_mmvq_batched(
15028 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15029 )
15030 }
15031
15032 #[allow(clippy::too_many_arguments)]
15035 pub fn qmatvec_e4m3_blk_mmvq_raw(
15036 &self,
15037 bytes: &CudaSlice<u8>,
15038 x: &CudaSlice<f32>,
15039 scales: &CudaSlice<f32>,
15040 m: usize,
15041 in_f: usize,
15042 out_f: usize,
15043 row_bytes: usize,
15044 scale_cols: usize,
15045 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15046 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15047 self.qmatvec_e4m3_blk_mmvq(
15048 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15049 )
15050 }
15051
15052 #[allow(clippy::too_many_arguments)]
15055 pub fn qmatvec_e4m3_fused2_raw(
15056 &self,
15057 b0: &CudaSlice<u8>,
15058 b1: &CudaSlice<u8>,
15059 x: &CudaSlice<f32>,
15060 in_f: usize,
15061 out0: usize,
15062 out1: usize,
15063 row_bytes: usize,
15064 ws0: f32,
15065 ws1: f32,
15066 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15067 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15068 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15069 }
15070
15071 #[allow(clippy::too_many_arguments)]
15072 pub fn qmatvec_e4m3_fused3_raw(
15073 &self,
15074 b0: &CudaSlice<u8>,
15075 b1: &CudaSlice<u8>,
15076 b2: &CudaSlice<u8>,
15077 x: &CudaSlice<f32>,
15078 in_f: usize,
15079 out0: usize,
15080 out1: usize,
15081 out2: usize,
15082 row_bytes: usize,
15083 ws0: f32,
15084 ws1: f32,
15085 ws2: f32,
15086 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15087 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15088 self.e4m3_fused3_core(
15089 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15090 )
15091 }
15092
15093 #[allow(clippy::too_many_arguments)]
15094 pub fn qmatvec_e4m3_fused2_t_raw(
15095 &self,
15096 b0: &CudaSlice<u8>,
15097 b1: &CudaSlice<u8>,
15098 x: &CudaSlice<f32>,
15099 m: usize,
15100 in_f: usize,
15101 out0: usize,
15102 out1: usize,
15103 row_bytes: usize,
15104 ws0: f32,
15105 ws1: f32,
15106 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15107 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15108 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15109 }
15110
15111 #[allow(clippy::too_many_arguments)]
15112 pub fn qmatvec_e4m3_fused3_t_raw(
15113 &self,
15114 b0: &CudaSlice<u8>,
15115 b1: &CudaSlice<u8>,
15116 b2: &CudaSlice<u8>,
15117 x: &CudaSlice<f32>,
15118 m: usize,
15119 in_f: usize,
15120 out0: usize,
15121 out1: usize,
15122 out2: usize,
15123 row_bytes: usize,
15124 ws0: f32,
15125 ws1: f32,
15126 ws2: f32,
15127 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15128 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15129 self.e4m3_fused3_t_core(
15130 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15131 )
15132 }
15133
15134 fn try_e4m3_blk_pre(
15145 &self,
15146 w: &crate::model::GpuTensor,
15147 aq: &CudaSlice<i8>,
15148 ad: &CudaSlice<f32>,
15149 m: usize,
15150 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15151 use crate::model::GpuTensor;
15152 if let GpuTensor::Quant {
15153 bytes,
15154 qtype,
15155 row_bytes,
15156 blk: Some(g),
15157 ..
15158 } = w
15159 {
15160 if *qtype == QT_F8_E4M3_BLK {
15161 if (2..=16).contains(&m)
15167 && std::env::var("MEMRA_NO_BATCHED").is_err()
15168 && (m <= 4 || Self::b8_enabled())
15169 {
15170 let mcols = Self::batched_mcols(m);
15171 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15172 bytes,
15173 aq,
15174 ad,
15175 &g.scales,
15176 m,
15177 w.in_features(),
15178 w.out_features(),
15179 *row_bytes,
15180 g.cols,
15181 mcols,
15182 )?));
15183 }
15184 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15185 bytes,
15186 aq,
15187 ad,
15188 &g.scales,
15189 m,
15190 w.in_features(),
15191 w.out_features(),
15192 *row_bytes,
15193 g.cols,
15194 )?));
15195 }
15196 }
15197 Ok(None)
15198 }
15199
15200 fn try_e4m3_blk_prefill(
15247 &self,
15248 w: &crate::model::GpuTensor,
15249 x: &CudaSlice<f32>,
15250 m: usize,
15251 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15252 use crate::model::GpuTensor;
15253 let GpuTensor::Quant {
15254 bytes,
15255 qtype,
15256 blk: Some(g),
15257 ..
15258 } = w
15259 else {
15260 return Ok(None);
15261 };
15262 if *qtype != QT_F8_E4M3_BLK {
15263 return Ok(None);
15264 }
15265 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
15270 return Ok(Some(y));
15271 }
15272 let (in_f, out_f) = (w.in_features(), w.out_features());
15273 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
15274 let tmp = GpuTensor::Quant {
15275 bytes: slab,
15276 qtype: QT_Q8_0,
15277 row_bytes: in_f / 32 * 34,
15278 ne: vec![in_f as u64, out_f as u64],
15279 scale: 1.0,
15280 rp: false,
15281 #[cfg(memra_cutlass)]
15282 cutlass: None,
15283 fp8: None,
15284 blk: None,
15285 f16: None,
15286 rp4: None,
15287 };
15288 Ok(Some(self.matmul(&tmp, x, m)?))
15290 }
15291
15292 pub fn matmul_pre_noscale(
15293 &self,
15294 w: &crate::model::GpuTensor,
15295 aq: &CudaSlice<i8>,
15296 ad: &CudaSlice<f32>,
15297 m: usize,
15298 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
15299 use crate::model::GpuTensor;
15300 if m == 1 {
15304 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
15305 return Ok(Some((y, 1.0)));
15306 }
15307 }
15308 if m != 1 || !self.uses_q8_1_fast(w) {
15310 return Ok(None);
15311 }
15312 let in_f = w.in_features();
15313 let out_f = w.out_features();
15314 let (bytes, qtype, row_bytes, scale, rp) = match w {
15315 GpuTensor::Quant {
15316 bytes,
15317 qtype,
15318 row_bytes,
15319 scale,
15320 rp,
15321 ..
15322 } => (bytes, *qtype, *row_bytes, *scale, *rp),
15323 _ => return Ok(None),
15324 };
15325 if self.mmvq_supports(qtype) {
15327 let (mbytes, mrp) = match w {
15329 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
15330 _ => (bytes, rp),
15331 };
15332 let y = self.qmatvec_mmvq(
15333 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
15334 )?;
15335 return Ok(Some((y, scale)));
15336 }
15337 let name = match qtype {
15339 QT_Q8_0 => "qmatvec_q8_0_dp4a",
15340 QT_Q4_K => "qmatvec_q4_K_dp4a",
15341 QT_Q6_K => "qmatvec_q6_K_dp4a",
15342 QT_Q5_K => "qmatvec_q5_K_dp4a",
15343 QT_Q3_K => "qmatvec_q3_K_dp4a",
15344 QT_NVFP4 => {
15345 if rp {
15346 "qmatvec_nvfp4_dp4a_rp"
15347 } else {
15348 "qmatvec_nvfp4_dp4a"
15349 }
15350 }
15351 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
15352 _ => return Ok(None),
15353 };
15354 let f = self.func(name);
15355 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15356 let cfg = LaunchConfig {
15357 grid_dim: (out_f as u32, m as u32, 1),
15358 block_dim: (128, 1, 1),
15359 shared_mem_bytes: 0,
15360 };
15361 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15362 let __s_b = self.gpu.stream();
15363 let mut b = __s_b.launch_builder(&f);
15364 b.arg(bytes)
15365 .arg(aq)
15366 .arg(ad)
15367 .arg(&mut y)
15368 .arg(&inf)
15369 .arg(&outf)
15370 .arg(&mi)
15371 .arg(&rb);
15372 unsafe {
15373 b.launch(cfg)?;
15374 }
15375 Ok(Some((y, scale)))
15376 }
15377
15378 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15381 if qtype == QT_F8_E4M3 {
15386 return true;
15387 }
15388 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15389 return false;
15390 }
15391 matches!(
15392 qtype,
15393 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15394 )
15395 }
15396
15397 pub fn qmatvec_mmvq(
15402 &self,
15403 bytes: &CudaSlice<u8>,
15404 aq: &CudaSlice<i8>,
15405 ad: &CudaSlice<f32>,
15406 m: usize,
15407 in_f: usize,
15408 out_f: usize,
15409 qtype: i32,
15410 row_bytes: usize,
15411 scale: f32,
15412 rp: bool,
15413 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15414 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15416 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15417 )?;
15418 Ok(y)
15419 }
15420
15421 #[allow(clippy::too_many_arguments)]
15423 pub fn qmatvec_mmvq_into(
15424 &self,
15425 bytes: &CudaSlice<u8>,
15426 aq: &CudaSlice<i8>,
15427 ad: &CudaSlice<f32>,
15428 m: usize,
15429 in_f: usize,
15430 out_f: usize,
15431 qtype: i32,
15432 row_bytes: usize,
15433 scale: f32,
15434 rp: bool,
15435 y: &mut CudaSlice<f32>,
15436 ) -> Result<(), Box<dyn std::error::Error>> {
15437 debug_assert!(y.len() >= m * out_f);
15438 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15444 && rp
15445 && m == 1
15446 && out_f >= 64
15447 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15448 && {
15449 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15450 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15451 }
15452 {
15453 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15454 let cfg = LaunchConfig {
15455 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15456 block_dim: (32, 2, 1),
15457 shared_mem_bytes: 0,
15458 };
15459 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15460 let __s_b = self.gpu.stream();
15461 let mut b = __s_b.launch_builder(&f);
15462 b.arg(bytes)
15463 .arg(aq)
15464 .arg(ad)
15465 .arg(&mut *y)
15466 .arg(&inf)
15467 .arg(&outf)
15468 .arg(&mi)
15469 .arg(&rb);
15470 unsafe {
15471 b.launch(cfg)?;
15472 }
15473 if scale != 1.0 {
15474 self.scale_inplace(y, scale, out_f)?;
15475 }
15476 return Ok(());
15477 }
15478 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15487 2
15488 } else {
15489 1
15490 };
15491 if m == 1 && qtype == QT_Q4_0 {
15496 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15497 mr = *Q40MR.get_or_init(|| {
15500 std::env::var("MEMRA_Q40_MR")
15501 .ok()
15502 .and_then(|v| v.parse().ok())
15503 .unwrap_or(1)
15504 });
15505 }
15506 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15517 let q5_force = q5_mode.as_deref() == Some("2");
15518 let q5_il = qtype == QT_Q5_K
15521 && m == 1
15522 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15523 if q5_il && !q5_force && out_f > 65536 {
15524 mr = 1;
15525 }
15526 if qtype == QT_Q4_0 && rp && mr != 1 {
15529 mr = 2;
15530 }
15531 if qtype == QT_Q8_0 && rp {
15535 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15536 mr = *Q80MR.get_or_init(|| {
15537 std::env::var("MEMRA_Q80_MR")
15538 .ok()
15539 .and_then(|v| v.parse().ok())
15540 .unwrap_or(1)
15541 });
15542 }
15543 let name = match (qtype, mr, rp) {
15544 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15545 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15546 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15547 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15548 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15549 (QT_Q5_K, 2, _) => {
15550 if q5_il {
15551 "qmatvec_q5_K_mmvq_mr2_il"
15552 } else {
15553 "qmatvec_q5_K_mmvq_mr2"
15554 }
15555 }
15556 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15557 (QT_Q8_0, _, true)
15562 if in_f % 1024 == 0 && {
15563 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15564 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15565 } =>
15566 {
15567 "qmatvec_q8_0_mmvq_rpca"
15568 }
15569 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15570 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15571 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15575 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15576 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15577 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15578 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15579 (QT_Q5_K, _, _) => {
15580 if q5_il {
15581 "qmatvec_q5_K_mmvq_il"
15582 } else {
15583 "qmatvec_q5_K_mmvq"
15584 }
15585 }
15586 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15587 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15588 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15589 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15590 };
15591 let f = self.func(name);
15592 let rows_per_block = ROWS_PER_BLOCK * mr;
15594 let cfg = LaunchConfig {
15595 grid_dim: (
15596 (out_f as u32 + rows_per_block - 1) / rows_per_block,
15597 m as u32,
15598 1,
15599 ),
15600 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15603 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15604 let __s_b = self.gpu.stream();
15605 let mut b = __s_b.launch_builder(&f);
15606 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15611 if Self::pdl_on()
15614 && Self::pdl_mmvq_on()
15615 && Self::pdl_nvfp4q8_on()
15616 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15617 {
15618 use cudarc::driver::{DevicePtr, DevicePtrMut};
15619 let s = &self.gpu.stream();
15620 let (pw, _g0) = bytes.device_ptr(s);
15621 let (paq, _g1) = aq.device_ptr(s);
15622 let (pad, _g2) = ad.device_ptr(s);
15623 let (py, _g3) = y.device_ptr_mut(s);
15624 let mut ps = [
15625 &pw as *const _ as *mut std::ffi::c_void,
15626 &paq as *const _ as *mut _,
15627 &pad as *const _ as *mut _,
15628 &py as *const _ as *mut _,
15629 &inf as *const _ as *mut _,
15630 &outf as *const _ as *mut _,
15631 &mi as *const _ as *mut _,
15632 &rb as *const _ as *mut _,
15633 &scale as *const _ as *mut _,
15634 ];
15635 unsafe {
15636 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15637 }
15638 return Ok(());
15639 }
15640 b.arg(bytes)
15641 .arg(aq)
15642 .arg(ad)
15643 .arg(&mut *y)
15644 .arg(&inf)
15645 .arg(&outf)
15646 .arg(&mi)
15647 .arg(&rb)
15648 .arg(&scale);
15649 unsafe {
15650 b.launch(cfg)?;
15651 }
15652 } else if Self::pdl_on()
15653 && Self::pdl_mmvq_on()
15654 && (matches!(
15655 name,
15656 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15657 ) || (Self::pdl_nvfp4q8_on()
15658 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15659 {
15660 {
15664 use cudarc::driver::{DevicePtr, DevicePtrMut};
15665 let s = &self.gpu.stream();
15666 let (pw, _g0) = bytes.device_ptr(s);
15667 let (paq, _g1) = aq.device_ptr(s);
15668 let (pad, _g2) = ad.device_ptr(s);
15669 let (py, _g3) = y.device_ptr_mut(s);
15670 let mut ps = [
15671 &pw as *const _ as *mut std::ffi::c_void,
15672 &paq as *const _ as *mut _,
15673 &pad as *const _ as *mut _,
15674 &py as *const _ as *mut _,
15675 &inf as *const _ as *mut _,
15676 &outf as *const _ as *mut _,
15677 &mi as *const _ as *mut _,
15678 &rb as *const _ as *mut _,
15679 ];
15680 unsafe {
15681 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15682 }
15683 }
15684 if scale != 1.0 {
15685 self.scale_inplace(y, scale, m * out_f)?;
15686 }
15687 } else {
15688 b.arg(bytes)
15689 .arg(aq)
15690 .arg(ad)
15691 .arg(&mut *y)
15692 .arg(&inf)
15693 .arg(&outf)
15694 .arg(&mi)
15695 .arg(&rb);
15696 unsafe {
15697 b.launch(cfg)?;
15698 }
15699 if scale != 1.0 {
15700 self.scale_inplace(y, scale, m * out_f)?;
15701 }
15702 }
15703 Ok(())
15704 }
15705
15706 pub fn qmatvec_mmvq_raw(
15710 &self,
15711 bytes: &CudaSlice<u8>,
15712 x: &CudaSlice<f32>,
15713 m: usize,
15714 in_f: usize,
15715 out_f: usize,
15716 qtype: i32,
15717 row_bytes: usize,
15718 rp: bool,
15719 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15720 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15721 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15722 }
15723
15724 pub fn batched_supports(&self, qtype: i32) -> bool {
15728 matches!(
15729 qtype,
15730 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15731 )
15732 }
15733
15734 pub fn iq_fast_enabled() -> bool {
15742 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15743 *ON.get_or_init(|| {
15744 std::env::var("MEMRA_IQ_FAST")
15745 .map(|v| v != "0")
15746 .unwrap_or(true)
15747 })
15748 }
15749
15750 pub fn b8_enabled() -> bool {
15753 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15754 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15755 }
15756
15757 pub fn batched_mcols(m: usize) -> usize {
15759 if m == 2 {
15760 2
15761 } else if m <= 4 {
15762 4
15763 } else if m <= 8 {
15764 8
15765 } else {
15766 16
15767 }
15768 }
15769
15770 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15775 Some(match (qtype, mcols) {
15776 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15777 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15778 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15779 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15785 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15786 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15787 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15788 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15791 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15792 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15793 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15794 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15797 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15798 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15799 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15800 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15801 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15802 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15803 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15804 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15808 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15809 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15810 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15811 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15815 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15816 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15817 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15818 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15819 _ => return None,
15820 })
15821 }
15822
15823 pub fn sm_count(&self) -> i32 {
15858 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15859 *SMS.get_or_init(|| {
15860 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15861 self.gpu
15862 .ctx
15863 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15864 .unwrap_or(82)
15865 })
15866 }
15867
15868 pub fn batched_variant(
15869 &self,
15870 _m: usize,
15871 in_f: usize,
15872 out_f: usize,
15873 qtype: i32,
15874 row_bytes: usize,
15875 mcols: usize,
15876 rp: bool,
15877 ) -> &'static str {
15878 if qtype == QT_Q8_0 {
15883 return if rp { "rp" } else { "base" };
15884 }
15885 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15886 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15887 Ok("base") => "base",
15888 Ok("pf") => "pf",
15889 Ok("r2") => "r2",
15890 Ok("r2w8") => "r2w8",
15891 Ok("pfr2") => "pfr2",
15892 Ok("ca") => "ca",
15893 Ok("car2") => "car2",
15894 Ok("rp") => "rp",
15897 Ok("rpr2") => "rpr2",
15898 Ok("rpr2w8") => "rpr2w8",
15899 Ok("rpca") => "rpca",
15902 Ok("rpcar2") => "rpcar2",
15903 Ok("rpsc") => "rpsc",
15910 Ok("rpms") => "rpms",
15911 Ok("rpmsc") => "rpmsc",
15912 Ok("rpks") => "rpks",
15913 Ok("rpksc") => "rpksc",
15914 _ => "auto",
15915 });
15916 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15920 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15925 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15926 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15927 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15928 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15929 let sms = *SMS.get_or_init(|| {
15930 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15931 self.gpu
15932 .ctx
15933 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15934 .unwrap_or(82)
15935 });
15936 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15956 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15959 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15960 Ok("base") => "base",
15961 Ok("r2") => "r2",
15962 Ok("r2w8") => "r2w8",
15963 _ => "auto",
15964 });
15965 let variant: &'static str = if qtype == QT_Q4_0 {
15966 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15970 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15971 Ok("base") => "base",
15977 Ok("r2") => "r2",
15978 Ok("ms") => "ms",
15979 Ok("sm") => "sm",
15980 Ok("la") => "la",
15981 _ => "auto",
15982 });
15983 let v = if q40 != "auto" {
15984 q40
15985 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15986 "r2"
15987 } else {
15988 "base"
15989 };
15990 if rp {
15995 match v {
15996 "ms" => "r2ms_rp",
15997 "sm" => "r2sm_rp",
15998 "la" => "r2la_rp",
15999 "r2" => "r2_rp",
16000 _ => "rp",
16001 }
16002 } else if matches!(v, "ms" | "sm" | "la") {
16003 "r2"
16004 } else {
16005 v
16006 }
16007 } else if qtype != QT_NVFP4 && !kq_r2 {
16008 "base"
16009 } else if kq_r2 && rp {
16010 "rp"
16014 } else if kq_r2 {
16015 if kq_bv != "auto" {
16018 if kq_bv == "r2w8" && mcols != 4 {
16019 "r2"
16020 } else {
16021 kq_bv
16022 }
16023 } else if bv != "auto" {
16024 match bv {
16025 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16026 "r2w8" | "rpr2w8" => {
16027 if mcols != 4 {
16028 "r2"
16029 } else {
16030 "r2w8"
16031 }
16032 }
16033 _ => "base", }
16035 } else {
16036 let blocks = (out_f + 7) / 8;
16037 let waves = blocks as f64 / (7 * sms as usize) as f64;
16038 let filled = blocks >= 4 * sms as usize;
16039 let use_r2 = if qtype == QT_Q4_K {
16040 filled
16041 } else {
16042 waves >= 2.0
16043 };
16044 if use_r2 { "r2" } else { "base" }
16045 }
16046 } else if bv != "auto" {
16047 let v = if bv == "r2w8" && mcols == 2 {
16052 "r2"
16053 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16054 "pf"
16055 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16056 "r2"
16057 } else if bv == "pfr2" && mcols == 8 {
16058 "r2"
16059 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16060 "rpr2"
16061 }
16062 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16064 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16065 } else if bv == "rpcar2" && mcols == 2 {
16066 "rpca"
16067 }
16068 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16071 "rpr2"
16072 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16073 "rpr2"
16074 } else {
16075 bv
16076 };
16077 if rp {
16078 match v {
16079 "base" | "pf" | "ca" | "rp" => "rp",
16080 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16081 "r2w8" | "rpr2w8" => {
16082 if mcols == 2 {
16083 "rpr2"
16084 } else {
16085 "rpr2w8"
16086 }
16087 }
16088 other => other, }
16090 } else {
16091 v
16092 }
16093 } else if mcols == 8 {
16094 if rp {
16105 if sc_ok { "rpsc" } else { "rpr2w8" }
16106 } else {
16107 "r2w8"
16108 }
16109 } else if mcols >= 4 {
16110 let blocks = (out_f + 7) / 8;
16114 let r7 = 7 * sms as usize;
16115 let r8 = 8 * sms as usize;
16116 let waves = blocks as f64 / r7 as f64;
16117 let filled = blocks >= 4 * sms as usize;
16118 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16122 if rp { "rpr2w8" } else { "r2w8" }
16126 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16127 if rp { "rpr2" } else { "r2" }
16130 } else {
16131 if rp { "rp" } else { "pf" }
16135 }
16136 } else if in_f >= 6144 {
16137 if rp { "rpr2" } else { "r2" }
16141 } else if rp {
16142 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16147 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16148 "rpsc"
16149 } else {
16150 "rp"
16151 }
16152 } else {
16153 "base"
16154 };
16155 variant
16156 }
16157
16158 pub fn qmatvec_mmvq_batched(
16159 &self,
16160 bytes: &CudaSlice<u8>,
16161 aq: &CudaSlice<i8>,
16162 ad: &CudaSlice<f32>,
16163 m: usize,
16164 in_f: usize,
16165 out_f: usize,
16166 qtype: i32,
16167 row_bytes: usize,
16168 mcols: usize,
16169 scale: f32,
16170 rp: bool,
16171 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16172 const ROWS_PER_BLOCK: u32 = 4;
16173 let forced: Option<&'static str> = {
16178 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16179 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16180 .as_deref()
16181 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16182 };
16183 let variant = match forced {
16184 Some(v) if !rp || v.contains("rp") => v,
16185 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16186 };
16187 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16188 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16189 })?;
16190 let variant = if mcols == 16 {
16194 if rp { "rp" } else { "base" }
16195 } else {
16196 variant
16197 };
16198 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16205 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16206 if b567
16207 && qtype == QT_NVFP4
16208 && rp
16209 && mcols == 8
16210 && (5..=7).contains(&m)
16211 && matches!(variant, "rpsc" | "rpr2w8")
16212 {
16213 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16214 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16216 let cfg = LaunchConfig {
16217 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16218 block_dim: (32, ROWS_PER_BLOCK, 1),
16219 shared_mem_bytes: 0,
16220 };
16221 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16222 let __s_b = self.gpu.stream();
16223 let mut b = __s_b.launch_builder(&f);
16224 b.arg(bytes)
16225 .arg(aq)
16226 .arg(ad)
16227 .arg(&mut y)
16228 .arg(&inf)
16229 .arg(&outf)
16230 .arg(&mi)
16231 .arg(&rb);
16232 unsafe {
16233 b.launch(cfg)?;
16234 }
16235 if scale != 1.0 {
16236 self.scale_inplace(&mut y, scale, m * out_f)?;
16237 }
16238 return Ok(y);
16239 }
16240 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16241 "base" => (base_name.into(), ROWS_PER_BLOCK),
16242 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16243 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16244 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16245 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16249 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16250 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16251 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16252 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16253 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16254 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16255 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16257 debug_assert!(
16258 !rp || name.contains("_rp"),
16259 "rp weight dispatched to a GGUF-layout kernel"
16260 );
16261 let f = self.func(&name);
16262 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16263 let smem = if name.contains("_r2sm_rp") {
16265 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
16266 } else {
16267 0
16268 };
16269 let cfg = LaunchConfig {
16270 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16271 block_dim: (32, ROWS_PER_BLOCK, 1),
16272 shared_mem_bytes: smem,
16273 };
16274 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16275 let __s_b = self.gpu.stream();
16276 let mut b = __s_b.launch_builder(&f);
16277 b.arg(bytes)
16278 .arg(aq)
16279 .arg(ad)
16280 .arg(&mut y)
16281 .arg(&inf)
16282 .arg(&outf)
16283 .arg(&mi)
16284 .arg(&rb);
16285 unsafe {
16286 b.launch(cfg)?;
16287 }
16288 if scale != 1.0 {
16289 self.scale_inplace(&mut y, scale, m * out_f)?;
16290 }
16291 Ok(y)
16292 }
16293
16294 pub fn qmatvec_batched_raw(
16298 &self,
16299 bytes: &CudaSlice<u8>,
16300 x: &CudaSlice<f32>,
16301 m: usize,
16302 in_f: usize,
16303 out_f: usize,
16304 qtype: i32,
16305 row_bytes: usize,
16306 mcols: usize,
16307 rp: bool,
16308 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16309 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16310 self.qmatvec_mmvq_batched(
16311 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
16312 )
16313 }
16314
16315 pub fn qmatvec_nvfp4_batched_raw(
16317 &self,
16318 bytes: &CudaSlice<u8>,
16319 x: &CudaSlice<f32>,
16320 m: usize,
16321 in_f: usize,
16322 out_f: usize,
16323 row_bytes: usize,
16324 mcols: usize,
16325 rp: bool,
16326 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16327 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
16328 }
16329
16330 fn try_fp4_gemm(
16334 &self,
16335 w: &crate::model::GpuTensor,
16336 x: &CudaSlice<f32>,
16337 m: usize,
16338 in_f: usize,
16339 out_f: usize,
16340 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16341 use crate::model::GpuTensor;
16342 if cfg!(memra_portable_cuda) {
16343 return Ok(None);
16344 }
16345 if std::env::var("MEMRA_FP4").is_ok() {
16348 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
16349 }
16350 if std::env::var("MEMRA_FP4").is_err() {
16351 return Ok(None);
16352 }
16353 #[cfg(memra_cutlass)]
16362 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
16363 if let GpuTensor::Quant {
16364 bytes,
16365 qtype,
16366 scale,
16367 row_bytes,
16368 cutlass,
16369 ..
16370 } = w
16371 {
16372 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16373 if let Some(cw) = cutlass {
16374 let y = self.cutlass_fp4_gemm(
16376 &cw.b_packed,
16377 &cw.sfb_swizzled,
16378 x,
16379 *scale,
16380 m,
16381 out_f,
16382 in_f,
16383 )?;
16384 return Ok(Some(y));
16385 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16386 let (b_packed, sfb_sw) =
16391 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16392 let y =
16393 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16394 return Ok(Some(y));
16395 }
16396 }
16397 }
16398 }
16399 if let GpuTensor::Quant {
16400 bytes,
16401 qtype,
16402 row_bytes,
16403 scale,
16404 rp,
16405 ..
16406 } = w
16407 {
16408 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16411 let y =
16412 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16413 return Ok(Some(y));
16414 }
16415 }
16416 Ok(None)
16417 }
16418
16419 pub fn rms_norm_f16out(
16423 &self,
16424 x: &CudaSlice<f32>,
16425 w: &CudaSlice<f32>,
16426 dst: &mut CudaSlice<f32>,
16427 dst16: &mut CudaSlice<u8>,
16428 ncols: usize,
16429 nrows: usize,
16430 eps: f32,
16431 ) -> Result<(), Box<dyn std::error::Error>> {
16432 let f = self.func("rms_norm_f16out_f32");
16433 let cfg = LaunchConfig {
16434 grid_dim: (nrows as u32, 1, 1),
16435 block_dim: (rms_block(), 1, 1),
16436 shared_mem_bytes: 0,
16437 };
16438 let (nc, e) = (ncols as i32, eps);
16439 let __s_b = self.gpu.stream();
16440 let mut b = __s_b.launch_builder(&f);
16441 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16442 unsafe {
16443 b.launch(cfg)?;
16444 }
16445 Ok(())
16446 }
16447
16448 #[allow(clippy::too_many_arguments)]
16451 pub fn add_rms_norm_f16out(
16452 &self,
16453 a: &CudaSlice<f32>,
16454 b: &CudaSlice<f32>,
16455 w: &CudaSlice<f32>,
16456 res: &mut CudaSlice<f32>,
16457 dst: &mut CudaSlice<f32>,
16458 dst16: &mut CudaSlice<u8>,
16459 ncols: usize,
16460 nrows: usize,
16461 eps: f32,
16462 ) -> Result<(), Box<dyn std::error::Error>> {
16463 let f = self.func("add_rms_norm_f16out_f32");
16464 let cfg = LaunchConfig {
16465 grid_dim: (nrows as u32, 1, 1),
16466 block_dim: (rms_block(), 1, 1),
16467 shared_mem_bytes: 0,
16468 };
16469 let (nc, e) = (ncols as i32, eps);
16470 let __s_lb = self.gpu.stream();
16471 let mut lb = __s_lb.launch_builder(&f);
16472 lb.arg(a)
16473 .arg(b)
16474 .arg(w)
16475 .arg(res)
16476 .arg(dst)
16477 .arg(dst16)
16478 .arg(&nc)
16479 .arg(&e);
16480 unsafe {
16481 lb.launch(cfg)?;
16482 }
16483 Ok(())
16484 }
16485
16486 pub fn matmul_group_xh(
16489 &self,
16490 ws: &[&crate::model::GpuTensor],
16491 x: &CudaSlice<f32>,
16492 xh: &CudaSlice<u8>,
16493 m: usize,
16494 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16495 let mut out = Vec::with_capacity(ws.len());
16496 let in_f = ws[0].in_features();
16497 for w in ws {
16498 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16499 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16500 out.push(y);
16501 continue;
16502 }
16503 }
16504 out.push(self.matmul(w, x, m)?);
16505 }
16506 Ok(out)
16507 }
16508
16509 pub fn gdn_pad_mask(
16512 &self,
16513 beta: &mut CudaSlice<f32>,
16514 g_log: &mut CudaSlice<f32>,
16515 len_d: &CudaSlice<i32>,
16516 h: usize,
16517 t: usize,
16518 ) -> Result<(), Box<dyn std::error::Error>> {
16519 let f = self.func("gdn_pad_mask_f32");
16520 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16521 let (hi, ti) = (h as i32, t as i32);
16522 let __s_b = self.gpu.stream();
16523 let mut b = __s_b.launch_builder(&f);
16524 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16525 unsafe {
16526 b.launch(cfg)?;
16527 }
16528 Ok(())
16529 }
16530
16531 pub fn row_gather_dev(
16534 &self,
16535 src: &CudaSlice<f32>,
16536 dst: &mut CudaSlice<f32>,
16537 len_d: &CudaSlice<i32>,
16538 ncols: usize,
16539 ) -> Result<(), Box<dyn std::error::Error>> {
16540 let f = self.func("row_gather_dev_f32");
16541 let cfg = LaunchConfig::for_num_elems(ncols as u32);
16542 let nc = ncols as i32;
16543 let __s_b = self.gpu.stream();
16544 let mut b = __s_b.launch_builder(&f);
16545 b.arg(src).arg(dst).arg(len_d).arg(&nc);
16546 unsafe {
16547 b.launch(cfg)?;
16548 }
16549 Ok(())
16550 }
16551
16552 pub fn matmul_group(
16559 &self,
16560 ws: &[&crate::model::GpuTensor],
16561 x: &CudaSlice<f32>,
16562 m: usize,
16563 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16564 use crate::model::GpuTensor;
16565 let mut out = Vec::with_capacity(ws.len());
16566 let any_mirror = ws
16567 .iter()
16568 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16569 if m >= 16 && any_mirror && !self.verify_exact_on() {
16570 let in_f = ws[0].in_features();
16571 let xh = self.f16_act(x, m * in_f, in_f)?;
16572 for w in ws {
16573 if w.in_features() == in_f {
16574 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16575 out.push(y);
16576 continue;
16577 }
16578 }
16579 out.push(self.matmul(w, x, m)?);
16580 }
16581 return Ok(out);
16582 }
16583 for w in ws {
16584 out.push(self.matmul(w, x, m)?);
16585 }
16586 Ok(out)
16587 }
16588
16589 pub fn matmul_group_multi(
16596 &self,
16597 ws: &[&crate::model::GpuTensor],
16598 xs: &[&CudaSlice<f32>],
16599 ms: &[usize],
16600 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16601 assert_eq!(xs.len(), ms.len());
16602 let in_f = ws[0].in_features();
16603 let total: usize = ms.iter().sum();
16604 let mut xcat = self.uninit(total * in_f)?;
16605 let mut off = 0usize;
16606 for (x, &m) in xs.iter().zip(ms) {
16607 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16608 off += m;
16609 }
16610 let ys = self.matmul_group(ws, &xcat, total)?;
16611 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16612 for (w, y) in ws.iter().zip(ys) {
16613 let out_f = w.out_features();
16614 let mut off = 0usize;
16615 for (s, &m) in ms.iter().enumerate() {
16616 let mut ys_s = self.uninit(m * out_f)?;
16617 let src = y.slice(off * out_f..(off + m) * out_f);
16618 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16619 out[s].push(ys_s);
16620 off += m;
16621 }
16622 }
16623 Ok(out)
16624 }
16625
16626 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16636 use crate::model::GpuTensor;
16637 if !legacy_quant_gemm_allowed(
16638 cfg!(memra_portable_cuda),
16639 cfg!(memra_hopper_mma),
16640 std::env::var_os("MEMRA_NO_GEMM").is_some(),
16641 ) {
16642 return false;
16643 }
16644 match w {
16645 GpuTensor::Quant { qtype, .. } => {
16646 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16647 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16648 }
16649 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16650 }
16651 }
16652
16653 pub fn qmatvec_gemm(
16660 &self,
16661 w: &crate::model::GpuTensor,
16662 aq: &CudaSlice<i8>,
16663 ad: &CudaSlice<f32>,
16664 m: usize,
16665 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16666 use crate::model::GpuTensor;
16667 let in_f = w.in_features();
16668 let out_f = w.out_features();
16669 let (bytes, qtype, row_bytes, scale, rp) = match w {
16670 GpuTensor::Quant {
16671 bytes,
16672 qtype,
16673 row_bytes,
16674 scale,
16675 rp,
16676 ..
16677 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16678 _ => unreachable!("gemm_supports guaranteed Quant"),
16679 };
16680 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16686 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16687 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16688 if scale != 1.0 {
16689 self.scale_inplace(&mut y, scale, m * out_f)?;
16690 }
16691 return Ok(y);
16692 }
16693 }
16694 let name = match qtype {
16695 QT_Q8_0 => "qmatvec_gemm_q8_0",
16696 QT_Q4_K => "qmatvec_gemm_q4_K",
16697 QT_Q4_0 => {
16698 if rp {
16699 "qmatvec_gemm_q4_0_rp"
16700 } else {
16701 "qmatvec_gemm_q4_0"
16702 }
16703 }
16704 QT_Q5_K => "qmatvec_gemm_q5_K",
16705 QT_Q6_K => "qmatvec_gemm_q6_K",
16706 QT_NVFP4 => {
16707 if rp {
16708 "qmatvec_gemm_nvfp4_rp"
16709 } else {
16710 "qmatvec_gemm_nvfp4"
16711 }
16712 }
16713 _ => unreachable!(),
16714 };
16715 let f = self.func(name);
16716 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);
16721 let k1_tile = if is_k1 {
16723 k1_launch_override().unwrap_or((128, 128, 8))
16724 } else {
16725 (128, 128, 8)
16726 };
16727 let (bm, bn): (u32, u32) = if is_k1 {
16728 (k1_tile.0, k1_tile.1)
16729 } else {
16730 (64, 256)
16731 };
16732 let warps: u32 = if is_k1 {
16733 k1_tile.2
16734 } else {
16735 match qtype {
16736 QT_NVFP4 => 8,
16737 _ => 4,
16738 }
16739 };
16740 let cfg = LaunchConfig {
16741 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16742 block_dim: (32, warps, 1),
16743 shared_mem_bytes: 0,
16744 };
16745 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16746 let __s_b = self.gpu.stream();
16747 let mut b = __s_b.launch_builder(&f);
16748 b.arg(bytes)
16749 .arg(aq)
16750 .arg(ad)
16751 .arg(&mut y)
16752 .arg(&inf)
16753 .arg(&outf)
16754 .arg(&mi)
16755 .arg(&rb);
16756 unsafe {
16757 b.launch(cfg)?;
16758 }
16759 if scale != 1.0 {
16760 self.scale_inplace(&mut y, scale, m * out_f)?;
16761 }
16762 Ok(y)
16763 }
16764
16765 pub fn qmatvec_gemm_raw(
16770 &self,
16771 bytes: &CudaSlice<u8>,
16772 x: &CudaSlice<f32>,
16773 m: usize,
16774 in_f: usize,
16775 out_f: usize,
16776 qtype: i32,
16777 row_bytes: usize,
16778 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16779 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16780 let name = match qtype {
16781 QT_Q8_0 => "qmatvec_gemm_q8_0",
16782 QT_Q4_K => "qmatvec_gemm_q4_K",
16783 QT_Q4_0 => "qmatvec_gemm_q4_0",
16784 QT_Q5_K => "qmatvec_gemm_q5_K",
16785 QT_Q6_K => "qmatvec_gemm_q6_K",
16786 QT_NVFP4 => "qmatvec_gemm_nvfp4",
16787 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16788 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16789 };
16790 let f = self.func(name);
16791 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);
16795 let k1_tile = if is_k1 {
16797 k1_launch_override().unwrap_or((128, 128, 8))
16798 } else {
16799 (128, 128, 8)
16800 };
16801 let (bm, bn): (u32, u32) = if is_k1 {
16802 (k1_tile.0, k1_tile.1)
16803 } else {
16804 (64, 256)
16805 };
16806 let warps: u32 = if is_k1 {
16807 k1_tile.2
16808 } else {
16809 match qtype {
16810 QT_NVFP4 | QT_NVFP4_RP => 8,
16811 _ => 4,
16812 }
16813 };
16814 let cfg = LaunchConfig {
16815 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16816 block_dim: (32, warps, 1),
16817 shared_mem_bytes: 0,
16818 };
16819 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16820 let __s_b = self.gpu.stream();
16821 let mut b = __s_b.launch_builder(&f);
16822 b.arg(bytes)
16823 .arg(&aq)
16824 .arg(&ad)
16825 .arg(&mut y)
16826 .arg(&inf)
16827 .arg(&outf)
16828 .arg(&mi)
16829 .arg(&rb);
16830 unsafe {
16831 b.launch(cfg)?;
16832 }
16833 Ok(y)
16834 }
16835
16836 pub fn qmatvec_gemm_q8_0_wgmma_raw(
16843 &self,
16844 rp4: &CudaSlice<u8>,
16845 aq: &CudaSlice<i8>,
16846 ad: &CudaSlice<f32>,
16847 m: usize,
16848 in_f: usize,
16849 out_f: usize,
16850 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16851 assert!(
16852 out_f % 64 == 0 && in_f % 32 == 0,
16853 "wgmma GEMM needs out_f%64==0, in_f%32==0"
16854 );
16855 let f = self.func("qmatvec_gemm_q8_0_wgmma");
16856 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
16858 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16859 block_dim: (128, 1, 1),
16860 shared_mem_bytes: 0,
16861 };
16862 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16863 let __s_b = self.gpu.stream();
16864 let mut b = __s_b.launch_builder(&f);
16865 b.arg(rp4)
16866 .arg(aq)
16867 .arg(ad)
16868 .arg(&mut y)
16869 .arg(&inf)
16870 .arg(&outf)
16871 .arg(&mi);
16872 unsafe {
16873 b.launch(cfg)?;
16874 }
16875 Ok(y)
16876 }
16877
16878 pub fn scale_inplace(
16880 &self,
16881 y: &mut CudaSlice<f32>,
16882 s: f32,
16883 n: usize,
16884 ) -> Result<(), Box<dyn std::error::Error>> {
16885 let f = self.func("scale_f32");
16886 let cfg = LaunchConfig::for_num_elems(n as u32);
16887 let (sf, ni) = (s, n as i32);
16888 let __s_b = self.gpu.stream();
16889 let mut b = __s_b.launch_builder(&f);
16890 b.arg(y).arg(&sf).arg(&ni);
16891 unsafe {
16892 b.launch(cfg)?;
16893 }
16894 Ok(())
16895 }
16896
16897 pub fn bf16_to_f32(
16902 &self,
16903 data: &cudarc::driver::CudaView<'_, u8>,
16904 n: usize,
16905 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16906 let mut out = self.alloc_uninit::<f32>(n)?;
16907 let f = self.func("bf16_to_f32");
16908 let cfg = LaunchConfig::for_num_elems(n as u32);
16909 let ni = n as i32;
16910 let __s_b = self.gpu.stream();
16911 let mut b = __s_b.launch_builder(&f);
16912 b.arg(data).arg(&mut out).arg(&ni);
16913 unsafe {
16914 b.launch(cfg)?;
16915 }
16916 Ok(out)
16917 }
16918
16919 fn linear_bf16_chunked(
16926 &self,
16927 x: &CudaSlice<f32>,
16928 data: &CudaSlice<u8>,
16929 m: usize,
16930 in_f: usize,
16931 out_f: usize,
16932 exact: bool,
16933 canonical_chunk_rows: Option<usize>,
16934 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16935 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16939 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16940 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16941 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16942 let started = timing.then(std::time::Instant::now);
16943 let result =
16944 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16945 if let Some(started) = started {
16946 use std::sync::atomic::Ordering;
16947 self.stream().synchronize()?;
16948 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16949 + started.elapsed().as_nanos() as u64;
16950 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16951 + (in_f * out_f * 2) as u64;
16952 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16953 if calls % 1024 == 0 {
16954 eprintln!(
16955 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16956 weight_gb={:.2}",
16957 ns as f64 / 1.0e6,
16958 ns as f64 / calls as f64 / 1.0e3,
16959 wb as f64 / 1.0e9,
16960 );
16961 }
16962 }
16963 result
16964 }
16965
16966 pub(crate) fn bf16_mmv_on() -> bool {
16971 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16972 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16973 }
16974
16975 fn matvec_bf16(
16978 &self,
16979 data: &CudaSlice<u8>,
16980 x: &CudaSlice<f32>,
16981 in_f: usize,
16982 out_f: usize,
16983 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16984 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16985 return Err(format!(
16986 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16987 data.len(),
16988 x.len()
16989 )
16990 .into());
16991 }
16992 let mut y = self.alloc_uninit::<f32>(out_f)?;
16993 let f = self.func("matvec_bf16_f32acc");
16994 let cfg = LaunchConfig {
16995 grid_dim: (out_f as u32, 1, 1),
16996 block_dim: (mmv_block(), 1, 1),
16997 shared_mem_bytes: 0,
16998 };
16999 let ini = in_f as i32;
17000 let __s_bld = self.gpu.stream();
17001 let mut bld = __s_bld.launch_builder(&f);
17002 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17003 unsafe {
17004 bld.launch(cfg)?;
17005 }
17006 Ok(y)
17007 }
17008
17009 #[allow(clippy::too_many_arguments)]
17013 #[allow(clippy::too_many_arguments)]
17018 pub fn qk_norm_rope_append_inc_dcw(
17019 &self,
17020 q_raw: &CudaSlice<f32>,
17021 k_raw: &CudaSlice<f32>,
17022 v_raw: &CudaSlice<f32>,
17023 qw: &CudaSlice<f32>,
17024 kw: &CudaSlice<f32>,
17025 q_out: &mut CudaSlice<f32>,
17026 k_out: &mut CudaSlice<f32>,
17027 pos: &CudaSlice<i32>,
17028 k_plane: &mut CudaSlice<u8>,
17029 v_plane: &mut CudaSlice<u8>,
17030 len_dev: &CudaSlice<i32>,
17033 base_dev: Option<&CudaSlice<i32>>,
17034 done_ctr: &mut CudaSlice<u32>,
17035 kv_dim_k: usize,
17036 kv_dim_v: usize,
17037 k_tok_bytes: usize,
17038 v_tok_bytes: usize,
17039 head_dim: usize,
17040 n_dims: usize,
17041 nh_q: usize,
17042 nh_k: usize,
17043 eps: f32,
17044 freq_base: f32,
17045 freq_scale: f32,
17046 ff: Option<&CudaSlice<f32>>,
17047 ) -> Result<(), Box<dyn std::error::Error>> {
17048 if head_dim != 128
17049 || kv_dim_v != kv_dim_k
17050 || kv_dim_k != nh_k * head_dim
17051 || q_raw.len() < nh_q * head_dim
17052 || k_raw.len() < nh_k * head_dim
17053 || v_raw.len() < kv_dim_v
17054 || q_out.len() < nh_q * head_dim
17055 || k_out.len() < nh_k * head_dim
17056 || pos.is_empty()
17057 || done_ctr.is_empty()
17058 {
17059 return Err(format!(
17060 "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}"
17061 )
17062 .into());
17063 }
17064 let f = self.func("qk_norm_rope_append_inc_dcw");
17065 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17066 let cfg = LaunchConfig {
17067 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17068 block_dim: (128, 1, 1),
17069 shared_mem_bytes: 0,
17070 };
17071 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17072 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17073 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17074 let null: u64 = 0;
17075 let __s_b = self.gpu.stream();
17076 let mut b = __s_b.launch_builder(&f);
17077 b.arg(q_raw)
17078 .arg(k_raw)
17079 .arg(v_raw)
17080 .arg(qw)
17081 .arg(kw)
17082 .arg(q_out)
17083 .arg(k_out)
17084 .arg(pos)
17085 .arg(&mut *k_plane)
17086 .arg(&mut *v_plane)
17087 .arg(len_dev);
17088 match base_dev {
17089 Some(base) => {
17090 b.arg(base);
17091 }
17092 None => {
17093 b.arg(&null);
17094 }
17095 }
17096 b.arg(&mut *done_ctr)
17097 .arg(&kvk)
17098 .arg(&kvv)
17099 .arg(&ktb)
17100 .arg(&vtb)
17101 .arg(&hd)
17102 .arg(&nd)
17103 .arg(&nq)
17104 .arg(&eps)
17105 .arg(&theta_scale)
17106 .arg(&freq_scale);
17107 match ff {
17108 Some(freqs) => {
17109 b.arg(freqs);
17110 }
17111 None => {
17112 b.arg(&null);
17113 }
17114 }
17115 unsafe {
17116 b.launch(cfg)?;
17117 }
17118 Ok(())
17119 }
17120
17121 pub fn qk_norm_rope_into(
17122 &self,
17123 q_raw: &CudaSlice<f32>,
17124 k_raw: &CudaSlice<f32>,
17125 qw: &CudaSlice<f32>,
17126 kw: &CudaSlice<f32>,
17127 q_out: &mut CudaSlice<f32>,
17128 k_out: &mut CudaSlice<f32>,
17129 pos: &CudaSlice<i32>,
17130 head_dim: usize,
17131 n_dims: usize,
17132 nh_q: usize,
17133 nh_k: usize,
17134 eps: f32,
17135 freq_base: f32,
17136 freq_scale: f32,
17137 ff: Option<&CudaSlice<f32>>,
17138 ) -> Result<(), Box<dyn std::error::Error>> {
17139 if head_dim > 512
17140 || q_raw.len() < nh_q * head_dim
17141 || k_raw.len() < nh_k * head_dim
17142 || q_out.len() < nh_q * head_dim
17143 || k_out.len() < nh_k * head_dim
17144 || qw.len() < head_dim
17145 || kw.len() < head_dim
17146 || pos.is_empty()
17147 {
17148 return Err(format!(
17149 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17150 )
17151 .into());
17152 }
17153 let f = self.func("qk_norm_rope_f32");
17154 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17155 let cfg = LaunchConfig {
17156 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17157 block_dim: (128, 1, 1),
17158 shared_mem_bytes: 0,
17159 };
17160 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17161 let __s_b = self.gpu.stream();
17162 let mut b = __s_b.launch_builder(&f);
17163 b.arg(q_raw)
17164 .arg(k_raw)
17165 .arg(qw)
17166 .arg(kw)
17167 .arg(q_out)
17168 .arg(k_out)
17169 .arg(pos)
17170 .arg(&hd)
17171 .arg(&nd)
17172 .arg(&nq)
17173 .arg(&eps)
17174 .arg(&theta_scale)
17175 .arg(&freq_scale);
17176 match ff {
17177 Some(ffv) => {
17178 b.arg(ffv);
17179 unsafe {
17180 b.launch(cfg)?;
17181 }
17182 }
17183 None => {
17184 let null: u64 = 0;
17185 b.arg(&null);
17186 unsafe {
17187 b.launch(cfg)?;
17188 }
17189 }
17190 }
17191 Ok(())
17192 }
17193
17194 #[allow(clippy::too_many_arguments)]
17197 pub fn matvec_f32_b4_into(
17198 &self,
17199 w: [&CudaSlice<f32>; 4],
17200 x: &CudaSlice<f32>,
17201 y: &mut CudaSlice<f32>,
17202 block_cols: usize,
17203 out_f: usize,
17204 ) -> Result<(), Box<dyn std::error::Error>> {
17205 if block_cols % 4 != 0
17206 || x.len() < 4 * block_cols
17207 || y.len() < out_f
17208 || w.iter().any(|w| w.len() != out_f * block_cols)
17209 {
17210 return Err(format!(
17211 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
17212 x.len()
17213 )
17214 .into());
17215 }
17216 let f = self.func("matvec_f32_b4");
17217 let cfg = LaunchConfig {
17218 grid_dim: (out_f as u32, 1, 1),
17219 block_dim: (128, 1, 1),
17220 shared_mem_bytes: 0,
17221 };
17222 let (bc, of) = (block_cols as i32, out_f as i32);
17223 let __s_b = self.gpu.stream();
17224 let mut b = __s_b.launch_builder(&f);
17225 b.arg(w[0])
17226 .arg(w[1])
17227 .arg(w[2])
17228 .arg(w[3])
17229 .arg(x)
17230 .arg(y)
17231 .arg(&bc)
17232 .arg(&of);
17233 unsafe {
17234 b.launch(cfg)?;
17235 }
17236 Ok(())
17237 }
17238
17239 pub fn axpy_rows_seq_into(
17242 &self,
17243 x: &CudaSlice<f32>,
17244 w: &CudaSlice<f32>,
17245 y: &mut CudaSlice<f32>,
17246 width: usize,
17247 n_rows: usize,
17248 ) -> Result<(), Box<dyn std::error::Error>> {
17249 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
17250 return Err(format!(
17251 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
17252 x.len(),
17253 w.len(),
17254 y.len()
17255 )
17256 .into());
17257 }
17258 let f = self.func("axpy_rows_seq_f32");
17259 let cfg = LaunchConfig::for_num_elems(width as u32);
17260 let (wi, nr) = (width as i32, n_rows as i32);
17261 let __s_b = self.gpu.stream();
17262 let mut b = __s_b.launch_builder(&f);
17263 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
17264 unsafe {
17265 b.launch(cfg)?;
17266 }
17267 Ok(())
17268 }
17269
17270 #[allow(clippy::too_many_arguments)]
17274 pub fn axpy_rows_seq_md_off_into(
17275 &self,
17276 x: &CudaSlice<f32>,
17277 w_route: &CudaSlice<f32>,
17278 md: &CudaSlice<f32>,
17279 sel: &CudaSlice<i32>,
17280 y: &mut CudaSlice<f32>,
17281 width: usize,
17282 n_rows: usize,
17283 row0: usize,
17284 ) -> Result<(), Box<dyn std::error::Error>> {
17285 if x.len() < (row0 + n_rows) * width
17286 || w_route.len() < row0 + n_rows
17287 || sel.len() < row0 + n_rows
17288 || y.len() < width
17289 {
17290 return Err(format!(
17291 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
17292 rows={n_rows} row0={row0}",
17293 x.len(),
17294 w_route.len(),
17295 sel.len(),
17296 y.len()
17297 )
17298 .into());
17299 }
17300 let f = self.func("axpy_rows_seq_md_off_f32");
17301 let cfg = LaunchConfig::for_num_elems(width as u32);
17302 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
17303 let __s_b = self.gpu.stream();
17304 let mut b = __s_b.launch_builder(&f);
17305 b.arg(x)
17306 .arg(w_route)
17307 .arg(md)
17308 .arg(sel)
17309 .arg(y)
17310 .arg(&wi)
17311 .arg(&nr)
17312 .arg(&r0);
17313 unsafe {
17314 b.launch(cfg)?;
17315 }
17316 Ok(())
17317 }
17318
17319 #[allow(clippy::too_many_arguments)]
17324 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
17325 &self,
17326 gate_bank: &CudaSlice<u8>,
17327 up_bank: &CudaSlice<u8>,
17328 sel: &CudaSlice<i32>,
17329 aq: &CudaSlice<i8>,
17330 ad: &CudaSlice<f32>,
17331 yg: &mut CudaSlice<f32>,
17332 yu: &mut CudaSlice<f32>,
17333 n_sel: usize,
17334 n_sel_col: usize,
17335 in_f: usize,
17336 out_f: usize,
17337 row_bytes: usize,
17338 expert_stride: usize,
17339 act_row_stride: usize,
17340 ad_row_stride: usize,
17341 ) -> Result<(), Box<dyn std::error::Error>> {
17342 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
17343 if yg.len() < n_sel * out_f
17344 || yu.len() < n_sel * out_f
17345 || sel.len() < n_sel
17346 || n_sel_col == 0
17347 || n_sel % n_sel_col != 0
17348 {
17349 return Err("NVFP4 gu tcol geometry".into());
17350 }
17351 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
17352 let cfg = LaunchConfig {
17353 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
17354 block_dim: (128, 1, 1),
17355 shared_mem_bytes: 0,
17356 };
17357 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
17358 let (rb, es) = (row_bytes as i64, expert_stride as i64);
17359 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
17360 let __s_b = self.gpu.stream();
17361 let mut b = __s_b.launch_builder(&f);
17362 b.arg(gate_bank)
17363 .arg(up_bank)
17364 .arg(sel)
17365 .arg(aq)
17366 .arg(ad)
17367 .arg(yg)
17368 .arg(yu)
17369 .arg(&inf)
17370 .arg(&outf)
17371 .arg(&ns)
17372 .arg(&rb)
17373 .arg(&es)
17374 .arg(&ars)
17375 .arg(&adrs)
17376 .arg(&nsc);
17377 unsafe {
17378 b.launch(cfg)?;
17379 }
17380 Ok(())
17381 }
17382
17383 #[allow(clippy::too_many_arguments)]
17386 pub fn axpy_rows_seq_md_into(
17387 &self,
17388 x: &CudaSlice<f32>,
17389 w_route: &CudaSlice<f32>,
17390 md: &CudaSlice<f32>,
17391 sel: &CudaSlice<i32>,
17392 y: &mut CudaSlice<f32>,
17393 width: usize,
17394 n_rows: usize,
17395 ) -> Result<(), Box<dyn std::error::Error>> {
17396 if x.len() < n_rows * width
17397 || w_route.len() < n_rows
17398 || sel.len() < n_rows
17399 || y.len() < width
17400 {
17401 return Err(format!(
17402 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17403 x.len(),
17404 w_route.len(),
17405 sel.len(),
17406 y.len()
17407 )
17408 .into());
17409 }
17410 let f = self.func("axpy_rows_seq_md_f32");
17411 let cfg = LaunchConfig::for_num_elems(width as u32);
17412 let (wi, nr) = (width as i32, n_rows as i32);
17413 let __s_b = self.gpu.stream();
17414 let mut b = __s_b.launch_builder(&f);
17415 b.arg(x)
17416 .arg(w_route)
17417 .arg(md)
17418 .arg(sel)
17419 .arg(y)
17420 .arg(&wi)
17421 .arg(&nr);
17422 unsafe {
17423 b.launch(cfg)?;
17424 }
17425 Ok(())
17426 }
17427
17428 #[allow(clippy::too_many_arguments)]
17430 #[allow(clippy::too_many_arguments)]
17434 pub fn matvec_bf16_qkvg_tcol_into(
17435 &self,
17436 wq: &CudaSlice<u8>,
17437 wk: &CudaSlice<u8>,
17438 wv: &CudaSlice<u8>,
17439 wg: &CudaSlice<u8>,
17440 x_t: &CudaSlice<f32>,
17441 yq: &mut CudaSlice<f32>,
17442 yk: &mut CudaSlice<f32>,
17443 yv: &mut CudaSlice<f32>,
17444 yg: &mut CudaSlice<f32>,
17445 in_f: usize,
17446 out_q: usize,
17447 out_kv: usize,
17448 out_g: usize,
17449 t: usize,
17450 ) -> Result<(), Box<dyn std::error::Error>> {
17451 if t == 0
17452 || t > 8
17453 || in_f % 8 != 0
17454 || x_t.len() < t * in_f
17455 || yq.len() < t * out_q
17456 || yk.len() < t * out_kv
17457 || yv.len() < t * out_kv
17458 || (out_g > 0 && yg.len() < t * out_g)
17459 {
17460 return Err("matvec_bf16_qkvg_tcol geometry".into());
17461 }
17462 let f = self.func("matvec_bf16_qkvg_tcol");
17463 let grid = out_q + 2 * out_kv + out_g;
17464 let cfg = LaunchConfig {
17465 grid_dim: (grid as u32, 1, 1),
17466 block_dim: (mmv_block(), 1, 1),
17467 shared_mem_bytes: 0,
17468 };
17469 let (ini, oq, okv, og, ti) = (
17470 in_f as i32,
17471 out_q as i32,
17472 out_kv as i32,
17473 out_g as i32,
17474 t as i32,
17475 );
17476 let __s_b = self.gpu.stream();
17477 let mut b = __s_b.launch_builder(&f);
17478 b.arg(wq)
17479 .arg(wk)
17480 .arg(wv)
17481 .arg(wg)
17482 .arg(x_t)
17483 .arg(yq)
17484 .arg(yk)
17485 .arg(yv)
17486 .arg(yg)
17487 .arg(&ini)
17488 .arg(&oq)
17489 .arg(&okv)
17490 .arg(&og)
17491 .arg(&ti);
17492 unsafe {
17493 b.launch(cfg)?;
17494 }
17495 Ok(())
17496 }
17497
17498 pub fn matvec_bf16_qkvg_into(
17499 &self,
17500 wq: &CudaSlice<u8>,
17501 wk: &CudaSlice<u8>,
17502 wv: &CudaSlice<u8>,
17503 wg: &CudaSlice<u8>,
17504 x: &CudaSlice<f32>,
17505 yq: &mut CudaSlice<f32>,
17506 yk: &mut CudaSlice<f32>,
17507 yv: &mut CudaSlice<f32>,
17508 yg: &mut CudaSlice<f32>,
17509 in_f: usize,
17510 out_q: usize,
17511 out_kv: usize,
17512 out_g: usize,
17513 ) -> Result<(), Box<dyn std::error::Error>> {
17514 if in_f % 8 != 0
17515 || wq.len() != out_q * in_f * 2
17516 || wk.len() != out_kv * in_f * 2
17517 || wv.len() != out_kv * in_f * 2
17518 || wg.len() < out_g * in_f * 2
17519 || x.len() < in_f
17520 || yq.len() < out_q
17521 || yk.len() < out_kv
17522 || yv.len() < out_kv
17523 || (out_g > 0 && yg.len() < out_g)
17524 {
17525 return Err(format!(
17526 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17527 )
17528 .into());
17529 }
17530 let f = self.func("matvec_bf16_qkvg");
17531 let cfg = LaunchConfig {
17532 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17533 block_dim: (mmv_block(), 1, 1),
17534 shared_mem_bytes: 0,
17535 };
17536 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17537 let __s_b = self.gpu.stream();
17538 let mut b = __s_b.launch_builder(&f);
17539 b.arg(wq)
17540 .arg(wk)
17541 .arg(wv)
17542 .arg(wg)
17543 .arg(x)
17544 .arg(yq)
17545 .arg(yk)
17546 .arg(yv)
17547 .arg(yg)
17548 .arg(&inf)
17549 .arg(&oq)
17550 .arg(&okv)
17551 .arg(&og);
17552 unsafe {
17553 b.launch(cfg)?;
17554 }
17555 Ok(())
17556 }
17557
17558 pub fn matvec_bf16_b4_into(
17560 &self,
17561 w: [&CudaSlice<u8>; 4],
17562 x: &CudaSlice<f32>,
17563 y: &mut CudaSlice<f32>,
17564 block_cols: usize,
17565 out_f: usize,
17566 ) -> Result<(), Box<dyn std::error::Error>> {
17567 if block_cols % 8 != 0
17568 || x.len() < 4 * block_cols
17569 || y.len() < out_f
17570 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17571 {
17572 return Err(format!(
17573 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17574 x.len()
17575 )
17576 .into());
17577 }
17578 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17581 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17582 let f = self.func(if x2 {
17583 "matvec_bf16_b4_x2"
17584 } else {
17585 "matvec_bf16_b4"
17586 });
17587 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17588 let cfg = LaunchConfig {
17589 grid_dim: (grid as u32, 1, 1),
17590 block_dim: (mmv_block(), 1, 1),
17591 shared_mem_bytes: 0,
17592 };
17593 let (bc, of) = (block_cols as i32, out_f as i32);
17594 let __s_b = self.gpu.stream();
17595 let mut b = __s_b.launch_builder(&f);
17596 b.arg(w[0])
17597 .arg(w[1])
17598 .arg(w[2])
17599 .arg(w[3])
17600 .arg(x)
17601 .arg(y)
17602 .arg(&bc)
17603 .arg(&of);
17604 unsafe {
17605 b.launch(cfg)?;
17606 }
17607 Ok(())
17608 }
17609
17610 pub fn matvec_bf16_b4_tcol_into(
17616 &self,
17617 w: [&CudaSlice<u8>; 4],
17618 x_t: &CudaSlice<f32>,
17619 y_t: &mut CudaSlice<f32>,
17620 block_cols: usize,
17621 out_f: usize,
17622 t: usize,
17623 ) -> Result<(), Box<dyn std::error::Error>> {
17624 if block_cols % 8 != 0
17625 || t == 0
17626 || t > 8
17627 || x_t.len() < t * 4 * block_cols
17628 || y_t.len() < t * out_f
17629 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17630 {
17631 return Err(format!(
17632 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17633 x_t.len()
17634 )
17635 .into());
17636 }
17637 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17638 return Err(
17639 "b4 tcol verify is qualified against the plain b4 kernel only \
17640 (MEMRA_B4_X2=1 is a different t=1 program)"
17641 .into(),
17642 );
17643 }
17644 let f = self.func("matvec_bf16_b4_tcol");
17645 let cfg = LaunchConfig {
17646 grid_dim: (out_f as u32, 1, 1),
17647 block_dim: (mmv_block(), 1, 1),
17648 shared_mem_bytes: 0,
17649 };
17650 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17651 let __s_b = self.gpu.stream();
17652 let mut b = __s_b.launch_builder(&f);
17653 b.arg(w[0])
17654 .arg(w[1])
17655 .arg(w[2])
17656 .arg(w[3])
17657 .arg(x_t)
17658 .arg(y_t)
17659 .arg(&bc)
17660 .arg(&of)
17661 .arg(&ti);
17662 unsafe {
17663 b.launch(cfg)?;
17664 }
17665 Ok(())
17666 }
17667
17668 pub fn matvec_bf16_into(
17670 &self,
17671 data: &CudaSlice<u8>,
17672 x: &CudaSlice<f32>,
17673 y: &mut CudaSlice<f32>,
17674 in_f: usize,
17675 out_f: usize,
17676 ) -> Result<(), Box<dyn std::error::Error>> {
17677 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17678 return Err(format!(
17679 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17680 data.len(),
17681 x.len(),
17682 y.len()
17683 )
17684 .into());
17685 }
17686 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17690 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17691 && in_f <= 2048;
17692 if x4 {
17693 let f = self.func("matvec_bf16_f32acc_x4");
17694 let cfg = LaunchConfig {
17695 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17696 block_dim: (mmv_block(), 1, 1),
17697 shared_mem_bytes: 0,
17698 };
17699 let (ini, outi) = (in_f as i32, out_f as i32);
17700 let __s_b = self.gpu.stream();
17701 let mut b = __s_b.launch_builder(&f);
17702 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17703 unsafe {
17704 b.launch(cfg)?;
17705 }
17706 return Ok(());
17707 }
17708 let f = self.func("matvec_bf16_f32acc");
17709 let cfg = LaunchConfig {
17710 grid_dim: (out_f as u32, 1, 1),
17711 block_dim: (mmv_block(), 1, 1),
17712 shared_mem_bytes: 0,
17713 };
17714 let ini = in_f as i32;
17715 let __s_b = self.gpu.stream();
17716 let mut b = __s_b.launch_builder(&f);
17717 b.arg(data).arg(x).arg(y).arg(&ini);
17718 unsafe {
17719 b.launch(cfg)?;
17720 }
17721 Ok(())
17722 }
17723
17724 pub fn matvec_bf16_view_into(
17727 &self,
17728 data: &cudarc::driver::CudaView<'_, u8>,
17729 x: &CudaSlice<f32>,
17730 y: &mut CudaSlice<f32>,
17731 in_f: usize,
17732 out_f: usize,
17733 ) -> Result<(), Box<dyn std::error::Error>> {
17734 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17735 return Err(format!(
17736 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17737 data.len(),
17738 x.len(),
17739 y.len()
17740 )
17741 .into());
17742 }
17743 let f = self.func("matvec_bf16_f32acc");
17744 let cfg = LaunchConfig {
17745 grid_dim: (out_f as u32, 1, 1),
17746 block_dim: (mmv_block(), 1, 1),
17747 shared_mem_bytes: 0,
17748 };
17749 let ini = in_f as i32;
17750 let __s_b = self.gpu.stream();
17751 let mut b = __s_b.launch_builder(&f);
17752 b.arg(data).arg(x).arg(y).arg(&ini);
17753 unsafe {
17754 b.launch(cfg)?;
17755 }
17756 Ok(())
17757 }
17758
17759 pub fn matvec_bf16_raw_out(
17762 &self,
17763 w: &CudaSlice<u8>,
17764 x: &CudaSlice<f32>,
17765 y_raw: u64,
17766 in_f: usize,
17767 out_f: usize,
17768 ) -> Result<(), Box<dyn std::error::Error>> {
17769 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17770 return Err("matvec_bf16_raw_out geometry".into());
17771 }
17772 let f = self.func("matvec_bf16_f32acc");
17773 let cfg = LaunchConfig {
17774 grid_dim: (out_f as u32, 1, 1),
17775 block_dim: (mmv_block(), 1, 1),
17776 shared_mem_bytes: 0,
17777 };
17778 let ini = in_f as i32;
17779 let __s_b = self.gpu.stream();
17780 let mut b = __s_b.launch_builder(&f);
17781 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17782 unsafe {
17783 b.launch(cfg)?;
17784 }
17785 Ok(())
17786 }
17787
17788 pub fn add3_raw(
17792 &self,
17793 a: &CudaSlice<f32>,
17794 b: &CudaSlice<f32>,
17795 sh_raw: u64,
17796 scale_raw: u64,
17797 dst: &mut CudaSlice<f32>,
17798 n: usize,
17799 ) -> Result<(), Box<dyn std::error::Error>> {
17800 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17801 return Err("add3_raw geometry".into());
17802 }
17803 let f = self.func("add3_f32");
17804 let cfg = LaunchConfig {
17805 grid_dim: ((n as u32).div_ceil(256), 1, 1),
17806 block_dim: (256, 1, 1),
17807 shared_mem_bytes: 0,
17808 };
17809 let ni = n as i32;
17810 let __s_b = self.gpu.stream();
17811 let mut bld = __s_b.launch_builder(&f);
17812 bld.arg(a)
17813 .arg(b)
17814 .arg(&sh_raw)
17815 .arg(&scale_raw)
17816 .arg(dst)
17817 .arg(&ni);
17818 unsafe {
17819 bld.launch(cfg)?;
17820 }
17821 Ok(())
17822 }
17823
17824 pub fn matvec_bf16_down_addscale_into(
17827 &self,
17828 w: &CudaSlice<u8>,
17829 x: &CudaSlice<f32>,
17830 scale: &CudaSlice<f32>,
17831 dst: &mut CudaSlice<f32>,
17832 in_f: usize,
17833 out_f: usize,
17834 ) -> Result<(), Box<dyn std::error::Error>> {
17835 if w.len() != in_f * out_f * 2
17836 || x.len() < in_f
17837 || in_f % 8 != 0
17838 || dst.len() < out_f
17839 || scale.is_empty()
17840 {
17841 return Err("matvec_bf16_down_addscale geometry".into());
17842 }
17843 let f = self.func("matvec_bf16_down_addscale");
17844 let cfg = LaunchConfig {
17845 grid_dim: (out_f as u32, 1, 1),
17846 block_dim: (mmv_block(), 1, 1),
17847 shared_mem_bytes: 0,
17848 };
17849 let ini = in_f as i32;
17850 let __s_b = self.gpu.stream();
17851 let mut b = __s_b.launch_builder(&f);
17852 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17853 unsafe {
17854 b.launch(cfg)?;
17855 }
17856 Ok(())
17857 }
17858
17859 pub fn matvec_bf16_dual_silu_into(
17862 &self,
17863 wg: &CudaSlice<u8>,
17864 wu: &CudaSlice<u8>,
17865 x: &CudaSlice<f32>,
17866 act: &mut CudaSlice<f32>,
17867 in_f: usize,
17868 out_f: usize,
17869 limit: Option<f32>,
17870 ) -> Result<(), Box<dyn std::error::Error>> {
17871 if wg.len() != in_f * out_f * 2
17872 || wu.len() != in_f * out_f * 2
17873 || x.len() < in_f
17874 || in_f % 8 != 0
17875 || act.len() < out_f
17876 {
17877 return Err("matvec_bf16_dual_silu geometry".into());
17878 }
17879 let f = self.func("matvec_bf16_dual_silu");
17880 let cfg = LaunchConfig {
17881 grid_dim: (out_f as u32, 1, 1),
17882 block_dim: (mmv_block(), 1, 1),
17883 shared_mem_bytes: 0,
17884 };
17885 let (ini, outi) = (in_f as i32, out_f as i32);
17886 let lim = limit.unwrap_or(0.0);
17887 let __s_b = self.gpu.stream();
17888 let mut b = __s_b.launch_builder(&f);
17889 b.arg(wg)
17890 .arg(wu)
17891 .arg(x)
17892 .arg(act)
17893 .arg(&ini)
17894 .arg(&outi)
17895 .arg(&lim);
17896 unsafe {
17897 b.launch(cfg)?;
17898 }
17899 Ok(())
17900 }
17901
17902 #[allow(clippy::too_many_arguments)]
17905 pub fn matvec_bf16_dual_view_into(
17906 &self,
17907 wg: &cudarc::driver::CudaView<'_, u8>,
17908 wu: &cudarc::driver::CudaView<'_, u8>,
17909 x: &CudaSlice<f32>,
17910 yg: &mut CudaSlice<f32>,
17911 yu: &mut CudaSlice<f32>,
17912 in_f: usize,
17913 out_f: usize,
17914 ) -> Result<(), Box<dyn std::error::Error>> {
17915 if wg.len() != in_f * out_f * 2
17916 || wu.len() != in_f * out_f * 2
17917 || x.len() < in_f
17918 || in_f % 8 != 0
17919 || yg.len() < out_f
17920 || yu.len() < out_f
17921 {
17922 return Err(format!(
17923 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17924 wg.len(),
17925 wu.len(),
17926 x.len()
17927 )
17928 .into());
17929 }
17930 let f = self.func("matvec_bf16_dual");
17931 let cfg = LaunchConfig {
17932 grid_dim: ((2 * out_f) as u32, 1, 1),
17933 block_dim: (mmv_block(), 1, 1),
17934 shared_mem_bytes: 0,
17935 };
17936 let (ini, outi) = (in_f as i32, out_f as i32);
17937 let __s_b = self.gpu.stream();
17938 let mut b = __s_b.launch_builder(&f);
17939 b.arg(wg)
17940 .arg(wu)
17941 .arg(x)
17942 .arg(yg)
17943 .arg(yu)
17944 .arg(&ini)
17945 .arg(&outi);
17946 unsafe {
17947 b.launch(cfg)?;
17948 }
17949 Ok(())
17950 }
17951
17952 #[allow(clippy::too_many_arguments)]
17954 pub fn matvec_bf16_dual_into(
17955 &self,
17956 wg: &CudaSlice<u8>,
17957 wu: &CudaSlice<u8>,
17958 x: &CudaSlice<f32>,
17959 yg: &mut CudaSlice<f32>,
17960 yu: &mut CudaSlice<f32>,
17961 in_f: usize,
17962 out_f: usize,
17963 ) -> Result<(), Box<dyn std::error::Error>> {
17964 if wg.len() != in_f * out_f * 2
17965 || wu.len() != in_f * out_f * 2
17966 || x.len() < in_f
17967 || in_f % 8 != 0
17968 || yg.len() < out_f
17969 || yu.len() < out_f
17970 {
17971 return Err(format!(
17972 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17973 wg.len(),
17974 wu.len(),
17975 x.len()
17976 )
17977 .into());
17978 }
17979 let f = self.func("matvec_bf16_dual");
17980 let cfg = LaunchConfig {
17981 grid_dim: ((2 * out_f) as u32, 1, 1),
17982 block_dim: (mmv_block(), 1, 1),
17983 shared_mem_bytes: 0,
17984 };
17985 let (ini, outi) = (in_f as i32, out_f as i32);
17986 let __s_b = self.gpu.stream();
17987 let mut b = __s_b.launch_builder(&f);
17988 b.arg(wg)
17989 .arg(wu)
17990 .arg(x)
17991 .arg(yg)
17992 .arg(yu)
17993 .arg(&ini)
17994 .arg(&outi);
17995 unsafe {
17996 b.launch(cfg)?;
17997 }
17998 Ok(())
17999 }
18000
18001 pub(crate) fn matvec_bf16_dual(
18004 &self,
18005 wg: &CudaSlice<u8>,
18006 wu: &CudaSlice<u8>,
18007 x: &CudaSlice<f32>,
18008 in_f: usize,
18009 out_f: usize,
18010 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
18011 if wg.len() != in_f * out_f * 2
18012 || wu.len() != in_f * out_f * 2
18013 || x.len() < in_f
18014 || in_f % 8 != 0
18015 {
18016 return Err(format!(
18017 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
18018 wg.len(),
18019 wu.len(),
18020 x.len()
18021 )
18022 .into());
18023 }
18024 let mut yg = self.alloc_uninit::<f32>(out_f)?;
18025 let mut yu = self.alloc_uninit::<f32>(out_f)?;
18026 let f = self.func("matvec_bf16_dual");
18027 let cfg = LaunchConfig {
18028 grid_dim: ((2 * out_f) as u32, 1, 1),
18029 block_dim: (mmv_block(), 1, 1),
18030 shared_mem_bytes: 0,
18031 };
18032 let (ini, outi) = (in_f as i32, out_f as i32);
18033 let __s_b = self.gpu.stream();
18034 let mut b = __s_b.launch_builder(&f);
18035 b.arg(wg)
18036 .arg(wu)
18037 .arg(x)
18038 .arg(&mut yg)
18039 .arg(&mut yu)
18040 .arg(&ini)
18041 .arg(&outi);
18042 unsafe {
18043 b.launch(cfg)?;
18044 }
18045 Ok((yg, yu))
18046 }
18047
18048 #[allow(clippy::too_many_arguments)]
18049 fn linear_bf16_chunked_inner(
18050 &self,
18051 x: &CudaSlice<f32>,
18052 data: &CudaSlice<u8>,
18053 m: usize,
18054 in_f: usize,
18055 out_f: usize,
18056 exact: bool,
18057 canonical_chunk_rows: Option<usize>,
18058 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18059 const CHUNK_BYTES: usize = 256 << 20;
18060 if m == 1
18063 && !exact
18064 && canonical_chunk_rows.is_none()
18065 && in_f % 8 == 0
18066 && Self::bf16_mmv_on()
18067 {
18068 return self.matvec_bf16(data, x, in_f, out_f);
18069 }
18070 let row_bytes = in_f
18071 .checked_mul(std::mem::size_of::<f32>())
18072 .ok_or("BF16 chunk row byte count overflow")?;
18073 if row_bytes == 0 || out_f == 0 {
18074 return Err("BF16 chunk dimensions must be nonzero".into());
18075 }
18076 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
18077 let chunk_rows = match canonical_chunk_rows {
18078 Some(rows) if rows == 0 => {
18079 return Err("canonical BF16 chunk rows must be nonzero".into());
18080 }
18081 Some(rows) if rows > max_chunk_rows => {
18082 return Err(format!(
18083 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
18084 )
18085 .into());
18086 }
18087 Some(rows) if out_f % rows != 0 => {
18088 return Err(format!(
18089 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
18090 )
18091 .into());
18092 }
18093 Some(rows) => rows,
18094 None => max_chunk_rows,
18095 };
18096 if chunk_rows >= out_f {
18097 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
18098 return if exact {
18099 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
18100 } else {
18101 self.linear(x, &wf32, m, in_f, out_f)
18102 };
18103 }
18104 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18105 let mut r0 = 0usize;
18106 while r0 < out_f {
18107 let rows = chunk_rows.min(out_f - r0);
18108 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
18109 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
18110 let yc = if exact {
18111 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
18112 } else {
18113 self.linear(x, &wf32, m, in_f, rows)?
18114 };
18115 for mi in 0..m {
18117 let src = yc.slice(mi * rows..(mi + 1) * rows);
18118 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
18119 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
18120 }
18121 r0 += rows;
18122 }
18123 Ok(y)
18124 }
18125
18126 pub fn linear_bf16_resident(
18130 &self,
18131 x: &CudaSlice<f32>,
18132 data: &CudaSlice<u8>,
18133 m: usize,
18134 in_f: usize,
18135 out_f: usize,
18136 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18137 if data.len() != in_f * out_f * 2 {
18138 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18139 }
18140 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
18141 }
18142
18143 pub fn linear_bf16_resident_canonical_rows(
18149 &self,
18150 x: &CudaSlice<f32>,
18151 data: &CudaSlice<u8>,
18152 m: usize,
18153 in_f: usize,
18154 out_f: usize,
18155 canonical_chunk_rows: usize,
18156 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18157 if data.len() != in_f * out_f * 2 {
18158 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
18159 }
18160 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
18161 }
18162
18163 pub fn linear_f32_resident_canonical_rows(
18168 &self,
18169 x: &CudaSlice<f32>,
18170 data: &CudaSlice<f32>,
18171 m: usize,
18172 in_f: usize,
18173 out_f: usize,
18174 canonical_chunk_rows: usize,
18175 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18176 self.linear_f32_resident_canonical_rows_inner(
18177 x,
18178 data,
18179 m,
18180 in_f,
18181 out_f,
18182 canonical_chunk_rows,
18183 false,
18184 )
18185 }
18186
18187 pub fn linear_f32_resident_canonical_rows_strided(
18193 &self,
18194 x: &CudaSlice<f32>,
18195 data: &CudaSlice<f32>,
18196 m: usize,
18197 in_f: usize,
18198 out_f: usize,
18199 canonical_chunk_rows: usize,
18200 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18201 self.linear_f32_resident_canonical_rows_inner(
18202 x,
18203 data,
18204 m,
18205 in_f,
18206 out_f,
18207 canonical_chunk_rows,
18208 true,
18209 )
18210 }
18211
18212 fn linear_f32_resident_canonical_rows_inner(
18213 &self,
18214 x: &CudaSlice<f32>,
18215 data: &CudaSlice<f32>,
18216 m: usize,
18217 in_f: usize,
18218 out_f: usize,
18219 canonical_chunk_rows: usize,
18220 strided_output: bool,
18221 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18222 if data.len() != in_f * out_f {
18223 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18224 }
18225 if canonical_chunk_rows == 0
18226 || canonical_chunk_rows > out_f
18227 || out_f % canonical_chunk_rows != 0
18228 {
18229 return Err(format!(
18230 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18231 )
18232 .into());
18233 }
18234 if canonical_chunk_rows == out_f {
18235 return self.linear(x, data, m, in_f, out_f);
18236 }
18237
18238 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
18239 let input = x.slice(0..x.len());
18240 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18241 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18242 if m == 1 {
18243 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18244 self.linear_device_into(
18245 &input,
18246 &weights,
18247 &mut destination,
18248 1,
18249 in_f,
18250 canonical_chunk_rows,
18251 )?;
18252 continue;
18253 }
18254 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
18255 if strided_output {
18256 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
18257 } else {
18258 for token in 0..m {
18259 let source = chunk
18260 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
18261 let mut destination =
18262 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
18263 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
18264 }
18265 }
18266 }
18267 Ok(y)
18268 }
18269
18270 pub fn linear_f32_resident_canonical_rows_t1_into(
18275 &self,
18276 x: &CudaSlice<f32>,
18277 data: &CudaSlice<f32>,
18278 y: &mut CudaSlice<f32>,
18279 in_f: usize,
18280 out_f: usize,
18281 canonical_chunk_rows: usize,
18282 ) -> Result<(), Box<dyn std::error::Error>> {
18283 if data.len() != in_f * out_f {
18284 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
18285 }
18286 if y.len() != out_f || x.len() != in_f {
18287 return Err(format!(
18288 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
18289 x.len(),
18290 y.len()
18291 )
18292 .into());
18293 }
18294 if canonical_chunk_rows == 0
18295 || canonical_chunk_rows > out_f
18296 || out_f % canonical_chunk_rows != 0
18297 {
18298 return Err(format!(
18299 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
18300 )
18301 .into());
18302 }
18303 let input = x.slice(0..x.len());
18304 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
18305 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
18306 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
18307 self.linear_device_into(
18308 &input,
18309 &weights,
18310 &mut destination,
18311 1,
18312 in_f,
18313 canonical_chunk_rows,
18314 )?;
18315 }
18316 Ok(())
18317 }
18318
18319 pub fn linear_t1_into(
18322 &self,
18323 x: &cudarc::driver::CudaView<'_, f32>,
18324 w: &cudarc::driver::CudaView<'_, f32>,
18325 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
18326 in_f: usize,
18327 out_f: usize,
18328 ) -> Result<(), Box<dyn std::error::Error>> {
18329 self.linear_device_into(x, w, y, 1, in_f, out_f)
18330 }
18331
18332 pub fn linear_decode_exact(
18339 &self,
18340 x: &CudaSlice<f32>,
18341 w: &CudaSlice<f32>,
18342 m_tokens: usize,
18343 in_f: usize,
18344 out_f: usize,
18345 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18346 if m_tokens == 1 {
18347 return self.linear(x, w, 1, in_f, out_f);
18348 }
18349 let xv = self.view(x, m_tokens * in_f);
18350 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
18351 for t in 0..m_tokens {
18352 let row = xv.slice(t * in_f..(t + 1) * in_f);
18353 let mut xr = self.alloc_uninit::<f32>(in_f)?;
18354 self.copy_view_into(&mut xr, 0, &row, in_f)?;
18355 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
18356 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
18357 }
18358 Ok(y)
18359 }
18360
18361 pub fn linear(
18362 &self,
18363 x: &CudaSlice<f32>,
18364 w: &CudaSlice<f32>,
18365 m_tokens: usize,
18366 in_f: usize,
18367 out_f: usize,
18368 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
18369 self.linear_device(x, w, m_tokens, in_f, out_f)
18370 }
18371
18372 fn linear_device<I>(
18373 &self,
18374 x: &I,
18375 w: &I,
18376 m_tokens: usize,
18377 in_f: usize,
18378 out_f: usize,
18379 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18380 where
18381 I: cudarc::driver::DevicePtr<f32>,
18382 {
18383 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)?;
18385 Ok(c)
18386 }
18387
18388 fn linear_device_into<I, O>(
18389 &self,
18390 x: &I,
18391 w: &I,
18392 c: &mut O,
18393 m_tokens: usize,
18394 in_f: usize,
18395 out_f: usize,
18396 ) -> Result<(), Box<dyn std::error::Error>>
18397 where
18398 I: cudarc::driver::DevicePtr<f32>,
18399 O: cudarc::driver::DevicePtrMut<f32>,
18400 {
18401 use cudarc::cublaslt::{Matmul, MatmulConfig};
18402 let cfg = MatmulConfig {
18403 transa: true,
18404 transb: false,
18405 transc: false,
18406 m: out_f as u64,
18407 n: m_tokens as u64,
18408 k: in_f as u64,
18409 alpha: 1.0,
18410 lda: in_f as i64,
18411 ldb: in_f as i64,
18412 beta: 0.0,
18413 ldc: out_f as i64,
18414 stride_a: None,
18415 stride_b: None,
18416 stride_c: None,
18417 stride_bias: None,
18418 batch_size: None,
18419 };
18420 let blas = self.gpu.blas();
18421 unsafe {
18422 blas.matmul(cfg, w, x, c, None, None)?;
18423 }
18424 Ok(())
18425 }
18426
18427 pub fn sdpa_naive(
18429 &self,
18430 q: &CudaSlice<f32>,
18431 k: &CudaSlice<f32>,
18432 v: &CudaSlice<f32>,
18433 o: &mut CudaSlice<f32>,
18434 head_dim: usize,
18435 n_head: usize,
18436 n_head_kv: usize,
18437 t: usize,
18438 t_kv: usize,
18439 scale: f32,
18440 causal: bool,
18441 ) -> Result<(), Box<dyn std::error::Error>> {
18442 let f = self.func("sdpa_naive_f32");
18443 let cfg = LaunchConfig {
18444 grid_dim: (n_head as u32, t as u32, 1),
18445 block_dim: (128, 1, 1),
18446 shared_mem_bytes: (t_kv * 4) as u32,
18447 };
18448 let (hd, nh, nhkv, ti, tkvi, cz) = (
18449 head_dim as i32,
18450 n_head as i32,
18451 n_head_kv as i32,
18452 t as i32,
18453 t_kv as i32,
18454 causal as i32,
18455 );
18456 let __s_b = self.gpu.stream();
18457 let mut b = __s_b.launch_builder(&f);
18458 b.arg(q)
18459 .arg(k)
18460 .arg(v)
18461 .arg(o)
18462 .arg(&hd)
18463 .arg(&nh)
18464 .arg(&nhkv)
18465 .arg(&ti)
18466 .arg(&tkvi)
18467 .arg(&scale)
18468 .arg(&cz);
18469 unsafe {
18470 b.launch(cfg)?;
18471 }
18472 Ok(())
18473 }
18474
18475 #[allow(clippy::too_many_arguments)]
18480 pub fn sdpa_naive_island(
18481 &self,
18482 q: &CudaSlice<f32>,
18483 k: &CudaSlice<f32>,
18484 v: &CudaSlice<f32>,
18485 o: &mut CudaSlice<f32>,
18486 span_id: &CudaSlice<i32>,
18487 head_dim: usize,
18488 n_head: usize,
18489 n_head_kv: usize,
18490 t: usize,
18491 t_kv: usize,
18492 scale: f32,
18493 window: usize,
18494 ) -> Result<(), Box<dyn std::error::Error>> {
18495 let f = self.func("sdpa_naive_island_f32");
18496 let cfg = LaunchConfig {
18497 grid_dim: (n_head as u32, t as u32, 1),
18498 block_dim: (128, 1, 1),
18499 shared_mem_bytes: (t_kv * 4) as u32,
18500 };
18501 let (hd, nh, nhkv, ti, tkvi, wi) = (
18502 head_dim as i32,
18503 n_head as i32,
18504 n_head_kv as i32,
18505 t as i32,
18506 t_kv as i32,
18507 window as i32,
18508 );
18509 let __s_b = self.gpu.stream();
18510 let mut b = __s_b.launch_builder(&f);
18511 b.arg(q)
18512 .arg(k)
18513 .arg(v)
18514 .arg(o)
18515 .arg(span_id)
18516 .arg(&hd)
18517 .arg(&nh)
18518 .arg(&nhkv)
18519 .arg(&ti)
18520 .arg(&tkvi)
18521 .arg(&scale)
18522 .arg(&wi);
18523 unsafe {
18524 b.launch(cfg)?;
18525 }
18526 Ok(())
18527 }
18528
18529 #[allow(clippy::too_many_arguments)]
18531 pub fn sdpa_naive_w(
18532 &self,
18533 q: &CudaSlice<f32>,
18534 k: &CudaSlice<f32>,
18535 v: &CudaSlice<f32>,
18536 o: &mut CudaSlice<f32>,
18537 head_dim: usize,
18538 n_head: usize,
18539 n_head_kv: usize,
18540 t: usize,
18541 t_kv: usize,
18542 scale: f32,
18543 causal: bool,
18544 window: usize,
18545 ) -> Result<(), Box<dyn std::error::Error>> {
18546 let f = self.func("sdpa_naive_w_f32");
18547 let cfg = LaunchConfig {
18548 grid_dim: (n_head as u32, t as u32, 1),
18549 block_dim: (128, 1, 1),
18550 shared_mem_bytes: (t_kv * 4) as u32,
18551 };
18552 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18553 head_dim as i32,
18554 n_head as i32,
18555 n_head_kv as i32,
18556 t as i32,
18557 t_kv as i32,
18558 causal as i32,
18559 window as i32,
18560 );
18561 let __s_b = self.gpu.stream();
18562 let mut b = __s_b.launch_builder(&f);
18563 b.arg(q)
18564 .arg(k)
18565 .arg(v)
18566 .arg(o)
18567 .arg(&hd)
18568 .arg(&nh)
18569 .arg(&nhkv)
18570 .arg(&ti)
18571 .arg(&tkvi)
18572 .arg(&scale)
18573 .arg(&cz)
18574 .arg(&wi);
18575 unsafe {
18576 b.launch(cfg)?;
18577 }
18578 Ok(())
18579 }
18580
18581 #[allow(clippy::too_many_arguments)]
18591 pub fn sdpa_naive_w_lo(
18592 &self,
18593 q: &CudaSlice<f32>,
18594 k: &CudaSlice<f32>,
18595 v: &CudaSlice<f32>,
18596 o: &mut CudaSlice<f32>,
18597 head_dim: usize,
18598 n_head: usize,
18599 n_head_kv: usize,
18600 t: usize,
18601 t_kv: usize,
18602 scale: f32,
18603 causal: bool,
18604 window: usize,
18605 ) -> Result<(), Box<dyn std::error::Error>> {
18606 let kv_lo = if window > 0 {
18607 (t_kv - t + 1).saturating_sub(window)
18608 } else {
18609 0
18610 };
18611 let smem = (t_kv - kv_lo) * 4;
18612 if smem > 48 * 1024 {
18613 return Err(format!(
18614 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18615 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18616 a window this wide needs the multi-pass long-ctx kernel"
18617 )
18618 .into());
18619 }
18620 let f = self.func("sdpa_naive_w_lo_f32");
18621 let cfg = LaunchConfig {
18622 grid_dim: (n_head as u32, t as u32, 1),
18623 block_dim: (128, 1, 1),
18624 shared_mem_bytes: smem as u32,
18625 };
18626 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18627 head_dim as i32,
18628 n_head as i32,
18629 n_head_kv as i32,
18630 t as i32,
18631 t_kv as i32,
18632 causal as i32,
18633 window as i32,
18634 kv_lo as i32,
18635 );
18636 let __s_b = self.gpu.stream();
18637 let mut b = __s_b.launch_builder(&f);
18638 b.arg(q)
18639 .arg(k)
18640 .arg(v)
18641 .arg(o)
18642 .arg(&hd)
18643 .arg(&nh)
18644 .arg(&nhkv)
18645 .arg(&ti)
18646 .arg(&tkvi)
18647 .arg(&scale)
18648 .arg(&cz)
18649 .arg(&wi)
18650 .arg(&lo);
18651 unsafe {
18652 b.launch(cfg)?;
18653 }
18654 Ok(())
18655 }
18656
18657 pub fn sdpa_naive_view(
18659 &self,
18660 q: &CudaSlice<f32>,
18661 k: &cudarc::driver::CudaView<f32>,
18662 v: &cudarc::driver::CudaView<f32>,
18663 o: &mut CudaSlice<f32>,
18664 head_dim: usize,
18665 n_head: usize,
18666 n_head_kv: usize,
18667 t: usize,
18668 t_kv: usize,
18669 scale: f32,
18670 causal: bool,
18671 ) -> Result<(), Box<dyn std::error::Error>> {
18672 let f = self.func("sdpa_naive_f32");
18673 let cfg = LaunchConfig {
18674 grid_dim: (n_head as u32, t as u32, 1),
18675 block_dim: (128, 1, 1),
18676 shared_mem_bytes: (t_kv * 4) as u32,
18677 };
18678 let (hd, nh, nhkv, ti, tkvi, cz) = (
18679 head_dim as i32,
18680 n_head as i32,
18681 n_head_kv as i32,
18682 t as i32,
18683 t_kv as i32,
18684 causal as i32,
18685 );
18686 let __s_b = self.gpu.stream();
18687 let mut b = __s_b.launch_builder(&f);
18688 b.arg(q)
18689 .arg(k)
18690 .arg(v)
18691 .arg(o)
18692 .arg(&hd)
18693 .arg(&nh)
18694 .arg(&nhkv)
18695 .arg(&ti)
18696 .arg(&tkvi)
18697 .arg(&scale)
18698 .arg(&cz);
18699 unsafe {
18700 b.launch(cfg)?;
18701 }
18702 Ok(())
18703 }
18704
18705 #[allow(clippy::too_many_arguments)]
18713 pub fn fa_dequant_kv_view_f32(
18714 &self,
18715 k: &cudarc::driver::CudaView<u8>,
18716 v: &cudarc::driver::CudaView<u8>,
18717 kf: &mut CudaSlice<f32>,
18718 vf: &mut CudaSlice<f32>,
18719 kv_dim_k: usize,
18720 kv_dim_v: usize,
18721 t_kv: usize,
18722 k_tok_bytes: usize,
18723 v_tok_bytes: usize,
18724 g: bool,
18725 ) -> Result<(), Box<dyn std::error::Error>> {
18726 let f = if g {
18727 self.func_g("fa_dequant_kv_ws_f32")
18728 } else {
18729 self.func("fa_dequant_kv_ws_f32")
18730 };
18731 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18732 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18733 let cfg = LaunchConfig {
18734 grid_dim: (nblk.max(1), 1, 1),
18735 block_dim: (256, 1, 1),
18736 shared_mem_bytes: 0,
18737 };
18738 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18739 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18740 let __s_b = self.gpu.stream();
18741 let mut b = __s_b.launch_builder(&f);
18742 b.arg(k)
18743 .arg(v)
18744 .arg(&mut *kf)
18745 .arg(&mut *vf)
18746 .arg(&kdk)
18747 .arg(&kdv)
18748 .arg(&tkvi)
18749 .arg(&ktb)
18750 .arg(&vtb);
18751 unsafe {
18752 b.launch(cfg)?;
18753 }
18754 Ok(())
18755 }
18756
18757 #[allow(clippy::too_many_arguments)]
18758 pub fn sdpa_naive_quantized_view(
18759 &self,
18760 q: &CudaSlice<f32>,
18761 k: &cudarc::driver::CudaView<u8>,
18762 v: &cudarc::driver::CudaView<u8>,
18763 o: &mut CudaSlice<f32>,
18764 head_dim: usize,
18765 n_head: usize,
18766 n_head_kv: usize,
18767 t: usize,
18768 t_kv: usize,
18769 scale: f32,
18770 causal: bool,
18771 k_tok_bytes: usize,
18772 v_tok_bytes: usize,
18773 ) -> Result<(), Box<dyn std::error::Error>> {
18774 let kv_dim = n_head_kv * head_dim;
18775 let mut kf = self.uninit(t_kv * kv_dim)?;
18776 let mut vf = self.uninit(t_kv * kv_dim)?;
18777 let f = self.func("fa_dequant_kv_ws_f32");
18778 let total = (2 * t_kv * kv_dim) as u64;
18779 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18780 let cfg = LaunchConfig {
18781 grid_dim: (nblk.max(1), 1, 1),
18782 block_dim: (256, 1, 1),
18783 shared_mem_bytes: 0,
18784 };
18785 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18786 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18787 let __s_b = self.gpu.stream();
18788 let mut b = __s_b.launch_builder(&f);
18789 b.arg(k)
18790 .arg(v)
18791 .arg(&mut kf)
18792 .arg(&mut vf)
18793 .arg(&kv_dim_i)
18794 .arg(&kv_dim_i)
18795 .arg(&t_kv_i)
18796 .arg(&k_tok_bytes_i)
18797 .arg(&v_tok_bytes_i);
18798 unsafe { b.launch(cfg)? };
18799 self.sdpa_naive(
18800 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18801 )
18802 }
18803
18804 #[allow(clippy::too_many_arguments)]
18816 pub fn sdpa_naive_w_quantized_view(
18817 &self,
18818 q: &CudaSlice<f32>,
18819 k: &cudarc::driver::CudaView<u8>,
18820 v: &cudarc::driver::CudaView<u8>,
18821 o: &mut CudaSlice<f32>,
18822 head_dim: usize,
18823 n_head: usize,
18824 n_head_kv: usize,
18825 t: usize,
18826 t_kv: usize,
18827 scale: f32,
18828 causal: bool,
18829 window: usize,
18830 k_tok_bytes: usize,
18831 v_tok_bytes: usize,
18832 ) -> Result<(), Box<dyn std::error::Error>> {
18833 let kv_dim = n_head_kv * head_dim;
18834 let mut kf = self.uninit(t_kv * kv_dim)?;
18835 let mut vf = self.uninit(t_kv * kv_dim)?;
18836 let f = self.func("fa_dequant_kv_ws_f32");
18837 let total = (2 * t_kv * kv_dim) as u64;
18838 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18839 let cfg = LaunchConfig {
18840 grid_dim: (nblk.max(1), 1, 1),
18841 block_dim: (256, 1, 1),
18842 shared_mem_bytes: 0,
18843 };
18844 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18845 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18846 let __s_b = self.gpu.stream();
18847 let mut b = __s_b.launch_builder(&f);
18848 b.arg(k)
18849 .arg(v)
18850 .arg(&mut kf)
18851 .arg(&mut vf)
18852 .arg(&kv_dim_i)
18853 .arg(&kv_dim_i)
18854 .arg(&t_kv_i)
18855 .arg(&k_tok_bytes_i)
18856 .arg(&v_tok_bytes_i);
18857 unsafe { b.launch(cfg)? };
18858 self.sdpa_naive_w(
18859 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18860 )
18861 }
18862
18863 pub fn fa_prefill(
18867 &self,
18868 q: &CudaSlice<f32>,
18869 k: &CudaSlice<f32>,
18870 v: &CudaSlice<f32>,
18871 o: &mut CudaSlice<f32>,
18872 head_dim: usize,
18873 n_head: usize,
18874 n_head_kv: usize,
18875 t: usize,
18876 t_kv: usize,
18877 scale: f32,
18878 causal: bool,
18879 ) -> Result<(), Box<dyn std::error::Error>> {
18880 if portable_mma_gated() {
18881 return self.sdpa_naive(
18882 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18883 );
18884 }
18885 let fa3_on = head_dim == 256
18893 && causal
18894 && t == t_kv
18895 && match std::env::var("MEMRA_FA3").as_deref() {
18896 Ok("0") => false,
18897 Ok("1") => {
18901 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
18902 true
18903 }
18904 _ => cfg!(memra_hopper_mma),
18905 };
18906 if fa3_on {
18907 let n = t * n_head * head_dim;
18908 let nkv = t * n_head_kv * head_dim;
18909 let mut q16 = self.alloc_u8_uninit(n * 2)?;
18910 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18911 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18912 self.f32_to_bf16_into(q, &mut q16, n)?;
18913 self.f32_to_bf16_into(k, &mut k16, nkv)?;
18914 self.f32_to_bf16_into(v, &mut v16, nkv)?;
18915 let rc = {
18916 use cudarc::driver::{DevicePtr, DevicePtrMut};
18917 let stream = self.gpu.stream();
18918 let (qp, _g1) = q16.device_ptr(&stream);
18919 let (kp, _g2) = k16.device_ptr(&stream);
18920 let (vp, _g3) = v16.device_ptr(&stream);
18921 let (op, _g4) = o.device_ptr_mut(&stream);
18922 unsafe {
18923 memra_fa3_prefill(
18924 qp as *const core::ffi::c_void,
18925 kp as *const core::ffi::c_void,
18926 vp as *const core::ffi::c_void,
18927 op as *mut f32,
18928 t as i32,
18929 n_head as i32,
18930 n_head_kv as i32,
18931 head_dim as i32,
18932 scale,
18933 stream.cu_stream() as *mut core::ffi::c_void,
18934 )
18935 }
18936 };
18937 if rc != 0 {
18938 return Err(format!("memra_fa3_prefill rc={rc}").into());
18939 }
18940 return Ok(());
18941 }
18942 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18947 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18948 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18949 const BLOCK_Q: usize = 64;
18950 const BKX: usize = 32;
18951 let f = self.func("fa_prefill_bf16_p1");
18952 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18953 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18954 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18955 f.set_attribute(
18956 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18957 shmem as i32,
18958 )?;
18959 let cfg = LaunchConfig {
18960 grid_dim: (
18961 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18962 n_head as u32,
18963 1,
18964 ),
18965 block_dim: (32, 4, 1),
18966 shared_mem_bytes: shmem,
18967 };
18968 let (hd, nh, nhkv, ti, tkvi, cz) = (
18969 head_dim as i32,
18970 n_head as i32,
18971 n_head_kv as i32,
18972 t as i32,
18973 t_kv as i32,
18974 causal as i32,
18975 );
18976 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18977 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18978 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18979 let __s_b = self.gpu.stream();
18980 let mut b = __s_b.launch_builder(&f);
18981 b.arg(&qb)
18982 .arg(&kb)
18983 .arg(&vb)
18984 .arg(o)
18985 .arg(&hd)
18986 .arg(&nh)
18987 .arg(&nhkv)
18988 .arg(&ti)
18989 .arg(&tkvi)
18990 .arg(&scale)
18991 .arg(&cz);
18992 unsafe {
18993 b.launch(cfg)?;
18994 }
18995 return Ok(());
18996 }
18997 const BK: usize = 32;
19003 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
19006 let (block_q, warps, w2_sfx): (usize, u32, &str) =
19007 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
19008 let hd_sfx = fa_hd_suffix(head_dim)?;
19012 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19013 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
19018 let (kb16, vb16) = if bf16kv {
19019 let n = t_kv * n_head_kv * head_dim;
19020 let mut kb = self.alloc_u8_uninit(n * 2)?;
19021 let mut vb = self.alloc_u8_uninit(n * 2)?;
19022 let fcv = self.func("f32_to_bf16_bulk");
19023 let ni = n as i64;
19024 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
19025 let __s_b = self.gpu.stream();
19026 let mut b = __s_b.launch_builder(&fcv);
19027 b.arg(k).arg(&mut kb).arg(&ni);
19028 unsafe {
19029 b.launch(cfgc)?;
19030 }
19031 let __s_b = self.gpu.stream();
19032 let mut b = __s_b.launch_builder(&fcv);
19033 b.arg(v).arg(&mut vb).arg(&ni);
19034 unsafe {
19035 b.launch(cfgc)?;
19036 }
19037 (Some(kb), Some(vb))
19038 } else {
19039 (None, None)
19040 };
19041 let f = self.func(&if bf16kv {
19042 format!("fa_prefill_bf16kv_pp{hd_sfx}")
19043 } else {
19044 format!(
19045 "fa_prefill_f32{}{}{hd_sfx}",
19046 if floor { "" } else { "_pp" },
19047 if floor { "" } else { w2_sfx }
19048 )
19049 });
19050 let kv_stages = if bf16kv { 2 } else { 1 };
19053 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19054 + 4 * (block_q * BK + 2 * block_q)) as u32;
19055 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19056 f.set_attribute(
19057 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19058 shmem as i32,
19059 )?;
19060 let cfg = LaunchConfig {
19061 grid_dim: (
19062 (t as u32 + block_q as u32 - 1) / block_q as u32,
19063 n_head as u32,
19064 1,
19065 ),
19066 block_dim: (32, warps, 1),
19067 shared_mem_bytes: shmem,
19068 };
19069 let (hd, nh, nhkv, ti, tkvi, cz) = (
19070 head_dim as i32,
19071 n_head as i32,
19072 n_head_kv as i32,
19073 t as i32,
19074 t_kv as i32,
19075 causal as i32,
19076 );
19077 let __s_b = self.gpu.stream();
19078 let mut b = __s_b.launch_builder(&f);
19079 b.arg(q);
19080 match (&kb16, &vb16) {
19081 (Some(kb), Some(vb)) => {
19082 b.arg(kb).arg(vb);
19083 }
19084 _ => {
19085 b.arg(k).arg(v);
19086 }
19087 }
19088 b.arg(o)
19089 .arg(&hd)
19090 .arg(&nh)
19091 .arg(&nhkv)
19092 .arg(&ti)
19093 .arg(&tkvi)
19094 .arg(&scale)
19095 .arg(&cz);
19096 unsafe {
19097 b.launch(cfg)?;
19098 }
19099 Ok(())
19100 }
19101
19102 #[allow(clippy::too_many_arguments)]
19106 pub fn fa_prefill_w(
19107 &self,
19108 q: &CudaSlice<f32>,
19109 k: &CudaSlice<f32>,
19110 v: &CudaSlice<f32>,
19111 o: &mut CudaSlice<f32>,
19112 head_dim: usize,
19113 n_head: usize,
19114 n_head_kv: usize,
19115 t: usize,
19116 t_kv: usize,
19117 scale: f32,
19118 causal: bool,
19119 window: usize,
19120 ) -> Result<(), Box<dyn std::error::Error>> {
19121 if portable_mma_gated() {
19124 return self.sdpa_naive_w(
19125 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
19126 );
19127 }
19128 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19132 let faw_f32 =
19133 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
19134 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
19135 self.fa_prefill_w_arm(
19136 q,
19137 k,
19138 v,
19139 o,
19140 head_dim,
19141 n_head,
19142 n_head_kv,
19143 t,
19144 t_kv,
19145 scale,
19146 causal,
19147 window,
19148 floor || faw_f32,
19149 floor,
19150 )
19151 }
19152
19153 #[allow(clippy::too_many_arguments)]
19156 pub fn fa_prefill_w_pre(
19157 &self,
19158 qb: &CudaSlice<u8>,
19159 kb: &CudaSlice<u8>,
19160 vb: &CudaSlice<u8>,
19161 o: &mut CudaSlice<f32>,
19162 head_dim: usize,
19163 n_head: usize,
19164 n_head_kv: usize,
19165 t: usize,
19166 t_kv: usize,
19167 scale: f32,
19168 causal: bool,
19169 window: usize,
19170 v_f16: bool,
19171 ) -> Result<(), Box<dyn std::error::Error>> {
19172 const BLOCK_Q: usize = 64;
19173 const BK: usize = 32;
19174 debug_assert_eq!(head_dim, 256);
19175 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19176 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
19177 if hp {
19178 const BLOCK_QH: usize = 32;
19179 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19182 let vh: &CudaSlice<u8> = if v_f16 {
19183 vb
19184 } else {
19185 let n = t_kv * n_head_kv * head_dim;
19186 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
19187 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
19188 }
19189 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
19190 vguard.as_ref().unwrap()
19191 };
19192 let f = self.func("fa_prefill_w_bf16_p1h2");
19193 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19194 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19195 f.set_attribute(
19196 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19197 shmem as i32,
19198 )?;
19199 let cfg = LaunchConfig {
19200 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19201 block_dim: (32, 4, 1),
19202 shared_mem_bytes: shmem,
19203 };
19204 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19205 head_dim as i32,
19206 n_head as i32,
19207 n_head_kv as i32,
19208 t as i32,
19209 t_kv as i32,
19210 causal as i32,
19211 window as i32,
19212 );
19213 let __s_b = self.gpu.stream();
19214 let mut b = __s_b.launch_builder(&f);
19215 b.arg(qb)
19216 .arg(kb)
19217 .arg(vh)
19218 .arg(o)
19219 .arg(&hd)
19220 .arg(&nh)
19221 .arg(&nhkv)
19222 .arg(&ti)
19223 .arg(&tkvi)
19224 .arg(&scale)
19225 .arg(&cz)
19226 .arg(&wi);
19227 unsafe {
19228 b.launch(cfg)?;
19229 }
19230 return Ok(());
19231 }
19232 let f = self.func("fa_prefill_w_bf16_p1");
19233 let shmem =
19234 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19235 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19236 f.set_attribute(
19237 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19238 shmem as i32,
19239 )?;
19240 let cfg = LaunchConfig {
19241 grid_dim: (
19242 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19243 n_head as u32,
19244 1,
19245 ),
19246 block_dim: (32, 4, 1),
19247 shared_mem_bytes: shmem,
19248 };
19249 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19250 head_dim as i32,
19251 n_head as i32,
19252 n_head_kv as i32,
19253 t as i32,
19254 t_kv as i32,
19255 causal as i32,
19256 window as i32,
19257 );
19258 let __s_b = self.gpu.stream();
19259 let mut b = __s_b.launch_builder(&f);
19260 b.arg(qb)
19261 .arg(kb)
19262 .arg(vb)
19263 .arg(o)
19264 .arg(&hd)
19265 .arg(&nh)
19266 .arg(&nhkv)
19267 .arg(&ti)
19268 .arg(&tkvi)
19269 .arg(&scale)
19270 .arg(&cz)
19271 .arg(&wi);
19272 unsafe {
19273 b.launch(cfg)?;
19274 }
19275 Ok(())
19276 }
19277
19278 #[allow(clippy::too_many_arguments)]
19280 pub fn fa_prefill_w_arm(
19281 &self,
19282 q: &CudaSlice<f32>,
19283 k: &CudaSlice<f32>,
19284 v: &CudaSlice<f32>,
19285 o: &mut CudaSlice<f32>,
19286 head_dim: usize,
19287 n_head: usize,
19288 n_head_kv: usize,
19289 t: usize,
19290 t_kv: usize,
19291 scale: f32,
19292 causal: bool,
19293 window: usize,
19294 f32_stage: bool,
19295 floor: bool,
19296 ) -> Result<(), Box<dyn std::error::Error>> {
19297 const BLOCK_Q: usize = 64;
19298 const BK: usize = 32;
19299 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
19300 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19304 let p1 = !floor
19305 && !f32_stage
19306 && *P1_ON.get_or_init(|| {
19307 std::env::var("MEMRA_FAW_P1")
19308 .map(|v| v != "0")
19309 .unwrap_or(true)
19310 });
19311 let hp =
19312 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19313 if hp {
19314 const BLOCK_QH: usize = 32;
19315 let f = self.func("fa_prefill_w_bf16_p1h2");
19316 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
19317 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19318 f.set_attribute(
19319 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19320 shmem as i32,
19321 )?;
19322 let cfg = LaunchConfig {
19323 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
19324 block_dim: (32, 4, 1),
19325 shared_mem_bytes: shmem,
19326 };
19327 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19328 head_dim as i32,
19329 n_head as i32,
19330 n_head_kv as i32,
19331 t as i32,
19332 t_kv as i32,
19333 causal as i32,
19334 window as i32,
19335 );
19336 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19337 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19338 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
19339 let __s_b = self.gpu.stream();
19340 let mut b = __s_b.launch_builder(&f);
19341 b.arg(&qb)
19342 .arg(&kb)
19343 .arg(&vh)
19344 .arg(o)
19345 .arg(&hd)
19346 .arg(&nh)
19347 .arg(&nhkv)
19348 .arg(&ti)
19349 .arg(&tkvi)
19350 .arg(&scale)
19351 .arg(&cz)
19352 .arg(&wi);
19353 unsafe {
19354 b.launch(cfg)?;
19355 }
19356 return Ok(());
19357 }
19358 if p1 {
19359 let f = self.func("fa_prefill_w_bf16_p1");
19360 let shmem =
19361 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19362 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19363 f.set_attribute(
19364 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19365 shmem as i32,
19366 )?;
19367 let cfg = LaunchConfig {
19368 grid_dim: (
19369 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19370 n_head as u32,
19371 1,
19372 ),
19373 block_dim: (32, 4, 1),
19374 shared_mem_bytes: shmem,
19375 };
19376 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19377 head_dim as i32,
19378 n_head as i32,
19379 n_head_kv as i32,
19380 t as i32,
19381 t_kv as i32,
19382 causal as i32,
19383 window as i32,
19384 );
19385 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19386 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19387 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19388 let __s_b = self.gpu.stream();
19389 let mut b = __s_b.launch_builder(&f);
19390 b.arg(&qb)
19391 .arg(&kb)
19392 .arg(&vb)
19393 .arg(o)
19394 .arg(&hd)
19395 .arg(&nh)
19396 .arg(&nhkv)
19397 .arg(&ti)
19398 .arg(&tkvi)
19399 .arg(&scale)
19400 .arg(&cz)
19401 .arg(&wi);
19402 unsafe {
19403 b.launch(cfg)?;
19404 }
19405 return Ok(());
19406 }
19407 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19410 let g4 = !floor
19411 && !f32_stage
19412 && n_head_kv == 1
19413 && n_head % 4 == 0
19414 && *G4_ON.get_or_init(|| {
19415 std::env::var("MEMRA_FAW_G4")
19416 .map(|v| v != "0")
19417 .unwrap_or(true)
19418 });
19419 if g4 {
19420 const SP_M: usize = 16;
19421 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19424 let o2 = *O2_ON.get_or_init(|| {
19425 std::env::var("MEMRA_FAW_O2")
19426 .map(|v| v != "0")
19427 .unwrap_or(true)
19428 });
19429 let f = self.func(if o2 {
19430 "fa_prefill_w_bf16_g4o2"
19431 } else {
19432 "fa_prefill_w_bf16_g4"
19433 });
19434 let shmem = if o2 {
19435 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19436 } else {
19437 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19438 as u32
19439 };
19440 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19441 f.set_attribute(
19442 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19443 shmem as i32,
19444 )?;
19445 let cfg = LaunchConfig {
19446 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19447 block_dim: (32, 4, 1),
19448 shared_mem_bytes: shmem,
19449 };
19450 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19451 head_dim as i32,
19452 n_head as i32,
19453 n_head_kv as i32,
19454 t as i32,
19455 t_kv as i32,
19456 causal as i32,
19457 window as i32,
19458 );
19459 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19460 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19461 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19462 let __s_b = self.gpu.stream();
19463 let mut b = __s_b.launch_builder(&f);
19464 b.arg(&qb)
19465 .arg(&kb)
19466 .arg(&vb)
19467 .arg(o)
19468 .arg(&hd)
19469 .arg(&nh)
19470 .arg(&nhkv)
19471 .arg(&ti)
19472 .arg(&tkvi)
19473 .arg(&scale)
19474 .arg(&cz)
19475 .arg(&wi);
19476 unsafe {
19477 b.launch(cfg)?;
19478 }
19479 return Ok(());
19480 }
19481 let f = self.func(if floor {
19482 "fa_prefill_w_f32"
19483 } else if f32_stage {
19484 "fa_prefill_w_f32_pp"
19485 } else {
19486 "fa_prefill_w_bf16_pp"
19487 });
19488 let shmem =
19489 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19490 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19491 f.set_attribute(
19492 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19493 shmem as i32,
19494 )?;
19495 let cfg = LaunchConfig {
19496 grid_dim: (
19497 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19498 n_head as u32,
19499 1,
19500 ),
19501 block_dim: (32, 4, 1),
19502 shared_mem_bytes: shmem,
19503 };
19504 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19505 head_dim as i32,
19506 n_head as i32,
19507 n_head_kv as i32,
19508 t as i32,
19509 t_kv as i32,
19510 causal as i32,
19511 window as i32,
19512 );
19513 if f32_stage {
19514 let __s_b = self.gpu.stream();
19515 let mut b = __s_b.launch_builder(&f);
19516 b.arg(q)
19517 .arg(k)
19518 .arg(v)
19519 .arg(o)
19520 .arg(&hd)
19521 .arg(&nh)
19522 .arg(&nhkv)
19523 .arg(&ti)
19524 .arg(&tkvi)
19525 .arg(&scale)
19526 .arg(&cz)
19527 .arg(&wi);
19528 unsafe {
19529 b.launch(cfg)?;
19530 }
19531 } else {
19532 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19533 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19534 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19535 let __s_b = self.gpu.stream();
19536 let mut b = __s_b.launch_builder(&f);
19537 b.arg(&qb)
19538 .arg(&kb)
19539 .arg(&vb)
19540 .arg(o)
19541 .arg(&hd)
19542 .arg(&nh)
19543 .arg(&nhkv)
19544 .arg(&ti)
19545 .arg(&tkvi)
19546 .arg(&scale)
19547 .arg(&cz)
19548 .arg(&wi);
19549 unsafe {
19550 b.launch(cfg)?;
19551 }
19552 }
19553 Ok(())
19554 }
19555
19556 #[allow(clippy::too_many_arguments)]
19560 pub fn fa_prefill_hd512(
19561 &self,
19562 q: &CudaSlice<f32>,
19563 k: &CudaSlice<f32>,
19564 v: &CudaSlice<f32>,
19565 o: &mut CudaSlice<f32>,
19566 head_dim: usize,
19567 n_head: usize,
19568 n_head_kv: usize,
19569 t: usize,
19570 t_kv: usize,
19571 scale: f32,
19572 causal: bool,
19573 ) -> Result<(), Box<dyn std::error::Error>> {
19574 if portable_mma_gated() {
19576 return self.sdpa_naive(
19577 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19578 );
19579 }
19580 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19586 let f32_stage =
19587 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19588 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19592 let sp = !f32_stage
19593 && *SP_ON.get_or_init(|| {
19594 std::env::var("MEMRA_FA512_SP")
19595 .map(|v| v != "0")
19596 .unwrap_or(true)
19597 });
19598 self.fa_prefill_hd512_arm(
19599 q,
19600 k,
19601 v,
19602 o,
19603 head_dim,
19604 n_head,
19605 n_head_kv,
19606 t,
19607 t_kv,
19608 scale,
19609 causal,
19610 f32_stage,
19611 sp,
19612 sp && fa_f16pv_on(),
19613 )
19614 }
19615
19616 #[allow(clippy::too_many_arguments)]
19618 pub fn fa_prefill_hd512_pre(
19619 &self,
19620 qb: &CudaSlice<u8>,
19621 kb: &CudaSlice<u8>,
19622 vb: &CudaSlice<u8>,
19623 o: &mut CudaSlice<f32>,
19624 head_dim: usize,
19625 n_head: usize,
19626 n_head_kv: usize,
19627 t: usize,
19628 t_kv: usize,
19629 scale: f32,
19630 causal: bool,
19631 v_f16: bool,
19632 ) -> Result<(), Box<dyn std::error::Error>> {
19633 debug_assert_eq!(head_dim, 512);
19634 const SP_M: usize = 16;
19635 const BKS: usize = 32;
19636 let f16pv = fa_f16pv_on();
19640 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19641 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19642 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19643 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19644 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19645 let n = t_kv * n_head_kv * head_dim;
19647 let need = n * 2;
19648 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19649 *vguard = Some(self.alloc_uninit::<u8>(need)?);
19650 }
19651 let dst = vguard.as_mut().unwrap();
19652 self.bf16_to_f16_into(vb, n, dst)?;
19653 vguard.as_ref().unwrap()
19654 } else {
19655 vb
19656 };
19657 let f = self.func(if hp {
19658 "fa_prefill_bf16_hd512_sp16h2"
19659 } else {
19660 match (f16pv, nw) {
19661 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19662 (true, _) => "fa_prefill_bf16_hd512_sp16",
19663 _ => "fa_prefill_bf16_hd512_sp",
19664 }
19665 });
19666 let (nwarp, npart) = if hp {
19667 (4usize, 4usize)
19668 } else if nw > 2 {
19669 (nw, nw)
19670 } else {
19671 (2, 1)
19672 };
19673 let shmem = if hp {
19675 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19676 as u32
19677 } else {
19678 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19679 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19680 };
19681 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19682 f.set_attribute(
19683 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19684 shmem as i32,
19685 )?;
19686 let grid_y = if hp {
19687 (n_head / 2) as u32
19688 } else {
19689 n_head as u32
19690 };
19691 let cfg = LaunchConfig {
19692 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19693 block_dim: (32, nwarp as u32, 1),
19694 shared_mem_bytes: shmem,
19695 };
19696 let (hd, nh, nhkv, ti, tkvi, cz) = (
19697 head_dim as i32,
19698 n_head as i32,
19699 n_head_kv as i32,
19700 t as i32,
19701 t_kv as i32,
19702 causal as i32,
19703 );
19704 let __s_b = self.gpu.stream();
19705 let mut b = __s_b.launch_builder(&f);
19706 b.arg(qb)
19707 .arg(kb)
19708 .arg(vref)
19709 .arg(o)
19710 .arg(&hd)
19711 .arg(&nh)
19712 .arg(&nhkv)
19713 .arg(&ti)
19714 .arg(&tkvi)
19715 .arg(&scale)
19716 .arg(&cz);
19717 unsafe {
19718 b.launch(cfg)?;
19719 }
19720 Ok(())
19721 }
19722
19723 #[allow(clippy::too_many_arguments)]
19726 pub fn fa_prefill_hd512_arm(
19727 &self,
19728 q: &CudaSlice<f32>,
19729 k: &CudaSlice<f32>,
19730 v: &CudaSlice<f32>,
19731 o: &mut CudaSlice<f32>,
19732 head_dim: usize,
19733 n_head: usize,
19734 n_head_kv: usize,
19735 t: usize,
19736 t_kv: usize,
19737 scale: f32,
19738 causal: bool,
19739 f32_stage: bool,
19740 sp: bool,
19741 f16pv: bool,
19742 ) -> Result<(), Box<dyn std::error::Error>> {
19743 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19744 if sp && !f32_stage {
19745 const SP_M: usize = 16;
19749 const BKS: usize = 32;
19750 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19751 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19752 let f = self.func(if hp {
19753 "fa_prefill_bf16_hd512_sp16h2"
19754 } else {
19755 match (f16pv, nw) {
19756 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19757 (true, _) => "fa_prefill_bf16_hd512_sp16",
19758 _ => "fa_prefill_bf16_hd512_sp",
19759 }
19760 });
19761 let (nwarp, npart) = if hp {
19762 (4usize, 4usize)
19763 } else if nw > 2 {
19764 (nw, nw)
19765 } else {
19766 (2, 1)
19767 };
19768 let shmem = if hp {
19769 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19770 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19771 } else {
19772 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19773 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19774 };
19775 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19776 f.set_attribute(
19777 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19778 shmem as i32,
19779 )?;
19780 let grid_y = if hp {
19781 (n_head / 2) as u32
19782 } else {
19783 n_head as u32
19784 };
19785 let cfg = LaunchConfig {
19786 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19787 block_dim: (32, nwarp as u32, 1),
19788 shared_mem_bytes: shmem,
19789 };
19790 let (hd, nh, nhkv, ti, tkvi, cz) = (
19791 head_dim as i32,
19792 n_head as i32,
19793 n_head_kv as i32,
19794 t as i32,
19795 t_kv as i32,
19796 causal as i32,
19797 );
19798 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19799 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19800 let vb = if f16pv {
19801 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19802 } else {
19803 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19804 };
19805 let __s_b = self.gpu.stream();
19806 let mut b = __s_b.launch_builder(&f);
19807 b.arg(&qb)
19808 .arg(&kb)
19809 .arg(&vb)
19810 .arg(o)
19811 .arg(&hd)
19812 .arg(&nh)
19813 .arg(&nhkv)
19814 .arg(&ti)
19815 .arg(&tkvi)
19816 .arg(&scale)
19817 .arg(&cz);
19818 unsafe {
19819 b.launch(cfg)?;
19820 }
19821 return Ok(());
19822 }
19823 const BLOCK_Q: usize = 32;
19824 const BK: usize = 32;
19825 const HALF: usize = 256;
19826 let f = self.func(if f32_stage {
19827 "fa_prefill_f32_hd512"
19828 } else {
19829 "fa_prefill_bf16_hd512"
19830 });
19831 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19833 + 4 * BLOCK_Q) as u32;
19834 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19835 f.set_attribute(
19836 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19837 shmem as i32,
19838 )?;
19839 let cfg = LaunchConfig {
19840 grid_dim: (
19841 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19842 n_head as u32,
19843 2,
19844 ),
19845 block_dim: (32, 2, 1),
19846 shared_mem_bytes: shmem,
19847 };
19848 let (hd, nh, nhkv, ti, tkvi, cz) = (
19849 head_dim as i32,
19850 n_head as i32,
19851 n_head_kv as i32,
19852 t as i32,
19853 t_kv as i32,
19854 causal as i32,
19855 );
19856 if f32_stage {
19857 let __s_b = self.gpu.stream();
19858 let mut b = __s_b.launch_builder(&f);
19859 b.arg(q)
19860 .arg(k)
19861 .arg(v)
19862 .arg(o)
19863 .arg(&hd)
19864 .arg(&nh)
19865 .arg(&nhkv)
19866 .arg(&ti)
19867 .arg(&tkvi)
19868 .arg(&scale)
19869 .arg(&cz);
19870 unsafe {
19871 b.launch(cfg)?;
19872 }
19873 } else {
19874 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19875 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19876 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19877 let __s_b = self.gpu.stream();
19878 let mut b = __s_b.launch_builder(&f);
19879 b.arg(&qb)
19880 .arg(&kb)
19881 .arg(&vb)
19882 .arg(o)
19883 .arg(&hd)
19884 .arg(&nh)
19885 .arg(&nhkv)
19886 .arg(&ti)
19887 .arg(&tkvi)
19888 .arg(&scale)
19889 .arg(&cz);
19890 unsafe {
19891 b.launch(cfg)?;
19892 }
19893 }
19894 Ok(())
19895 }
19896
19897 #[allow(clippy::too_many_arguments)]
19901 pub fn rope_neox2_bf16e(
19902 &self,
19903 q: &mut CudaSlice<f32>,
19904 k: &mut CudaSlice<f32>,
19905 qb: &mut CudaSlice<u8>,
19906 kb: &mut CudaSlice<u8>,
19907 pos: &CudaSlice<i32>,
19908 head_dim: usize,
19909 n_dims: usize,
19910 nh_q: usize,
19911 nh_k: usize,
19912 n_tokens: usize,
19913 base: f32,
19914 freq_scale: f32,
19915 ff: Option<&CudaSlice<f32>>,
19916 ) -> Result<(), Box<dyn std::error::Error>> {
19917 let f = self.func("rope_neox2_bf16e_f32");
19918 let rows = ((nh_q + nh_k) * n_tokens) as u32;
19919 let cfg = LaunchConfig {
19920 grid_dim: (rows, 1, 1),
19921 block_dim: ((head_dim / 2) as u32, 1, 1),
19922 shared_mem_bytes: 0,
19923 };
19924 let theta_scale = base.powf(-2.0 / n_dims as f32);
19925 let (hd, nd, nhq, nhk, nt) = (
19926 head_dim as i32,
19927 n_dims as i32,
19928 nh_q as i32,
19929 nh_k as i32,
19930 n_tokens as i32,
19931 );
19932 let __s_b = self.gpu.stream();
19933 let mut b = __s_b.launch_builder(&f);
19934 match ff {
19935 Some(t) => {
19936 b.arg(&mut *q)
19937 .arg(&mut *k)
19938 .arg(&mut *qb)
19939 .arg(&mut *kb)
19940 .arg(pos)
19941 .arg(&hd)
19942 .arg(&nd)
19943 .arg(&nhq)
19944 .arg(&nhk)
19945 .arg(&nt)
19946 .arg(&theta_scale)
19947 .arg(&freq_scale)
19948 .arg(t);
19949 unsafe {
19950 b.launch(cfg)?;
19951 }
19952 }
19953 None => {
19954 let null: u64 = 0;
19955 b.arg(&mut *q)
19956 .arg(&mut *k)
19957 .arg(&mut *qb)
19958 .arg(&mut *kb)
19959 .arg(pos)
19960 .arg(&hd)
19961 .arg(&nd)
19962 .arg(&nhq)
19963 .arg(&nhk)
19964 .arg(&nt)
19965 .arg(&theta_scale)
19966 .arg(&freq_scale)
19967 .arg(&null);
19968 unsafe {
19969 b.launch(cfg)?;
19970 }
19971 }
19972 }
19973 Ok(())
19974 }
19975
19976 pub fn f32_to_bf16(
19979 &self,
19980 x: &CudaSlice<f32>,
19981 n: usize,
19982 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19983 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19984 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19985 let f = self.func("f32_to_bf16_flat");
19986 let n_i = n as i64;
19987 let cfg = LaunchConfig {
19988 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19989 block_dim: (256, 1, 1),
19990 shared_mem_bytes: 0,
19991 };
19992 let __s_b = self.gpu.stream();
19993 let mut b = __s_b.launch_builder(&f);
19994 b.arg(x).arg(&mut y).arg(&n_i);
19995 unsafe {
19996 b.launch(cfg)?;
19997 }
19998 Ok(y)
19999 }
20000
20001 pub fn f32_to_f16(
20002 &self,
20003 x: &CudaSlice<f32>,
20004 n: usize,
20005 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20006 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
20007 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20008 let f = self.func("f32_to_f16_flat");
20009 let n_i = n as i64;
20010 let cfg = LaunchConfig {
20011 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
20012 block_dim: (256, 1, 1),
20013 shared_mem_bytes: 0,
20014 };
20015 let __s_b = self.gpu.stream();
20016 let mut b = __s_b.launch_builder(&f);
20017 b.arg(x).arg(&mut y).arg(&n_i);
20018 unsafe {
20019 b.launch(cfg)?;
20020 }
20021 Ok(y)
20022 }
20023
20024 pub fn bf16_to_f16(
20026 &self,
20027 xb: &CudaSlice<u8>,
20028 n: usize,
20029 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
20030 let mut y = self.alloc_uninit::<u8>(n * 2)?;
20031 self.bf16_to_f16_into(xb, n, &mut y)?;
20032 Ok(y)
20033 }
20034
20035 pub fn bf16_to_f16_into(
20037 &self,
20038 xb: &CudaSlice<u8>,
20039 n: usize,
20040 y: &mut CudaSlice<u8>,
20041 ) -> Result<(), Box<dyn std::error::Error>> {
20042 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
20043 assert!(y.len() >= n * 2);
20044 let f = self.func("bf16_to_f16_flat");
20045 let n2 = (n / 2) as i64;
20046 let cfg = LaunchConfig {
20047 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
20048 block_dim: (256, 1, 1),
20049 shared_mem_bytes: 0,
20050 };
20051 let __s_b = self.gpu.stream();
20052 let mut b = __s_b.launch_builder(&f);
20053 b.arg(xb).arg(y).arg(&n2);
20054 unsafe {
20055 b.launch(cfg)?;
20056 }
20057 Ok(())
20058 }
20059
20060 #[allow(clippy::too_many_arguments)]
20065 pub fn fa_prefill_vl8(
20066 &self,
20067 seqs: &[FaSeqVl],
20068 head_dim: usize,
20069 n_head: usize,
20070 n_head_kv: usize,
20071 scale: f32,
20072 ) -> Result<(), Box<dyn std::error::Error>> {
20073 const BK: usize = 32;
20074 let b = seqs.len();
20075 assert!(b >= 1 && b <= 8);
20076 let mut packed = [FaSeqVl::default(); 8];
20077 packed[..b].copy_from_slice(seqs);
20078 let v = FaVl8(packed);
20079 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20080 let ept = (n_head_kv * head_dim) as i32;
20081 {
20082 let f = self.func("fa_mirror_vl");
20083 let max_n = (max_t as i64) * ept as i64;
20084 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
20085 for which in 0..2i32 {
20086 let cfg = LaunchConfig {
20087 grid_dim: (blocks, 1, b as u32),
20088 block_dim: (256, 1, 1),
20089 shared_mem_bytes: 0,
20090 };
20091 let __s_lb = self.gpu.stream();
20092 let mut lb = __s_lb.launch_builder(&f);
20093 lb.arg(&v).arg(&ept).arg(&which);
20094 unsafe {
20095 lb.launch(cfg)?;
20096 }
20097 }
20098 }
20099 let hd_sfx = fa_hd_suffix(head_dim)?;
20100 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
20101 let block_q = 64usize;
20102 let kv_stages = 2usize;
20103 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20104 + 4 * (block_q * BK + 2 * block_q)) as u32;
20105 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20106 f.set_attribute(
20107 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20108 shmem as i32,
20109 )?;
20110 let cfg = LaunchConfig {
20111 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
20112 block_dim: (32, 4, 1),
20113 shared_mem_bytes: shmem,
20114 };
20115 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20116 let __s_lb = self.gpu.stream();
20117 let mut lb = __s_lb.launch_builder(&f);
20118 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
20119 unsafe {
20120 lb.launch(cfg)?;
20121 }
20122 Ok(())
20123 }
20124
20125 #[allow(clippy::too_many_arguments)]
20129 pub fn attn_pre_vl8(
20130 &self,
20131 seqs: &[AttnPreVl],
20132 wq: &CudaSlice<f32>,
20133 wk: &CudaSlice<f32>,
20134 head_dim: usize,
20135 rope_dims: usize,
20136 n_head: usize,
20137 n_head_kv: usize,
20138 eps: f32,
20139 freq_base: f32,
20140 freq_scale: f32,
20141 kv_dim_k: usize,
20142 kv_dim_v: usize,
20143 k_tok_bytes: usize,
20144 v_tok_bytes: usize,
20145 ) -> Result<(), Box<dyn std::error::Error>> {
20146 let b = seqs.len();
20147 assert!(b >= 1 && b <= 8);
20148 let mut packed = [AttnPreVl::default(); 8];
20149 packed[..b].copy_from_slice(seqs);
20150 let v = AttnPreVl8(packed);
20151 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
20152 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20153 {
20154 let f = self.func("q_gate_split_vl");
20155 let n = max_t * (n_head * head_dim) as u32;
20156 let cfg = LaunchConfig {
20157 grid_dim: (n.div_ceil(256), 1, b as u32),
20158 block_dim: (256, 1, 1),
20159 shared_mem_bytes: 0,
20160 };
20161 let __s_lb = self.gpu.stream();
20162 let mut lb = __s_lb.launch_builder(&f);
20163 lb.arg(&v).arg(&hd).arg(&nh);
20164 unsafe {
20165 lb.launch(cfg)?;
20166 }
20167 }
20168 {
20169 let f = self.func("attn_rms_vl");
20170 let cfg = LaunchConfig {
20171 grid_dim: (max_t * n_head as u32, 2, b as u32),
20172 block_dim: (rms_block(), 1, 1),
20173 shared_mem_bytes: 0,
20174 };
20175 let __s_lb = self.gpu.stream();
20176 let mut lb = __s_lb.launch_builder(&f);
20177 lb.arg(&v)
20178 .arg(wq)
20179 .arg(wk)
20180 .arg(&hd)
20181 .arg(&nh)
20182 .arg(&nhkv)
20183 .arg(&eps);
20184 unsafe {
20185 lb.launch(cfg)?;
20186 }
20187 }
20188 {
20189 let f = self.func("attn_rope_vl");
20190 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
20191 let nd = rope_dims as i32;
20192 let cfg = LaunchConfig {
20193 grid_dim: (max_t * n_head as u32, 2, b as u32),
20194 block_dim: ((head_dim / 2) as u32, 1, 1),
20195 shared_mem_bytes: 0,
20196 };
20197 let __s_lb = self.gpu.stream();
20198 let mut lb = __s_lb.launch_builder(&f);
20199 lb.arg(&v)
20200 .arg(&hd)
20201 .arg(&nd)
20202 .arg(&nh)
20203 .arg(&nhkv)
20204 .arg(&theta_scale)
20205 .arg(&freq_scale);
20206 unsafe {
20207 lb.launch(cfg)?;
20208 }
20209 }
20210 {
20211 let f = self.func("append_kv_vl");
20212 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20213 let cfg = LaunchConfig {
20214 grid_dim: (nblk, max_t, b as u32),
20215 block_dim: (32, 1, 1),
20216 shared_mem_bytes: 0,
20217 };
20218 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20219 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20220 let __s_lb = self.gpu.stream();
20221 let mut lb = __s_lb.launch_builder(&f);
20222 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
20223 unsafe {
20224 lb.launch(cfg)?;
20225 }
20226 }
20227 Ok(())
20228 }
20229
20230 pub fn fa_prefill_view(
20235 &self,
20236 q: &CudaSlice<f32>,
20237 k: &cudarc::driver::CudaView<u8>,
20238 v: &cudarc::driver::CudaView<u8>,
20239 o: &mut CudaSlice<f32>,
20240 head_dim: usize,
20241 n_head: usize,
20242 n_head_kv: usize,
20243 t: usize,
20244 t_kv: usize,
20245 scale: f32,
20246 causal: bool,
20247 k_tok_bytes: usize,
20248 v_tok_bytes: usize,
20249 g: bool,
20250 ) -> Result<(), Box<dyn std::error::Error>> {
20251 if portable_mma_gated() {
20252 return self.sdpa_naive_quantized_view(
20253 q,
20254 k,
20255 v,
20256 o,
20257 head_dim,
20258 n_head,
20259 n_head_kv,
20260 t,
20261 t_kv,
20262 scale,
20263 causal,
20264 k_tok_bytes,
20265 v_tok_bytes,
20266 );
20267 }
20268 const BLOCK_Q: usize = 64;
20269 const BK: usize = 32;
20270 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
20273 let f = if g {
20274 self.func_g(&name)
20275 } else {
20276 self.func(&name)
20277 };
20278 let shmem =
20279 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20280 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20281 f.set_attribute(
20282 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20283 shmem as i32,
20284 )?;
20285 let cfg = LaunchConfig {
20286 grid_dim: (
20287 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20288 n_head as u32,
20289 1,
20290 ),
20291 block_dim: (32, 4, 1),
20292 shared_mem_bytes: shmem,
20293 };
20294 let (hd, nh, nhkv, ti, tkvi, cz) = (
20295 head_dim as i32,
20296 n_head as i32,
20297 n_head_kv as i32,
20298 t as i32,
20299 t_kv as i32,
20300 causal as i32,
20301 );
20302 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20303 let __s_b = self.gpu.stream();
20304 let mut b = __s_b.launch_builder(&f);
20305 b.arg(q)
20306 .arg(k)
20307 .arg(v)
20308 .arg(o)
20309 .arg(&hd)
20310 .arg(&nh)
20311 .arg(&nhkv)
20312 .arg(&ti)
20313 .arg(&tkvi)
20314 .arg(&scale)
20315 .arg(&cz)
20316 .arg(&ktb)
20317 .arg(&vtb);
20318 unsafe {
20319 b.launch(cfg)?;
20320 }
20321 Ok(())
20322 }
20323
20324 #[allow(clippy::too_many_arguments)]
20334 pub fn fa_prefill_view_ws(
20335 &self,
20336 q: &CudaSlice<f32>,
20337 k: &cudarc::driver::CudaView<u8>,
20338 v: &cudarc::driver::CudaView<u8>,
20339 o: &mut CudaSlice<f32>,
20340 head_dim: usize,
20341 n_head: usize,
20342 n_head_kv: usize,
20343 t: usize,
20344 t_kv: usize,
20345 scale: f32,
20346 causal: bool,
20347 k_tok_bytes: usize,
20348 v_tok_bytes: usize,
20349 g: bool,
20350 ) -> Result<(), Box<dyn std::error::Error>> {
20351 if portable_mma_gated() {
20352 return self.sdpa_naive_quantized_view(
20353 q,
20354 k,
20355 v,
20356 o,
20357 head_dim,
20358 n_head,
20359 n_head_kv,
20360 t,
20361 t_kv,
20362 scale,
20363 causal,
20364 k_tok_bytes,
20365 v_tok_bytes,
20366 );
20367 }
20368 const BLOCK_Q: usize = 64;
20369 const BK: usize = 32;
20370 let kv_dim_k = n_head_kv * head_dim;
20371 let kv_dim_v = n_head_kv * head_dim;
20372 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20374 let mut guard = self.prime_deqw_ws.lock().unwrap();
20376 let need_grow = match guard.as_ref() {
20377 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20378 None => true,
20379 };
20380 if need_grow {
20381 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20382 let (ck, cv) = guard
20383 .as_ref()
20384 .map(|(a, b)| (a.len(), b.len()))
20385 .unwrap_or((0, 0));
20386 *guard = Some((
20387 self.alloc_u8(grow(ck, k_ws_bytes))?,
20388 self.alloc_u8(grow(cv, v_ws_bytes))?,
20389 ));
20390 }
20391 let (kw, vw) = guard.as_mut().unwrap();
20392 {
20394 let f = if g {
20396 self.func_g("fa_dequant_kv_ws_bf16")
20397 } else {
20398 self.func("fa_dequant_kv_ws_bf16")
20399 };
20400 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20401 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20402 let cfg = LaunchConfig {
20403 grid_dim: (nblk.max(1), 1, 1),
20404 block_dim: (256, 1, 1),
20405 shared_mem_bytes: 0,
20406 };
20407 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20408 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20409 let __s_b = self.gpu.stream();
20410 let mut b = __s_b.launch_builder(&f);
20411 b.arg(k)
20412 .arg(v)
20413 .arg(&mut *kw)
20414 .arg(&mut *vw)
20415 .arg(&kdk)
20416 .arg(&kdv)
20417 .arg(&tkvi)
20418 .arg(&ktb)
20419 .arg(&vtb);
20420 unsafe {
20421 b.launch(cfg)?;
20422 }
20423 }
20424 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20432 .map(|v| v != "0")
20433 .unwrap_or(true);
20434 {
20435 let hd_sfx = fa_hd_suffix(head_dim)?;
20436 let f = self.func(&format!(
20437 "fa_prefill_qw{}{hd_sfx}",
20438 if db { "_db" } else { "" }
20439 ));
20440 let shmem = if db {
20441 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20443 } else {
20444 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20445 };
20446 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20447 f.set_attribute(
20448 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20449 shmem as i32,
20450 )?;
20451 let cfg = LaunchConfig {
20452 grid_dim: (
20453 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20454 n_head as u32,
20455 1,
20456 ),
20457 block_dim: (32, 4, 1),
20458 shared_mem_bytes: shmem,
20459 };
20460 let (hd, nh, nhkv, ti, tkvi, cz) = (
20461 head_dim as i32,
20462 n_head as i32,
20463 n_head_kv as i32,
20464 t as i32,
20465 t_kv as i32,
20466 causal as i32,
20467 );
20468 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20469 let __s_b = self.gpu.stream();
20470 let mut b = __s_b.launch_builder(&f);
20471 b.arg(q)
20472 .arg(&*kw)
20473 .arg(&*vw)
20474 .arg(o)
20475 .arg(&hd)
20476 .arg(&nh)
20477 .arg(&nhkv)
20478 .arg(&ti)
20479 .arg(&tkvi)
20480 .arg(&scale)
20481 .arg(&cz)
20482 .arg(&kdk)
20483 .arg(&kdv);
20484 unsafe {
20485 b.launch(cfg)?;
20486 }
20487 }
20488 Ok(())
20489 }
20490
20491 #[allow(clippy::too_many_arguments)]
20507 pub fn fa_prefill_view_ws_w_hd128(
20508 &self,
20509 q: &CudaSlice<f32>,
20510 k: &cudarc::driver::CudaView<u8>,
20511 v: &cudarc::driver::CudaView<u8>,
20512 o: &mut CudaSlice<f32>,
20513 head_dim: usize,
20514 n_head: usize,
20515 n_head_kv: usize,
20516 t: usize,
20517 t_kv: usize,
20518 scale: f32,
20519 causal: bool,
20520 window: usize,
20521 k_tok_bytes: usize,
20522 v_tok_bytes: usize,
20523 ) -> Result<(), Box<dyn std::error::Error>> {
20524 assert_eq!(
20525 head_dim, 128,
20526 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20527 );
20528 if portable_mma_gated() {
20529 return self.sdpa_naive_w_quantized_view(
20530 q,
20531 k,
20532 v,
20533 o,
20534 head_dim,
20535 n_head,
20536 n_head_kv,
20537 t,
20538 t_kv,
20539 scale,
20540 causal,
20541 window,
20542 k_tok_bytes,
20543 v_tok_bytes,
20544 );
20545 }
20546 const BLOCK_Q: usize = 64;
20547 const BK: usize = 32;
20548 let kv_dim_k = n_head_kv * head_dim;
20549 let kv_dim_v = n_head_kv * head_dim;
20550 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20552 let mut guard = self.prime_deqw_ws.lock().unwrap();
20553 let need_grow = match guard.as_ref() {
20554 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20555 None => true,
20556 };
20557 if need_grow {
20558 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20559 let (ck, cv) = guard
20560 .as_ref()
20561 .map(|(a, b)| (a.len(), b.len()))
20562 .unwrap_or((0, 0));
20563 *guard = Some((
20564 self.alloc_u8(grow(ck, k_ws_bytes))?,
20565 self.alloc_u8(grow(cv, v_ws_bytes))?,
20566 ));
20567 }
20568 let (kw, vw) = guard.as_mut().unwrap();
20569 {
20572 let f = self.func("fa_dequant_kv_ws_bf16");
20573 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20574 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20575 let cfg = LaunchConfig {
20576 grid_dim: (nblk.max(1), 1, 1),
20577 block_dim: (256, 1, 1),
20578 shared_mem_bytes: 0,
20579 };
20580 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20581 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20582 let __s_b = self.gpu.stream();
20583 let mut b = __s_b.launch_builder(&f);
20584 b.arg(k)
20585 .arg(v)
20586 .arg(&mut *kw)
20587 .arg(&mut *vw)
20588 .arg(&kdk)
20589 .arg(&kdv)
20590 .arg(&tkvi)
20591 .arg(&ktb)
20592 .arg(&vtb);
20593 unsafe {
20594 b.launch(cfg)?;
20595 }
20596 }
20597 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20599 .map(|v| v != "0")
20600 .unwrap_or(true);
20601 {
20602 let f = self.func(if db {
20603 "fa_prefill_qw_db_w_hd128"
20604 } else {
20605 "fa_prefill_qw_w_hd128"
20606 });
20607 let shmem = if db {
20608 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20609 } else {
20610 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20611 };
20612 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20613 f.set_attribute(
20614 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20615 shmem as i32,
20616 )?;
20617 let cfg = LaunchConfig {
20618 grid_dim: (
20619 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20620 n_head as u32,
20621 1,
20622 ),
20623 block_dim: (32, 4, 1),
20624 shared_mem_bytes: shmem,
20625 };
20626 let (hd, nh, nhkv, ti, tkvi, cz) = (
20627 head_dim as i32,
20628 n_head as i32,
20629 n_head_kv as i32,
20630 t as i32,
20631 t_kv as i32,
20632 causal as i32,
20633 );
20634 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20635 let __s_b = self.gpu.stream();
20636 let mut b = __s_b.launch_builder(&f);
20637 b.arg(q)
20638 .arg(&*kw)
20639 .arg(&*vw)
20640 .arg(o)
20641 .arg(&hd)
20642 .arg(&nh)
20643 .arg(&nhkv)
20644 .arg(&ti)
20645 .arg(&tkvi)
20646 .arg(&scale)
20647 .arg(&cz)
20648 .arg(&kdk)
20649 .arg(&kdv)
20650 .arg(&wnd);
20651 unsafe {
20652 b.launch(cfg)?;
20653 }
20654 }
20655 Ok(())
20656 }
20657
20658 pub fn fa_decode(
20662 &self,
20663 q: &CudaSlice<f32>,
20664 k: &cudarc::driver::CudaView<u8>,
20665 v: &cudarc::driver::CudaView<u8>,
20666 o: &mut CudaSlice<f32>,
20667 head_dim: usize,
20668 n_head: usize,
20669 n_head_kv: usize,
20670 t_kv: usize,
20671 scale: f32,
20672 k_tok_bytes: usize,
20673 v_tok_bytes: usize,
20674 ) -> Result<(), Box<dyn std::error::Error>> {
20675 self.fa_decode_kvmod(
20676 q,
20677 k,
20678 v,
20679 o,
20680 head_dim,
20681 n_head,
20682 n_head_kv,
20683 t_kv,
20684 scale,
20685 k_tok_bytes,
20686 v_tok_bytes,
20687 false,
20688 )
20689 }
20690
20691 #[allow(clippy::too_many_arguments)]
20695 #[allow(clippy::too_many_arguments)]
20699 #[allow(clippy::too_many_arguments)]
20700 fn fa_decode_scalar_unified(
20701 &self,
20702 q: &cudarc::driver::CudaView<f32>,
20703 k: &cudarc::driver::CudaView<u8>,
20704 v: &cudarc::driver::CudaView<u8>,
20705 o: &mut cudarc::driver::CudaViewMut<f32>,
20706 head_dim: usize,
20707 n_head: usize,
20708 n_head_kv: usize,
20709 t_kv_host: usize,
20710 t_kv_dev: Option<&CudaSlice<i32>>,
20711 scale: f32,
20712 n_splits: usize,
20713 split_keys: usize,
20714 k_tok_bytes: usize,
20715 v_tok_bytes: usize,
20716 g: bool,
20717 part_o: &mut CudaSlice<f32>,
20718 part_m: &mut CudaSlice<f32>,
20719 part_l: &mut CudaSlice<f32>,
20720 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20721 ) -> Result<(), Box<dyn std::error::Error>> {
20722 let f = if g {
20723 self.func_g("fa_decode_f32")
20724 } else {
20725 self.fa_func("fa_decode_f32", head_dim)
20726 };
20727 let cfg = LaunchConfig {
20728 grid_dim: (n_head as u32, n_splits as u32, 1),
20729 block_dim: (head_dim as u32, 1, 1),
20730 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20731 };
20732 let (hd, nh, nhkv, nsp) = (
20733 head_dim as i32,
20734 n_head as i32,
20735 n_head_kv as i32,
20736 n_splits as i32,
20737 );
20738 let (ktb, vtb, tkvi, ski) = (
20739 k_tok_bytes as i64,
20740 v_tok_bytes as i64,
20741 t_kv_host as i32,
20742 split_keys as i32,
20743 );
20744 let __s_b = self.gpu.stream();
20745 let mut b = __s_b.launch_builder(&f);
20746 match t_kv_dev {
20747 Some(d) => {
20748 b.arg(q)
20749 .arg(k)
20750 .arg(v)
20751 .arg(&mut *part_o)
20752 .arg(&mut *part_m)
20753 .arg(&mut *part_l)
20754 .arg(&hd)
20755 .arg(&nh)
20756 .arg(&nhkv)
20757 .arg(&tkvi)
20758 .arg(d)
20759 .arg(&scale)
20760 .arg(&nsp)
20761 .arg(&ski)
20762 .arg(&ktb)
20763 .arg(&vtb);
20764 unsafe {
20765 b.launch(cfg)?;
20766 }
20767 }
20768 None => {
20769 let null: u64 = 0;
20770 b.arg(q)
20771 .arg(k)
20772 .arg(v)
20773 .arg(&mut *part_o)
20774 .arg(&mut *part_m)
20775 .arg(&mut *part_l)
20776 .arg(&hd)
20777 .arg(&nh)
20778 .arg(&nhkv)
20779 .arg(&tkvi)
20780 .arg(&null)
20781 .arg(&scale)
20782 .arg(&nsp)
20783 .arg(&ski)
20784 .arg(&ktb)
20785 .arg(&vtb);
20786 unsafe {
20787 b.launch(cfg)?;
20788 }
20789 }
20790 }
20791 let cfg2 = LaunchConfig {
20792 grid_dim: (n_head as u32, 1, 1),
20793 block_dim: (head_dim as u32, 1, 1),
20794 shared_mem_bytes: 0,
20795 };
20796 if let Some((oq, od)) = q8_out {
20797 let fc = if g {
20799 self.func_g("fa_decode_combine_q8_1")
20800 } else {
20801 self.fa_func("fa_decode_combine_q8_1", head_dim)
20802 };
20803 let __s_b2 = self.gpu.stream();
20804 let mut b2 = __s_b2.launch_builder(&fc);
20805 b2.arg(&*part_o)
20806 .arg(&*part_m)
20807 .arg(&*part_l)
20808 .arg(oq)
20809 .arg(od)
20810 .arg(&hd)
20811 .arg(&nh)
20812 .arg(&nsp);
20813 unsafe {
20814 b2.launch(cfg2)?;
20815 }
20816 return Ok(());
20817 }
20818 let fc = if g {
20819 self.func_g("fa_decode_combine_f32")
20820 } else {
20821 self.fa_func("fa_decode_combine_f32", head_dim)
20822 };
20823 let __s_b2 = self.gpu.stream();
20824 let mut b2 = __s_b2.launch_builder(&fc);
20825 b2.arg(&*part_o)
20826 .arg(&*part_m)
20827 .arg(&*part_l)
20828 .arg(o)
20829 .arg(&hd)
20830 .arg(&nh)
20831 .arg(&nsp);
20832 unsafe {
20833 b2.launch(cfg2)?;
20834 }
20835 Ok(())
20836 }
20837
20838 pub fn fa_decode_kvmod(
20839 &self,
20840 q: &CudaSlice<f32>,
20841 k: &cudarc::driver::CudaView<u8>,
20842 v: &cudarc::driver::CudaView<u8>,
20843 o: &mut CudaSlice<f32>,
20844 head_dim: usize,
20845 n_head: usize,
20846 n_head_kv: usize,
20847 t_kv: usize,
20848 scale: f32,
20849 k_tok_bytes: usize,
20850 v_tok_bytes: usize,
20851 g: bool,
20852 ) -> Result<(), Box<dyn std::error::Error>> {
20853 let q_view = q.as_view();
20854 let mut o_view = o.as_view_mut();
20855 self.fa_decode_kvmod_view(
20856 &q_view,
20857 k,
20858 v,
20859 &mut o_view,
20860 head_dim,
20861 n_head,
20862 n_head_kv,
20863 t_kv,
20864 scale,
20865 k_tok_bytes,
20866 v_tok_bytes,
20867 g,
20868 )
20869 }
20870
20871 #[allow(clippy::too_many_arguments)]
20876 pub fn fa_decode_kvmod_view(
20877 &self,
20878 q: &cudarc::driver::CudaView<f32>,
20879 k: &cudarc::driver::CudaView<u8>,
20880 v: &cudarc::driver::CudaView<u8>,
20881 o: &mut cudarc::driver::CudaViewMut<f32>,
20882 head_dim: usize,
20883 n_head: usize,
20884 n_head_kv: usize,
20885 t_kv: usize,
20886 scale: f32,
20887 k_tok_bytes: usize,
20888 v_tok_bytes: usize,
20889 g: bool,
20890 ) -> Result<(), Box<dyn std::error::Error>> {
20891 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20912 if g && head_dim == 256 && !fa_v4_at(t_kv) {
20916 fa_vec = false;
20917 }
20918 let sp = fa_split_keys(t_kv, n_head_kv);
20919 let n_splits = if fa_vec {
20920 ((t_kv + sp - 1) / sp).max(1)
20921 } else {
20922 ((t_kv + 255) / 256).max(1)
20923 };
20924 let o_len = n_head * n_splits * head_dim;
20925 let ml_len = n_head * n_splits;
20926 let mut part_guard = self.fa_part_pool.lock().unwrap();
20927 if part_guard
20928 .as_ref()
20929 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20930 .unwrap_or(true)
20931 {
20932 let old = part_guard.take();
20943 let (co, cm) = old
20944 .as_ref()
20945 .map(|pp| (pp.0.len(), pp.1.len()))
20946 .unwrap_or((0, 0));
20947 if let Some(old) = old {
20948 self.fa_part_retired.lock().unwrap().push(old);
20949 }
20950 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20951 eprintln!(
20952 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20953 co, o_len, cm, ml_len
20954 );
20955 }
20956 *part_guard = Some((
20957 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20958 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20959 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20960 ));
20961 }
20962 let pg = part_guard.as_mut().unwrap();
20963 self.gpu
20964 .stream()
20965 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20966 self.gpu
20967 .stream()
20968 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20969 self.gpu
20970 .stream()
20971 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20972 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20973 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20974 let (hd, nh, nhkv, tkvi, nsp) = (
20975 head_dim as i32,
20976 n_head as i32,
20977 n_head_kv as i32,
20978 t_kv as i32,
20979 n_splits as i32,
20980 );
20981 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20982 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20986 let fa512_min = fa512_min_tkv();
20991 let deep = fa_vec
20994 && head_dim == 256
20995 && fa_v4_at(t_kv)
20996 && !g
20997 && fa_deep_at(t_kv)
20998 && !matches!(fa_v4_mode(), "noB3" | "stage");
20999 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
21000 let gqa = (n_head / n_head_kv).max(1) as u32;
21003 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
21004 (
21005 fv,
21006 LaunchConfig {
21007 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21008 block_dim: (32, gqa, 1),
21009 shared_mem_bytes: 0,
21010 },
21011 )
21012 } else if fa_vec && head_dim <= 256 {
21013 let gqa = (n_head / n_head_kv).max(1) as u32;
21014 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21025 let smem_tkv = *SMEM_TKV.get_or_init(|| {
21026 std::env::var("MEMRA_FA_SMEM_TKV")
21027 .ok()
21028 .and_then(|v| v.parse().ok())
21029 .unwrap_or_else(|| {
21030 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21031 })
21032 });
21033 if fa_v4_at(t_kv) && head_dim == 256 {
21034 let v4name = match fa_v4_mode() {
21038 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
21041 _ => "fa_decode_vec_q_v4",
21042 };
21043 let fv = if g {
21044 self.func_g(v4name)
21045 } else {
21046 self.func(v4name)
21047 };
21048 let shmem = (if deep { 12160 } else { 11520 }
21051 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
21052 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21053 fv.set_attribute(
21054 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21055 shmem as i32,
21056 )?;
21057 (
21058 fv,
21059 LaunchConfig {
21060 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21061 block_dim: (32, gqa, 1),
21062 shared_mem_bytes: shmem,
21063 },
21064 )
21065 } else if fa_v3_active(head_dim) {
21066 let fv = if g {
21069 self.func_g("fa_decode_vec_q_v3")
21070 } else {
21071 self.func("fa_decode_vec_q_v3")
21072 };
21073 let shmem = (32 * head_dim * 2) as u32; (
21075 fv,
21076 LaunchConfig {
21077 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21078 block_dim: (32, gqa, 1),
21079 shared_mem_bytes: shmem,
21080 },
21081 )
21082 } else if fa_v2_on() {
21083 let fv = if g {
21087 self.func_g("fa_decode_vec_q_v2")
21088 } else {
21089 self.func("fa_decode_vec_q_v2")
21090 };
21091 let shmem = (2 * 32 * head_dim * 2) as u32; (
21093 fv,
21094 LaunchConfig {
21095 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21096 block_dim: (32, gqa, 1),
21097 shared_mem_bytes: shmem,
21098 },
21099 )
21100 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
21101 {
21102 let fv = if g {
21106 self.func_g("fa_decode_vec_q_smem")
21107 } else {
21108 self.func("fa_decode_vec_q_smem")
21109 };
21110 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
21112 fv.set_attribute(
21113 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21114 shmem as i32,
21115 )?;
21116 (
21117 fv,
21118 LaunchConfig {
21119 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21120 block_dim: (32, gqa, 1),
21121 shared_mem_bytes: shmem,
21122 },
21123 )
21124 } else {
21125 let fv = if g {
21128 self.func_g("fa_decode_vec_q")
21129 } else {
21130 self.func("fa_decode_vec_q")
21131 };
21132 (
21133 fv,
21134 LaunchConfig {
21135 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
21136 block_dim: (32, gqa, 1),
21137 shared_mem_bytes: 0,
21138 },
21139 )
21140 }
21141 } else {
21142 return self.fa_decode_scalar_unified(
21145 q,
21146 k,
21147 v,
21148 o,
21149 head_dim,
21150 n_head,
21151 n_head_kv,
21152 t_kv,
21153 None,
21154 scale,
21155 n_splits,
21156 if fa_vec { sp } else { 256 },
21157 k_tok_bytes,
21158 v_tok_bytes,
21159 g,
21160 part_o,
21161 part_m,
21162 part_l,
21163 None,
21164 );
21165 };
21166 let __s_b = self.gpu.stream();
21167 let mut b = __s_b.launch_builder(&f);
21168 b.arg(q)
21169 .arg(k)
21170 .arg(v)
21171 .arg(&mut *part_o)
21172 .arg(&mut *part_m)
21173 .arg(&mut *part_l)
21174 .arg(&hd)
21175 .arg(&nh)
21176 .arg(&nhkv)
21177 .arg(&tkvi)
21178 .arg(&scale)
21179 .arg(&nsp)
21180 .arg(&ktb)
21181 .arg(&vtb);
21182 unsafe {
21183 b.launch(cfg)?;
21184 }
21185 let (fc, cfg2) = (
21188 if g {
21189 self.func_g("fa_decode_combine_f32")
21190 } else {
21191 self.fa_func("fa_decode_combine_f32", head_dim)
21192 },
21193 LaunchConfig {
21194 grid_dim: (n_head as u32, 1, 1),
21195 block_dim: (head_dim as u32, 1, 1),
21196 shared_mem_bytes: 0,
21197 },
21198 );
21199 let __s_b2 = self.gpu.stream();
21200 let mut b2 = __s_b2.launch_builder(&fc);
21201 b2.arg(&*part_o)
21202 .arg(&*part_m)
21203 .arg(&*part_l)
21204 .arg(o)
21205 .arg(&hd)
21206 .arg(&nh)
21207 .arg(&nsp);
21208 unsafe {
21209 b2.launch(cfg2)?;
21210 }
21211 Ok(())
21212 }
21213
21214 #[allow(clippy::too_many_arguments)]
21225 pub fn fa_decode_batch_seqs_v4(
21226 &self,
21227 q: &CudaSlice<f32>,
21228 kv_ptrs: &cudarc::driver::CudaView<u64>,
21229 pos_seq: &CudaSlice<i32>,
21230 o: &mut CudaSlice<f32>,
21231 head_dim: usize,
21232 n_head: usize,
21233 n_head_kv: usize,
21234 b_n: usize,
21235 t_kv_max: usize,
21236 scale: f32,
21237 split_keys: usize,
21238 k_tok_bytes: usize,
21239 v_tok_bytes: usize,
21240 ) -> Result<(), Box<dyn std::error::Error>> {
21241 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
21242 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
21243 let o_len = b_n * n_head * n_splits_max * head_dim;
21244 let ml_len = b_n * n_head * n_splits_max;
21245 let mut part_guard = self.fa_part_pool.lock().unwrap();
21246 if part_guard
21247 .as_ref()
21248 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21249 .unwrap_or(true)
21250 {
21251 let old = part_guard.take();
21262 let (co, cm) = old
21263 .as_ref()
21264 .map(|pp| (pp.0.len(), pp.1.len()))
21265 .unwrap_or((0, 0));
21266 if let Some(old) = old {
21267 self.fa_part_retired.lock().unwrap().push(old);
21268 }
21269 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21270 eprintln!(
21271 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21272 co, o_len, cm, ml_len
21273 );
21274 }
21275 *part_guard = Some((
21276 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21277 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21278 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21279 ));
21280 }
21281 let pg = part_guard.as_mut().unwrap();
21282 self.gpu
21283 .stream()
21284 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21285 self.gpu
21286 .stream()
21287 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21288 self.gpu
21289 .stream()
21290 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21291 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21292 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21293 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
21294 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21295 let gqa = (n_head / n_head_kv).max(1) as u32;
21296 let f = self.func("fa_decode_vec_q_seqs_v4");
21297 let shmem = (11520 + 32 * head_dim * 2) as u32;
21299 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21300 f.set_attribute(
21301 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21302 shmem as i32,
21303 )?;
21304 let cfg = LaunchConfig {
21305 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
21306 block_dim: (32, gqa, 1),
21307 shared_mem_bytes: shmem,
21308 };
21309 {
21310 let __s_b = self.gpu.stream();
21311 let mut b = __s_b.launch_builder(&f);
21312 b.arg(q)
21313 .arg(kv_ptrs)
21314 .arg(pos_seq)
21315 .arg(&mut *part_o)
21316 .arg(&mut *part_m)
21317 .arg(&mut *part_l)
21318 .arg(&hd)
21319 .arg(&nh)
21320 .arg(&nhkv)
21321 .arg(&scale)
21322 .arg(&nspm)
21323 .arg(&spk)
21324 .arg(&ktb)
21325 .arg(&vtb);
21326 unsafe {
21327 b.launch(cfg)?;
21328 }
21329 }
21330 let fc = self.func("fa_decode_combine_seqs");
21331 let cfg2 = LaunchConfig {
21332 grid_dim: (n_head as u32, b_n as u32, 1),
21333 block_dim: (head_dim as u32, 1, 1),
21334 shared_mem_bytes: 0,
21335 };
21336 let __s_b2 = self.gpu.stream();
21337 let mut b2 = __s_b2.launch_builder(&fc);
21338 b2.arg(&*part_o)
21339 .arg(&*part_m)
21340 .arg(&*part_l)
21341 .arg(o)
21342 .arg(&hd)
21343 .arg(&nh)
21344 .arg(pos_seq)
21345 .arg(&nspm)
21346 .arg(&spk);
21347 unsafe {
21348 b2.launch(cfg2)?;
21349 }
21350 Ok(())
21351 }
21352
21353 #[allow(clippy::too_many_arguments)]
21360 pub fn append_kv_quantized_seqs(
21361 &self,
21362 k_rows: &CudaSlice<f32>,
21363 v_rows: &CudaSlice<f32>,
21364 kv_ptrs: &cudarc::driver::CudaView<u64>,
21365 pos_seq: &CudaSlice<i32>,
21366 b_n: usize,
21367 kv_dim_k: usize,
21368 kv_dim_v: usize,
21369 k_tok_bytes: usize,
21370 v_tok_bytes: usize,
21371 ) -> Result<(), Box<dyn std::error::Error>> {
21372 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
21373 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21374 let cfg = LaunchConfig {
21375 grid_dim: (nblk, b_n as u32, 1),
21376 block_dim: (32, 1, 1),
21377 shared_mem_bytes: 0,
21378 };
21379 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21380 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21381 let __s_b = self.gpu.stream();
21382 let mut b = __s_b.launch_builder(&f);
21383 b.arg(k_rows)
21384 .arg(v_rows)
21385 .arg(kv_ptrs)
21386 .arg(pos_seq)
21387 .arg(&kdk)
21388 .arg(&kdv)
21389 .arg(&ktb)
21390 .arg(&vtb);
21391 unsafe {
21392 b.launch(cfg)?;
21393 }
21394 Ok(())
21395 }
21396
21397 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21403 std::env::var("MEMRA_NO_FA_VEC").is_err()
21404 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21405 && base_len + 1 >= fa_vec_min_tkv()
21406 && head_dim <= 256
21407 && head_dim % 32 == 0
21408 }
21409
21410 #[allow(clippy::too_many_arguments)]
21419 pub fn fa_decode_rows(
21420 &self,
21421 q: &CudaSlice<f32>,
21422 k: &cudarc::driver::CudaView<u8>,
21423 v: &cudarc::driver::CudaView<u8>,
21424 o: &mut CudaSlice<f32>,
21425 head_dim: usize,
21426 n_head: usize,
21427 n_head_kv: usize,
21428 base_len: usize,
21429 t: usize,
21430 scale: f32,
21431 k_tok_bytes: usize,
21432 v_tok_bytes: usize,
21433 base_dev: Option<(&CudaSlice<i32>, i32)>,
21437 kv_shared: bool,
21440 g: bool,
21444 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21447 ) -> Result<(), Box<dyn std::error::Error>> {
21448 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21449 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
21456 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21457 let v = *SP512.get_or_init(|| {
21460 std::env::var("MEMRA_FA_SP512")
21461 .ok()
21462 .and_then(|x| x.parse().ok())
21463 .unwrap_or(0)
21464 });
21465 sp = if v >= 8 {
21466 v
21467 } else {
21468 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21469 };
21470 }
21471 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21472 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21473 let gqa = (n_head / n_head_kv).max(1) as u32;
21474 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21485 groups.push((0, t, sp));
21486 } else {
21487 let mut r0 = 0usize;
21488 while r0 < t {
21489 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21490 let mut r1 = r0 + 1;
21491 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21492 r1 += 1;
21493 }
21494 groups.push((r0, r1 - r0, sp_g));
21495 r0 = r1;
21496 }
21497 }
21498 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21502 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21503 std::env::var("MEMRA_FA_SMEM_TKV")
21504 .ok()
21505 .and_then(|v| v.parse().ok())
21506 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21507 });
21508 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21509 let v3 = fa_v3_active(head_dim);
21510 let smem_rows =
21511 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21512 let _ = kv_shared;
21517 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21520 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21534 let tb512 = head_dim == 512
21536 && sp <= 32
21537 && n_head / n_head_kv.max(1) <= 16
21538 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21539 let fname = if tb512 {
21540 "fa_decode_vec_q_rows_v4_512_tb"
21541 } else if i2 {
21542 "fa_decode_vec_q_rows_dpl16_i2"
21543 } else if head_dim == 512 {
21544 "fa_decode_vec_q_rows_dpl16"
21545 }
21546 else if v4 {
21548 "fa_decode_vec_q_rows_v4"
21549 } else if v3 {
21550 "fa_decode_vec_q_rows_v3"
21551 } else if fa_v2_on() {
21552 "fa_decode_vec_q_rows_v2"
21553 } else if smem_rows {
21554 "fa_decode_vec_q_rows_smem"
21555 } else {
21556 "fa_decode_vec_q_rows"
21557 };
21558 let f = if head_dim == 512 {
21559 self.fa_func(fname, head_dim)
21560 } else if g {
21561 self.func_g(if smem_rows {
21569 "fa_decode_vec_q_rows"
21570 } else {
21571 fname
21572 })
21573 } else {
21574 self.func(fname)
21575 };
21576 let shmem = if tb512 {
21577 let gk = Self::gkv_on();
21579 let sh =
21580 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21581 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21582 f.set_attribute(
21583 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21584 sh as i32,
21585 )?;
21586 sh
21587 } else if v4 || v3 || smem_rows || fa_v2_on() {
21588 let sh = (if v4 {
21590 11520 + 32 * head_dim * if g { 1 } else { 2 }
21591 } else if v3 {
21592 32 * head_dim * 2
21593 } else {
21594 2 * 32 * head_dim * 2
21595 }) as u32;
21596 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21597 f.set_attribute(
21598 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21599 sh as i32,
21600 )?;
21601 sh
21602 } else {
21603 0
21604 };
21605 for &(r0, t_g, sp_g) in &groups {
21609 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21610 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21611 let base_i = (base_len + r0) as i32;
21612 let o_len = t_g * n_head * n_splits_g * head_dim;
21613 let ml_len = t_g * n_head * n_splits_g;
21614 let mut part_guard = self.fa_part_pool.lock().unwrap();
21615 if part_guard
21616 .as_ref()
21617 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21618 .unwrap_or(true)
21619 {
21620 let old = part_guard.take();
21631 let (co, cm) = old
21632 .as_ref()
21633 .map(|pp| (pp.0.len(), pp.1.len()))
21634 .unwrap_or((0, 0));
21635 if let Some(old) = old {
21636 self.fa_part_retired.lock().unwrap().push(old);
21637 }
21638 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21639 eprintln!(
21640 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21641 co, o_len, cm, ml_len
21642 );
21643 }
21644 *part_guard = Some((
21645 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21646 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21647 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21648 ));
21649 }
21650 let pg = part_guard.as_mut().unwrap();
21651 self.gpu
21652 .stream()
21653 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21654 self.gpu
21655 .stream()
21656 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21657 self.gpu
21658 .stream()
21659 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21660 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21661 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21662 let qv = self.view(q, t * n_head * head_dim);
21663 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21664 let cfg = LaunchConfig {
21665 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21666 block_dim: (32, gqa, 1),
21667 shared_mem_bytes: shmem,
21668 };
21669 {
21670 let __s_b = self.gpu.stream();
21671 let mut b = __s_b.launch_builder(&f);
21672 if tb512 {
21673 let (bd, plus) =
21675 base_dev.expect("hd512 rows twin requires a device base counter");
21676 let plus_g = plus + r0 as i32;
21677 let nr = t_g as i32;
21678 if Self::pdl_on() && Self::pdl_wb_on() {
21679 use cudarc::driver::{DevicePtr, DevicePtrMut};
21681 let s = &self.gpu.stream();
21682 let (pq, _b0) = q_g.device_ptr(s);
21683 let (pk, _b1) = k.device_ptr(s);
21684 let (pv, _b2) = v.device_ptr(s);
21685 let (po, _b3) = part_o.device_ptr_mut(s);
21686 let (pm, _b4) = part_m.device_ptr_mut(s);
21687 let (pl, _b5) = part_l.device_ptr_mut(s);
21688 let (pb, _b6) = bd.device_ptr(s);
21689 let mut ps = [
21690 &pq as *const _ as *mut std::ffi::c_void,
21691 &pk as *const _ as *mut _,
21692 &pv as *const _ as *mut _,
21693 &po as *const _ as *mut _,
21694 &pm as *const _ as *mut _,
21695 &pl as *const _ as *mut _,
21696 &hd as *const _ as *mut _,
21697 &nh as *const _ as *mut _,
21698 &nhkv as *const _ as *mut _,
21699 &pb as *const _ as *mut _,
21700 &plus_g as *const _ as *mut _,
21701 &scale as *const _ as *mut _,
21702 &nspm as *const _ as *mut _,
21703 &spk as *const _ as *mut _,
21704 &ktb as *const _ as *mut _,
21705 &vtb as *const _ as *mut _,
21706 &nr as *const _ as *mut _,
21707 ];
21708 unsafe {
21709 self.launch_pdl_flash(
21710 Self::gkv_on(),
21711 "fa_decode_vec_q_rows_v4_512_tb",
21712 (n_head_kv as u32, n_splits_g as u32, 1),
21713 (32, gqa, 1),
21714 shmem,
21715 &mut ps,
21716 )?;
21717 }
21718 } else {
21719 let cfg_tb = LaunchConfig {
21720 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21721 block_dim: (32, gqa, 1),
21722 shared_mem_bytes: shmem,
21723 };
21724 b.arg(&q_g)
21725 .arg(k)
21726 .arg(v)
21727 .arg(&mut *part_o)
21728 .arg(&mut *part_m)
21729 .arg(&mut *part_l)
21730 .arg(&hd)
21731 .arg(&nh)
21732 .arg(&nhkv)
21733 .arg(bd)
21734 .arg(&plus_g)
21735 .arg(&scale)
21736 .arg(&nspm)
21737 .arg(&spk)
21738 .arg(&ktb)
21739 .arg(&vtb)
21740 .arg(&nr);
21741 unsafe {
21742 b.launch(cfg_tb)?;
21743 }
21744 }
21745 } else if head_dim == 512 {
21746 let (bd, plus) =
21747 base_dev.expect("hd512 rows twin requires a device base counter");
21748 let plus_g = plus + r0 as i32;
21749 b.arg(&q_g)
21750 .arg(k)
21751 .arg(v)
21752 .arg(&mut *part_o)
21753 .arg(&mut *part_m)
21754 .arg(&mut *part_l)
21755 .arg(&hd)
21756 .arg(&nh)
21757 .arg(&nhkv)
21758 .arg(bd)
21759 .arg(&plus_g)
21760 .arg(&scale)
21761 .arg(&nspm)
21762 .arg(&spk)
21763 .arg(&ktb)
21764 .arg(&vtb);
21765 unsafe {
21766 b.launch(cfg)?;
21767 }
21768 } else {
21769 b.arg(&q_g)
21770 .arg(k)
21771 .arg(v)
21772 .arg(&mut *part_o)
21773 .arg(&mut *part_m)
21774 .arg(&mut *part_l)
21775 .arg(&hd)
21776 .arg(&nh)
21777 .arg(&nhkv)
21778 .arg(&base_i)
21779 .arg(&scale)
21780 .arg(&nspm)
21781 .arg(&spk)
21782 .arg(&ktb)
21783 .arg(&vtb);
21784 unsafe {
21785 b.launch(cfg)?;
21786 }
21787 }
21788 }
21789 let cfg2 = LaunchConfig {
21790 grid_dim: (n_head as u32, t_g as u32, 1),
21791 block_dim: (head_dim as u32, 1, 1),
21792 shared_mem_bytes: 0,
21793 };
21794 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21795 if head_dim == 512 {
21796 let (bd, plus) = base_dev.unwrap();
21799 let plus_g = plus + r0 as i32;
21800 if let Some((oq, od)) = q8_out.as_mut() {
21801 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21803 if Self::pdl_on() && Self::pdl_wb_on() {
21804 use cudarc::driver::{DevicePtr, DevicePtrMut};
21806 let s = &self.gpu.stream();
21807 let (po, _g0) = part_o.device_ptr(s);
21808 let (pm, _g1) = part_m.device_ptr(s);
21809 let (pl, _g2) = part_l.device_ptr(s);
21810 let (pq, _g3) = oq.device_ptr_mut(s);
21811 let (pd, _g4) = od.device_ptr_mut(s);
21812 let (pb, _g5) = bd.device_ptr(s);
21813 let mut ps = [
21814 &po as *const _ as *mut std::ffi::c_void,
21815 &pm as *const _ as *mut _,
21816 &pl as *const _ as *mut _,
21817 &pq as *const _ as *mut _,
21818 &pd as *const _ as *mut _,
21819 &hd as *const _ as *mut _,
21820 &nh as *const _ as *mut _,
21821 &pb as *const _ as *mut _,
21822 &plus_g as *const _ as *mut _,
21823 &nspm as *const _ as *mut _,
21824 &spk as *const _ as *mut _,
21825 ];
21826 unsafe {
21827 self.launch_pdl_flash(
21828 Self::gkv_on(),
21829 "fa_decode_combine_rows_dc_q8_1",
21830 cfg2.grid_dim,
21831 cfg2.block_dim,
21832 0,
21833 &mut ps,
21834 )?;
21835 }
21836 continue;
21837 }
21838 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21839 let __s_b2 = self.gpu.stream();
21840 let mut b2 = __s_b2.launch_builder(&fc);
21841 b2.arg(&*part_o)
21842 .arg(&*part_m)
21843 .arg(&*part_l)
21844 .arg(&mut **oq)
21845 .arg(&mut **od)
21846 .arg(&hd)
21847 .arg(&nh)
21848 .arg(bd)
21849 .arg(&plus_g)
21850 .arg(&nspm)
21851 .arg(&spk);
21852 unsafe {
21853 b2.launch(cfg2)?;
21854 }
21855 continue;
21856 }
21857 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21858 let __s_b2 = self.gpu.stream();
21859 let mut b2 = __s_b2.launch_builder(&fc);
21860 b2.arg(&*part_o)
21861 .arg(&*part_m)
21862 .arg(&*part_l)
21863 .arg(&mut o_g)
21864 .arg(&hd)
21865 .arg(&nh)
21866 .arg(bd)
21867 .arg(&plus_g)
21868 .arg(&nspm)
21869 .arg(&spk);
21870 unsafe {
21871 b2.launch(cfg2)?;
21872 }
21873 } else {
21874 assert!(
21877 q8_out.is_none(),
21878 "rows q8 emit requires the hd512 dc combine"
21879 );
21880 let fc = self.func("fa_decode_combine_rows");
21881 let __s_b2 = self.gpu.stream();
21882 let mut b2 = __s_b2.launch_builder(&fc);
21883 b2.arg(&*part_o)
21884 .arg(&*part_m)
21885 .arg(&*part_l)
21886 .arg(&mut o_g)
21887 .arg(&hd)
21888 .arg(&nh)
21889 .arg(&base_i)
21890 .arg(&nspm)
21891 .arg(&spk);
21892 unsafe {
21893 b2.launch(cfg2)?;
21894 }
21895 }
21896 }
21897 Ok(())
21898 }
21899
21900 #[allow(clippy::too_many_arguments)]
21904 pub fn fa_decode_rows_w(
21905 &self,
21906 q: &CudaSlice<f32>,
21907 k: &cudarc::driver::CudaView<u8>,
21908 v: &cudarc::driver::CudaView<u8>,
21909 o: &mut CudaSlice<f32>,
21910 head_dim: usize,
21911 n_head: usize,
21912 n_head_kv: usize,
21913 base_dev: &CudaSlice<i32>,
21914 base_plus: i32,
21915 t: usize,
21916 scale: f32,
21917 window: usize,
21918 k_tok_bytes: usize,
21919 v_tok_bytes: usize,
21920 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21921 ) -> Result<(), Box<dyn std::error::Error>> {
21922 debug_assert!(head_dim == 256);
21927 let sp = {
21935 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21936 let v = *SPW.get_or_init(|| {
21937 std::env::var("MEMRA_FA_SPW")
21938 .ok()
21939 .and_then(|x| x.parse().ok())
21940 .unwrap_or(0)
21941 });
21942 if v >= 8 {
21943 v
21944 } else {
21945 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21946 }
21947 };
21948 let n_splits_max = (window + sp - 1) / sp;
21949 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21950 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21951 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21952 let gqa = (n_head / n_head_kv).max(1) as u32;
21953 let o_len = t * n_head * n_splits_max * head_dim;
21954 let ml_len = t * n_head * n_splits_max;
21955 let mut part_guard = self.fa_part_pool.lock().unwrap();
21956 if part_guard
21957 .as_ref()
21958 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21959 .unwrap_or(true)
21960 {
21961 let old = part_guard.take();
21972 let (co, cm) = old
21973 .as_ref()
21974 .map(|pp| (pp.0.len(), pp.1.len()))
21975 .unwrap_or((0, 0));
21976 if let Some(old) = old {
21977 self.fa_part_retired.lock().unwrap().push(old);
21978 }
21979 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21980 eprintln!(
21981 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21982 co, o_len, cm, ml_len
21983 );
21984 }
21985 *part_guard = Some((
21986 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21987 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21988 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21989 ));
21990 }
21991 let pg = part_guard.as_mut().unwrap();
21992 self.gpu
21993 .stream()
21994 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21995 self.gpu
21996 .stream()
21997 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21998 self.gpu
21999 .stream()
22000 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22001 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22002 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22008 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
22009 std::env::var("MEMRA_FA_SMEM_TKV")
22010 .ok()
22011 .and_then(|v| v.parse().ok())
22012 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
22013 });
22014 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22020 let wg = Self::wkv_on();
22025 let sp2 =
22028 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
22029 if sp2 {
22030 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22031 if Self::pdl_on() && Self::pdl_wb_on() {
22032 use cudarc::driver::{DevicePtr, DevicePtrMut};
22034 let s = &self.gpu.stream();
22035 let (pq, _b0) = q.device_ptr(s);
22036 let (pk, _b1) = k.device_ptr(s);
22037 let (pv, _b2) = v.device_ptr(s);
22038 let (po, _b3) = part_o.device_ptr_mut(s);
22039 let (pm, _b4) = part_m.device_ptr_mut(s);
22040 let (pl, _b5) = part_l.device_ptr_mut(s);
22041 let (pb, _b6) = base_dev.device_ptr(s);
22042 let mut ps = [
22043 &pq as *const _ as *mut std::ffi::c_void,
22044 &pk as *const _ as *mut _,
22045 &pv as *const _ as *mut _,
22046 &po as *const _ as *mut _,
22047 &pm as *const _ as *mut _,
22048 &pl as *const _ as *mut _,
22049 &hd as *const _ as *mut _,
22050 &nh as *const _ as *mut _,
22051 &nhkv as *const _ as *mut _,
22052 &pb as *const _ as *mut _,
22053 &base_plus as *const _ as *mut _,
22054 &scale as *const _ as *mut _,
22055 &nspm as *const _ as *mut _,
22056 &spk as *const _ as *mut _,
22057 &ktb as *const _ as *mut _,
22058 &vtb as *const _ as *mut _,
22059 &wini as *const _ as *mut _,
22060 ];
22061 unsafe {
22062 self.launch_pdl_flash(
22063 wg,
22064 "fa_decode_vec_q_rows_v4_w_sp",
22065 (n_head_kv as u32, n_splits_max as u32, t as u32),
22066 (32, gqa + 1, 1),
22067 sh,
22068 &mut ps,
22069 )?;
22070 }
22071 } else {
22072 let f = if wg {
22073 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
22074 } else {
22075 self.func("fa_decode_vec_q_rows_v4_w_sp")
22076 };
22077 f.set_attribute(
22078 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22079 sh as i32,
22080 )?;
22081 let cfg = LaunchConfig {
22082 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22083 block_dim: (32, gqa + 1, 1),
22084 shared_mem_bytes: sh,
22085 };
22086 let __s_b = self.gpu.stream();
22087 let mut b = __s_b.launch_builder(&f);
22088 b.arg(q)
22089 .arg(k)
22090 .arg(v)
22091 .arg(&mut *part_o)
22092 .arg(&mut *part_m)
22093 .arg(&mut *part_l)
22094 .arg(&hd)
22095 .arg(&nh)
22096 .arg(&nhkv)
22097 .arg(base_dev)
22098 .arg(&base_plus)
22099 .arg(&scale)
22100 .arg(&nspm)
22101 .arg(&spk)
22102 .arg(&ktb)
22103 .arg(&vtb)
22104 .arg(&wini);
22105 unsafe {
22106 b.launch(cfg)?;
22107 }
22108 }
22109 } else {
22110 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
22111 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
22113 use cudarc::driver::{DevicePtr, DevicePtrMut};
22114 let s = &self.gpu.stream();
22115 let (pq, _b0) = q.device_ptr(s);
22116 let (pk, _b1) = k.device_ptr(s);
22117 let (pv, _b2) = v.device_ptr(s);
22118 let (po, _b3) = part_o.device_ptr_mut(s);
22119 let (pm, _b4) = part_m.device_ptr_mut(s);
22120 let (pl, _b5) = part_l.device_ptr_mut(s);
22121 let (pb, _b6) = base_dev.device_ptr(s);
22122 let mut ps = [
22123 &pq as *const _ as *mut std::ffi::c_void,
22124 &pk as *const _ as *mut _,
22125 &pv as *const _ as *mut _,
22126 &po as *const _ as *mut _,
22127 &pm as *const _ as *mut _,
22128 &pl as *const _ as *mut _,
22129 &hd as *const _ as *mut _,
22130 &nh as *const _ as *mut _,
22131 &nhkv as *const _ as *mut _,
22132 &pb as *const _ as *mut _,
22133 &base_plus as *const _ as *mut _,
22134 &scale as *const _ as *mut _,
22135 &nspm as *const _ as *mut _,
22136 &spk as *const _ as *mut _,
22137 &ktb as *const _ as *mut _,
22138 &vtb as *const _ as *mut _,
22139 &wini as *const _ as *mut _,
22140 ];
22141 unsafe {
22142 self.launch_pdl_flash(
22143 wg,
22144 "fa_decode_vec_q_rows_v4_w",
22145 (n_head_kv as u32, n_splits_max as u32, t as u32),
22146 (32, gqa, 1),
22147 sh,
22148 &mut ps,
22149 )?;
22150 }
22151 } else {
22152 let pick = |name: &str| {
22153 if wg {
22154 self.func_g(name)
22155 } else {
22156 self.func(name)
22157 }
22158 };
22159 let (f, sh) = if fa_v4_at(window) {
22160 let f = pick("fa_decode_vec_q_rows_v4_w");
22161 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
22162 } else if smem_tkv > 0 && window >= smem_tkv {
22163 (
22166 pick("fa_decode_vec_q_rows_smem_w"),
22167 (2 * 32 * head_dim * 2) as u32,
22168 )
22169 } else {
22170 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
22171 };
22172 f.set_attribute(
22173 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22174 sh as i32,
22175 )?;
22176 let cfg = LaunchConfig {
22177 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22178 block_dim: (32, gqa, 1),
22179 shared_mem_bytes: sh,
22180 };
22181 let __s_b = self.gpu.stream();
22182 let mut b = __s_b.launch_builder(&f);
22183 b.arg(q)
22184 .arg(k)
22185 .arg(v)
22186 .arg(&mut *part_o)
22187 .arg(&mut *part_m)
22188 .arg(&mut *part_l)
22189 .arg(&hd)
22190 .arg(&nh)
22191 .arg(&nhkv)
22192 .arg(base_dev)
22193 .arg(&base_plus)
22194 .arg(&scale)
22195 .arg(&nspm)
22196 .arg(&spk)
22197 .arg(&ktb)
22198 .arg(&vtb)
22199 .arg(&wini);
22200 unsafe {
22201 b.launch(cfg)?;
22202 }
22203 }
22204 }
22205 let cfg2 = LaunchConfig {
22206 grid_dim: (n_head as u32, t as u32, 1),
22207 block_dim: (head_dim as u32, 1, 1),
22208 shared_mem_bytes: 0,
22209 };
22210 if let Some((oq, od)) = q8_out {
22211 if Self::pdl_on() && Self::pdl_wb_on() {
22214 use cudarc::driver::{DevicePtr, DevicePtrMut};
22216 let s = &self.gpu.stream();
22217 let (po, _g0) = part_o.device_ptr(s);
22218 let (pm, _g1) = part_m.device_ptr(s);
22219 let (pl, _g2) = part_l.device_ptr(s);
22220 let (pq, _g3) = oq.device_ptr_mut(s);
22221 let (pd, _g4) = od.device_ptr_mut(s);
22222 let mut ps = [
22223 &po as *const _ as *mut std::ffi::c_void,
22224 &pm as *const _ as *mut _,
22225 &pl as *const _ as *mut _,
22226 &pq as *const _ as *mut _,
22227 &pd as *const _ as *mut _,
22228 &hd as *const _ as *mut _,
22229 &nh as *const _ as *mut _,
22230 &nspm as *const _ as *mut _,
22231 &spk as *const _ as *mut _,
22232 &wini as *const _ as *mut _,
22233 ];
22234 unsafe {
22235 self.launch_pdl_flash(
22236 wg,
22237 "fa_decode_combine_rows_w_q8_1",
22238 cfg2.grid_dim,
22239 cfg2.block_dim,
22240 0,
22241 &mut ps,
22242 )?;
22243 }
22244 return Ok(());
22245 }
22246 let fc = if wg {
22247 self.func_g("fa_decode_combine_rows_w_q8_1")
22248 } else {
22249 self.func("fa_decode_combine_rows_w_q8_1")
22250 };
22251 let __s_b2 = self.gpu.stream();
22252 let mut b2 = __s_b2.launch_builder(&fc);
22253 b2.arg(&*part_o)
22254 .arg(&*part_m)
22255 .arg(&*part_l)
22256 .arg(oq)
22257 .arg(od)
22258 .arg(&hd)
22259 .arg(&nh)
22260 .arg(&nspm)
22261 .arg(&spk)
22262 .arg(&wini);
22263 unsafe {
22264 b2.launch(cfg2)?;
22265 }
22266 return Ok(());
22267 }
22268 let fc = if wg {
22269 self.func_g("fa_decode_combine_rows_w")
22270 } else {
22271 self.func("fa_decode_combine_rows_w")
22272 };
22273 let __s_b2 = self.gpu.stream();
22274 let mut b2 = __s_b2.launch_builder(&fc);
22275 b2.arg(&*part_o)
22276 .arg(&*part_m)
22277 .arg(&*part_l)
22278 .arg(o)
22279 .arg(&hd)
22280 .arg(&nh)
22281 .arg(&nspm)
22282 .arg(&spk)
22283 .arg(&wini);
22284 unsafe {
22285 b2.launch(cfg2)?;
22286 }
22287 Ok(())
22288 }
22289
22290 #[allow(clippy::too_many_arguments)]
22296 pub fn fa_decode_rows_dc(
22297 &self,
22298 q: &CudaSlice<f32>,
22299 k: &cudarc::driver::CudaView<u8>,
22300 v: &cudarc::driver::CudaView<u8>,
22301 o: &mut CudaSlice<f32>,
22302 head_dim: usize,
22303 n_head: usize,
22304 n_head_kv: usize,
22305 base_dev: &CudaSlice<i32>,
22306 t_kv_upper: usize,
22307 t: usize,
22308 scale: f32,
22309 k_tok_bytes: usize,
22310 v_tok_bytes: usize,
22311 base_plus: i32,
22312 g: bool,
22313 ) -> Result<(), Box<dyn std::error::Error>> {
22314 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
22315 assert!(
22316 v4 || fa_v3_active(head_dim),
22317 "stream fa rows requires the v3 or v4 lane"
22318 );
22319 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
22320 if v4 {
22321 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22322 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22323 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22324 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22325 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22326 let gqa = (n_head / n_head_kv).max(1) as u32;
22327 let o_len = t * n_head * n_splits_max * head_dim;
22328 let ml_len = t * n_head * n_splits_max;
22329 let mut part_guard = self.fa_part_pool.lock().unwrap();
22330 if part_guard
22331 .as_ref()
22332 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22333 .unwrap_or(true)
22334 {
22335 let old = part_guard.take();
22346 let (co, cm) = old
22347 .as_ref()
22348 .map(|pp| (pp.0.len(), pp.1.len()))
22349 .unwrap_or((0, 0));
22350 if let Some(old) = old {
22351 self.fa_part_retired.lock().unwrap().push(old);
22352 }
22353 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22354 eprintln!(
22355 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22356 co, o_len, cm, ml_len
22357 );
22358 }
22359 *part_guard = Some((
22360 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22361 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22362 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22363 ));
22364 }
22365 let pg = part_guard.as_mut().unwrap();
22366 self.gpu
22367 .stream()
22368 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22369 self.gpu
22370 .stream()
22371 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22372 self.gpu
22373 .stream()
22374 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22375 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22376 let f = if g {
22377 self.func_g("fa_decode_vec_q_rows_v4_dc")
22378 } else {
22379 self.func("fa_decode_vec_q_rows_v4_dc")
22380 };
22381 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22382 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22383 f.set_attribute(
22384 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22385 sh as i32,
22386 )?;
22387 let cfg = LaunchConfig {
22388 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22389 block_dim: (32, gqa, 1),
22390 shared_mem_bytes: sh,
22391 };
22392 let __s_b = self.gpu.stream();
22393 let mut b = __s_b.launch_builder(&f);
22394 b.arg(q)
22395 .arg(k)
22396 .arg(v)
22397 .arg(&mut *part_o)
22398 .arg(&mut *part_m)
22399 .arg(&mut *part_l)
22400 .arg(&hd)
22401 .arg(&nh)
22402 .arg(&nhkv)
22403 .arg(base_dev)
22404 .arg(&base_plus)
22405 .arg(&scale)
22406 .arg(&nspm)
22407 .arg(&spk)
22408 .arg(&ktb)
22409 .arg(&vtb);
22410 unsafe {
22411 b.launch(cfg)?;
22412 }
22413 let fc = self.func("fa_decode_combine_rows_dc");
22414 let cfg2 = LaunchConfig {
22415 grid_dim: (n_head as u32, t as u32, 1),
22416 block_dim: (head_dim as u32, 1, 1),
22417 shared_mem_bytes: 0,
22418 };
22419 let __s_b2 = self.gpu.stream();
22420 let mut b2 = __s_b2.launch_builder(&fc);
22421 b2.arg(&*part_o)
22422 .arg(&*part_m)
22423 .arg(&*part_l)
22424 .arg(o)
22425 .arg(&hd)
22426 .arg(&nh)
22427 .arg(base_dev)
22428 .arg(&base_plus)
22429 .arg(&nspm)
22430 .arg(&spk);
22431 unsafe {
22432 b2.launch(cfg2)?;
22433 }
22434 return Ok(());
22435 }
22436 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22437 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22438 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22439 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22440 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22441 let gqa = (n_head / n_head_kv).max(1) as u32;
22442 let o_len = t * n_head * n_splits_max * head_dim;
22443 let ml_len = t * n_head * n_splits_max;
22444 let mut part_guard = self.fa_part_pool.lock().unwrap();
22445 if part_guard
22446 .as_ref()
22447 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22448 .unwrap_or(true)
22449 {
22450 let old = part_guard.take();
22461 let (co, cm) = old
22462 .as_ref()
22463 .map(|pp| (pp.0.len(), pp.1.len()))
22464 .unwrap_or((0, 0));
22465 if let Some(old) = old {
22466 self.fa_part_retired.lock().unwrap().push(old);
22467 }
22468 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22469 eprintln!(
22470 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22471 co, o_len, cm, ml_len
22472 );
22473 }
22474 *part_guard = Some((
22475 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22476 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22477 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22478 ));
22479 }
22480 let pg = part_guard.as_mut().unwrap();
22481 self.gpu
22482 .stream()
22483 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22484 self.gpu
22485 .stream()
22486 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22487 self.gpu
22488 .stream()
22489 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22490 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22491 let f = self.func("fa_decode_vec_q_rows_v3_dc");
22492 let sh = (32 * head_dim * 2) as u32;
22493 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22494 f.set_attribute(
22495 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22496 sh as i32,
22497 )?;
22498 let cfg = LaunchConfig {
22499 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22500 block_dim: (32, gqa, 1),
22501 shared_mem_bytes: sh,
22502 };
22503 let __s_b = self.gpu.stream();
22504 let mut b = __s_b.launch_builder(&f);
22505 b.arg(q)
22506 .arg(k)
22507 .arg(v)
22508 .arg(&mut *part_o)
22509 .arg(&mut *part_m)
22510 .arg(&mut *part_l)
22511 .arg(&hd)
22512 .arg(&nh)
22513 .arg(&nhkv)
22514 .arg(base_dev)
22515 .arg(&scale)
22516 .arg(&nspm)
22517 .arg(&spk)
22518 .arg(&ktb)
22519 .arg(&vtb);
22520 unsafe {
22521 b.launch(cfg)?;
22522 }
22523 let fc = self.func("fa_decode_combine_rows_dc");
22524 let cfg2 = LaunchConfig {
22525 grid_dim: (n_head as u32, t as u32, 1),
22526 block_dim: (head_dim as u32, 1, 1),
22527 shared_mem_bytes: 0,
22528 };
22529 let plus0 = 0i32;
22530 let __s_b2 = self.gpu.stream();
22531 let mut b2 = __s_b2.launch_builder(&fc);
22532 b2.arg(&*part_o)
22533 .arg(&*part_m)
22534 .arg(&*part_l)
22535 .arg(o)
22536 .arg(&hd)
22537 .arg(&nh)
22538 .arg(base_dev)
22539 .arg(&plus0)
22540 .arg(&nspm)
22541 .arg(&spk);
22542 unsafe {
22543 b2.launch(cfg2)?;
22544 }
22545 Ok(())
22546 }
22547
22548 pub fn fa_decode_dc(
22559 &self,
22560 q: &CudaSlice<f32>,
22561 k: &cudarc::driver::CudaView<u8>,
22562 v: &cudarc::driver::CudaView<u8>,
22563 o: &mut CudaSlice<f32>,
22564 head_dim: usize,
22565 n_head: usize,
22566 n_head_kv: usize,
22567 t_kv_dev: &CudaSlice<i32>,
22568 bucket_max: usize,
22569 scale: f32,
22570 k_tok_bytes: usize,
22571 v_tok_bytes: usize,
22572 g: bool,
22573 ) -> Result<(), Box<dyn std::error::Error>> {
22574 self.fa_decode_dc_q8(
22575 q,
22576 k,
22577 v,
22578 o,
22579 head_dim,
22580 n_head,
22581 n_head_kv,
22582 t_kv_dev,
22583 bucket_max,
22584 scale,
22585 k_tok_bytes,
22586 v_tok_bytes,
22587 g,
22588 None,
22589 )
22590 }
22591
22592 #[allow(clippy::too_many_arguments)]
22595 pub fn fa_decode_dc_q8(
22596 &self,
22597 q: &CudaSlice<f32>,
22598 k: &cudarc::driver::CudaView<u8>,
22599 v: &cudarc::driver::CudaView<u8>,
22600 o: &mut CudaSlice<f32>,
22601 head_dim: usize,
22602 n_head: usize,
22603 n_head_kv: usize,
22604 t_kv_dev: &CudaSlice<i32>,
22605 bucket_max: usize,
22606 scale: f32,
22607 k_tok_bytes: usize,
22608 v_tok_bytes: usize,
22609 g: bool,
22610 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22611 ) -> Result<(), Box<dyn std::error::Error>> {
22612 let mut fa_vec =
22620 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22621 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22622 fa_vec = false;
22623 } let sp = fa_split_keys(bucket_max, n_head_kv);
22625 let n_splits = if fa_vec {
22626 ((bucket_max + sp - 1) / sp).max(1)
22627 } else {
22628 ((bucket_max + 255) / 256).max(1)
22629 };
22630 let o_len = n_head * n_splits * head_dim;
22631 let ml_len = n_head * n_splits;
22632 let mut part_guard = self.fa_part_pool.lock().unwrap();
22633 if part_guard
22634 .as_ref()
22635 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22636 .unwrap_or(true)
22637 {
22638 let old = part_guard.take();
22649 let (co, cm) = old
22650 .as_ref()
22651 .map(|pp| (pp.0.len(), pp.1.len()))
22652 .unwrap_or((0, 0));
22653 if let Some(old) = old {
22654 self.fa_part_retired.lock().unwrap().push(old);
22655 }
22656 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22657 eprintln!(
22658 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22659 co, o_len, cm, ml_len
22660 );
22661 }
22662 *part_guard = Some((
22663 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22664 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22665 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22666 ));
22667 }
22668 let pg = part_guard.as_mut().unwrap();
22669 self.gpu
22670 .stream()
22671 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22672 self.gpu
22673 .stream()
22674 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22675 self.gpu
22676 .stream()
22677 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22678 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22679 let (hd, nh, nhkv, nsp) = (
22680 head_dim as i32,
22681 n_head as i32,
22682 n_head_kv as i32,
22683 n_splits as i32,
22684 );
22685 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22686 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22687 let deep = fa_vec
22690 && head_dim == 256
22691 && fa_v4_at(bucket_max)
22692 && !g
22693 && fa_deep_at(bucket_max)
22694 && !matches!(fa_v4_mode(), "noB3" | "stage");
22695 let (f, cfg) = if fa_vec
22696 && head_dim == 512
22697 && bucket_max >= {
22698 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22699 *FA512_MIN_DC.get_or_init(|| {
22700 std::env::var("MEMRA_FA512_MIN")
22701 .ok()
22702 .and_then(|v| v.parse().ok())
22703 .unwrap_or(512)
22704 })
22705 } {
22706 let gqa = (n_head / n_head_kv).max(1) as u32;
22708 (
22709 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22710 LaunchConfig {
22711 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22712 block_dim: (32, gqa, 1),
22713 shared_mem_bytes: 0,
22714 },
22715 )
22716 } else if fa_vec && head_dim == 512 {
22717 let q_view = q.as_view();
22720 let mut o_view = o.as_view_mut();
22721 return self.fa_decode_scalar_unified(
22722 &q_view,
22723 k,
22724 v,
22725 &mut o_view,
22726 head_dim,
22727 n_head,
22728 n_head_kv,
22729 0,
22730 Some(t_kv_dev),
22731 scale,
22732 n_splits,
22733 sp,
22734 k_tok_bytes,
22735 v_tok_bytes,
22736 g,
22737 &mut *part_o,
22738 &mut *part_m,
22739 &mut *part_l,
22740 q8_out,
22741 );
22742 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22743 let gqa = (n_head / n_head_kv).max(1) as u32;
22746 let fv = if g {
22747 self.func_g("fa_decode_vec_q_v4_dc")
22748 } else if deep {
22749 self.func("fa_decode_vec_q_v4_deep_dc")
22750 } else {
22751 self.func("fa_decode_vec_q_v4_dc")
22752 };
22753 let shmem =
22754 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22755 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22756 fv.set_attribute(
22757 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22758 shmem as i32,
22759 )?;
22760 (
22761 fv,
22762 LaunchConfig {
22763 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22764 block_dim: (32, gqa, 1),
22765 shared_mem_bytes: shmem,
22766 },
22767 )
22768 } else if fa_vec && fa_v3_active(head_dim) {
22769 let gqa = (n_head / n_head_kv).max(1) as u32;
22772 let fv = if g {
22773 self.func_g("fa_decode_vec_q_v3_dc")
22774 } else {
22775 self.func("fa_decode_vec_q_v3_dc")
22776 };
22777 let shmem = (32 * head_dim * 2) as u32; (
22779 fv,
22780 LaunchConfig {
22781 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22782 block_dim: (32, gqa, 1),
22783 shared_mem_bytes: shmem,
22784 },
22785 )
22786 } else if fa_vec && fa_v2_on() {
22787 let gqa = (n_head / n_head_kv).max(1) as u32;
22791 let fv = if g {
22792 self.func_g("fa_decode_vec_q_v2_dc")
22793 } else {
22794 self.func("fa_decode_vec_q_v2_dc")
22795 };
22796 let shmem = (2 * 32 * head_dim * 2) as u32; (
22798 fv,
22799 LaunchConfig {
22800 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22801 block_dim: (32, gqa, 1),
22802 shared_mem_bytes: shmem,
22803 },
22804 )
22805 } else if fa_vec {
22806 let gqa = (n_head / n_head_kv).max(1) as u32;
22807 let fv = if g {
22809 self.func_g("fa_decode_vec_q_dc")
22810 } else {
22811 self.func("fa_decode_vec_q_dc")
22812 };
22813 (
22814 fv,
22815 LaunchConfig {
22816 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22817 block_dim: (32, gqa, 1),
22818 shared_mem_bytes: 0,
22819 },
22820 )
22821 } else {
22822 let q_view = q.as_view();
22823 let mut o_view = o.as_view_mut();
22824 return self.fa_decode_scalar_unified(
22825 &q_view,
22826 k,
22827 v,
22828 &mut o_view,
22829 head_dim,
22830 n_head,
22831 n_head_kv,
22832 0,
22833 Some(t_kv_dev),
22834 scale,
22835 n_splits,
22836 if fa_vec { sp } else { 256 },
22837 k_tok_bytes,
22838 v_tok_bytes,
22839 g,
22840 &mut *part_o,
22841 &mut *part_m,
22842 &mut *part_l,
22843 q8_out,
22844 );
22845 };
22846 let ski = sp as i32; let __s_b = self.gpu.stream();
22848 let mut b = __s_b.launch_builder(&f);
22849 b.arg(q)
22850 .arg(k)
22851 .arg(v)
22852 .arg(&mut *part_o)
22853 .arg(&mut *part_m)
22854 .arg(&mut *part_l)
22855 .arg(&hd)
22856 .arg(&nh)
22857 .arg(&nhkv)
22858 .arg(t_kv_dev)
22859 .arg(&scale)
22860 .arg(&nsp)
22861 .arg(&ski)
22862 .arg(&ktb)
22863 .arg(&vtb);
22864 unsafe {
22865 b.launch(cfg)?;
22866 }
22867 let cfg2 = LaunchConfig {
22868 grid_dim: (n_head as u32, 1, 1),
22869 block_dim: (head_dim as u32, 1, 1),
22870 shared_mem_bytes: 0,
22871 };
22872 if let Some((oq, od)) = q8_out {
22873 let fc = if g {
22874 self.func_g("fa_decode_combine_q8_1")
22875 } else {
22876 self.fa_func("fa_decode_combine_q8_1", head_dim)
22877 };
22878 let __s_b2 = self.gpu.stream();
22879 let mut b2 = __s_b2.launch_builder(&fc);
22880 b2.arg(&*part_o)
22881 .arg(&*part_m)
22882 .arg(&*part_l)
22883 .arg(oq)
22884 .arg(od)
22885 .arg(&hd)
22886 .arg(&nh)
22887 .arg(&nsp);
22888 unsafe {
22889 b2.launch(cfg2)?;
22890 }
22891 return Ok(());
22892 }
22893 let fc = if g {
22894 self.func_g("fa_decode_combine_f32")
22895 } else {
22896 self.fa_func("fa_decode_combine_f32", head_dim)
22897 };
22898 let __s_b2 = self.gpu.stream();
22899 let mut b2 = __s_b2.launch_builder(&fc);
22900 b2.arg(&*part_o)
22901 .arg(&*part_m)
22902 .arg(&*part_l)
22903 .arg(o)
22904 .arg(&hd)
22905 .arg(&nh)
22906 .arg(&nsp);
22907 unsafe {
22908 b2.launch(cfg2)?;
22909 }
22910 Ok(())
22911 }
22912
22913 #[allow(clippy::too_many_arguments)]
22917 pub fn append_kv_quantized_dcw(
22918 &self,
22919 k_row: &CudaSlice<f32>,
22920 v_row: &CudaSlice<f32>,
22921 kc: &mut CudaSlice<u8>,
22922 vc: &mut CudaSlice<u8>,
22923 len_dev: &CudaSlice<i32>,
22924 base_dev: Option<&CudaSlice<i32>>,
22925 kv_dim_k: usize,
22926 kv_dim_v: usize,
22927 k_tok_bytes: usize,
22928 v_tok_bytes: usize,
22929 ) -> Result<(), Box<dyn std::error::Error>> {
22930 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22931 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22932 let cfg = LaunchConfig {
22933 grid_dim: (nblk, 1, 1),
22934 block_dim: (32, 1, 1),
22935 shared_mem_bytes: 0,
22936 };
22937 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22938 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22939 let null: u64 = 0;
22940 let __s_b = self.gpu.stream();
22941 let mut b = __s_b.launch_builder(&f);
22942 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22943 match base_dev {
22944 Some(base) => {
22945 b.arg(base);
22946 }
22947 None => {
22948 b.arg(&null);
22949 }
22950 }
22951 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22952 unsafe {
22953 b.launch(cfg)?;
22954 }
22955 Ok(())
22956 }
22957
22958 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22960 let f = self.func("inc_i32");
22961 let cfg = LaunchConfig {
22962 grid_dim: (1, 1, 1),
22963 block_dim: (1, 1, 1),
22964 shared_mem_bytes: 0,
22965 };
22966 let __s_b = self.gpu.stream();
22967 let mut b = __s_b.launch_builder(&f);
22968 b.arg(counter);
22969 unsafe {
22970 b.launch(cfg)?;
22971 }
22972 Ok(())
22973 }
22974
22975 #[allow(clippy::too_many_arguments)]
22984 fn fa_part_pool_grow(
22988 &self,
22989 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22990 o_len: usize,
22991 ml_len: usize,
22992 ) -> Result<(), Box<dyn std::error::Error>> {
22993 if part_guard
22994 .as_ref()
22995 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22996 .unwrap_or(true)
22997 {
22998 let old = part_guard.take();
22999 let (co, cm) = old
23000 .as_ref()
23001 .map(|pp| (pp.0.len(), pp.1.len()))
23002 .unwrap_or((0, 0));
23003 if let Some(old) = old {
23004 self.fa_part_retired.lock().unwrap().push(old);
23005 }
23006 *part_guard = Some((
23007 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
23008 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23009 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
23010 ));
23011 }
23012 Ok(())
23013 }
23014
23015 pub fn fa_dcw_pool_ensure(
23018 &self,
23019 head_dim: usize,
23020 n_head: usize,
23021 n_head_kv: usize,
23022 bucket_max: usize,
23023 ) -> Result<(), Box<dyn std::error::Error>> {
23024 let sp = fa_split_keys(bucket_max, n_head_kv);
23025 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23026 let o_len = n_head * n_splits * head_dim;
23027 let ml_len = n_head * n_splits;
23028 let mut part_guard = self.fa_part_pool.lock().unwrap();
23029 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
23030 }
23031
23032 #[allow(clippy::too_many_arguments)]
23040 pub fn fa_decode_dcw2(
23041 &self,
23042 q2: &CudaSlice<f32>,
23043 k_ring: &cudarc::driver::CudaView<u8>,
23044 v_ring: &cudarc::driver::CudaView<u8>,
23045 o2: &mut CudaSlice<f32>,
23046 head_dim: usize,
23047 n_head: usize,
23048 n_head_kv: usize,
23049 len_dev: &CudaSlice<i32>,
23050 base_dev: Option<&CudaSlice<i32>>,
23051 window: usize,
23052 bucket_max: usize,
23053 scale: f32,
23054 k_tok_bytes: usize,
23055 v_tok_bytes: usize,
23056 gate2: &CudaSlice<f32>,
23057 ) -> Result<(), Box<dyn std::error::Error>> {
23058 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23059 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23060 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
23061 }
23062 let sp = fa_split_keys(bucket_max, n_head_kv);
23063 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23064 let o_len = 2 * n_head * n_splits * head_dim;
23066 let ml_len = 2 * n_head * n_splits;
23067 let mut part_guard = self.fa_part_pool.lock().unwrap();
23068 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23069 let pg = part_guard.as_mut().unwrap();
23070 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23071 let (hd, nh, nhkv, nsp) = (
23072 head_dim as i32,
23073 n_head as i32,
23074 n_head_kv as i32,
23075 n_splits as i32,
23076 );
23077 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23078 let (ski, win) = (sp as i32, window as i32);
23079 let gqa = (n_head / n_head_kv).max(1) as u32;
23080 let smem = (32 * head_dim * 2) as u32;
23081 let f = self.func("fa_decode_vec_q_v3_dcw2");
23082 let cfg = LaunchConfig {
23083 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
23084 block_dim: (32, gqa, 1),
23085 shared_mem_bytes: smem,
23086 };
23087 let null: u64 = 0;
23088 {
23089 let __s_b = self.gpu.stream();
23090 let mut b = __s_b.launch_builder(&f);
23091 b.arg(q2)
23092 .arg(k_ring)
23093 .arg(v_ring)
23094 .arg(&mut *part_o)
23095 .arg(&mut *part_m)
23096 .arg(&mut *part_l)
23097 .arg(&hd)
23098 .arg(&nh)
23099 .arg(&nhkv)
23100 .arg(len_dev);
23101 match base_dev {
23102 Some(base) => {
23103 b.arg(base);
23104 }
23105 None => {
23106 b.arg(&null);
23107 }
23108 }
23109 b.arg(&win)
23110 .arg(&scale)
23111 .arg(&nsp)
23112 .arg(&ski)
23113 .arg(&ktb)
23114 .arg(&vtb);
23115 unsafe {
23116 b.launch(cfg)?;
23117 }
23118 }
23119 let fc = {
23123 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23124 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23125 self.func("fa_decode_combine_gate_f32_s")
23126 } else {
23127 self.func("fa_decode_combine_gate_f32")
23128 }
23129 };
23130 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
23131 let nh2 = (2 * n_head) as i32;
23132 let cfg2 = LaunchConfig {
23133 grid_dim: ((2 * n_head) as u32, 1, 1),
23134 block_dim: (head_dim as u32, 1, 1),
23135 shared_mem_bytes: if combine_shared {
23136 (2 * n_splits * 4) as u32
23137 } else {
23138 0
23139 },
23140 };
23141 let __s_b2 = self.gpu.stream();
23142 let mut b2 = __s_b2.launch_builder(&fc);
23143 b2.arg(&*part_o)
23144 .arg(&*part_m)
23145 .arg(&*part_l)
23146 .arg(gate2)
23147 .arg(o2)
23148 .arg(&hd)
23149 .arg(&nh2)
23150 .arg(&nsp);
23151 unsafe {
23152 b2.launch(cfg2)?;
23153 }
23154 Ok(())
23155 }
23156
23157 pub fn fa_decode_dcw(
23158 &self,
23159 q: &CudaSlice<f32>,
23160 k_ring: &cudarc::driver::CudaView<u8>,
23161 v_ring: &cudarc::driver::CudaView<u8>,
23162 o: &mut CudaSlice<f32>,
23163 head_dim: usize,
23164 n_head: usize,
23165 n_head_kv: usize,
23166 len_dev: &CudaSlice<i32>,
23167 base_dev: Option<&CudaSlice<i32>>,
23168 window: usize,
23169 bucket_max: usize,
23170 scale: f32,
23171 k_tok_bytes: usize,
23172 v_tok_bytes: usize,
23173 fused_gate: Option<&CudaSlice<f32>>,
23177 ) -> Result<(), Box<dyn std::error::Error>> {
23178 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
23179 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
23180 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"
23181 .into());
23182 }
23183 let sp = fa_split_keys(bucket_max, n_head_kv);
23184 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
23185 let o_len = n_head * n_splits * head_dim;
23186 let ml_len = n_head * n_splits;
23187 let mut part_guard = self.fa_part_pool.lock().unwrap();
23188 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
23189 let pg = part_guard.as_mut().unwrap();
23190 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23195 let memset_on = *MEMSET_ON
23200 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
23201 || crate::tp::token_graph_building();
23202 if memset_on {
23203 self.gpu
23204 .stream()
23205 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23206 self.gpu
23207 .stream()
23208 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23209 self.gpu
23210 .stream()
23211 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23212 }
23213 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23214 let (hd, nh, nhkv, nsp) = (
23215 head_dim as i32,
23216 n_head as i32,
23217 n_head_kv as i32,
23218 n_splits as i32,
23219 );
23220 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23221 let (ski, win) = (sp as i32, window as i32);
23222 let gqa = (n_head / n_head_kv).max(1) as u32;
23223 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23227 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
23228 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
23229 Ok("2") => 2,
23230 Ok("1") => 1,
23231 _ => 0,
23232 });
23233 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23238 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
23239 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
23240 std::sync::Mutex::new(None);
23241 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23245 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
23246 && (n_head / n_head_kv) % 2 == 0
23247 && (n_head / n_head_kv) >= 2;
23248 let f = if fprof {
23249 self.func("fa_decode_vec_q_v3_dcw_prof")
23250 } else if hs2 {
23251 self.func("fa_decode_vec_q_v3_dcw_hs2")
23252 } else if hoist == 2 {
23253 self.func("fa_decode_vec_q_v3_dcw_hc")
23255 } else if hoist == 1 {
23256 self.func("fa_decode_vec_q_v3_dcw_h")
23258 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
23259 self.func("fa_decode_vec_q_v3_dcw_u8")
23260 } else {
23261 self.func("fa_decode_vec_q_v3_dcw")
23262 };
23263 let cfg = LaunchConfig {
23264 grid_dim: if hs2 {
23265 ((2 * n_head_kv) as u32, n_splits as u32, 1)
23266 } else {
23267 (n_head_kv as u32, n_splits as u32, 1)
23268 },
23269 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
23270 shared_mem_bytes: smem,
23271 };
23272 let null: u64 = 0;
23273 let __s_b = self.gpu.stream();
23274 let mut b = __s_b.launch_builder(&f);
23275 b.arg(q)
23276 .arg(k_ring)
23277 .arg(v_ring)
23278 .arg(&mut *part_o)
23279 .arg(&mut *part_m)
23280 .arg(&mut *part_l)
23281 .arg(&hd)
23282 .arg(&nh)
23283 .arg(&nhkv)
23284 .arg(len_dev);
23285 match base_dev {
23286 Some(base) => {
23287 b.arg(base);
23288 }
23289 None => {
23290 b.arg(&null);
23291 }
23292 }
23293 b.arg(&win)
23294 .arg(&scale)
23295 .arg(&nsp)
23296 .arg(&ski)
23297 .arg(&ktb)
23298 .arg(&vtb);
23299 if fprof {
23300 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
23301 if guard
23302 .as_ref()
23303 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
23304 {
23305 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
23306 }
23307 let (_, buf) = guard.as_mut().expect("armed above");
23308 b.arg(&*buf);
23309 unsafe {
23310 b.launch(cfg)?;
23311 }
23312 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
23313 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
23314 if n % 430 == 0 {
23315 self.stream().synchronize()?;
23316 let h = self.dtoh_u64(buf)?;
23317 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
23318 let tot: u64 = h[..6].iter().sum();
23319 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
23320 for (i, name) in phases.iter().enumerate() {
23321 let pct = if tot > 0 {
23322 h[i] as f64 / tot as f64 * 100.0
23323 } else {
23324 0.0
23325 };
23326 line.push_str(&format!(" {name}={pct:.1}%"));
23327 }
23328 if h[6] > 0 {
23329 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
23330 }
23331 eprintln!("{line}");
23332 }
23333 } else {
23334 unsafe {
23335 b.launch(cfg)?;
23336 }
23337 }
23338 let mut combine_shared = false;
23339 let fc = if fused_gate.is_some() {
23340 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23343 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
23344 combine_shared = true;
23345 self.func("fa_decode_combine_gate_f32_s")
23346 } else {
23347 self.func("fa_decode_combine_gate_f32")
23348 }
23349 } else {
23350 self.fa_func("fa_decode_combine_f32", head_dim)
23351 };
23352 let cfg2 = LaunchConfig {
23353 grid_dim: (n_head as u32, 1, 1),
23354 block_dim: (head_dim as u32, 1, 1),
23355 shared_mem_bytes: if combine_shared {
23356 (2 * n_splits * 4) as u32
23357 } else {
23358 0
23359 },
23360 };
23361 let __s_b2 = self.gpu.stream();
23362 let mut b2 = __s_b2.launch_builder(&fc);
23363 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
23364 if let Some(gate_row) = fused_gate {
23365 b2.arg(gate_row);
23366 }
23367 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
23368 unsafe {
23369 b2.launch(cfg2)?;
23370 }
23371 Ok(())
23372 }
23373
23374 pub fn fa_geom_eager(
23380 &self,
23381 t_kv: usize,
23382 head_dim: usize,
23383 n_head_kv: usize,
23384 g: bool,
23385 ) -> (bool, usize) {
23386 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
23390 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
23396 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
23397 if g && head_dim == 256 && !fa_v4_at(t_kv) {
23403 fa_vec = false;
23404 }
23405 let sp = fa_split_keys(t_kv, n_head_kv);
23406 let n_splits = if fa_vec {
23407 ((t_kv + sp - 1) / sp).max(1)
23408 } else {
23409 ((t_kv + 255) / 256).max(1)
23410 };
23411 (fa_vec, n_splits)
23412 }
23413
23414 pub fn fa_bucket_key(
23420 &self,
23421 t_kv: usize,
23422 head_dim: usize,
23423 n_head_kv: usize,
23424 g: bool,
23425 ) -> (bool, usize) {
23426 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
23427 }
23428
23429 pub fn capture_graph_retained<F>(
23441 &self,
23442 step: F,
23443 ) -> Result<
23444 (
23445 cudarc::driver::CudaGraph,
23446 Vec<Box<dyn std::any::Any + Send>>,
23447 ),
23448 Box<dyn std::error::Error>,
23449 >
23450 where
23451 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23452 {
23453 use cudarc::driver::sys::CUgraphInstantiate_flags;
23454 self.capture_graph_retained_flags(
23455 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23456 step,
23457 )
23458 }
23459
23460 pub fn capture_graph_retained_flags<F>(
23465 &self,
23466 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
23467 mut step: F,
23468 ) -> Result<
23469 (
23470 cudarc::driver::CudaGraph,
23471 Vec<Box<dyn std::any::Any + Send>>,
23472 ),
23473 Box<dyn std::error::Error>,
23474 >
23475 where
23476 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23477 {
23478 use cudarc::driver::sys::CUstreamCaptureMode;
23479 self.capture_keep.lock().unwrap().clear();
23487 let was_tracking = self.gpu.ctx.is_event_tracking();
23488 if was_tracking {
23489 unsafe {
23490 self.gpu.ctx.disable_event_tracking();
23491 }
23492 }
23493 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23494 self.capture_keep_on
23495 .store(true, std::sync::atomic::Ordering::Relaxed);
23496 let w = (|| {
23497 step(self)?;
23498 step(self)
23499 })();
23500 self.capture_keep_on
23501 .store(false, std::sync::atomic::Ordering::Relaxed);
23502 w?;
23503 self.gpu.stream().synchronize()?;
23504 self.gpu
23505 .stream()
23506 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23507 let r = step(self);
23508 let g = self.gpu.stream().end_capture(flags);
23509 r?;
23510 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23511 graph.upload()?;
23512 Ok(graph)
23513 };
23514 let result = run();
23515 self.capture_keep_on
23516 .store(false, std::sync::atomic::Ordering::Relaxed);
23517 if was_tracking {
23518 unsafe {
23519 self.gpu.ctx.enable_event_tracking();
23520 }
23521 }
23522 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
23523 Ok((result?, keeper))
23524 }
23525
23526 pub fn capture_graph_retained_nowarm<F>(
23532 &self,
23533 mut step: F,
23534 ) -> Result<
23535 (
23536 cudarc::driver::CudaGraph,
23537 Vec<Box<dyn std::any::Any + Send>>,
23538 ),
23539 Box<dyn std::error::Error>,
23540 >
23541 where
23542 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23543 {
23544 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23545 let was_tracking = self.gpu.ctx.is_event_tracking();
23546 if was_tracking {
23547 unsafe {
23548 self.gpu.ctx.disable_event_tracking();
23549 }
23550 }
23551 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23552 self.gpu.stream().synchronize()?;
23553 self.gpu
23554 .stream()
23555 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23556 let r = step(self);
23557 let g = self.gpu.stream().end_capture(
23558 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23559 );
23560 r?;
23561 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23562 graph.upload()?;
23563 Ok(graph)
23564 };
23565 let result = run();
23566 if was_tracking {
23567 unsafe {
23568 self.gpu.ctx.enable_event_tracking();
23569 }
23570 }
23571 Ok((result?, Vec::new()))
23572 }
23573
23574 pub fn capture_graph<F>(
23575 &self,
23576 mut step: F,
23577 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
23578 where
23579 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
23580 {
23581 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
23582 let was_tracking = self.gpu.ctx.is_event_tracking();
23590 if was_tracking {
23591 unsafe {
23592 self.gpu.ctx.disable_event_tracking();
23593 }
23594 }
23595 let iflag = {
23602 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23603 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23604 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23607 Ok("priority") => {
23608 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23609 }
23610 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23611 })
23612 };
23613 let ct = {
23620 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23621 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23622 };
23623 let warmups = {
23646 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23647 *W.get_or_init(|| {
23648 std::env::var("MEMRA_GRAPH_WARMUPS")
23649 .ok()
23650 .and_then(|v| v.parse().ok())
23651 .filter(|n| *n >= 1)
23652 .unwrap_or(1)
23653 })
23654 };
23655 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23656 let t_w = std::time::Instant::now();
23657 for _ in 0..warmups {
23659 step(self)?;
23660 }
23661 self.gpu.stream().synchronize()?;
23662 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23663 let t_c = std::time::Instant::now();
23665 self.gpu
23666 .stream()
23667 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23668 let r = step(self);
23671 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23672 let t_i = std::time::Instant::now();
23673 let g = self.gpu.stream().end_capture(iflag);
23674 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23675 r?;
23676 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23677 let t_u = std::time::Instant::now();
23678 graph.upload()?;
23679 if ct {
23680 println!(
23681 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
23682 instantiate {ms_inst:.2} ms upload {:.2} ms",
23683 t_u.elapsed().as_secs_f64() * 1e3
23684 );
23685 }
23686 Ok(graph)
23687 };
23688 let result = run();
23689 if was_tracking {
23690 unsafe {
23691 self.gpu.ctx.enable_event_tracking();
23692 }
23693 }
23694 result
23695 }
23696
23697 pub fn gdn_scan_s128_view(
23699 &self,
23700 q: &CudaSlice<f32>,
23701 k: &CudaSlice<f32>,
23702 v: &CudaSlice<f32>,
23703 g: &CudaSlice<f32>,
23704 beta: &CudaSlice<f32>,
23705 state_in: &cudarc::driver::CudaView<f32>,
23706 state_out: &mut cudarc::driver::CudaViewMut<f32>,
23707 o: &mut CudaSlice<f32>,
23708 n_head: usize,
23709 t: usize,
23710 scale: f32,
23711 ) -> Result<(), Box<dyn std::error::Error>> {
23712 let f = self.func("gdn_scan_s128");
23713 const S_V: u32 = 128;
23714 const WARP: u32 = 32;
23715 const COLS: u32 = 4;
23716 let cfg = LaunchConfig {
23717 grid_dim: (n_head as u32, 1, S_V / COLS),
23718 block_dim: (WARP, COLS, 1),
23719 shared_mem_bytes: 0,
23720 };
23721 let (h, ti) = (n_head as i32, t as i32);
23722 let __s_b = self.gpu.stream();
23723 let mut b = __s_b.launch_builder(&f);
23724 b.arg(q)
23725 .arg(k)
23726 .arg(v)
23727 .arg(g)
23728 .arg(beta)
23729 .arg(state_in)
23730 .arg(state_out)
23731 .arg(o)
23732 .arg(&h)
23733 .arg(&ti)
23734 .arg(&scale);
23735 unsafe {
23736 b.launch(cfg)?;
23737 }
23738 Ok(())
23739 }
23740
23741 pub fn ssm_conv1d_view(
23743 &self,
23744 x: &cudarc::driver::CudaView<f32>,
23745 w: &CudaSlice<f32>,
23746 y: &mut CudaSlice<f32>,
23747 conv_dim: usize,
23748 t: usize,
23749 d_conv: usize,
23750 silu: bool,
23751 ) -> Result<(), Box<dyn std::error::Error>> {
23752 let f = self.func("ssm_conv1d_silu_f32");
23753 let cfg = LaunchConfig {
23755 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23756 block_dim: (256, 1, 1),
23757 shared_mem_bytes: 0,
23758 };
23759 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23760 let __s_b = self.gpu.stream();
23761 let mut b = __s_b.launch_builder(&f);
23762 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23763 unsafe {
23764 b.launch(cfg)?;
23765 }
23766 Ok(())
23767 }
23768
23769 pub fn ssm_conv1d_tm(
23776 &self,
23777 qkv_tm: &CudaSlice<f32>,
23778 w: &CudaSlice<f32>,
23779 y: &mut CudaSlice<f32>,
23780 conv_dim: usize,
23781 t: usize,
23782 d_conv: usize,
23783 ) -> Result<(), Box<dyn std::error::Error>> {
23784 let f = self.func("ssm_conv1d_tm_f32");
23785 let cfg = LaunchConfig {
23786 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23787 block_dim: (256, 1, 1),
23788 shared_mem_bytes: 0,
23789 };
23790 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23791 let __s_b = self.gpu.stream();
23792 let mut b = __s_b.launch_builder(&f);
23793 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23794 unsafe {
23795 b.launch(cfg)?;
23796 }
23797 Ok(())
23798 }
23799
23800 pub fn ssm_conv1d_tm_state(
23808 &self,
23809 qkv_tm: &CudaSlice<f32>,
23810 conv_state: &mut CudaSlice<f32>,
23811 w: &CudaSlice<f32>,
23812 y: &mut CudaSlice<f32>,
23813 conv_dim: usize,
23814 t: usize,
23815 d_conv: usize,
23816 ) -> Result<(), Box<dyn std::error::Error>> {
23817 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23818 }
23819
23820 #[allow(clippy::too_many_arguments)]
23823 pub fn ssm_conv1d_tm_state_pad(
23824 &self,
23825 qkv_tm: &CudaSlice<f32>,
23826 conv_state: &mut CudaSlice<f32>,
23827 w: &CudaSlice<f32>,
23828 y: &mut CudaSlice<f32>,
23829 conv_dim: usize,
23830 t: usize,
23831 d_conv: usize,
23832 pad_len: Option<&CudaSlice<i32>>,
23833 ) -> Result<(), Box<dyn std::error::Error>> {
23834 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23835 let ring_old = if t < d_conv - 1 {
23839 Some(self.clone_dtod(conv_state)?)
23840 } else {
23841 None
23842 };
23843 {
23844 let f = self.func("ssm_conv1d_tm_state_f32");
23845 let cfg = LaunchConfig {
23846 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23847 block_dim: (256, 1, 1),
23848 shared_mem_bytes: 0,
23849 };
23850 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23851 let __s_b = self.gpu.stream();
23852 let mut b = __s_b.launch_builder(&f);
23853 b.arg(qkv_tm)
23854 .arg(&*conv_state)
23855 .arg(w)
23856 .arg(y)
23857 .arg(&cd)
23858 .arg(&ti)
23859 .arg(&dc);
23860 unsafe {
23861 b.launch(cfg)?;
23862 }
23863 }
23864 match (ring_old, pad_len) {
23865 (None, Some(len_d)) => {
23866 let f = self.func("ssm_conv_ring_update_dev_f32");
23867 let n = conv_dim * (d_conv - 1);
23868 let cfg = LaunchConfig::for_num_elems(n as u32);
23869 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23870 let __s_b = self.gpu.stream();
23871 let mut b = __s_b.launch_builder(&f);
23872 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23873 unsafe {
23874 b.launch(cfg)?;
23875 }
23876 }
23877 (None, None) => {
23878 let f = self.func("ssm_conv_ring_update_f32");
23879 let n = conv_dim * (d_conv - 1);
23880 let cfg = LaunchConfig::for_num_elems(n as u32);
23881 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23882 let __s_b = self.gpu.stream();
23883 let mut b = __s_b.launch_builder(&f);
23884 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23885 unsafe {
23886 b.launch(cfg)?;
23887 }
23888 }
23889 (Some(old), _) => {
23890 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23891 }
23892 }
23893 Ok(())
23894 }
23895
23896 pub fn ssm_conv1d_tm_state_pad_v(
23898 &self,
23899 qkv_tm: &cudarc::driver::CudaView<f32>,
23900 conv_state: &mut CudaSlice<f32>,
23901 w: &CudaSlice<f32>,
23902 y: &mut CudaSlice<f32>,
23903 conv_dim: usize,
23904 t: usize,
23905 d_conv: usize,
23906 pad_len: Option<&CudaSlice<i32>>,
23907 ) -> Result<(), Box<dyn std::error::Error>> {
23908 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23909 let ring_old = if t < d_conv - 1 {
23913 Some(self.clone_dtod(conv_state)?)
23914 } else {
23915 None
23916 };
23917 {
23918 let f = self.func("ssm_conv1d_tm_state_f32");
23919 let cfg = LaunchConfig {
23920 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23921 block_dim: (256, 1, 1),
23922 shared_mem_bytes: 0,
23923 };
23924 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23925 let __s_b = self.gpu.stream();
23926 let mut b = __s_b.launch_builder(&f);
23927 b.arg(qkv_tm)
23928 .arg(&*conv_state)
23929 .arg(w)
23930 .arg(y)
23931 .arg(&cd)
23932 .arg(&ti)
23933 .arg(&dc);
23934 unsafe {
23935 b.launch(cfg)?;
23936 }
23937 }
23938 match (ring_old, pad_len) {
23939 (None, Some(len_d)) => {
23940 let f = self.func("ssm_conv_ring_update_dev_f32");
23941 let n = conv_dim * (d_conv - 1);
23942 let cfg = LaunchConfig::for_num_elems(n as u32);
23943 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23944 let __s_b = self.gpu.stream();
23945 let mut b = __s_b.launch_builder(&f);
23946 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23947 unsafe {
23948 b.launch(cfg)?;
23949 }
23950 }
23951 (None, None) => {
23952 let f = self.func("ssm_conv_ring_update_f32");
23953 let n = conv_dim * (d_conv - 1);
23954 let cfg = LaunchConfig::for_num_elems(n as u32);
23955 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23956 let __s_b = self.gpu.stream();
23957 let mut b = __s_b.launch_builder(&f);
23958 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23959 unsafe {
23960 b.launch(cfg)?;
23961 }
23962 }
23963 (Some(_), _) => unreachable!(
23964 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23965 ),
23966 }
23967 Ok(())
23968 }
23969
23970 pub fn ssm_conv_ring_rebuild(
23975 &self,
23976 qkv_tm: &CudaSlice<f32>,
23977 ring_old: &CudaSlice<f32>,
23978 conv_state: &mut CudaSlice<f32>,
23979 conv_dim: usize,
23980 tc: usize,
23981 d_conv: usize,
23982 ) -> Result<(), Box<dyn std::error::Error>> {
23983 let f = self.func("ssm_conv_ring_rebuild_f32");
23984 let n = conv_dim * (d_conv - 1);
23985 let cfg = LaunchConfig::for_num_elems(n as u32);
23986 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23987 let __s_b = self.gpu.stream();
23988 let mut b = __s_b.launch_builder(&f);
23989 b.arg(qkv_tm)
23990 .arg(ring_old)
23991 .arg(conv_state)
23992 .arg(&cd)
23993 .arg(&ti)
23994 .arg(&dc);
23995 unsafe {
23996 b.launch(cfg)?;
23997 }
23998 Ok(())
23999 }
24000
24001 #[allow(clippy::too_many_arguments)]
24006 pub fn gdn_prep_decode(
24007 &self,
24008 conv_out: &CudaSlice<f32>,
24009 beta_raw: &CudaSlice<f32>,
24010 alpha: &CudaSlice<f32>,
24011 dt_bias: &CudaSlice<f32>,
24012 a: &CudaSlice<f32>,
24013 q_l2: &mut CudaSlice<f32>,
24014 k_l2: &mut CudaSlice<f32>,
24015 v_g: &mut CudaSlice<f32>,
24016 beta: &mut CudaSlice<f32>,
24017 g_log: &mut CudaSlice<f32>,
24018 d_state: usize,
24019 num_v: usize,
24020 num_k: usize,
24021 key_dim: usize,
24022 eps: f32,
24023 ) -> Result<(), Box<dyn std::error::Error>> {
24024 let f = self.func("gdn_prep_decode_f32");
24025 let cfg = LaunchConfig {
24026 grid_dim: (num_v as u32, 1, 1),
24027 block_dim: (32, 4, 1),
24028 shared_mem_bytes: 0,
24029 };
24030 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24031 let __s_b = self.gpu.stream();
24032 let mut b = __s_b.launch_builder(&f);
24033 b.arg(conv_out)
24034 .arg(beta_raw)
24035 .arg(alpha)
24036 .arg(dt_bias)
24037 .arg(a)
24038 .arg(q_l2)
24039 .arg(k_l2)
24040 .arg(v_g)
24041 .arg(beta)
24042 .arg(g_log)
24043 .arg(&ds)
24044 .arg(&nv)
24045 .arg(&nk)
24046 .arg(&kd)
24047 .arg(&eps);
24048 unsafe {
24049 b.launch(cfg)?;
24050 }
24051 Ok(())
24052 }
24053
24054 #[allow(clippy::too_many_arguments)]
24058 pub fn ssm_conv1d_gdn(
24059 &self,
24060 qkv_tm: &CudaSlice<f32>,
24061 w: &CudaSlice<f32>,
24062 q_g: &mut CudaSlice<f32>,
24063 k_g: &mut CudaSlice<f32>,
24064 v_g: &mut CudaSlice<f32>,
24065 conv_dim: usize,
24066 t: usize,
24067 d_conv: usize,
24068 d_state: usize,
24069 num_v: usize,
24070 num_k: usize,
24071 key_dim: usize,
24072 ) -> Result<(), Box<dyn std::error::Error>> {
24073 let f = self.func("ssm_conv1d_gdn_f32");
24074 let cfg = LaunchConfig {
24075 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24076 block_dim: (256, 1, 1),
24077 shared_mem_bytes: 0,
24078 };
24079 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24080 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24081 let __s_b = self.gpu.stream();
24082 let mut b = __s_b.launch_builder(&f);
24083 b.arg(qkv_tm)
24084 .arg(w)
24085 .arg(q_g)
24086 .arg(k_g)
24087 .arg(v_g)
24088 .arg(&cd)
24089 .arg(&ti)
24090 .arg(&dc)
24091 .arg(&ds)
24092 .arg(&nv)
24093 .arg(&nk)
24094 .arg(&kd);
24095 unsafe {
24096 b.launch(cfg)?;
24097 }
24098 Ok(())
24099 }
24100
24101 pub fn ssm_conv1d(
24102 &self,
24103 x: &CudaSlice<f32>,
24104 w: &CudaSlice<f32>,
24105 y: &mut CudaSlice<f32>,
24106 conv_dim: usize,
24107 t: usize,
24108 d_conv: usize,
24109 silu: bool,
24110 ) -> Result<(), Box<dyn std::error::Error>> {
24111 let f = self.func("ssm_conv1d_silu_f32");
24112 let cfg = LaunchConfig {
24113 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
24114 block_dim: (256, 1, 1),
24115 shared_mem_bytes: 0,
24116 };
24117 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
24118 let __s_b = self.gpu.stream();
24119 let mut b = __s_b.launch_builder(&f);
24120 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
24121 unsafe {
24122 b.launch(cfg)?;
24123 }
24124 Ok(())
24125 }
24126
24127 pub fn gdn_scan_s128(
24130 &self,
24131 q: &CudaSlice<f32>,
24132 k: &CudaSlice<f32>,
24133 v: &CudaSlice<f32>,
24134 g: &CudaSlice<f32>,
24135 beta: &CudaSlice<f32>,
24136 state_in: &CudaSlice<f32>,
24137 state_out: &mut CudaSlice<f32>,
24138 o: &mut CudaSlice<f32>,
24139 n_head: usize,
24140 t: usize,
24141 scale: f32,
24142 ) -> Result<(), Box<dyn std::error::Error>> {
24143 let f = self.func("gdn_scan_s128");
24144 const S_V: u32 = 128;
24145 const WARP: u32 = 32;
24146 const COLS_PER_BLOCK: u32 = 4;
24147 let cfg = LaunchConfig {
24148 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
24149 block_dim: (WARP, COLS_PER_BLOCK, 1),
24150 shared_mem_bytes: 0,
24151 };
24152 let (h, ti) = (n_head as i32, t as i32);
24153 let __s_b = self.gpu.stream();
24154 let mut b = __s_b.launch_builder(&f);
24155 b.arg(q)
24156 .arg(k)
24157 .arg(v)
24158 .arg(g)
24159 .arg(beta)
24160 .arg(state_in)
24161 .arg(state_out)
24162 .arg(o)
24163 .arg(&h)
24164 .arg(&ti)
24165 .arg(&scale);
24166 unsafe {
24167 b.launch(cfg)?;
24168 }
24169 Ok(())
24170 }
24171
24172 #[allow(clippy::too_many_arguments)]
24177 pub fn ssm_conv1d_fused_decode_b(
24178 &self,
24179 qkv_cols: &CudaSlice<f32>,
24180 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24181 w: &CudaSlice<f32>,
24182 conv_outs: &mut CudaSlice<f32>,
24183 conv_dim: usize,
24184 d_conv: usize,
24185 b_n: usize,
24186 ) -> Result<(), Box<dyn std::error::Error>> {
24187 let f = self.func("ssm_conv1d_fused_decode_b_f32");
24188 let cfg = LaunchConfig {
24189 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24190 block_dim: (256, 1, 1),
24191 shared_mem_bytes: 0,
24192 };
24193 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24194 let __s_b = self.gpu.stream();
24195 let mut b = __s_b.launch_builder(&f);
24196 b.arg(qkv_cols)
24197 .arg(conv_state_ptrs)
24198 .arg(w)
24199 .arg(conv_outs)
24200 .arg(&cd)
24201 .arg(&dc);
24202 unsafe {
24203 b.launch(cfg)?;
24204 }
24205 Ok(())
24206 }
24207
24208 #[allow(clippy::too_many_arguments)]
24209 pub fn gdn_prep_decode_b(
24210 &self,
24211 conv_outs: &CudaSlice<f32>,
24212 beta_raws: &CudaSlice<f32>,
24213 alphas: &CudaSlice<f32>,
24214 dt_bias: &CudaSlice<f32>,
24215 a: &CudaSlice<f32>,
24216 q_l2: &mut CudaSlice<f32>,
24217 k_l2: &mut CudaSlice<f32>,
24218 v_g: &mut CudaSlice<f32>,
24219 beta: &mut CudaSlice<f32>,
24220 g_log: &mut CudaSlice<f32>,
24221 d_state: usize,
24222 num_v: usize,
24223 num_k: usize,
24224 key_dim: usize,
24225 eps: f32,
24226 conv_dim: usize,
24227 b_n: usize,
24228 ) -> Result<(), Box<dyn std::error::Error>> {
24229 let f = self.func("gdn_prep_decode_b_f32");
24230 let cfg = LaunchConfig {
24231 grid_dim: (num_v as u32, 1, b_n as u32),
24232 block_dim: (32, 4, 1),
24233 shared_mem_bytes: 0,
24234 };
24235 let (ds, nv, nk, kd, cd) = (
24236 d_state as i32,
24237 num_v as i32,
24238 num_k as i32,
24239 key_dim as i32,
24240 conv_dim as i32,
24241 );
24242 let __s_b = self.gpu.stream();
24243 let mut b = __s_b.launch_builder(&f);
24244 b.arg(conv_outs)
24245 .arg(beta_raws)
24246 .arg(alphas)
24247 .arg(dt_bias)
24248 .arg(a)
24249 .arg(q_l2)
24250 .arg(k_l2)
24251 .arg(v_g)
24252 .arg(beta)
24253 .arg(g_log)
24254 .arg(&ds)
24255 .arg(&nv)
24256 .arg(&nk)
24257 .arg(&kd)
24258 .arg(&eps)
24259 .arg(&cd);
24260 unsafe {
24261 b.launch(cfg)?;
24262 }
24263 Ok(())
24264 }
24265
24266 #[allow(clippy::too_many_arguments)]
24267 pub fn gdn_scan_s128_batched(
24268 &self,
24269 q: &CudaSlice<f32>,
24270 k: &CudaSlice<f32>,
24271 v: &CudaSlice<f32>,
24272 g: &CudaSlice<f32>,
24273 beta: &CudaSlice<f32>,
24274 state_in_ptrs: &cudarc::driver::CudaView<u64>,
24275 state_out_ptrs: &cudarc::driver::CudaView<u64>,
24276 o: &mut CudaSlice<f32>,
24277 n_head: usize,
24278 b_n: usize,
24279 scale: f32,
24280 ) -> Result<(), Box<dyn std::error::Error>> {
24281 let f = self.func("gdn_scan_s128_b");
24282 const S_V: u32 = 128;
24283 const WARP: u32 = 32;
24284 const COLS_PER_BLOCK: u32 = 4;
24285 let cfg = LaunchConfig {
24286 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24287 block_dim: (WARP, COLS_PER_BLOCK, 1),
24288 shared_mem_bytes: 0,
24289 };
24290 let h = n_head as i32;
24291 let __s_b = self.gpu.stream();
24292 let mut b = __s_b.launch_builder(&f);
24293 b.arg(q)
24294 .arg(k)
24295 .arg(v)
24296 .arg(g)
24297 .arg(beta)
24298 .arg(state_in_ptrs)
24299 .arg(state_out_ptrs)
24300 .arg(o)
24301 .arg(&h)
24302 .arg(&scale);
24303 unsafe {
24304 b.launch(cfg)?;
24305 }
24306 Ok(())
24307 }
24308
24309 #[allow(clippy::too_many_arguments)]
24315 pub fn ssm_conv1d_fused_decode_b_view(
24316 &self,
24317 qkv_cols: &cudarc::driver::CudaView<f32>,
24318 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
24319 w: &CudaSlice<f32>,
24320 conv_outs: &mut CudaSlice<f32>,
24321 conv_dim: usize,
24322 d_conv: usize,
24323 b_n: usize,
24324 ) -> Result<(), Box<dyn std::error::Error>> {
24325 let f = self.func("ssm_conv1d_fused_decode_b_f32");
24326 let cfg = LaunchConfig {
24327 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
24328 block_dim: (256, 1, 1),
24329 shared_mem_bytes: 0,
24330 };
24331 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24332 let __s_b = self.gpu.stream();
24333 let mut b = __s_b.launch_builder(&f);
24334 b.arg(qkv_cols)
24335 .arg(conv_state_ptrs)
24336 .arg(w)
24337 .arg(conv_outs)
24338 .arg(&cd)
24339 .arg(&dc);
24340 unsafe {
24341 b.launch(cfg)?;
24342 }
24343 Ok(())
24344 }
24345
24346 #[allow(clippy::too_many_arguments)]
24347 pub fn gdn_prep_decode_b_view(
24348 &self,
24349 conv_outs: &CudaSlice<f32>,
24350 beta_raws: &cudarc::driver::CudaView<f32>,
24351 alphas: &cudarc::driver::CudaView<f32>,
24352 dt_bias: &CudaSlice<f32>,
24353 a: &CudaSlice<f32>,
24354 q_l2: &mut CudaSlice<f32>,
24355 k_l2: &mut CudaSlice<f32>,
24356 v_g: &mut CudaSlice<f32>,
24357 beta: &mut CudaSlice<f32>,
24358 g_log: &mut CudaSlice<f32>,
24359 d_state: usize,
24360 num_v: usize,
24361 num_k: usize,
24362 key_dim: usize,
24363 eps: f32,
24364 conv_dim: usize,
24365 b_n: usize,
24366 ) -> Result<(), Box<dyn std::error::Error>> {
24367 let f = self.func("gdn_prep_decode_b_f32");
24368 let cfg = LaunchConfig {
24369 grid_dim: (num_v as u32, 1, b_n as u32),
24370 block_dim: (32, 4, 1),
24371 shared_mem_bytes: 0,
24372 };
24373 let (ds, nv, nk, kd, cd) = (
24374 d_state as i32,
24375 num_v as i32,
24376 num_k as i32,
24377 key_dim as i32,
24378 conv_dim as i32,
24379 );
24380 let __s_b = self.gpu.stream();
24381 let mut b = __s_b.launch_builder(&f);
24382 b.arg(conv_outs)
24383 .arg(beta_raws)
24384 .arg(alphas)
24385 .arg(dt_bias)
24386 .arg(a)
24387 .arg(q_l2)
24388 .arg(k_l2)
24389 .arg(v_g)
24390 .arg(beta)
24391 .arg(g_log)
24392 .arg(&ds)
24393 .arg(&nv)
24394 .arg(&nk)
24395 .arg(&kd)
24396 .arg(&eps)
24397 .arg(&cd);
24398 unsafe {
24399 b.launch(cfg)?;
24400 }
24401 Ok(())
24402 }
24403
24404 #[allow(clippy::too_many_arguments)]
24405 pub fn gdn_scan_s128_batched_view(
24406 &self,
24407 q: &CudaSlice<f32>,
24408 k: &CudaSlice<f32>,
24409 v: &CudaSlice<f32>,
24410 g: &CudaSlice<f32>,
24411 beta: &CudaSlice<f32>,
24412 state_in_ptrs: &cudarc::driver::CudaView<u64>,
24413 state_out_ptrs: &cudarc::driver::CudaView<u64>,
24414 o: &mut cudarc::driver::CudaViewMut<f32>,
24415 n_head: usize,
24416 b_n: usize,
24417 scale: f32,
24418 ) -> Result<(), Box<dyn std::error::Error>> {
24419 let f = self.func("gdn_scan_s128_b");
24420 const S_V: u32 = 128;
24421 const WARP: u32 = 32;
24422 const COLS_PER_BLOCK: u32 = 4;
24423 let cfg = LaunchConfig {
24424 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
24425 block_dim: (WARP, COLS_PER_BLOCK, 1),
24426 shared_mem_bytes: 0,
24427 };
24428 let h = n_head as i32;
24429 let __s_b = self.gpu.stream();
24430 let mut b = __s_b.launch_builder(&f);
24431 b.arg(q)
24432 .arg(k)
24433 .arg(v)
24434 .arg(g)
24435 .arg(beta)
24436 .arg(state_in_ptrs)
24437 .arg(state_out_ptrs)
24438 .arg(o)
24439 .arg(&h)
24440 .arg(&scale);
24441 unsafe {
24442 b.launch(cfg)?;
24443 }
24444 Ok(())
24445 }
24446
24447 pub fn gdn_chunked_enabled() -> bool {
24456 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24457 *E.get_or_init(|| {
24458 std::env::var("MEMRA_GDN_CHUNKED")
24459 .map(|v| v != "0")
24460 .unwrap_or(true)
24461 })
24462 }
24463
24464 pub fn gdn_chunk_size() -> usize {
24469 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24470 *C.get_or_init(|| {
24471 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
24472 .ok()
24473 .and_then(|v| v.parse().ok())
24474 .unwrap_or(32);
24475 c.clamp(32, 128) / 32 * 32
24476 })
24477 }
24478
24479 #[allow(clippy::too_many_arguments)]
24484 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
24487 #[allow(clippy::too_many_arguments)]
24488 pub fn gdn_chunk_k123(
24489 &self,
24490 q: &CudaSlice<f32>,
24491 k: &CudaSlice<f32>,
24492 v: &CudaSlice<f32>,
24493 g: &CudaSlice<f32>,
24494 beta: &CudaSlice<f32>,
24495 wb16: Option<&mut CudaSlice<u8>>,
24496 n_head: usize,
24497 t: usize,
24498 c: usize,
24499 hk: usize,
24500 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
24501 ) -> Result<
24502 (
24503 CudaSlice<f32>,
24504 CudaSlice<f32>,
24505 CudaSlice<f32>,
24506 CudaSlice<f32>,
24507 ),
24508 Box<dyn std::error::Error>,
24509 > {
24510 const D: usize = 128;
24511 let h = n_head;
24512 let nc = (t + c - 1) / c;
24513 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24514 let mut gcum = self.uninit(t * h)?;
24515 let mut a = self.uninit(nc * h * c * c)?;
24516 let mut p = self.uninit(nc * h * c * c)?;
24517 let mut u = self.uninit(nc * h * c * D)?;
24518 let mut w = self.uninit(nc * h * c * D)?;
24519 {
24520 let f = self.func("gdn_chunk_cumgate_f32");
24522 let cfg = LaunchConfig {
24523 grid_dim: (nc as u32, h as u32, 1),
24524 block_dim: (32, 1, 1),
24525 shared_mem_bytes: 0,
24526 };
24527 let __s_b = self.gpu.stream();
24528 let mut b = __s_b.launch_builder(&f);
24529 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
24530 unsafe {
24531 b.launch(cfg)?;
24532 }
24533 }
24534 if let Some((qb, kb, pb)) = k2w {
24535 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
24538 let f = self.func("gdn_k2_wgmma");
24539 let cfg = LaunchConfig {
24540 grid_dim: (nc as u32, h as u32, 1),
24541 block_dim: (128, 1, 1),
24542 shared_mem_bytes: 0,
24543 };
24544 let hki = hk as i32;
24545 let __s_b = self.gpu.stream();
24546 let mut b = __s_b.launch_builder(&f);
24547 b.arg(qb)
24548 .arg(kb)
24549 .arg(&gcum)
24550 .arg(beta)
24551 .arg(&mut a)
24552 .arg(&mut *pb)
24553 .arg(&hi)
24554 .arg(&ti)
24555 .arg(&ci)
24556 .arg(&hki);
24557 unsafe {
24558 b.launch(cfg)?;
24559 }
24560 } else if c <= 64 && !portable_mma_gated() {
24561 let f = self.func("gdn_chunk_attn_f32");
24563 f.set_attribute(
24564 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24565 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24566 )?;
24567 let jt = ((c + 31) / 32) as u32;
24568 let cfg = LaunchConfig {
24569 grid_dim: (nc as u32, h as u32, jt),
24570 block_dim: (256, 1, 1),
24571 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24572 };
24573 let hki = hk as i32;
24574 let __s_b = self.gpu.stream();
24575 let mut b = __s_b.launch_builder(&f);
24576 b.arg(q)
24577 .arg(k)
24578 .arg(&gcum)
24579 .arg(beta)
24580 .arg(&mut a)
24581 .arg(&mut p)
24582 .arg(&hi)
24583 .arg(&ti)
24584 .arg(&ci)
24585 .arg(&hki);
24586 unsafe {
24587 b.launch(cfg)?;
24588 }
24589 } else {
24590 assert!(
24592 hk == h,
24593 "generic K2 is broadcast-only (de-broadcast rides C==32)"
24594 );
24595 let f = self.func("gdn_chunk_attn_g_f32");
24596 let cfg = LaunchConfig {
24597 grid_dim: (nc as u32, h as u32, 1),
24598 block_dim: (32, 8, 1),
24599 shared_mem_bytes: 0,
24600 };
24601 let __s_b = self.gpu.stream();
24602 let mut b = __s_b.launch_builder(&f);
24603 b.arg(q)
24604 .arg(k)
24605 .arg(&gcum)
24606 .arg(beta)
24607 .arg(&mut a)
24608 .arg(&mut p)
24609 .arg(&hi)
24610 .arg(&ti)
24611 .arg(&ci);
24612 unsafe {
24613 b.launch(cfg)?;
24614 }
24615 }
24616 {
24617 let cfg = LaunchConfig {
24619 grid_dim: (nc as u32, h as u32, 1),
24620 block_dim: (256, 1, 1),
24621 shared_mem_bytes: 0,
24622 };
24623 match c {
24624 32 | 64 => {
24625 let f = self.func(if c == 32 {
24626 "gdn_chunk_solve32_f32"
24627 } else {
24628 "gdn_chunk_solve64_f32"
24629 });
24630 let wb: u64 = match wb16 {
24632 Some(d) => self.addr_u8(d),
24633 None => 0,
24634 };
24635 let hki = hk as i32;
24636 let __s_b = self.gpu.stream();
24637 let mut b = __s_b.launch_builder(&f);
24638 b.arg(v)
24639 .arg(k)
24640 .arg(&a)
24641 .arg(&gcum)
24642 .arg(&mut u)
24643 .arg(&mut w)
24644 .arg(&wb)
24645 .arg(&hi)
24646 .arg(&ti)
24647 .arg(&hki);
24648 unsafe {
24649 b.launch(cfg)?;
24650 }
24651 }
24652 _ => {
24653 assert!(hk == h, "generic K3 is broadcast-only");
24654 let f = self.func("gdn_chunk_solve_f32");
24655 let __s_b = self.gpu.stream();
24656 let mut b = __s_b.launch_builder(&f);
24657 b.arg(v)
24658 .arg(k)
24659 .arg(&a)
24660 .arg(&gcum)
24661 .arg(&mut u)
24662 .arg(&mut w)
24663 .arg(&hi)
24664 .arg(&ti)
24665 .arg(&ci);
24666 unsafe {
24667 b.launch(cfg)?;
24668 }
24669 }
24670 }
24671 }
24672 Ok((gcum, p, u, w))
24673 }
24674
24675 pub fn gdn_db_on() -> bool {
24679 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24680 }
24681
24682 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24691 !portable_mma_gated()
24692 && c == 32
24693 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24694 Ok("1") => true,
24695 Ok("0") => false,
24696 _ => gdn_mma_default_on(),
24697 }
24698 }
24699
24700 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24707 cfg!(memra_hopper_mma)
24708 && self.gdn_mma_enabled(c)
24709 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24710 }
24711
24712 #[allow(clippy::too_many_arguments)]
24717 pub fn ssm_conv1d_gdn_state_pad(
24718 &self,
24719 qkv_tm: &cudarc::driver::CudaView<f32>,
24720 conv_state: &mut CudaSlice<f32>,
24721 w: &CudaSlice<f32>,
24722 q_g: &mut CudaSlice<f32>,
24723 k_g: &mut CudaSlice<f32>,
24724 v_g: &mut CudaSlice<f32>,
24725 conv_dim: usize,
24726 t: usize,
24727 d_conv: usize,
24728 d_state: usize,
24729 num_v: usize,
24730 num_k: usize,
24731 key_dim: usize,
24732 hk: usize,
24733 pad_len: Option<&CudaSlice<i32>>,
24734 ) -> Result<(), Box<dyn std::error::Error>> {
24735 assert!(
24736 t >= d_conv - 1,
24737 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24738 );
24739 {
24740 let f = self.func("ssm_conv1d_gdn_state_f32");
24741 let cfg = LaunchConfig {
24742 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24743 block_dim: (256, 1, 1),
24744 shared_mem_bytes: 0,
24745 };
24746 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24747 let (ds, nv, nk, kd, hki) = (
24748 d_state as i32,
24749 num_v as i32,
24750 num_k as i32,
24751 key_dim as i32,
24752 hk as i32,
24753 );
24754 let __s_b = self.gpu.stream();
24755 let mut b = __s_b.launch_builder(&f);
24756 b.arg(qkv_tm)
24757 .arg(&*conv_state)
24758 .arg(w)
24759 .arg(q_g)
24760 .arg(k_g)
24761 .arg(v_g)
24762 .arg(&cd)
24763 .arg(&ti)
24764 .arg(&dc)
24765 .arg(&ds)
24766 .arg(&nv)
24767 .arg(&nk)
24768 .arg(&kd)
24769 .arg(&hki);
24770 unsafe {
24771 b.launch(cfg)?;
24772 }
24773 }
24774 match pad_len {
24775 Some(len_d) => {
24776 let f = self.func("ssm_conv_ring_update_dev_f32");
24777 let n = conv_dim * (d_conv - 1);
24778 let cfg = LaunchConfig::for_num_elems(n as u32);
24779 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24780 let __s_b = self.gpu.stream();
24781 let mut b = __s_b.launch_builder(&f);
24782 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24783 unsafe {
24784 b.launch(cfg)?;
24785 }
24786 }
24787 None => {
24788 let f = self.func("ssm_conv_ring_update_f32");
24789 let n = conv_dim * (d_conv - 1);
24790 let cfg = LaunchConfig::for_num_elems(n as u32);
24791 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24792 let __s_b = self.gpu.stream();
24793 let mut b = __s_b.launch_builder(&f);
24794 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24795 unsafe {
24796 b.launch(cfg)?;
24797 }
24798 }
24799 }
24800 Ok(())
24801 }
24802
24803 pub fn gdn_chunk_alloc(
24807 &self,
24808 n_head: usize,
24809 t: usize,
24810 c: usize,
24811 hk: usize,
24812 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24813 const D: usize = 128;
24814 assert!(
24815 c == 32,
24816 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24817 );
24818 let h = n_head;
24819 let nc = (t + c - 1) / c;
24820 Ok(GdnChunkBufs {
24821 gcum: self.uninit(t * h)?,
24822 a: self.uninit(nc * h * c * c)?,
24823 p: self.uninit(nc * h * c * c)?,
24824 u: self.uninit(nc * h * c * D)?,
24825 w: self.uninit(nc * h * c * D)?,
24826 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24827 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24828 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24829 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24830 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24831 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24832 o: self.uninit(D * h * t)?,
24833 t,
24834 nc,
24835 })
24836 }
24837
24838 pub fn f32_to_bf16_v(
24840 &self,
24841 x: &cudarc::driver::CudaView<f32>,
24842 dst: &mut CudaSlice<u8>,
24843 n: usize,
24844 ) -> Result<(), Box<dyn std::error::Error>> {
24845 let f = self.func("f32_to_bf16_bulk");
24846 let ni = n as i64;
24847 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24848 let __s_b = self.gpu.stream();
24849 let mut b = __s_b.launch_builder(&f);
24850 b.arg(x).arg(dst).arg(&ni);
24851 unsafe {
24852 b.launch(cfg)?;
24853 }
24854 Ok(())
24855 }
24856
24857 pub fn f32_to_bf16_into(
24859 &self,
24860 x: &CudaSlice<f32>,
24861 dst: &mut CudaSlice<u8>,
24862 n: usize,
24863 ) -> Result<(), Box<dyn std::error::Error>> {
24864 let f = self.func("f32_to_bf16_bulk");
24865 let ni = n as i64;
24866 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24867 let __s_b = self.gpu.stream();
24868 let mut b = __s_b.launch_builder(&f);
24869 b.arg(x).arg(dst).arg(&ni);
24870 unsafe {
24871 b.launch(cfg)?;
24872 }
24873 Ok(())
24874 }
24875
24876 pub fn gdn_chunk_k123_vl8(
24879 &self,
24880 seqs: &[GdnSeqVl],
24881 n_head: usize,
24882 hk: usize,
24883 wq: Option<&GdnWVl8>,
24884 ) -> Result<(), Box<dyn std::error::Error>> {
24885 let b = seqs.len();
24886 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24887 let mut packed = [GdnSeqVl::default(); 8];
24888 packed[..b].copy_from_slice(seqs);
24889 let v = GdnVl8(packed);
24890 let (hi, ci) = (n_head as i32, 32i32);
24891 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24892 {
24893 let f = self.func("gdn_chunk_cumgate_vl");
24894 let cfg = LaunchConfig {
24895 grid_dim: (max_nc, n_head as u32, b as u32),
24896 block_dim: (32, 1, 1),
24897 shared_mem_bytes: 0,
24898 };
24899 let __s_lb = self.gpu.stream();
24900 let mut lb = __s_lb.launch_builder(&f);
24901 lb.arg(&v).arg(&hi).arg(&ci);
24902 unsafe {
24903 lb.launch(cfg)?;
24904 }
24905 }
24906 let hki = hk as i32;
24907 if let Some(w) = wq {
24908 let f = self.func("gdn_k2_wgmma_vl");
24910 let cfg = LaunchConfig {
24911 grid_dim: (max_nc, n_head as u32, b as u32),
24912 block_dim: (128, 1, 1),
24913 shared_mem_bytes: 0,
24914 };
24915 let __s_lb = self.gpu.stream();
24916 let mut lb = __s_lb.launch_builder(&f);
24917 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24918 unsafe {
24919 lb.launch(cfg)?;
24920 }
24921 } else {
24922 let f = self.func("gdn_chunk_attn_vl");
24923 f.set_attribute(
24924 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24925 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24926 )?;
24927 let cfg = LaunchConfig {
24928 grid_dim: (max_nc, n_head as u32, b as u32),
24929 block_dim: (256, 1, 1),
24930 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24931 };
24932 let __s_lb = self.gpu.stream();
24933 let mut lb = __s_lb.launch_builder(&f);
24934 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24935 unsafe {
24936 lb.launch(cfg)?;
24937 }
24938 }
24939 {
24940 let f = self.func("gdn_chunk_solve32_vl");
24941 let cfg = LaunchConfig {
24942 grid_dim: (max_nc, n_head as u32, b as u32),
24943 block_dim: (256, 1, 1),
24944 shared_mem_bytes: 0,
24945 };
24946 let __s_lb = self.gpu.stream();
24947 let mut lb = __s_lb.launch_builder(&f);
24948 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24949 unsafe {
24950 lb.launch(cfg)?;
24951 }
24952 }
24953 Ok(())
24954 }
24955
24956 #[allow(clippy::too_many_arguments)]
24960 pub fn gdn_prep_vl8(
24961 &self,
24962 seqs: &[GdnPrepVl],
24963 conv_w: &CudaSlice<f32>,
24964 dt_bias: &CudaSlice<f32>,
24965 a: &CudaSlice<f32>,
24966 conv_dim: usize,
24967 d_conv: usize,
24968 d_state: usize,
24969 num_v: usize,
24970 num_k: usize,
24971 key_dim: usize,
24972 hk: usize,
24973 eps: f32,
24974 ) -> Result<(), Box<dyn std::error::Error>> {
24975 let b = seqs.len();
24976 assert!(b >= 1 && b <= 8);
24977 let mut packed = [GdnPrepVl::default(); 8];
24978 packed[..b].copy_from_slice(seqs);
24979 let v = GdnPrepVl8(packed);
24980 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24981 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24982 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24983 assert!(
24984 conv_fuse || hk == num_v,
24985 "de-broadcast requires the fused conv"
24986 );
24987 if conv_fuse {
24988 let f = self.func("ssm_conv1d_gdn_state_vl");
24989 let cfg = LaunchConfig {
24990 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24991 block_dim: (256, 1, 1),
24992 shared_mem_bytes: 0,
24993 };
24994 let (dsi, nvi, nki, kdi, hki) = (
24995 d_state as i32,
24996 num_v as i32,
24997 num_k as i32,
24998 key_dim as i32,
24999 hk as i32,
25000 );
25001 let __s_lb = self.gpu.stream();
25002 let mut lb = __s_lb.launch_builder(&f);
25003 lb.arg(&v)
25004 .arg(conv_w)
25005 .arg(&cdi)
25006 .arg(&dci)
25007 .arg(&dsi)
25008 .arg(&nvi)
25009 .arg(&nki)
25010 .arg(&kdi)
25011 .arg(&hki);
25012 unsafe {
25013 lb.launch(cfg)?;
25014 }
25015 } else {
25016 let f = self.func("ssm_conv1d_tm_state_vl");
25017 let cfg = LaunchConfig {
25018 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
25019 block_dim: (256, 1, 1),
25020 shared_mem_bytes: 0,
25021 };
25022 let __s_lb = self.gpu.stream();
25023 let mut lb = __s_lb.launch_builder(&f);
25024 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
25025 unsafe {
25026 lb.launch(cfg)?;
25027 }
25028 }
25029 {
25030 let f = self.func("ssm_conv_ring_update_vl");
25031 let n = (conv_dim * (d_conv - 1)) as u32;
25032 let cfg = LaunchConfig {
25033 grid_dim: (n.div_ceil(256), 1, b as u32),
25034 block_dim: (256, 1, 1),
25035 shared_mem_bytes: 0,
25036 };
25037 let __s_lb = self.gpu.stream();
25038 let mut lb = __s_lb.launch_builder(&f);
25039 lb.arg(&v).arg(&cdi).arg(&dci);
25040 unsafe {
25041 lb.launch(cfg)?;
25042 }
25043 }
25044 if !conv_fuse {
25045 let f = self.func("qkv_to_gdn_repack_vl");
25046 let n = max_t * (num_v * d_state) as u32;
25047 let cfg = LaunchConfig {
25048 grid_dim: (n.div_ceil(256), 1, b as u32),
25049 block_dim: (256, 1, 1),
25050 shared_mem_bytes: 0,
25051 };
25052 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25053 let __s_lb = self.gpu.stream();
25054 let mut lb = __s_lb.launch_builder(&f);
25055 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
25056 unsafe {
25057 lb.launch(cfg)?;
25058 }
25059 }
25060 if Self::l2_v2_on(d_state) {
25061 let f = self.func("gdn_l2_v2_vl");
25062 let cfg = LaunchConfig {
25063 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
25064 block_dim: (256, 1, 1),
25065 shared_mem_bytes: 0,
25066 };
25067 let (dsi, nvi) = (d_state as i32, hk as i32);
25068 let __s_lb = self.gpu.stream();
25069 let mut lb = __s_lb.launch_builder(&f);
25070 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25071 unsafe {
25072 lb.launch(cfg)?;
25073 }
25074 } else {
25075 let f = self.func("gdn_l2_vl");
25076 let cfg = LaunchConfig {
25077 grid_dim: (max_t * hk as u32, 2, b as u32),
25078 block_dim: (256, 1, 1),
25079 shared_mem_bytes: 0,
25080 };
25081 let (dsi, nvi) = (d_state as i32, hk as i32);
25082 let __s_lb = self.gpu.stream();
25083 let mut lb = __s_lb.launch_builder(&f);
25084 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
25085 unsafe {
25086 lb.launch(cfg)?;
25087 }
25088 }
25089 {
25090 let f = self.func("gdn_gate_prep_vl");
25091 let n = max_t * num_v as u32;
25092 let cfg = LaunchConfig {
25093 grid_dim: (n.div_ceil(256), 1, b as u32),
25094 block_dim: (256, 1, 1),
25095 shared_mem_bytes: 0,
25096 };
25097 let nvi = num_v as i32;
25098 let __s_lb = self.gpu.stream();
25099 let mut lb = __s_lb.launch_builder(&f);
25100 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
25101 unsafe {
25102 lb.launch(cfg)?;
25103 }
25104 }
25105 Ok(())
25106 }
25107
25108 pub fn gdn_mirror_vl8(
25110 &self,
25111 seqs: &[GdnSeqVl],
25112 n_head: usize,
25113 which: i32,
25114 hk: usize,
25115 ) -> Result<(), Box<dyn std::error::Error>> {
25116 let b = seqs.len();
25117 assert!(b >= 1 && b <= 8);
25118 let mut packed = [GdnSeqVl::default(); 8];
25119 packed[..b].copy_from_slice(seqs);
25120 let v = GdnVl8(packed);
25121 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
25122 let max_n = seqs
25123 .iter()
25124 .map(|s| {
25125 if which == 0 {
25126 s.t as i64 * ept as i64
25127 } else {
25128 s.nc as i64 * ept as i64 * 32
25129 }
25130 })
25131 .max()
25132 .unwrap();
25133 let f = self.func("gdn_mirror_vl");
25134 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
25135 let cfg = LaunchConfig {
25136 grid_dim: (blocks, 1, b as u32),
25137 block_dim: (256, 1, 1),
25138 shared_mem_bytes: 0,
25139 };
25140 let __s_lb = self.gpu.stream();
25141 let mut lb = __s_lb.launch_builder(&f);
25142 lb.arg(&v).arg(&ept).arg(&which);
25143 unsafe {
25144 lb.launch(cfg)?;
25145 }
25146 Ok(())
25147 }
25148
25149 pub fn gdn_tail_vl8(
25151 &self,
25152 seqs: &[GdnPrepVl],
25153 norm_w: &CudaSlice<f32>,
25154 d_state: usize,
25155 num_v: usize,
25156 eps: f32,
25157 ) -> Result<(), Box<dyn std::error::Error>> {
25158 let b = seqs.len();
25159 assert!(b >= 1 && b <= 8);
25160 let mut packed = [GdnPrepVl::default(); 8];
25161 packed[..b].copy_from_slice(seqs);
25162 let v = GdnPrepVl8(packed);
25163 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
25164 let f = self.func("gated_rmsnorm_f16out_vl");
25165 let cfg = LaunchConfig {
25167 grid_dim: (max_t * num_v as u32, 1, b as u32),
25168 block_dim: (128, 1, 1),
25169 shared_mem_bytes: 0,
25170 };
25171 let (dsi, nvi) = (d_state as i32, num_v as i32);
25172 let __s_lb = self.gpu.stream();
25173 let mut lb = __s_lb.launch_builder(&f);
25174 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
25175 unsafe {
25176 lb.launch(cfg)?;
25177 }
25178 Ok(())
25179 }
25180
25181 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
25184 use cudarc::driver::DevicePtr;
25185 let s = self.gpu.stream();
25186 let (p, _g) = x.device_ptr(&s);
25187 p as u64
25188 }
25189 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
25190 use cudarc::driver::DevicePtrMut;
25191 let s = self.gpu.stream();
25192 let (p, _g) = x.device_ptr_mut(&s);
25193 p as u64
25194 }
25195 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
25196 use cudarc::driver::DevicePtr;
25197 let s = self.gpu.stream();
25198 let (p, _g) = x.device_ptr(&s);
25199 p as u64
25200 }
25201 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
25202 use cudarc::driver::DevicePtr;
25203 let s = self.gpu.stream();
25204 let (p, _g) = x.device_ptr(&s);
25205 p as u64
25206 }
25207
25208 pub fn gdn_chunk_vl8(
25212 &self,
25213 seqs: &[GdnSeqVl],
25214 n_head: usize,
25215 scale: f32,
25216 hk: usize,
25217 wq: Option<&GdnWVl8>,
25218 ) -> Result<(), Box<dyn std::error::Error>> {
25219 const NSPLIT: u32 = 4;
25220 let b = seqs.len();
25221 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
25222 let mut packed = [GdnSeqVl::default(); 8];
25223 packed[..b].copy_from_slice(seqs);
25224 let v = GdnVl8(packed);
25225 let (hi, ci) = (n_head as i32, 32i32);
25226 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
25227 let hki = hk as i32;
25228 if let Some(w) = wq {
25229 let f = self.func("gdn_k45_wgmma_vl");
25231 let cfg = LaunchConfig {
25232 grid_dim: (n_head as u32, NSPLIT, b as u32),
25233 block_dim: (256, 1, 1),
25234 shared_mem_bytes: 0,
25235 };
25236 let __s_lb = self.gpu.stream();
25237 let mut lb = __s_lb.launch_builder(&f);
25238 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
25239 unsafe {
25240 lb.launch(cfg)?;
25241 }
25242 let _ = max_nc;
25243 return Ok(());
25244 }
25245 {
25246 let f = self.func("gdn_chunk_state_mma_vl");
25247 let cfg = LaunchConfig {
25248 grid_dim: (n_head as u32, NSPLIT, b as u32),
25249 block_dim: (256, 1, 1),
25250 shared_mem_bytes: 0,
25251 };
25252 let __s_lb = self.gpu.stream();
25253 let mut lb = __s_lb.launch_builder(&f);
25254 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
25255 unsafe {
25256 lb.launch(cfg)?;
25257 }
25258 }
25259 {
25260 let f = self.func("gdn_chunk_output_mma_vl");
25261 let cfg = LaunchConfig {
25262 grid_dim: (max_nc, n_head as u32, b as u32),
25263 block_dim: (256, 1, 1),
25264 shared_mem_bytes: 0,
25265 };
25266 let __s_lb = self.gpu.stream();
25267 let mut lb = __s_lb.launch_builder(&f);
25268 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
25269 unsafe {
25270 lb.launch(cfg)?;
25271 }
25272 }
25273 Ok(())
25274 }
25275 pub fn gdn_scan_chunked(
25276 &self,
25277 q: &CudaSlice<f32>,
25278 k: &CudaSlice<f32>,
25279 v: &CudaSlice<f32>,
25280 g: &CudaSlice<f32>,
25281 beta: &CudaSlice<f32>,
25282 kb16_pre: Option<&CudaSlice<u8>>,
25283 qb16_pre: Option<&CudaSlice<u8>>,
25284 state_in: &CudaSlice<f32>,
25285 state_out: &mut CudaSlice<f32>,
25286 o: &mut CudaSlice<f32>,
25287 n_head: usize,
25288 t: usize,
25289 scale: f32,
25290 c: usize,
25291 hk: usize,
25292 ) -> Result<(), Box<dyn std::error::Error>> {
25293 const D: usize = 128;
25294 const NSPLIT: u32 = 4;
25295 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
25296 let h = n_head;
25297 let nc = (t + c - 1) / c;
25298 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
25299 let gdn_mma_pre = !portable_mma_gated()
25304 && c == 32
25305 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25306 Ok("1") => true,
25307 Ok("0") => false,
25308 _ => gdn_mma_default_on(),
25309 };
25310 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
25311 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
25312 } else {
25313 None
25314 };
25315 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
25320 && gdn_mma_pre
25321 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
25322 let nk = t * hk * D;
25323 let mut kb16_local: Option<CudaSlice<u8>> = None;
25324 if gdn_mma_pre && kb16_pre.is_none() {
25325 let mut kb = self.alloc_u8_uninit(nk * 2)?;
25326 let f = self.func("f32_to_bf16_bulk");
25327 let n2 = nk as i64;
25328 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25329 let __s_b = self.gpu.stream();
25330 let mut b = __s_b.launch_builder(&f);
25331 b.arg(k).arg(&mut kb).arg(&n2);
25332 unsafe {
25333 b.launch(cfg2)?;
25334 }
25335 kb16_local = Some(kb);
25336 }
25337 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
25338 if let Some(kb) = kb16_pre {
25339 assert!(kb.len() >= nk * 2, "kb16_pre too small");
25340 }
25341 let mut qb16: Option<CudaSlice<u8>> = None;
25342 let mut pb16: Option<CudaSlice<u8>> = None;
25343 if gdn_wgmma_pre {
25344 if qb16_pre.is_none() {
25347 let mut qb = self.alloc_u8_uninit(nk * 2)?;
25348 let f = self.func("f32_to_bf16_bulk");
25349 let n2 = nk as i64;
25350 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
25351 let __s_b = self.gpu.stream();
25352 let mut b = __s_b.launch_builder(&f);
25353 b.arg(q).arg(&mut qb).arg(&n2);
25354 unsafe {
25355 b.launch(cfg2)?;
25356 }
25357 qb16 = Some(qb);
25358 } else if let Some(qb) = qb16_pre {
25359 assert!(qb.len() >= nk * 2, "qb16_pre too small");
25360 }
25361 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
25362 }
25363 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
25364 let k2w = if gdn_wgmma_pre {
25365 Some((
25366 *qb16_ref0.as_ref().unwrap(),
25367 *kb16_ref0.as_ref().unwrap(),
25368 pb16.as_mut().unwrap(),
25369 ))
25370 } else {
25371 None
25372 };
25373 let (gcum, p, u, w) =
25374 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
25375 let _ = &w;
25376 let mut y = self.uninit(nc * h * c * D)?;
25377 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
25391 && c == 32
25392 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
25393 Ok("1") => true,
25394 Ok("0") => false,
25395 _ => gdn_mma_default_on(),
25396 };
25397 if gdn_mma {
25398 let wb16 = wb16_pre
25399 .take()
25400 .expect("mma path pre-allocates wb16 (K3 store fold)");
25401 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
25402 if gdn_wgmma_pre {
25414 let qb16 = qb16_ref0.unwrap();
25416 let pb16 = pb16.as_ref().unwrap();
25417 {
25418 let f = self.func("gdn_k45_wgmma");
25419 let cfg = LaunchConfig {
25420 grid_dim: (h as u32, 4, 1),
25421 block_dim: (256, 1, 1),
25422 shared_mem_bytes: 0,
25423 };
25424 let hki = hk as i32;
25425 let __s_b = self.gpu.stream();
25426 let mut b = __s_b.launch_builder(&f);
25427 b.arg(kb16_ref)
25428 .arg(&gcum)
25429 .arg(beta)
25430 .arg(&u)
25431 .arg(&wb16)
25432 .arg(qb16)
25433 .arg(pb16)
25434 .arg(o)
25435 .arg(&scale)
25436 .arg(state_in)
25437 .arg(&mut *state_out)
25438 .arg(&hi)
25439 .arg(&ti)
25440 .arg(&ci)
25441 .arg(&hki);
25442 unsafe {
25443 b.launch(cfg)?;
25444 }
25445 }
25446 return Ok(());
25447 }
25448 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
25452 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
25453 {
25454 let f = self.func("gdn_chunk_state_mma");
25455 let cfg = LaunchConfig {
25456 grid_dim: (h as u32, NSPLIT, 1),
25457 block_dim: (256, 1, 1),
25458 shared_mem_bytes: 0,
25459 };
25460 let hki = hk as i32;
25461 let __s_b = self.gpu.stream();
25462 let mut b = __s_b.launch_builder(&f);
25463 b.arg(kb16_ref)
25464 .arg(&gcum)
25465 .arg(beta)
25466 .arg(&u)
25467 .arg(&wb16)
25468 .arg(&mut y16)
25469 .arg(&mut ssnap16)
25470 .arg(state_in)
25471 .arg(&mut *state_out)
25472 .arg(&hi)
25473 .arg(&ti)
25474 .arg(&ci)
25475 .arg(&hki);
25476 unsafe {
25477 b.launch(cfg)?;
25478 }
25479 }
25480 {
25481 let f = self.func("gdn_chunk_output_mma");
25483 let jt = ((c + 31) / 32) as u32;
25484 let cfg = LaunchConfig {
25485 grid_dim: (nc as u32, h as u32, jt),
25486 block_dim: (256, 1, 1),
25487 shared_mem_bytes: 0,
25488 };
25489 let hki = hk as i32;
25490 let __s_b = self.gpu.stream();
25491 let mut b = __s_b.launch_builder(&f);
25492 b.arg(q)
25493 .arg(&gcum)
25494 .arg(&p)
25495 .arg(&y16)
25496 .arg(&ssnap16)
25497 .arg(o)
25498 .arg(&hi)
25499 .arg(&ti)
25500 .arg(&ci)
25501 .arg(&scale)
25502 .arg(&hki);
25503 unsafe {
25504 b.launch(cfg)?;
25505 }
25506 }
25507 return Ok(());
25508 }
25509 {
25510 let f = self.func("gdn_chunk_state_f32");
25512 let cfg = LaunchConfig {
25513 grid_dim: (h as u32, NSPLIT, 1),
25514 block_dim: (256, 1, 1),
25515 shared_mem_bytes: 0,
25516 };
25517 let __s_b = self.gpu.stream();
25518 let mut b = __s_b.launch_builder(&f);
25519 b.arg(k)
25520 .arg(&gcum)
25521 .arg(beta)
25522 .arg(&u)
25523 .arg(&w)
25524 .arg(&mut y)
25525 .arg(&mut ssnap)
25526 .arg(state_in)
25527 .arg(&mut *state_out)
25528 .arg(&hi)
25529 .arg(&ti)
25530 .arg(&ci);
25531 unsafe {
25532 b.launch(cfg)?;
25533 }
25534 }
25535 {
25536 let f = self.func("gdn_chunk_output_f32");
25538 let jt = ((c + 31) / 32) as u32;
25539 let cfg = LaunchConfig {
25540 grid_dim: (nc as u32, h as u32, jt),
25541 block_dim: (256, 1, 1),
25542 shared_mem_bytes: 0,
25543 };
25544 let __s_b = self.gpu.stream();
25545 let mut b = __s_b.launch_builder(&f);
25546 b.arg(q)
25547 .arg(&gcum)
25548 .arg(&p)
25549 .arg(&y)
25550 .arg(&ssnap)
25551 .arg(o)
25552 .arg(&hi)
25553 .arg(&ti)
25554 .arg(&ci)
25555 .arg(&scale);
25556 unsafe {
25557 b.launch(cfg)?;
25558 }
25559 }
25560 Ok(())
25561 }
25562
25563 #[allow(clippy::too_many_arguments)]
25572 #[allow(clippy::too_many_arguments)]
25573 pub fn gdn_scan_prefill(
25574 &self,
25575 q: &CudaSlice<f32>,
25576 k: &CudaSlice<f32>,
25577 v: &CudaSlice<f32>,
25578 g: &CudaSlice<f32>,
25579 beta: &CudaSlice<f32>,
25580 kb16_pre: Option<&CudaSlice<u8>>,
25581 qb16_pre: Option<&CudaSlice<u8>>,
25582 state_in: &CudaSlice<f32>,
25583 state_out: &mut CudaSlice<f32>,
25584 o: &mut CudaSlice<f32>,
25585 n_head: usize,
25586 t: usize,
25587 scale: f32,
25588 hk: usize,
25589 ) -> Result<(), Box<dyn std::error::Error>> {
25590 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
25591 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
25592 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25593 }
25594 if Self::gdn_chunked_enabled() && t >= 16 {
25595 self.gdn_scan_chunked(
25596 q,
25597 k,
25598 v,
25599 g,
25600 beta,
25601 kb16_pre,
25602 qb16_pre,
25603 state_in,
25604 state_out,
25605 o,
25606 n_head,
25607 t,
25608 scale,
25609 Self::gdn_chunk_size(),
25610 hk,
25611 )
25612 } else {
25613 assert!(
25614 hk == n_head,
25615 "s128 scan is broadcast-only (prep guarantees by predicate)"
25616 );
25617 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25618 }
25619 }
25620
25621 #[allow(clippy::too_many_arguments)]
25623 fn gdn_scan_diff(
25624 &self,
25625 q: &CudaSlice<f32>,
25626 k: &CudaSlice<f32>,
25627 v: &CudaSlice<f32>,
25628 g: &CudaSlice<f32>,
25629 beta: &CudaSlice<f32>,
25630 state_in: &CudaSlice<f32>,
25631 state_out: &mut CudaSlice<f32>,
25632 o: &mut CudaSlice<f32>,
25633 n_head: usize,
25634 t: usize,
25635 scale: f32,
25636 ) -> Result<(), Box<dyn std::error::Error>> {
25637 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25638 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25639 let mut o_c = self.uninit(o.len())?;
25640 let mut st_c = self.uninit(state_out.len())?;
25641 self.gdn_scan_chunked(
25642 q,
25643 k,
25644 v,
25645 g,
25646 beta,
25647 None,
25648 None,
25649 state_in,
25650 &mut st_c,
25651 &mut o_c,
25652 n_head,
25653 t,
25654 scale,
25655 Self::gdn_chunk_size(),
25656 n_head,
25657 )?;
25658 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25659 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25660 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25661 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25662 let mut max_abs = 0f32;
25663 let mut max_rel = 0f32;
25664 let mut sum_rel = 0f64;
25665 for (x, y) in a.iter().zip(b) {
25666 let ad = (x - y).abs();
25667 let rel = ad / x.abs().max(y.abs()).max(1e-3);
25668 if ad > max_abs {
25669 max_abs = ad;
25670 }
25671 if rel > max_rel {
25672 max_rel = rel;
25673 }
25674 sum_rel += rel as f64;
25675 }
25676 (max_abs, max_rel, sum_rel / a.len() as f64)
25677 };
25678 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25679 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25680 println!(
25681 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25682 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25683 Self::gdn_chunk_size()
25684 );
25685 Ok(())
25686 }
25687
25688 pub fn gdn_glog(
25690 &self,
25691 alpha: &CudaSlice<f32>,
25692 dt_bias: &CudaSlice<f32>,
25693 a: &CudaSlice<f32>,
25694 g_log: &mut CudaSlice<f32>,
25695 n_head: usize,
25696 t: usize,
25697 ) -> Result<(), Box<dyn std::error::Error>> {
25698 let f = self.func("gdn_glog_f32");
25699 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25700 let (h, ti) = (n_head as i32, t as i32);
25701 let __s_b = self.gpu.stream();
25702 let mut b = __s_b.launch_builder(&f);
25703 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25704 unsafe {
25705 b.launch(cfg)?;
25706 }
25707 Ok(())
25708 }
25709
25710 pub fn sigmoid_v(
25713 &self,
25714 x: &cudarc::driver::CudaView<f32>,
25715 y: &mut CudaSlice<f32>,
25716 n: usize,
25717 ) -> Result<(), Box<dyn std::error::Error>> {
25718 let f = self.func("sigmoid_f32");
25719 let cfg = LaunchConfig::for_num_elems(n as u32);
25720 let ni = n as i32;
25721 let __s_b = self.gpu.stream();
25722 let mut b = __s_b.launch_builder(&f);
25723 b.arg(x).arg(y).arg(&ni);
25724 unsafe {
25725 b.launch(cfg)?;
25726 }
25727 Ok(())
25728 }
25729
25730 pub fn gdn_glog_v(
25731 &self,
25732 alpha: &cudarc::driver::CudaView<f32>,
25733 dt_bias: &CudaSlice<f32>,
25734 a: &CudaSlice<f32>,
25735 g_log: &mut CudaSlice<f32>,
25736 n_head: usize,
25737 t: usize,
25738 ) -> Result<(), Box<dyn std::error::Error>> {
25739 let f = self.func("gdn_glog_f32");
25740 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25741 let (h, ti) = (n_head as i32, t as i32);
25742 let __s_b = self.gpu.stream();
25743 let mut b = __s_b.launch_builder(&f);
25744 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25745 unsafe {
25746 b.launch(cfg)?;
25747 }
25748 Ok(())
25749 }
25750
25751 pub fn sigmoid(
25752 &self,
25753 x: &CudaSlice<f32>,
25754 y: &mut CudaSlice<f32>,
25755 n: usize,
25756 ) -> Result<(), Box<dyn std::error::Error>> {
25757 let f = self.func("sigmoid_f32");
25758 let cfg = LaunchConfig::for_num_elems(n as u32);
25759 let ni = n as i32;
25760 let __s_b = self.gpu.stream();
25761 let mut b = __s_b.launch_builder(&f);
25762 b.arg(x).arg(y).arg(&ni);
25763 unsafe {
25764 b.launch(cfg)?;
25765 }
25766 Ok(())
25767 }
25768
25769 pub fn sig_mul_f16out(
25772 &self,
25773 a: &CudaSlice<f32>,
25774 g: &CudaSlice<f32>,
25775 dst: &mut CudaSlice<f32>,
25776 dst16: &mut CudaSlice<u8>,
25777 n: usize,
25778 ) -> Result<(), Box<dyn std::error::Error>> {
25779 let f = self.func("sig_mul_f16out_f32");
25780 let cfg = LaunchConfig::for_num_elems(n as u32);
25781 let ni = n as i32;
25782 let __s_b = self.gpu.stream();
25783 let mut b = __s_b.launch_builder(&f);
25784 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25785 unsafe {
25786 b.launch(cfg)?;
25787 }
25788 Ok(())
25789 }
25790
25791 #[allow(clippy::too_many_arguments)]
25800 pub fn attn_head_gate(
25801 &self,
25802 a: &CudaSlice<f32>,
25803 g: &CudaSlice<f32>,
25804 dst: &mut CudaSlice<f32>,
25805 dst16: Option<&mut CudaSlice<u8>>,
25806 head_dim: usize,
25807 n_head: usize,
25808 t: usize,
25809 ) -> Result<(), Box<dyn std::error::Error>> {
25810 let f = self.func("attn_head_gate_f32");
25811 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25812 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25813 let d16: u64 = match dst16 {
25815 Some(d) => self.addr_u8(d),
25816 None => 0,
25817 };
25818 let __s_b = self.gpu.stream();
25819 let mut b = __s_b.launch_builder(&f);
25820 b.arg(a)
25821 .arg(g)
25822 .arg(dst)
25823 .arg(&d16)
25824 .arg(&hd)
25825 .arg(&nh)
25826 .arg(&ti);
25827 unsafe {
25828 b.launch(cfg)?;
25829 }
25830 Ok(())
25831 }
25832
25833 #[allow(clippy::too_many_arguments)]
25842 pub fn swiglu_clamped_mul_scaled(
25843 &self,
25844 gate: &CudaSlice<f32>,
25845 up: &CudaSlice<f32>,
25846 gs: f32,
25847 us: f32,
25848 limit: f32,
25849 dst: &mut CudaSlice<f32>,
25850 n: usize,
25851 ) -> Result<(), Box<dyn std::error::Error>> {
25852 debug_assert!(
25853 limit > 1e-6,
25854 "swiglu_clamped needs a live limit; use silu_mul_scaled"
25855 );
25856 let f = self.func("swiglu_clamped_mul_scaled_f32");
25857 let cfg = LaunchConfig::for_num_elems(n as u32);
25858 let ni = n as i32;
25859 let __s_b = self.gpu.stream();
25860 let mut b = __s_b.launch_builder(&f);
25861 b.arg(gate)
25862 .arg(up)
25863 .arg(&gs)
25864 .arg(&us)
25865 .arg(&limit)
25866 .arg(dst)
25867 .arg(&ni);
25868 unsafe {
25869 b.launch(cfg)?;
25870 }
25871 Ok(())
25872 }
25873
25874 pub fn gated_rmsnorm(
25876 &self,
25877 o: &CudaSlice<f32>,
25878 w: &CudaSlice<f32>,
25879 z: &CudaSlice<f32>,
25880 dst: &mut CudaSlice<f32>,
25881 ncols: usize,
25882 nrows: usize,
25883 eps: f32,
25884 ) -> Result<(), Box<dyn std::error::Error>> {
25885 let f = self.func("gated_rmsnorm_f32");
25886 let cfg = LaunchConfig {
25887 grid_dim: (nrows as u32, 1, 1),
25888 block_dim: (128, 1, 1),
25889 shared_mem_bytes: 0,
25890 };
25891 let (nc, e) = (ncols as i32, eps);
25892 let __s_b = self.gpu.stream();
25893 let mut b = __s_b.launch_builder(&f);
25894 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25895 unsafe {
25896 b.launch(cfg)?;
25897 }
25898 Ok(())
25899 }
25900
25901 pub fn gated_rmsnorm_f16out(
25904 &self,
25905 o: &CudaSlice<f32>,
25906 w: &CudaSlice<f32>,
25907 z: &CudaSlice<f32>,
25908 dst: &mut CudaSlice<f32>,
25909 dst16: &mut CudaSlice<u8>,
25910 ncols: usize,
25911 nrows: usize,
25912 eps: f32,
25913 ) -> Result<(), Box<dyn std::error::Error>> {
25914 let f = self.func("gated_rmsnorm_f16out_f32");
25915 let cfg = LaunchConfig {
25917 grid_dim: (nrows as u32, 1, 1),
25918 block_dim: (128, 1, 1),
25919 shared_mem_bytes: 0,
25920 };
25921 let (nc, e) = (ncols as i32, eps);
25922 let __s_b = self.gpu.stream();
25923 let mut b = __s_b.launch_builder(&f);
25924 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25925 unsafe {
25926 b.launch(cfg)?;
25927 }
25928 Ok(())
25929 }
25930
25931 #[allow(clippy::too_many_arguments)]
25935 pub fn add_rms_norm_zq8(
25936 &self,
25937 a: &CudaSlice<f32>,
25938 b_in: &CudaSlice<f32>,
25939 w: &CudaSlice<f32>,
25940 res: &mut CudaSlice<f32>,
25941 z: &mut CudaSlice<f32>,
25942 ncols: usize,
25943 nrows: usize,
25944 eps: f32,
25945 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25946 assert!(ncols % 32 == 0);
25947 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25948 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25949 let f = self.func("add_rms_norm_zq8");
25950 let cfg = LaunchConfig {
25951 grid_dim: (nrows as u32, 1, 1),
25952 block_dim: (1024, 1, 1),
25953 shared_mem_bytes: 0,
25954 };
25955 let (nc, ep) = (ncols as i32, eps);
25956 let __s_b = self.gpu.stream();
25957 let mut b = __s_b.launch_builder(&f);
25958 b.arg(a)
25959 .arg(b_in)
25960 .arg(w)
25961 .arg(res)
25962 .arg(z)
25963 .arg(&mut q)
25964 .arg(&mut d)
25965 .arg(&nc)
25966 .arg(&ep);
25967 unsafe {
25968 b.launch(cfg)?;
25969 }
25970 Ok((q, d))
25971 }
25972
25973 pub fn gated_rmsnorm_zv(
25978 &self,
25979 o: &CudaSlice<f32>,
25980 w: &CudaSlice<f32>,
25981 z: &cudarc::driver::CudaView<f32>,
25982 dst: &mut CudaSlice<f32>,
25983 ncols: usize,
25984 nrows: usize,
25985 eps: f32,
25986 ) -> Result<(), Box<dyn std::error::Error>> {
25987 let f = self.func("gated_rmsnorm_f32");
25988 let cfg = LaunchConfig {
25989 grid_dim: (nrows as u32, 1, 1),
25990 block_dim: (128, 1, 1),
25991 shared_mem_bytes: 0,
25992 };
25993 let (nc, e) = (ncols as i32, eps);
25994 let __s_b = self.gpu.stream();
25995 let mut b = __s_b.launch_builder(&f);
25996 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25997 unsafe {
25998 b.launch(cfg)?;
25999 }
26000 Ok(())
26001 }
26002
26003 pub fn gated_rmsnorm_f16out_zv(
26004 &self,
26005 o: &CudaSlice<f32>,
26006 w: &CudaSlice<f32>,
26007 z: &cudarc::driver::CudaView<f32>,
26008 dst: &mut CudaSlice<f32>,
26009 dst16: &mut CudaSlice<u8>,
26010 ncols: usize,
26011 nrows: usize,
26012 eps: f32,
26013 ) -> Result<(), Box<dyn std::error::Error>> {
26014 let f = self.func("gated_rmsnorm_f16out_f32");
26015 let cfg = LaunchConfig {
26017 grid_dim: (nrows as u32, 1, 1),
26018 block_dim: (128, 1, 1),
26019 shared_mem_bytes: 0,
26020 };
26021 let (nc, e) = (ncols as i32, eps);
26022 let __s_b = self.gpu.stream();
26023 let mut b = __s_b.launch_builder(&f);
26024 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
26025 unsafe {
26026 b.launch(cfg)?;
26027 }
26028 Ok(())
26029 }
26030
26031 pub fn gated_rmsnorm_q8_1(
26032 &self,
26033 o: &CudaSlice<f32>,
26034 w: &CudaSlice<f32>,
26035 z: &CudaSlice<f32>,
26036 ncols: usize,
26037 nrows: usize,
26038 eps: f32,
26039 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
26040 assert!(ncols % 32 == 0);
26041 let f = self.func("gated_rmsnorm_q8_1");
26042 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
26043 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
26044 let cfg = LaunchConfig {
26045 grid_dim: (nrows as u32, 1, 1),
26046 block_dim: (128, 1, 1),
26047 shared_mem_bytes: 0,
26048 };
26049 let (nc, ep) = (ncols as i32, eps);
26050 let __s_b = self.gpu.stream();
26051 let mut b = __s_b.launch_builder(&f);
26052 b.arg(o)
26053 .arg(w)
26054 .arg(z)
26055 .arg(&mut out_q)
26056 .arg(&mut out_d)
26057 .arg(&nc)
26058 .arg(&ep);
26059 unsafe {
26060 b.launch(cfg)?;
26061 }
26062 Ok((out_q, out_d))
26063 }
26064
26065 pub fn transpose(
26067 &self,
26068 inp: &CudaSlice<f32>,
26069 rows: usize,
26070 cols: usize,
26071 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26072 let f = self.func("transpose_f32");
26073 let mut out = self.zeros(rows * cols)?;
26074 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
26075 let (r, c) = (rows as i32, cols as i32);
26076 let __s_b = self.gpu.stream();
26077 let mut b = __s_b.launch_builder(&f);
26078 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
26079 unsafe {
26080 b.launch(cfg)?;
26081 }
26082 Ok(out)
26083 }
26084
26085 pub fn repeat_heads(
26087 &self,
26088 inp: &CudaSlice<f32>,
26089 out: &mut CudaSlice<f32>,
26090 head_dim: usize,
26091 n_in: usize,
26092 n_out: usize,
26093 t: usize,
26094 ) -> Result<(), Box<dyn std::error::Error>> {
26095 let f = self.func("repeat_heads_f32");
26096 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
26097 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
26098 let __s_b = self.gpu.stream();
26099 let mut b = __s_b.launch_builder(&f);
26100 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
26101 unsafe {
26102 b.launch(cfg)?;
26103 }
26104 Ok(())
26105 }
26106
26107 pub fn q_gate_split(
26114 &self,
26115 qf: &CudaSlice<f32>,
26116 q_out: &mut CudaSlice<f32>,
26117 gate_out: &mut CudaSlice<f32>,
26118 head_dim: usize,
26119 n_head: usize,
26120 t: usize,
26121 ) -> Result<(), Box<dyn std::error::Error>> {
26122 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
26123 let out_need = head_dim * n_head * t;
26124 if q_out.len() < out_need || gate_out.len() < out_need {
26125 return Err(format!(
26126 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
26127 q_out.len(),
26128 gate_out.len()
26129 )
26130 .into());
26131 }
26132 let f = self.func("q_gate_split_f32");
26133 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
26134 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
26135 let __s_b = self.gpu.stream();
26136 let mut b = __s_b.launch_builder(&f);
26137 b.arg(qf)
26138 .arg(q_out)
26139 .arg(gate_out)
26140 .arg(&hd)
26141 .arg(&nh)
26142 .arg(&ti);
26143 unsafe {
26144 b.launch(cfg)?;
26145 }
26146 Ok(())
26147 }
26148
26149 pub fn qkv_to_gdn_repack(
26153 &self,
26154 conv_out: &CudaSlice<f32>,
26155 q_g: &mut CudaSlice<f32>,
26156 k_g: &mut CudaSlice<f32>,
26157 v_g: &mut CudaSlice<f32>,
26158 d_state: usize,
26159 num_v: usize,
26160 num_k: usize,
26161 key_dim: usize,
26162 t: usize,
26163 ) -> Result<(), Box<dyn std::error::Error>> {
26164 let f = self.func("qkv_to_gdn_repack_f32");
26165 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
26166 let (ds, nv, nk, kd, ti) = (
26167 d_state as i32,
26168 num_v as i32,
26169 num_k as i32,
26170 key_dim as i32,
26171 t as i32,
26172 );
26173 let __s_b = self.gpu.stream();
26174 let mut b = __s_b.launch_builder(&f);
26175 b.arg(conv_out)
26176 .arg(q_g)
26177 .arg(k_g)
26178 .arg(v_g)
26179 .arg(&ds)
26180 .arg(&nv)
26181 .arg(&nk)
26182 .arg(&kd)
26183 .arg(&ti);
26184 unsafe {
26185 b.launch(cfg)?;
26186 }
26187 Ok(())
26188 }
26189
26190 pub fn conv_left_pad(
26193 &self,
26194 src: &CudaSlice<f32>,
26195 dst: &mut CudaSlice<f32>,
26196 conv_dim: usize,
26197 t: usize,
26198 pad: usize,
26199 ) -> Result<(), Box<dyn std::error::Error>> {
26200 let f = self.func("conv_left_pad_f32");
26201 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
26202 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
26203 let __s_b = self.gpu.stream();
26204 let mut b = __s_b.launch_builder(&f);
26205 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
26206 unsafe {
26207 b.launch(cfg)?;
26208 }
26209 Ok(())
26210 }
26211
26212 pub fn conv_assemble_and_roll(
26216 &self,
26217 qkv_col: &CudaSlice<f32>,
26218 conv_state: &mut CudaSlice<f32>,
26219 conv_in: &mut CudaSlice<f32>,
26220 conv_dim: usize,
26221 pad: usize,
26222 ) -> Result<(), Box<dyn std::error::Error>> {
26223 let f = self.func("conv_assemble_and_roll_f32");
26224 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26225 let (cd, p) = (conv_dim as i32, pad as i32);
26226 let __s_b = self.gpu.stream();
26227 let mut b = __s_b.launch_builder(&f);
26228 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
26229 unsafe {
26230 b.launch(cfg)?;
26231 }
26232 Ok(())
26233 }
26234
26235 pub fn ssm_conv1d_fused_decode(
26241 &self,
26242 qkv_col: &CudaSlice<f32>,
26243 conv_state: &mut CudaSlice<f32>,
26244 w: &CudaSlice<f32>,
26245 conv_out: &mut CudaSlice<f32>,
26246 conv_dim: usize,
26247 d_conv: usize,
26248 ) -> Result<(), Box<dyn std::error::Error>> {
26249 let f = self.func("ssm_conv1d_fused_decode_f32");
26250 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
26251 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26252 let __s_b = self.gpu.stream();
26253 let mut b = __s_b.launch_builder(&f);
26254 b.arg(qkv_col)
26255 .arg(conv_state)
26256 .arg(w)
26257 .arg(conv_out)
26258 .arg(&cd)
26259 .arg(&dc);
26260 unsafe {
26261 b.launch(cfg)?;
26262 }
26263 Ok(())
26264 }
26265
26266 pub fn slice_range(
26269 &self,
26270 src: &CudaSlice<f32>,
26271 start: usize,
26272 len: usize,
26273 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26274 let host = self.gpu.stream().clone_dtoh(src)?;
26275 self.gpu.stream().synchronize()?;
26276 Ok(self.htod(&host[start..start + len])?)
26277 }
26278}
26279
26280#[cfg(test)]
26281mod target_dispatch_tests {
26282 use super::legacy_quant_gemm_allowed;
26283
26284 #[test]
26285 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
26286 assert!(legacy_quant_gemm_allowed(false, false, false));
26288 assert!(!legacy_quant_gemm_allowed(false, false, true));
26289 assert!(!legacy_quant_gemm_allowed(true, false, false));
26291 assert!(!legacy_quant_gemm_allowed(true, false, true));
26292 assert!(legacy_quant_gemm_allowed(true, true, false));
26294 assert!(!legacy_quant_gemm_allowed(true, true, true));
26295 }
26296
26297 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
26298 #[test]
26299 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
26300 assert!(!legacy_quant_gemm_allowed(
26301 cfg!(memra_portable_cuda),
26302 cfg!(memra_hopper_mma),
26303 false
26304 ));
26305 }
26306
26307 #[cfg(memra_hopper_mma)]
26308 #[test]
26309 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
26310 assert!(legacy_quant_gemm_allowed(
26311 cfg!(memra_portable_cuda),
26312 cfg!(memra_hopper_mma),
26313 false
26314 ));
26315 assert!(super::portable_mma_gated() == false);
26316 }
26317}
26318
26319impl memra_kv::KvDev for Engine {
26322 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26323 Engine::zeros(self, n)
26324 }
26325 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26326 Engine::uninit(self, n)
26327 }
26328 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
26329 Engine::alloc_u8(self, n)
26330 }
26331 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
26332 Engine::htod_i32(self, v)
26333 }
26334 fn clone_dtod(
26335 &self,
26336 src: &CudaSlice<f32>,
26337 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
26338 Engine::clone_dtod(self, src)
26339 }
26340 fn copy_into(
26341 &self,
26342 dst: &mut CudaSlice<f32>,
26343 off: usize,
26344 src: &CudaSlice<f32>,
26345 len: usize,
26346 ) -> Result<(), Box<dyn std::error::Error>> {
26347 Engine::copy_into(self, dst, off, src, len)
26348 }
26349 fn set_i32_one(
26350 &self,
26351 d: &mut CudaSlice<i32>,
26352 v: i32,
26353 ) -> Result<(), Box<dyn std::error::Error>> {
26354 Engine::set_i32_one(self, d, v)
26355 }
26356}
26357
26358#[cfg(test)]
26359mod fused_gate_bounds_tests {
26360 use super::*;
26361
26362 #[test]
26375 #[ignore = "requires a CUDA GPU"]
26376 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
26377 let e = Engine::new(0).unwrap();
26378 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
26379 let fused = 2 * head_dim * n_head * t;
26380 let out_n = head_dim * n_head * t;
26381
26382 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
26384 let mut q = e.uninit(out_n).unwrap();
26385 let mut gate = e.uninit(out_n).unwrap();
26386 let err = e
26387 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
26388 .expect_err("half-width wq must be refused, not read past")
26389 .to_string();
26390 assert!(err.contains("NO fused gate"), "{err}");
26391 assert!(err.contains(&format!("{fused}")), "{err}");
26392
26393 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
26396 let wide = e.htod(&host).unwrap();
26397 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
26398 .expect("full-width wq splits");
26399 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
26400 for tok in 0..t {
26401 for hh in 0..n_head {
26402 for d in 0..head_dim {
26403 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
26404 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
26405 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
26406 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
26407 }
26408 }
26409 }
26410
26411 let mut small = e.uninit(out_n - 1).unwrap();
26413 assert!(
26414 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
26415 .is_err()
26416 );
26417 }
26418}
26419
26420#[cfg(test)]
26424mod fused_rope_width_tests {
26425 use super::Engine;
26426
26427 #[test]
26430 fn full_width_is_accepted() {
26431 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
26432 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
26433 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
26434 }
26435
26436 #[test]
26449 fn gemma4_official_artifact_widths_pass() {
26450 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
26451 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
26452 }
26453
26454 #[test]
26457 fn partial_rotary_is_refused_with_the_geometry_named() {
26458 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
26460 .expect_err("partial rotary must refuse");
26461 let msg = err.to_string();
26462 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
26463 assert!(msg.contains("n_rot 64"), "{msg}");
26464 assert!(msg.contains("head_dim 256"), "{msg}");
26465 assert!(
26466 msg.contains("64..256"),
26467 "names the band it would corrupt: {msg}"
26468 );
26469 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
26471 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
26473 }
26474}