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
317pub(crate) const fn gdn_mma_default_on() -> bool {
326 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
327}
328
329const fn konst_eq(a: &str, b: &str) -> bool {
331 let (a, b) = (a.as_bytes(), b.as_bytes());
332 if a.len() != b.len() {
333 return false;
334 }
335 let mut i = 0;
336 while i < a.len() {
337 if a[i] != b[i] {
338 return false;
339 }
340 i += 1;
341 }
342 true
343}
344
345const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
350 (!portable_cuda || hopper_mma) && !no_gemm
351}
352
353const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
361const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
362const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
363const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
364const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
365
366pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
369
370fn flash_fatbin_bytes() -> &'static [u8] {
372 match kv_cache_formats() {
373 ("q8_0", "q5_1") => FLASH_FATBIN,
374 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
375 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
376 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
377 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
378 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
379 other => unreachable!("kv_cache_formats returned {other:?}"),
380 }
381}
382
383fn k1_launch_override() -> Option<(u32, u32, u32)> {
390 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
391 *K1.get_or_init(|| {
392 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
393 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
394 match p.as_slice() {
395 [bm, bn, w] => Some((*bm, *bn, *w)),
396 _ => None,
397 }
398 })
399}
400
401pub(crate) fn wgmma_gemm_enabled() -> bool {
408 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
409 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
410}
411
412pub const FA_VEC_MIN_TKV: usize = 96;
427pub fn fa_vec_min_tkv() -> usize {
431 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
432 *V.get_or_init(|| {
433 std::env::var("MEMRA_FA_VEC_MIN")
434 .ok()
435 .and_then(|v| v.parse().ok())
436 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
437 })
438}
439
440pub fn fa_f16pv_on() -> bool {
451 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
452 *ON.get_or_init(|| {
453 std::env::var("MEMRA_FA_F16PV")
454 .map(|v| v != "0")
455 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
456 })
457}
458
459pub fn fa512_hp_on() -> bool {
463 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
464 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
465}
466
467pub fn faw_hp_on() -> bool {
471 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
472 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
473}
474
475pub fn fa512_wide_warps() -> usize {
479 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
480 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
481 Ok("1") => 4,
482 _ => 2,
483 })
484}
485
486pub fn fa512_min_tkv() -> usize {
489 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
490 *FA512_MIN.get_or_init(|| {
491 std::env::var("MEMRA_FA512_MIN")
492 .ok()
493 .and_then(|v| v.parse().ok())
494 .unwrap_or(512)
495 })
496}
497pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
501 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
502pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
506pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
512 std::sync::atomic::AtomicBool::new(false);
513pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
520 std::sync::atomic::AtomicBool::new(true);
521pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
522 std::sync::atomic::AtomicUsize::new(16);
523pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
528pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
530pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
535pub use memra_kv::KV_FP8_FORCE;
538pub(crate) fn mmv_block() -> u32 {
543 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
544 *V.get_or_init(|| {
545 std::env::var("MEMRA_MMV_BLOCK")
546 .ok()
547 .and_then(|v| v.parse().ok())
548 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
549 .unwrap_or(128)
550 })
551}
552
553pub(crate) fn sig_expf_dev_on() -> bool {
558 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
559 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
560}
561
562pub(crate) fn topk_fast_on() -> bool {
563 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
564 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
565}
566
567pub(crate) fn rms_block() -> u32 {
568 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
569 *V.get_or_init(|| {
570 std::env::var("MEMRA_RMS_BLOCK")
571 .ok()
572 .and_then(|v| v.parse().ok())
573 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
574 })
575}
576
577pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
578 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
579 if let Some(forced) = *S.get_or_init(|| {
580 std::env::var("MEMRA_FA_SPLIT")
581 .ok()
582 .and_then(|v| v.parse().ok())
583 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
584 }) {
585 return forced;
586 }
587 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
605 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
606 {
607 return if t_kv <= 8192 {
608 16
609 } else if t_kv <= 16384 {
610 64
611 } else {
612 128
613 };
614 }
615 let big_rig = fa_sm_count() >= 128;
616 if big_rig {
617 let _ = n_head_kv;
618 if t_kv <= 2048 {
619 16
620 } else if t_kv <= 16384 {
621 64
622 } else {
623 128
624 }
625 } else if n_head_kv <= 4 {
626 if t_kv <= 512 {
647 8
648 } else if t_kv <= 16384 {
649 64
650 } else {
651 128
652 }
653 } else {
654 if t_kv <= 8192 {
655 32
656 } else if t_kv <= 16384 {
657 64
658 } else {
659 128
660 }
661 }
662}
663
664fn fa_sm_count() -> i32 {
667 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
668 *N.get_or_init(|| {
669 cudarc::driver::result::init().ok();
670 cudarc::driver::result::device::get(0)
671 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
672 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
673 .unwrap_or(82)
674 })
675}
676
677fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
681 match head_dim {
682 256 => Ok(""),
683 128 => Ok("_hd128"),
684 d => Err(format!(
685 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
686 callers must gate to sdpa_naive"
687 )
688 .into()),
689 }
690}
691
692pub const QT_Q8_0: i32 = 0;
694pub const QT_Q4_K: i32 = 1;
695pub const QT_Q6_K: i32 = 2;
696pub const QT_Q5_K: i32 = 3;
697pub const QT_Q3_K: i32 = 4;
698pub const QT_IQ4_XS: i32 = 5;
699pub const QT_IQ3_S: i32 = 6;
700pub const QT_NVFP4: i32 = 7;
701pub const QT_F8_E4M3: i32 = 10;
707pub const QT_NVFP4_RP: i32 = 9;
710pub const QT_F32: i32 = 8;
712pub const QT_BF16: i32 = 11;
713pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
718pub const QT_F8_E4M3_BLK: i32 = 14;
734
735pub struct Engine {
737 pub gpu: memra_runtime::Gpu,
738 module: Arc<CudaModule>,
739 hybrid: Arc<CudaModule>,
740 qmatvec: Arc<CudaModule>,
741 flash: Arc<CudaModule>,
742 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
746 gemm: Arc<CudaModule>,
747 router: Arc<CudaModule>,
748 sample: Arc<CudaModule>,
750 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
754 moe_cache_layout: Mutex<Option<Vec<usize>>>,
758 capture_keep_on: std::sync::atomic::AtomicBool,
764 verify_exact: std::sync::atomic::AtomicBool,
769 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
770 pub copy_stream: Arc<CudaStream>,
772 #[cfg(memra_cutlass)]
779 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
780 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
784 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
787 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
791 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
795 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
797 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
798 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
803 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
808 router_stage: Mutex<Option<PinnedStage>>,
812}
813
814fn fa_v2_on() -> bool {
824 std::env::var("MEMRA_FA_V2")
830 .map(|v| v != "0")
831 .unwrap_or(true)
832}
833
834fn fa_v3_on() -> bool {
842 std::env::var("MEMRA_FA_V3")
846 .map(|v| v != "0")
847 .unwrap_or(true)
848}
849
850fn fa_v4_mode() -> &'static str {
855 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
856 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
857}
858fn fa_v4_on() -> bool {
859 fa_v4_mode() != "0"
860} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
869 std::sync::atomic::AtomicUsize::new(1024);
870pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
871 std::sync::atomic::AtomicUsize::new(usize::MAX);
872pub fn fa_v4_at_pub(t_kv: usize) -> bool {
873 fa_v4_at(t_kv)
874}
875fn fa_v4_at(t_kv: usize) -> bool {
876 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
877 let mx = *M.get_or_init(|| {
878 std::env::var("MEMRA_FA_V4_MAX")
879 .ok()
880 .and_then(|v| v.parse().ok())
881 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
882 });
883 fa_v4_on() && t_kv < mx
884}
885pub const FA_DEEP_MIN_DEFAULT: usize = 0;
899fn fa_deep_at(t_kv: usize) -> bool {
900 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
901 return false;
902 }
903 let min = std::env::var("MEMRA_FA_DEEP_MIN")
904 .ok()
905 .and_then(|v| v.parse().ok())
906 .unwrap_or(FA_DEEP_MIN_DEFAULT);
907 t_kv >= min
908}
909pub fn fa_deep_at_pub(t_kv: usize) -> bool {
911 fa_deep_at(t_kv)
912}
913
914fn fa_v3_active(head_dim: usize) -> bool {
915 fa_v3_on()
918 && head_dim % 128 == 0
919 && kv_cache_formats() == ("q8_0", "q5_1")
920 && !Engine::kv_fp8_on()
921}
922
923pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
931 std::env::var("MEMRA_NO_FA_VEC").is_err()
932 && t_kv >= fa_vec_min_tkv()
933 && head_dim == 256
934 && fa_v4_at(t_kv)
935 && !matches!(fa_v4_mode(), "noB3" | "stage")
936 && !Engine::kv_fp8_on()
937}
938pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
940 fa_split_keys(t_kv, n_head_kv)
941}
942
943struct PinnedStage {
948 ptr: *mut u8,
949 cap: usize,
950}
951unsafe impl Send for PinnedStage {}
952impl PinnedStage {
953 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
954 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
955 Ok(PinnedStage { ptr, cap })
956 }
957}
958impl Drop for PinnedStage {
959 fn drop(&mut self) {
960 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
961 }
962}
963
964pub const ARGMAX_NB: usize = 256;
967
968pub(crate) use memra_fa3_vl as fa3_vl_raw;
970
971unsafe extern "C" {
972 fn memra_fa3_prefill(
974 q16: *const core::ffi::c_void,
975 k16: *const core::ffi::c_void,
976 v16: *const core::ffi::c_void,
977 o: *mut f32,
978 t: i32,
979 h: i32,
980 hkv: i32,
981 d: i32,
982 scale: f32,
983 stream: *mut core::ffi::c_void,
984 ) -> i32;
985 pub(crate) fn memra_fa3_vl(
987 q16s: *const *const core::ffi::c_void,
988 k16s: *const *const core::ffi::c_void,
989 v16s: *const *const core::ffi::c_void,
990 os: *const *mut f32,
991 ts: *const i32,
992 b: i32,
993 h: i32,
994 hkv: i32,
995 d: i32,
996 scale: f32,
997 stream: *mut core::ffi::c_void,
998 ) -> i32;
999}
1000
1001#[repr(C)]
1006#[derive(Clone, Copy)]
1007pub struct WPtr8(pub [u64; 8]);
1008unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1009
1010#[repr(C)]
1015#[derive(Clone, Copy, Default)]
1016pub struct GdnSeqVl {
1017 pub kb16: u64,
1018 pub gcum: u64,
1019 pub beta: u64,
1020 pub u: u64,
1021 pub wb16: u64,
1022 pub y: u64,
1023 pub ssnap: u64,
1024 pub state_in: u64,
1025 pub state_out: u64,
1026 pub q: u64,
1027 pub p: u64,
1028 pub o: u64,
1029 pub k: u64,
1030 pub v: u64,
1031 pub g: u64,
1032 pub a: u64,
1033 pub w: u64,
1034 pub t: i32,
1035 pub nc: i32,
1036}
1037unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1038#[repr(C)]
1039#[derive(Clone, Copy)]
1040pub struct GdnVl8(pub [GdnSeqVl; 8]);
1041unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1042
1043#[repr(C)]
1046#[derive(Clone, Copy, Default)]
1047pub struct GdnWVl {
1048 pub qb16: u64,
1049 pub pb16: u64,
1050}
1051unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1052#[repr(C)]
1053#[derive(Clone, Copy)]
1054pub struct GdnWVl8(pub [GdnWVl; 8]);
1055unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1056
1057#[repr(C)]
1059#[derive(Clone, Copy, Default)]
1060pub struct GdnPrepVl {
1061 pub qkv: u64,
1062 pub conv_state: u64,
1063 pub conv_out: u64,
1064 pub q_g: u64,
1065 pub k_g: u64,
1066 pub v_g: u64,
1067 pub q_l2: u64,
1068 pub k_l2: u64,
1069 pub beta_raw: u64,
1070 pub alpha: u64,
1071 pub beta: u64,
1072 pub g_log: u64,
1073 pub o: u64,
1074 pub z: u64,
1075 pub gn: u64,
1076 pub gn16: u64,
1077 pub kb16: u64,
1078 pub qb16: u64,
1079 pub t: i32,
1080 pub pad: i32,
1081}
1082unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1083#[repr(C)]
1084#[derive(Clone, Copy)]
1085pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1086unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1087
1088#[repr(C)]
1090#[derive(Clone, Copy, Default)]
1091pub struct FaSeqVl {
1092 pub q: u64,
1093 pub k16: u64,
1094 pub v16: u64,
1095 pub o: u64,
1096 pub kf: u64,
1097 pub vf: u64,
1098 pub t: i32,
1099 pub pad: i32,
1100}
1101unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1102#[repr(C)]
1103#[derive(Clone, Copy)]
1104pub struct FaVl8(pub [FaSeqVl; 8]);
1105unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1106
1107#[repr(C)]
1109#[derive(Clone, Copy, Default)]
1110pub struct AttnPreVl {
1111 pub qf: u64,
1112 pub kf: u64,
1113 pub vf: u64,
1114 pub q: u64,
1115 pub gate: u64,
1116 pub qn: u64,
1117 pub kn: u64,
1118 pub kc: u64,
1119 pub vc: u64,
1120 pub t: i32,
1121 pub pad: i32,
1122}
1123unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1124#[repr(C)]
1125#[derive(Clone, Copy)]
1126pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1127unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1128
1129pub struct GdnChunkBufs {
1132 pub gcum: CudaSlice<f32>,
1133 pub a: CudaSlice<f32>,
1134 pub p: CudaSlice<f32>,
1135 pub u: CudaSlice<f32>,
1136 pub w: CudaSlice<f32>,
1137 pub kb16: CudaSlice<u8>,
1138 pub wb16: CudaSlice<u8>,
1139 pub y16: CudaSlice<u8>,
1140 pub ssnap16: CudaSlice<u8>,
1141 pub qb16: CudaSlice<u8>,
1142 pub pb16: CudaSlice<u8>,
1143 pub o: CudaSlice<f32>,
1144 pub t: usize,
1145 pub nc: usize,
1146}
1147
1148#[repr(C)]
1150#[derive(Clone, Copy)]
1151pub struct F32x8(pub [f32; 8]);
1152unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1153
1154pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1158
1159impl Engine {
1160 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1161 let gpu = memra_runtime::Gpu::new(ordinal)?;
1162 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1166 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1167 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1168 .and_then(|d| unsafe {
1169 Ok((
1170 cudarc::driver::result::device::get_attribute(
1171 d,
1172 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1173 )?,
1174 cudarc::driver::result::device::get_attribute(
1175 d,
1176 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1177 )?,
1178 ))
1179 })
1180 .unwrap_or((0, 0));
1181 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1182 let ok = matches!(
1183 (built, maj, min),
1184 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1185 );
1186 if !ok {
1187 return Err(format!(
1188 "memra was built for sm_{built} but device {ordinal} reports compute \
1189 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1190 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1191 )
1192 .into());
1193 }
1194 }
1195 unsafe {
1200 use cudarc::driver::sys;
1201 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1202 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1203 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1204 let mut thresh: u64 = u64::MAX;
1205 let _ = sys::cuMemPoolSetAttribute(
1206 pool,
1207 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1208 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1209 );
1210 }
1211 }
1212 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1213 let hybrid = gpu
1214 .ctx
1215 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1216 let qmatvec = gpu
1217 .ctx
1218 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1219 let flash = gpu
1220 .ctx
1221 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1222 let gemm = gpu
1223 .ctx
1224 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1225 let router = gpu
1226 .ctx
1227 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1228 let sample = gpu
1229 .ctx
1230 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1231 let copy_stream = gpu.ctx.new_stream()?;
1232 if std::env::var("MEMRA_EVT")
1248 .map(|v| v == "1")
1249 .unwrap_or(false)
1250 {
1251 } else {
1253 unsafe {
1254 gpu.ctx.disable_event_tracking();
1255 }
1256 }
1257 Ok(Self {
1258 gpu,
1259 module,
1260 hybrid,
1261 qmatvec,
1262 flash,
1263 flash_g: std::sync::OnceLock::new(),
1264 gemm,
1265 router,
1266 sample,
1267 moe_cache: Mutex::new(None),
1268 moe_cache_layout: Mutex::new(None),
1269 copy_stream,
1270 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1271 verify_exact: std::sync::atomic::AtomicBool::new(false),
1272 capture_keep: Mutex::new(Vec::new()),
1273 argmax_partials: Mutex::new(None),
1274 prime_deqw_ws: Mutex::new(None),
1275 router_stage: Mutex::new(None),
1276 fp8_scratch: Mutex::new(None),
1277 fa_vf16_scratch: Mutex::new(None),
1278 fa_part_pool: Mutex::new(None),
1279 fa_part_retired: Mutex::new(Vec::new()),
1280 fn_cache: Mutex::new(Default::default()),
1281 f16_scratch: Mutex::new(None),
1282 #[cfg(memra_cutlass)]
1283 cutlass_scratch: Mutex::new(None),
1284 })
1285 }
1286
1287 pub fn ctx(&self) -> &Arc<CudaContext> {
1288 &self.gpu.ctx
1289 }
1290
1291 pub fn pool_cached_bytes(&self) -> usize {
1309 let (reserved, used) = self.pool_reserved_used();
1310 reserved.saturating_sub(used)
1311 }
1312
1313 pub fn pool_reserved_used(&self) -> (usize, usize) {
1320 use cudarc::driver::sys;
1321 unsafe {
1322 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1323 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1324 != sys::CUresult::CUDA_SUCCESS
1325 {
1326 return (0, 0);
1327 }
1328 let (mut reserved, mut used) = (0u64, 0u64);
1329 if sys::cuMemPoolGetAttribute(
1330 pool,
1331 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1332 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1333 ) != sys::CUresult::CUDA_SUCCESS
1334 {
1335 return (0, 0);
1336 }
1337 if sys::cuMemPoolGetAttribute(
1338 pool,
1339 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1340 &mut used as *mut u64 as *mut core::ffi::c_void,
1341 ) != sys::CUresult::CUDA_SUCCESS
1342 {
1343 return (0, 0);
1344 }
1345 (reserved as usize, used as usize)
1346 }
1347 }
1348
1349 pub fn stream(&self) -> Arc<CudaStream> {
1352 self.gpu.stream()
1353 }
1354 pub fn gkv_on() -> bool {
1357 memra_kv::gkv_on()
1358 }
1359
1360 pub fn wkv_on() -> bool {
1372 memra_kv::wkv_on()
1373 }
1374
1375 pub fn kv_fp8_on() -> bool {
1381 memra_kv::kv_fp8_on()
1382 }
1383
1384 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1387 if head_dim == 512 && Self::gkv_on() {
1388 self.func_g(name)
1389 } else {
1390 self.func(name)
1391 }
1392 }
1393
1394 fn func_g(&self, name: &str) -> CudaFunction {
1398 let m = self.flash_g.get_or_init(|| {
1399 self.gpu
1400 .ctx
1401 .load_module(cudarc::nvrtc::Ptx::from_binary(
1402 FLASH_FATBIN_KF8VF8.to_vec(),
1403 ))
1404 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1405 });
1406 let key = format!("g:{name}");
1407 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1408 return f.clone();
1409 }
1410 let f = match m.load_function(name) {
1411 Ok(f) => f,
1412 Err(_) => self.func(name),
1413 };
1414 self.fn_cache.lock().unwrap().insert(key, f.clone());
1415 f
1416 }
1417
1418 fn func(&self, name: &str) -> CudaFunction {
1419 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1422 return f.clone();
1423 }
1424 let f = self
1425 .module
1426 .load_function(name)
1427 .or_else(|_| self.hybrid.load_function(name))
1428 .or_else(|_| self.qmatvec.load_function(name))
1429 .or_else(|_| self.flash.load_function(name))
1430 .or_else(|_| self.gemm.load_function(name))
1431 .or_else(|_| self.router.load_function(name))
1432 .or_else(|_| self.sample.load_function(name))
1433 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1434 self.fn_cache
1435 .lock()
1436 .unwrap()
1437 .insert(name.to_string(), f.clone());
1438 f
1439 }
1440
1441 pub fn scatter_trim_logits(
1444 &self,
1445 src: &CudaSlice<f32>,
1446 d2t: &CudaSlice<u32>,
1447 dst: &mut CudaSlice<f32>,
1448 d_vocab: usize,
1449 n_vocab: usize,
1450 ) -> Result<(), Box<dyn std::error::Error>> {
1451 let f1 = self.func("scatter_trim_logits_f32");
1452 let f2 = self.func("scatter_trim_logits_pass2_f32");
1453 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1454 let cfg1 = LaunchConfig {
1455 grid_dim: (256, 1, 1),
1456 block_dim: (256, 1, 1),
1457 shared_mem_bytes: 0,
1458 };
1459 let __s_b1 = self.gpu.stream();
1460 let mut b1 = __s_b1.launch_builder(&f1);
1461 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1462 unsafe {
1463 b1.launch(cfg1)?;
1464 }
1465 let cfg2 = LaunchConfig {
1466 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1467 block_dim: (256, 1, 1),
1468 shared_mem_bytes: 0,
1469 };
1470 let __s_b2 = self.gpu.stream();
1471 let mut b2 = __s_b2.launch_builder(&f2);
1472 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1473 unsafe {
1474 b2.launch(cfg2)?;
1475 }
1476 Ok(())
1477 }
1478
1479 #[allow(clippy::too_many_arguments)]
1485 pub fn filter_stats(
1486 &self,
1487 x: &CudaSlice<f32>,
1488 row_stride: usize,
1489 rows: &CudaSlice<i32>,
1490 out_th: &mut CudaSlice<f32>,
1491 out_z: &mut CudaSlice<f32>,
1492 out_max: &mut CudaSlice<f32>,
1493 n: usize,
1494 nrow: usize,
1495 temp: f32,
1496 top_k: i32,
1497 top_p: f32,
1498 min_p: f32,
1499 ) -> Result<(), Box<dyn std::error::Error>> {
1500 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1515 let coop_on =
1516 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1517 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1518 if coop_on && 16 * nrow <= self.sm_count() as usize {
1519 let f = self.func("filter_stats_coop_f32");
1520 let mut ws = self.alloc_uninit::<f32>(nrow * (2 * 16 + 2))?;
1521 let cfg = LaunchConfig {
1522 grid_dim: (16, nrow as u32, 1),
1523 block_dim: (512, 1, 1),
1524 shared_mem_bytes: 0,
1525 };
1526 let __s_b = self.gpu.stream();
1527 let mut b = __s_b.launch_builder(&f);
1528 b.arg(x)
1529 .arg(&rs)
1530 .arg(rows)
1531 .arg(&mut *out_th)
1532 .arg(&mut *out_z)
1533 .arg(&mut *out_max)
1534 .arg(&mut ws)
1535 .arg(&ni)
1536 .arg(&nr)
1537 .arg(&temp)
1538 .arg(&top_k)
1539 .arg(&top_p)
1540 .arg(&min_p);
1541 unsafe {
1542 b.launch_cooperative(cfg)?;
1543 }
1544 return Ok(());
1545 }
1546 let f = self.func("filter_stats_f32");
1547 let cfg = LaunchConfig {
1548 grid_dim: (nrow as u32, 1, 1),
1549 block_dim: (1024, 1, 1),
1550 shared_mem_bytes: 0,
1551 };
1552 let __s_b = self.gpu.stream();
1553 let mut b = __s_b.launch_builder(&f);
1554 b.arg(x)
1555 .arg(&rs)
1556 .arg(rows)
1557 .arg(&mut *out_th)
1558 .arg(&mut *out_z)
1559 .arg(&mut *out_max)
1560 .arg(&ni)
1561 .arg(&nr)
1562 .arg(&temp)
1563 .arg(&top_k)
1564 .arg(&top_p)
1565 .arg(&min_p);
1566 unsafe {
1567 b.launch(cfg)?;
1568 }
1569 Ok(())
1570 }
1571
1572 #[allow(clippy::too_many_arguments)]
1574 pub fn softmax_gather_filtered(
1575 &self,
1576 x: &CudaSlice<f32>,
1577 row_stride: usize,
1578 ids: &CudaSlice<u32>,
1579 rows: &CudaSlice<i32>,
1580 th: &CudaSlice<f32>,
1581 z: &CudaSlice<f32>,
1582 out: &mut CudaSlice<f32>,
1583 n: usize,
1584 npair: usize,
1585 temp: f32,
1586 ) -> Result<(), Box<dyn std::error::Error>> {
1587 let f = self.func("softmax_gather_filtered_f32");
1588 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1589 let cfg = LaunchConfig {
1590 grid_dim: (npair as u32, 1, 1),
1591 block_dim: (256, 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(ids)
1599 .arg(rows)
1600 .arg(th)
1601 .arg(z)
1602 .arg(&mut *out)
1603 .arg(&ni)
1604 .arg(&np)
1605 .arg(&temp);
1606 unsafe {
1607 b.launch(cfg)?;
1608 }
1609 Ok(())
1610 }
1611
1612 #[allow(clippy::too_many_arguments)]
1614 pub fn residual_sample_filtered(
1615 &self,
1616 p: &CudaSlice<f32>,
1617 q: Option<&CudaSlice<f32>>,
1618 n: usize,
1619 temp: f32,
1620 seed: u64,
1621 stream_pos: u32,
1622 p_stats: (f32, f32, f32),
1623 q_stats: (f32, f32, f32),
1624 out_tok: &mut CudaSlice<u32>,
1625 ) -> Result<(), Box<dyn std::error::Error>> {
1626 let f = self.func("residual_sample_filtered_f32");
1627 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1628 let has_q: i32 = q.is_some() as i32;
1629 let qbuf = q.unwrap_or(p);
1630 let (pm, pth, pz) = p_stats;
1631 let (qm, qth, qz) = q_stats;
1632 let cfg = LaunchConfig {
1633 grid_dim: (1, 1, 1),
1634 block_dim: (1024, 1, 1),
1635 shared_mem_bytes: 0,
1636 };
1637 let __s_b = self.gpu.stream();
1638 let mut b = __s_b.launch_builder(&f);
1639 b.arg(p)
1640 .arg(qbuf)
1641 .arg(&has_q)
1642 .arg(&ni)
1643 .arg(&temp)
1644 .arg(&slo)
1645 .arg(&shi)
1646 .arg(&stream_pos)
1647 .arg(&pm)
1648 .arg(&pth)
1649 .arg(&pz)
1650 .arg(&qm)
1651 .arg(&qth)
1652 .arg(&qz)
1653 .arg(&mut *out_tok);
1654 unsafe {
1655 b.launch(cfg)?;
1656 }
1657 Ok(())
1658 }
1659
1660 #[allow(clippy::too_many_arguments)]
1666 pub fn residual_sample_sparse_q(
1667 &self,
1668 p: &CudaSlice<f32>,
1669 cand_ids: &CudaSlice<u32>,
1670 q_probs: &CudaSlice<f32>,
1671 n_cand: usize,
1672 n: usize,
1673 temp: f32,
1674 seed: u64,
1675 stream_pos: u32,
1676 p_stats: (f32, f32, f32),
1677 out_tok: &mut CudaSlice<u32>,
1678 ) -> Result<(), Box<dyn std::error::Error>> {
1679 assert!(
1680 n_cand >= 1 && n_cand <= 32,
1681 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
1682 );
1683 let f = self.func("residual_sample_sparse_q_f32");
1684 let (ni, nc) = (n as i32, n_cand as i32);
1685 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1686 let (pm, pth, pz) = p_stats;
1687 let cfg = LaunchConfig {
1688 grid_dim: (1, 1, 1),
1689 block_dim: (1024, 1, 1),
1690 shared_mem_bytes: 0,
1691 };
1692 let __s_b = self.gpu.stream();
1693 let mut b = __s_b.launch_builder(&f);
1694 b.arg(p)
1695 .arg(cand_ids)
1696 .arg(q_probs)
1697 .arg(&nc)
1698 .arg(&ni)
1699 .arg(&temp)
1700 .arg(&slo)
1701 .arg(&shi)
1702 .arg(&stream_pos)
1703 .arg(&pm)
1704 .arg(&pth)
1705 .arg(&pz)
1706 .arg(&mut *out_tok);
1707 unsafe {
1708 b.launch(cfg)?;
1709 }
1710 Ok(())
1711 }
1712
1713 #[allow(clippy::too_many_arguments)]
1715 pub fn gumbel_perturb_filtered(
1716 &self,
1717 x: &CudaSlice<f32>,
1718 y: &mut CudaSlice<f32>,
1719 n: usize,
1720 seed: u64,
1721 stream_pos: u32,
1722 temp: f32,
1723 row_max: f32,
1724 th: f32,
1725 ) -> Result<(), Box<dyn std::error::Error>> {
1726 let f = self.func("gumbel_perturb_filtered_f32");
1727 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1728 let cfg = LaunchConfig {
1729 grid_dim: (n.div_ceil(256) as u32, 1, 1),
1730 block_dim: (256, 1, 1),
1731 shared_mem_bytes: 0,
1732 };
1733 let __s_b = self.gpu.stream();
1734 let mut b = __s_b.launch_builder(&f);
1735 b.arg(x)
1736 .arg(&mut *y)
1737 .arg(&ni)
1738 .arg(&slo)
1739 .arg(&shi)
1740 .arg(&stream_pos)
1741 .arg(&temp)
1742 .arg(&row_max)
1743 .arg(&th);
1744 unsafe {
1745 b.launch(cfg)?;
1746 }
1747 Ok(())
1748 }
1749
1750 #[allow(clippy::too_many_arguments)]
1754 pub fn penalize_logits(
1755 &self,
1756 x: &mut CudaSlice<f32>,
1757 hist: &CudaSlice<u32>,
1758 n_hist: usize,
1759 rep: f32,
1760 freq: f32,
1761 present: f32,
1762 n: usize,
1763 ) -> Result<(), Box<dyn std::error::Error>> {
1764 if n_hist == 0 {
1765 return Ok(());
1766 }
1767 let f = self.func("penalize_logits_f32");
1768 let (nh, ni) = (n_hist as i32, n as i32);
1769 let cfg = LaunchConfig {
1770 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
1771 block_dim: (128, 1, 1),
1772 shared_mem_bytes: 0,
1773 };
1774 let __s_b = self.gpu.stream();
1775 let mut b = __s_b.launch_builder(&f);
1776 b.arg(&mut *x)
1777 .arg(hist)
1778 .arg(&nh)
1779 .arg(&rep)
1780 .arg(&freq)
1781 .arg(&present)
1782 .arg(&ni);
1783 unsafe {
1784 b.launch(cfg)?;
1785 }
1786 Ok(())
1787 }
1788
1789 #[allow(clippy::too_many_arguments)]
1791 pub fn penalize_logits_rows(
1792 &self,
1793 x: &mut CudaSlice<f32>,
1794 hist: &CudaSlice<u32>,
1795 n_hist: usize,
1796 rep: f32,
1797 freq: f32,
1798 present: f32,
1799 n: usize,
1800 nrow: usize,
1801 ) -> Result<(), Box<dyn std::error::Error>> {
1802 if n_hist == 0 || nrow == 0 {
1803 return Ok(());
1804 }
1805 let f = self.func("penalize_logits_rows_f32");
1806 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1807 let cfg = LaunchConfig {
1808 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
1809 block_dim: (128, 1, 1),
1810 shared_mem_bytes: 0,
1811 };
1812 let __s_b = self.gpu.stream();
1813 let mut b = __s_b.launch_builder(&f);
1814 b.arg(&mut *x)
1815 .arg(hist)
1816 .arg(&nh)
1817 .arg(&rep)
1818 .arg(&freq)
1819 .arg(&present)
1820 .arg(&ni)
1821 .arg(&nr);
1822 unsafe {
1823 b.launch(cfg)?;
1824 }
1825 Ok(())
1826 }
1827
1828 #[allow(clippy::too_many_arguments)]
1836 pub fn penalize_logits_rows_inc(
1837 &self,
1838 x: &mut CudaSlice<f32>,
1839 hist: &CudaSlice<u32>,
1840 n_hist0: usize,
1841 rep: f32,
1842 freq: f32,
1843 present: f32,
1844 n: usize,
1845 nrow: usize,
1846 win: usize,
1847 ) -> Result<(), Box<dyn std::error::Error>> {
1848 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
1849 return Ok(());
1850 }
1851 debug_assert!(
1852 hist.len() >= n_hist0 + nrow - 1,
1853 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
1854 );
1855 let f = self.func("penalize_logits_rows_inc_f32");
1856 let max_len = win.min(n_hist0 + nrow - 1).max(1);
1857 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
1858 let cfg = LaunchConfig {
1859 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
1860 block_dim: (128, 1, 1),
1861 shared_mem_bytes: 0,
1862 };
1863 let __s_b = self.gpu.stream();
1864 let mut b = __s_b.launch_builder(&f);
1865 b.arg(&mut *x)
1866 .arg(hist)
1867 .arg(&nh)
1868 .arg(&rep)
1869 .arg(&freq)
1870 .arg(&present)
1871 .arg(&ni)
1872 .arg(&nr)
1873 .arg(&wi);
1874 unsafe {
1875 b.launch(cfg)?;
1876 }
1877 Ok(())
1878 }
1879
1880 pub fn wpf_level() -> u32 {
1888 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1889 *ON.get_or_init(|| {
1890 std::env::var("MEMRA_WPF")
1891 .ok()
1892 .and_then(|v| v.parse().ok())
1893 .unwrap_or(1)
1894 })
1895 }
1896
1897 pub fn set_verify_exact(&self, on: bool) {
1909 self.verify_exact
1910 .store(on, std::sync::atomic::Ordering::Relaxed);
1911 }
1912 pub(crate) fn verify_exact_on(&self) -> bool {
1913 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1914 }
1915
1916 pub fn qkv_append_on() -> bool {
1919 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1920 *ON.get_or_init(|| {
1921 std::env::var("MEMRA_QKV_APPEND")
1922 .map(|v| v != "0")
1923 .unwrap_or(true)
1924 })
1925 }
1926
1927 pub fn pdl_wb_on() -> bool {
1930 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1931 *ON.get_or_init(|| {
1932 std::env::var("MEMRA_PDL_WB")
1933 .map(|v| v != "0")
1934 .unwrap_or(true)
1935 })
1936 }
1937
1938 pub fn norm_ilp_on() -> bool {
1946 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1947 *ON.get_or_init(|| {
1948 std::env::var("MEMRA_NORM_ILP")
1949 .map(|v| v != "0")
1950 .unwrap_or(true)
1951 })
1952 }
1953
1954 pub fn tk_ffn_dual_on() -> bool {
1962 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1963 *ON.get_or_init(|| {
1964 std::env::var("MEMRA_TK_FFN_DUAL")
1965 .map(|v| v != "0")
1966 .unwrap_or(true)
1967 })
1968 }
1969
1970 pub fn pdl_mmvq_on() -> bool {
1974 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1975 *ON.get_or_init(|| {
1976 std::env::var("MEMRA_PDL_MMVQ")
1977 .map(|v| v != "0")
1978 .unwrap_or(true)
1979 })
1980 }
1981
1982 pub fn pdl_on() -> bool {
1983 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1984 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1985 }
1986
1987 pub fn pdl_nvfp4q8_on() -> bool {
1993 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1994 *ON.get_or_init(|| {
1995 std::env::var("MEMRA_PDL_NVFP4")
1996 .map(|v| v != "0")
1997 .unwrap_or(true)
1998 })
1999 }
2000
2001 fn q40_mr1_on() -> bool {
2007 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2008 match *Q40MR.get_or_init(|| {
2009 std::env::var("MEMRA_Q40_MR")
2010 .ok()
2011 .and_then(|v| v.parse().ok())
2012 }) {
2013 Some(v) => v == 1,
2014 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2015 }
2016 }
2017
2018 fn pdl_func_flash(
2023 &self,
2024 g: bool,
2025 name: &'static str,
2026 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2027 use cudarc::driver::sys as cu;
2028 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2035 std::sync::Mutex::new(None);
2036 static FNS: std::sync::Mutex<
2037 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2038 > = std::sync::Mutex::new(None);
2039 let ctx_key = self.ctx().cu_ctx() as usize;
2040 if let Some(&f) = FNS
2041 .lock()
2042 .unwrap()
2043 .get_or_insert_with(Default::default)
2044 .get(&(ctx_key, g, name))
2045 {
2046 return Ok(f as cu::CUfunction);
2047 }
2048 let module = {
2049 let mut mods = MODS.lock().unwrap();
2050 let map = mods.get_or_insert_with(Default::default);
2051 match map.get(&(ctx_key, g)) {
2052 Some(&m) => m,
2053 None => {
2054 let m = self.pdl_load_module_in_ctx(if g {
2055 FLASH_FATBIN_KF8VF8
2056 } else {
2057 FLASH_FATBIN
2058 })?;
2059 map.insert((ctx_key, g), m);
2060 m
2061 }
2062 }
2063 };
2064 let cname = std::ffi::CString::new(name)?;
2065 let mut f: cu::CUfunction = std::ptr::null_mut();
2066 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2067 if r != cu::CUresult::CUDA_SUCCESS {
2068 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2069 }
2070 FNS.lock()
2071 .unwrap()
2072 .get_or_insert_with(Default::default)
2073 .insert((ctx_key, g, name), f as usize);
2074 Ok(f)
2075 }
2076
2077 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2082 use cudarc::driver::sys as cu;
2083 let mut prev: cu::CUcontext = std::ptr::null_mut();
2084 unsafe {
2085 cu::cuCtxGetCurrent(&mut prev).result()?;
2086 }
2087 self.ctx().bind_to_thread()?;
2088 let mut m: cu::CUmodule = std::ptr::null_mut();
2089 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2090 let restore = if prev.is_null() {
2091 cu::CUresult::CUDA_SUCCESS
2092 } else {
2093 unsafe { cu::cuCtxSetCurrent(prev) }
2094 };
2095 if r != cu::CUresult::CUDA_SUCCESS {
2096 return Err(format!("pdl module load: {r:?}").into());
2097 }
2098 if restore != cu::CUresult::CUDA_SUCCESS {
2099 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2100 }
2101 Ok(m as usize)
2102 }
2103
2104 pub fn raw_kernel_function(
2107 &self,
2108 name: &'static str,
2109 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2110 self.pdl_func(name)
2111 }
2112
2113 fn pdl_func(
2114 &self,
2115 name: &'static str,
2116 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2117 use cudarc::driver::sys as cu;
2118 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2121 std::sync::Mutex::new(None);
2122 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2125 std::sync::Mutex::new(None);
2126 static FNS: std::sync::Mutex<
2127 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2128 > = std::sync::Mutex::new(None);
2129 let ctx_key = self.ctx().cu_ctx() as usize;
2130 if let Some(&f) = FNS
2131 .lock()
2132 .unwrap()
2133 .get_or_insert_with(Default::default)
2134 .get(&(ctx_key, name))
2135 {
2136 return Ok(f as cu::CUfunction);
2137 }
2138 let module = {
2139 let mut mods = MODULES.lock().unwrap();
2140 let map = mods.get_or_insert_with(Default::default);
2141 match map.get(&ctx_key) {
2142 Some(&m) => m,
2143 None => {
2144 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2145 map.insert(ctx_key, m);
2146 m
2147 }
2148 }
2149 };
2150 let cname = std::ffi::CString::new(name)?;
2151 let mut f: cu::CUfunction = std::ptr::null_mut();
2152 let mut r =
2153 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2154 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2155 let qmodule = {
2156 let mut mods = QMODULES.lock().unwrap();
2157 let map = mods.get_or_insert_with(Default::default);
2158 match map.get(&ctx_key) {
2159 Some(&m) => m,
2160 None => {
2161 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2162 map.insert(ctx_key, m);
2163 m
2164 }
2165 }
2166 };
2167 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2168 }
2169 if r != cu::CUresult::CUDA_SUCCESS {
2170 return Err(format!("pdl_func {name}: {r:?}").into());
2171 }
2172 FNS.lock()
2173 .unwrap()
2174 .get_or_insert_with(Default::default)
2175 .insert((ctx_key, name), f as usize);
2176 Ok(f)
2177 }
2178
2179 unsafe fn launch_pdl_flash(
2191 &self,
2192 g: bool,
2193 name: &'static str,
2194 grid: (u32, u32, u32),
2195 block: (u32, u32, u32),
2196 smem: u32,
2197 params: &mut [*mut std::ffi::c_void],
2198 ) -> Result<(), Box<dyn std::error::Error>> {
2199 use cudarc::driver::sys as cu;
2200 let f = self.pdl_func_flash(g, name)?;
2201 if smem > 0 {
2202 let r =
2204 unsafe {
2205 cu::cuFuncSetAttribute(f,
2206 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2207 smem as i32)
2208 };
2209 if r != cu::CUresult::CUDA_SUCCESS {
2210 return Err(format!("pdl smem attr {name}: {r:?}").into());
2211 }
2212 }
2213 let mut attr = cu::CUlaunchAttribute {
2214 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2215 pad: [0; 4],
2216 value: cu::CUlaunchAttributeValue {
2217 programmaticStreamSerializationAllowed: 1,
2218 },
2219 };
2220 let cfg = cu::CUlaunchConfig {
2221 gridDimX: grid.0,
2222 gridDimY: grid.1,
2223 gridDimZ: grid.2,
2224 blockDimX: block.0,
2225 blockDimY: block.1,
2226 blockDimZ: block.2,
2227 sharedMemBytes: smem,
2228 hStream: self.gpu.stream().cu_stream(),
2229 attrs: &mut attr,
2230 numAttrs: 1,
2231 };
2232 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2233 if r != cu::CUresult::CUDA_SUCCESS {
2234 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2235 }
2236 Ok(())
2237 }
2238
2239 unsafe fn launch_pdl(
2240 &self,
2241 name: &'static str,
2242 grid: (u32, u32, u32),
2243 block: (u32, u32, u32),
2244 params: &mut [*mut std::ffi::c_void],
2245 ) -> Result<(), Box<dyn std::error::Error>> {
2246 use cudarc::driver::sys as cu;
2247 let f = self.pdl_func(name)?;
2248 let mut attr = cu::CUlaunchAttribute {
2249 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2250 pad: [0; 4],
2251 value: cu::CUlaunchAttributeValue {
2252 programmaticStreamSerializationAllowed: 1,
2253 },
2254 };
2255 let cfg = cu::CUlaunchConfig {
2256 gridDimX: grid.0,
2257 gridDimY: grid.1,
2258 gridDimZ: grid.2,
2259 blockDimX: block.0,
2260 blockDimY: block.1,
2261 blockDimZ: block.2,
2262 sharedMemBytes: 0,
2263 hStream: self.gpu.stream().cu_stream(),
2264 attrs: &mut attr,
2265 numAttrs: 1,
2266 };
2267 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2268 if r != cu::CUresult::CUDA_SUCCESS {
2269 return Err(format!("launch_pdl {name}: {r:?}").into());
2270 }
2271 Ok(())
2272 }
2273
2274 pub fn prefetch_weight_l2(
2277 &self,
2278 w: &crate::model::GpuTensor,
2279 ) -> Result<(), Box<dyn std::error::Error>> {
2280 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2281 let p = rp4.as_ref().unwrap_or(bytes);
2282 self.prefetch_l2(p, p.len())?;
2283 }
2284 Ok(())
2285 }
2286
2287 pub fn gather_row_bf16(
2290 &self,
2291 table: &CudaSlice<u8>,
2292 tok: &CudaSlice<u32>,
2293 idx: usize,
2294 dst: &mut CudaSlice<f32>,
2295 ncols: usize,
2296 ) -> Result<(), Box<dyn std::error::Error>> {
2297 let f = self.func("gather_row_bf16_f32");
2298 let cfg = LaunchConfig {
2299 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2300 block_dim: (256, 1, 1),
2301 shared_mem_bytes: 0,
2302 };
2303 let (nc, ix) = (ncols as i32, idx as i32);
2304 let __s_b = self.gpu.stream();
2305 let mut b = __s_b.launch_builder(&f);
2306 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2307 unsafe {
2308 b.launch(cfg)?;
2309 }
2310 Ok(())
2311 }
2312
2313 #[allow(clippy::too_many_arguments)]
2319 pub fn dflash2_dynconv(
2320 &self,
2321 x: &CudaSlice<f32>,
2322 dyn_: &CudaSlice<f32>,
2323 base: &CudaSlice<f32>,
2324 out: &mut CudaSlice<f32>,
2325 rows: usize,
2326 hidden: usize,
2327 group_size: usize,
2328 ksize: usize,
2329 half: usize,
2330 ) -> Result<(), Box<dyn std::error::Error>> {
2331 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2332 let f = self.func("dflash2_dynconv_f32");
2333 let n = rows * hidden;
2334 let cfg = LaunchConfig {
2335 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2336 block_dim: (256, 1, 1),
2337 shared_mem_bytes: 0,
2338 };
2339 let (ri, hi, gi, ki, hf) = (
2340 rows as i32,
2341 hidden as i32,
2342 group_size as i32,
2343 ksize as i32,
2344 half as i32,
2345 );
2346 let __s_b = self.gpu.stream();
2347 let mut b = __s_b.launch_builder(&f);
2348 b.arg(x)
2349 .arg(dyn_)
2350 .arg(base)
2351 .arg(out)
2352 .arg(&ri)
2353 .arg(&hi)
2354 .arg(&gi)
2355 .arg(&ki)
2356 .arg(&hf);
2357 unsafe {
2358 b.launch(cfg)?;
2359 }
2360 Ok(())
2361 }
2362
2363 pub fn topk_rows(
2367 &self,
2368 logits: &CudaSlice<f32>,
2369 n_rows: usize,
2370 n_cols: usize,
2371 k: usize,
2372 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2373 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2374 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2375 let f = self.func("topk_rows_f32");
2376 let nth = 256usize;
2377 let mut vals = self.uninit(n_rows * k)?;
2378 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2379 let cfg = LaunchConfig {
2380 grid_dim: (n_rows as u32, 1, 1),
2381 block_dim: (nth as u32, 1, 1),
2382 shared_mem_bytes: (nth * k * 8) as u32,
2383 };
2384 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2385 let __s_b = self.gpu.stream();
2386 let mut b = __s_b.launch_builder(&f);
2387 b.arg(logits)
2388 .arg(&nr)
2389 .arg(&nc)
2390 .arg(&ki)
2391 .arg(&mut vals)
2392 .arg(&mut idxs);
2393 unsafe {
2394 b.launch(cfg)?;
2395 }
2396 Ok((vals, idxs))
2397 }
2398
2399 pub fn add_row_inplace(
2401 &self,
2402 logits: &mut CudaSlice<f32>,
2403 bias: &CudaSlice<f32>,
2404 n: usize,
2405 row_off: usize,
2406 ) -> Result<(), Box<dyn std::error::Error>> {
2407 let f = self.func("add_row_inplace_f32");
2408 let cfg = LaunchConfig {
2409 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2410 block_dim: (256, 1, 1),
2411 shared_mem_bytes: 0,
2412 };
2413 let (ni, off) = (n as i32, row_off as i64);
2414 let __s_b = self.gpu.stream();
2415 let mut b = __s_b.launch_builder(&f);
2416 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2417 unsafe {
2418 b.launch(cfg)?;
2419 }
2420 Ok(())
2421 }
2422
2423 pub fn prefetch_l2(
2425 &self,
2426 p: &CudaSlice<u8>,
2427 n: usize,
2428 ) -> Result<(), Box<dyn std::error::Error>> {
2429 let f = self.func("prefetch_l2_bytes");
2430 let lines = n.div_ceil(128);
2431 let ni = n as i64;
2432 let cfg = LaunchConfig {
2433 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2434 block_dim: (256, 1, 1),
2435 shared_mem_bytes: 0,
2436 };
2437 let __s_b = self.gpu.stream();
2438 let mut b = __s_b.launch_builder(&f);
2439 b.arg(p).arg(&ni);
2440 unsafe {
2441 b.launch(cfg)?;
2442 }
2443 Ok(())
2444 }
2445
2446 pub fn router_gemv(
2449 &self,
2450 w: &CudaSlice<f32>,
2451 x: &CudaSlice<f32>,
2452 n_embd: usize,
2453 n_experts: usize,
2454 t: usize,
2455 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2456 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2462 Ok("0") => false,
2463 Ok(_) => true,
2464 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2465 };
2466 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
2476 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
2477 }
2478
2479 pub fn router_gemv_form(
2482 &self,
2483 w: &CudaSlice<f32>,
2484 x: &CudaSlice<f32>,
2485 n_embd: usize,
2486 n_experts: usize,
2487 t: usize,
2488 w8: bool,
2489 batch: bool,
2490 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2491 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
2492 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
2493 let f = if batch {
2494 self.func("router_gemv_f32_w8_batch")
2495 } else if w8 {
2496 self.func("router_gemv_f32_w8")
2497 } else {
2498 self.func("router_gemv_f32")
2499 };
2500 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2501 let cfg = if batch {
2502 LaunchConfig {
2503 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
2504 block_dim: (32, 8, 1),
2505 shared_mem_bytes: 0,
2506 }
2507 } else {
2508 LaunchConfig {
2509 grid_dim: (n_experts as u32, t as u32, 1),
2510 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2511 shared_mem_bytes: 0,
2512 }
2513 };
2514 let __s_b = self.gpu.stream();
2515 let mut b = __s_b.launch_builder(&f);
2516 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
2517 unsafe {
2518 b.launch(cfg)?;
2519 }
2520 Ok(y)
2521 }
2522
2523 pub fn router_gemv_into(
2526 &self,
2527 w: &CudaSlice<f32>,
2528 x: &CudaSlice<f32>,
2529 y: &mut CudaSlice<f32>,
2530 n_embd: usize,
2531 n_experts: usize,
2532 t: usize,
2533 ) -> Result<(), Box<dyn std::error::Error>> {
2534 if y.len() < t * n_experts {
2535 return Err("router_gemv_into output too small".into());
2536 }
2537 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
2538 Ok("0") => false,
2539 Ok(_) => true,
2540 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2541 };
2542 let f = if w8 {
2543 self.func("router_gemv_f32_w8")
2544 } else {
2545 self.func("router_gemv_f32")
2546 };
2547 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
2548 let cfg = LaunchConfig {
2549 grid_dim: (n_experts as u32, t as u32, 1),
2550 block_dim: (32, if w8 { 8 } else { 1 }, 1),
2551 shared_mem_bytes: 0,
2552 };
2553 let __s_b = self.gpu.stream();
2554 let mut b = __s_b.launch_builder(&f);
2555 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
2556 unsafe {
2557 b.launch(cfg)?;
2558 }
2559 Ok(())
2560 }
2561
2562 pub fn rows_permute(
2564 &self,
2565 src: &CudaSlice<f32>,
2566 idx: &CudaSlice<i32>,
2567 nrows: usize,
2568 ncols: usize,
2569 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2570 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
2571 let f = self.func("rows_permute_f32");
2572 let (nc, nr) = (ncols as i32, nrows as i32);
2573 let cfg = LaunchConfig {
2574 grid_dim: (nrows as u32, 1, 1),
2575 block_dim: (256, 1, 1),
2576 shared_mem_bytes: 0,
2577 };
2578 let __s_b = self.gpu.stream();
2579 let mut b = __s_b.launch_builder(&f);
2580 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
2581 unsafe {
2582 b.launch(cfg)?;
2583 }
2584 Ok(dst)
2585 }
2586
2587 pub fn sigmoid_dot_rows(
2592 &self,
2593 x: &CudaSlice<f32>,
2594 w: &CudaSlice<f32>,
2595 n_embd: usize,
2596 t: usize,
2597 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2598 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2601 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
2602 let gs = self.linear(x, w, t, n_embd, 1)?;
2603 let mut g = self.uninit(t)?;
2604 self.sigmoid(&gs, &mut g, t)?;
2605 return Ok(g);
2606 }
2607 let mut g = self.alloc_uninit::<f32>(t)?;
2613 let f = self.func("sigmoid_dot_rows_f32");
2614 let (ne, ti) = (n_embd as i32, t as i32);
2615 let cfg = LaunchConfig {
2616 grid_dim: (t as u32, 1, 1),
2617 block_dim: (32, 8, 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(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
2623 unsafe {
2624 b.launch(cfg)?;
2625 }
2626 Ok(g)
2627 }
2628
2629 pub fn sigmoid_dot_rows_into(
2631 &self,
2632 x: &CudaSlice<f32>,
2633 w: &CudaSlice<f32>,
2634 g: &mut CudaSlice<f32>,
2635 n_embd: usize,
2636 t: usize,
2637 ) -> Result<(), Box<dyn std::error::Error>> {
2638 if g.len() < t {
2639 return Err("sigmoid_dot_rows_into output too small".into());
2640 }
2641 let f = self.func("sigmoid_dot_rows_f32");
2642 let (ne, ti) = (n_embd as i32, t as i32);
2643 let cfg = LaunchConfig {
2644 grid_dim: (t as u32, 1, 1),
2645 block_dim: (32, 8, 1),
2646 shared_mem_bytes: 0,
2647 };
2648 let __s_b = self.gpu.stream();
2649 let mut b = __s_b.launch_builder(&f);
2650 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
2651 unsafe {
2652 b.launch(cfg)?;
2653 }
2654 Ok(())
2655 }
2656
2657 pub fn spec_rollback_stream(
2659 &self,
2660 len_ptrs: &CudaSlice<u64>,
2661 pos_start: &CudaSlice<i32>,
2662 acc: &CudaSlice<u32>,
2663 base: usize,
2664 n_rows: usize,
2665 ) -> Result<(), Box<dyn std::error::Error>> {
2666 let f = self.func("spec_rollback_stream");
2667 let (b, nr) = (base as i32, n_rows as i32);
2668 let cfg = LaunchConfig {
2669 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
2670 block_dim: (64, 1, 1),
2671 shared_mem_bytes: 0,
2672 };
2673 let __s_bl = self.gpu.stream();
2674 let mut bl = __s_bl.launch_builder(&f);
2675 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
2676 unsafe {
2677 bl.launch(cfg)?;
2678 }
2679 Ok(())
2680 }
2681
2682 pub fn plain_tok_ring(
2684 &self,
2685 vam: &CudaSlice<u32>,
2686 pos_start: &CudaSlice<i32>,
2687 base: usize,
2688 ring: &mut CudaSlice<u32>,
2689 ) -> Result<(), Box<dyn std::error::Error>> {
2690 let f = self.func("plain_tok_ring");
2691 let (b, cap) = (base as i32, ring.len() as i32);
2692 let cfg = LaunchConfig {
2693 grid_dim: (1, 1, 1),
2694 block_dim: (32, 1, 1),
2695 shared_mem_bytes: 0,
2696 };
2697 let __s_bl = self.gpu.stream();
2698 let mut bl = __s_bl.launch_builder(&f);
2699 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
2700 unsafe {
2701 bl.launch(cfg)?;
2702 }
2703 Ok(())
2704 }
2705
2706 pub fn spec_ring_commit(
2708 &self,
2709 vtok: &CudaSlice<u32>,
2710 acc: &CudaSlice<u32>,
2711 brk: &CudaSlice<u32>,
2712 ring: &mut CudaSlice<u32>,
2713 pend: &mut CudaSlice<u32>,
2714 ) -> Result<(), Box<dyn std::error::Error>> {
2715 let f = self.func("spec_ring_commit");
2716 let cfg = LaunchConfig {
2717 grid_dim: (1, 1, 1),
2718 block_dim: (32, 1, 1),
2719 shared_mem_bytes: 0,
2720 };
2721 let __s_b = self.gpu.stream();
2722 let mut b = __s_b.launch_builder(&f);
2723 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
2724 unsafe {
2725 b.launch(cfg)?;
2726 }
2727 Ok(())
2728 }
2729 pub fn i32_copy_add(
2730 &self,
2731 src: &CudaSlice<i32>,
2732 dst: &mut CudaSlice<i32>,
2733 delta: i32,
2734 ) -> Result<(), Box<dyn std::error::Error>> {
2735 let f = self.func("i32_copy_add");
2736 let cfg = LaunchConfig {
2737 grid_dim: (1, 1, 1),
2738 block_dim: (32, 1, 1),
2739 shared_mem_bytes: 0,
2740 };
2741 let __s_b = self.gpu.stream();
2742 let mut b = __s_b.launch_builder(&f);
2743 b.arg(src).arg(dst).arg(&delta);
2744 unsafe {
2745 b.launch(cfg)?;
2746 }
2747 Ok(())
2748 }
2749 pub fn u32_copy(
2750 &self,
2751 src: &CudaSlice<u32>,
2752 dst: &mut CudaSlice<u32>,
2753 ) -> Result<(), Box<dyn std::error::Error>> {
2754 let f = self.func("u32_copy");
2755 let cfg = LaunchConfig {
2756 grid_dim: (1, 1, 1),
2757 block_dim: (32, 1, 1),
2758 shared_mem_bytes: 0,
2759 };
2760 let __s_b = self.gpu.stream();
2761 let mut b = __s_b.launch_builder(&f);
2762 b.arg(src).arg(dst);
2763 unsafe {
2764 b.launch(cfg)?;
2765 }
2766 Ok(())
2767 }
2768
2769 pub fn spec_adapt_k(
2773 &self,
2774 acc: &CudaSlice<u32>,
2775 brk: &mut CudaSlice<u32>,
2776 floor: usize,
2777 cap: usize,
2778 ) -> Result<(), Box<dyn std::error::Error>> {
2779 let f = self.func("spec_adapt_k");
2780 let (fl, cp) = (floor as i32, cap as i32);
2781 let cfg = LaunchConfig {
2782 grid_dim: (1, 1, 1),
2783 block_dim: (32, 1, 1),
2784 shared_mem_bytes: 0,
2785 };
2786 let __s_b = self.gpu.stream();
2787 let mut b = __s_b.launch_builder(&f);
2788 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
2789 unsafe {
2790 b.launch(cfg)?;
2791 }
2792 Ok(())
2793 }
2794
2795 pub fn spec_accept_greedy_dc(
2797 &self,
2798 preds: &CudaSlice<u32>,
2799 vtok: &CudaSlice<u32>,
2800 last_pred: &CudaSlice<u32>,
2801 brk: &CudaSlice<u32>,
2802 out: &mut CudaSlice<u32>,
2803 ) -> Result<(), Box<dyn std::error::Error>> {
2804 let f = self.func("spec_accept_greedy_dc");
2805 let cfg = LaunchConfig {
2806 grid_dim: (1, 1, 1),
2807 block_dim: (32, 1, 1),
2808 shared_mem_bytes: 0,
2809 };
2810 let __s_b = self.gpu.stream();
2811 let mut b = __s_b.launch_builder(&f);
2812 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
2813 unsafe {
2814 b.launch(cfg)?;
2815 }
2816 Ok(())
2817 }
2818
2819 pub fn pos_iota(
2821 &self,
2822 pos0: &CudaSlice<i32>,
2823 out: &mut CudaSlice<i32>,
2824 t: usize,
2825 ) -> Result<(), Box<dyn std::error::Error>> {
2826 let f = self.func("pos_iota_i32");
2827 let ti = t as i32;
2828 let cfg = LaunchConfig {
2829 grid_dim: (1, 1, 1),
2830 block_dim: (t.max(1) as u32, 1, 1),
2831 shared_mem_bytes: 0,
2832 };
2833 let __s_b = self.gpu.stream();
2834 let mut b = __s_b.launch_builder(&f);
2835 b.arg(pos0).arg(out).arg(&ti);
2836 unsafe {
2837 b.launch(cfg)?;
2838 }
2839 Ok(())
2840 }
2841 #[allow(clippy::too_many_arguments)]
2842 pub fn append_kv_quantized_rows_dc(
2843 &self,
2844 k_rows: &CudaSlice<f32>,
2845 v_rows: &CudaSlice<f32>,
2846 kc: &mut CudaSlice<u8>,
2847 vc: &mut CudaSlice<u8>,
2848 t0_dev: &CudaSlice<i32>,
2849 t: usize,
2850 kv_dim_k: usize,
2851 kv_dim_v: usize,
2852 k_tok_bytes: usize,
2853 v_tok_bytes: usize,
2854 g: bool,
2855 ) -> Result<(), Box<dyn std::error::Error>> {
2856 let f = if g {
2857 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
2858 } else {
2859 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
2860 };
2861 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2862 let cfg = LaunchConfig {
2863 grid_dim: (nblk, t as u32, 1),
2864 block_dim: (32, 1, 1),
2865 shared_mem_bytes: 0,
2866 };
2867 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2868 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2869 let __s_b = self.gpu.stream();
2870 let mut b = __s_b.launch_builder(&f);
2871 b.arg(k_rows)
2872 .arg(v_rows)
2873 .arg(kc)
2874 .arg(vc)
2875 .arg(t0_dev)
2876 .arg(&kdk)
2877 .arg(&kdv)
2878 .arg(&ktb)
2879 .arg(&vtb);
2880 unsafe {
2881 b.launch(cfg)?;
2882 }
2883 Ok(())
2884 }
2885
2886 #[allow(clippy::too_many_arguments)]
2889 pub fn append_kv_quantized_row_dc_inc(
2890 &self,
2891 k_row: &CudaSlice<f32>,
2892 v_row: &CudaSlice<f32>,
2893 kc: &mut CudaSlice<u8>,
2894 vc: &mut CudaSlice<u8>,
2895 t0_dev: &mut CudaSlice<i32>,
2896 kv_dim_k: usize,
2897 kv_dim_v: usize,
2898 k_tok_bytes: usize,
2899 v_tok_bytes: usize,
2900 g: bool,
2901 ) -> Result<(), Box<dyn std::error::Error>> {
2902 let f = if g {
2903 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
2904 } else {
2905 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
2906 };
2907 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
2908 let cfg = LaunchConfig {
2909 grid_dim: (1, 1, 1),
2910 block_dim: (nthreads, 1, 1),
2911 shared_mem_bytes: 0,
2912 };
2913 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2914 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2915 let __s_b = self.gpu.stream();
2916 let mut b = __s_b.launch_builder(&f);
2917 b.arg(k_row)
2918 .arg(v_row)
2919 .arg(kc)
2920 .arg(vc)
2921 .arg(t0_dev)
2922 .arg(&kdk)
2923 .arg(&kdv)
2924 .arg(&ktb)
2925 .arg(&vtb);
2926 unsafe {
2927 b.launch(cfg)?;
2928 }
2929 Ok(())
2930 }
2931
2932 pub fn pack_tok_p(
2934 &self,
2935 tok: &CudaSlice<u32>,
2936 p: &CudaSlice<f32>,
2937 out: &mut CudaSlice<u32>,
2938 slot: usize,
2939 ) -> Result<(), Box<dyn std::error::Error>> {
2940 let f = self.func("pack_tok_p");
2941 let sl = slot as i32;
2942 let cfg = LaunchConfig {
2943 grid_dim: (1, 1, 1),
2944 block_dim: (32, 1, 1),
2945 shared_mem_bytes: 0,
2946 };
2947 let __s_b = self.gpu.stream();
2948 let mut b = __s_b.launch_builder(&f);
2949 b.arg(tok).arg(p).arg(out).arg(&sl);
2950 unsafe {
2951 b.launch(cfg)?;
2952 }
2953 Ok(())
2954 }
2955 pub fn tok_map_u32(
2956 &self,
2957 tok: &mut CudaSlice<u32>,
2958 map: &CudaSlice<u32>,
2959 ) -> Result<(), Box<dyn std::error::Error>> {
2960 let f = self.func("tok_map_u32");
2961 let cfg = LaunchConfig {
2962 grid_dim: (1, 1, 1),
2963 block_dim: (32, 1, 1),
2964 shared_mem_bytes: 0,
2965 };
2966 let __s_b = self.gpu.stream();
2967 let mut b = __s_b.launch_builder(&f);
2968 b.arg(tok).arg(map);
2969 unsafe {
2970 b.launch(cfg)?;
2971 }
2972 Ok(())
2973 }
2974
2975 #[allow(clippy::too_many_arguments)]
2977 pub fn spec_assemble_verify(
2978 &self,
2979 tokp: &CudaSlice<u32>,
2980 pend: &CudaSlice<u32>,
2981 d2t: Option<&CudaSlice<u32>>,
2982 vtok: &mut CudaSlice<u32>,
2983 brk: &mut CudaSlice<u32>,
2984 p_min: f32,
2985 k: usize,
2986 pmin0: bool,
2987 ) -> Result<(), Box<dyn std::error::Error>> {
2988 let f = self.func("spec_assemble_verify");
2989 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
2990 let cfg = LaunchConfig {
2991 grid_dim: (1, 1, 1),
2992 block_dim: (32, 1, 1),
2993 shared_mem_bytes: 0,
2994 };
2995 let __s_b = self.gpu.stream();
2996 let mut b = __s_b.launch_builder(&f);
2997 match d2t {
2998 Some(m) => {
2999 b.arg(tokp)
3000 .arg(pend)
3001 .arg(m)
3002 .arg(vtok)
3003 .arg(brk)
3004 .arg(&p_min)
3005 .arg(&ki)
3006 .arg(&pm);
3007 unsafe {
3008 b.launch(cfg)?;
3009 }
3010 }
3011 None => {
3012 let null: u64 = 0;
3013 b.arg(tokp)
3014 .arg(pend)
3015 .arg(&null)
3016 .arg(vtok)
3017 .arg(brk)
3018 .arg(&p_min)
3019 .arg(&ki)
3020 .arg(&pm);
3021 unsafe {
3022 b.launch(cfg)?;
3023 }
3024 }
3025 }
3026 Ok(())
3027 }
3028
3029 #[allow(clippy::too_many_arguments)]
3031 pub fn ssm_conv_ring_rebuild_dc(
3032 &self,
3033 qkv_tm: &CudaSlice<f32>,
3034 ring_old: &CudaSlice<f32>,
3035 conv_state: &mut CudaSlice<f32>,
3036 conv_dim: usize,
3037 acc: &CudaSlice<u32>,
3038 base: usize,
3039 t_v: usize,
3040 d_conv: usize,
3041 ) -> Result<(), Box<dyn std::error::Error>> {
3042 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3043 let n = conv_dim * (d_conv - 1);
3044 let cfg = LaunchConfig::for_num_elems(n as u32);
3045 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3046 let __s_b = self.gpu.stream();
3047 let mut b = __s_b.launch_builder(&f);
3048 b.arg(qkv_tm)
3049 .arg(ring_old)
3050 .arg(conv_state)
3051 .arg(&cd)
3052 .arg(acc)
3053 .arg(&b0)
3054 .arg(&tv)
3055 .arg(&dc);
3056 unsafe {
3057 b.launch(cfg)?;
3058 }
3059 Ok(())
3060 }
3061 #[allow(clippy::too_many_arguments)]
3062 pub fn gdn_scan_s128_dc(
3063 &self,
3064 q: &CudaSlice<f32>,
3065 k: &CudaSlice<f32>,
3066 v: &CudaSlice<f32>,
3067 g: &CudaSlice<f32>,
3068 beta: &CudaSlice<f32>,
3069 state_in: &CudaSlice<f32>,
3070 state_out: &mut CudaSlice<f32>,
3071 o: &mut CudaSlice<f32>,
3072 n_head: usize,
3073 acc: &CudaSlice<u32>,
3074 base: usize,
3075 t_v: usize,
3076 scale: f32,
3077 ) -> Result<(), Box<dyn std::error::Error>> {
3078 let f = self.func("gdn_scan_s128_dc");
3079 const S_V: u32 = 128;
3080 const WARP: u32 = 32;
3081 const COLS_PER_BLOCK: u32 = 4;
3082 let cfg = LaunchConfig {
3083 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3084 block_dim: (WARP, COLS_PER_BLOCK, 1),
3085 shared_mem_bytes: 0,
3086 };
3087 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3088 let __s_b = self.gpu.stream();
3089 let mut b = __s_b.launch_builder(&f);
3090 b.arg(q)
3091 .arg(k)
3092 .arg(v)
3093 .arg(g)
3094 .arg(beta)
3095 .arg(state_in)
3096 .arg(state_out)
3097 .arg(o)
3098 .arg(&h)
3099 .arg(acc)
3100 .arg(&b0)
3101 .arg(&tv)
3102 .arg(&scale);
3103 unsafe {
3104 b.launch(cfg)?;
3105 }
3106 Ok(())
3107 }
3108
3109 pub fn spec_rollback_kv(
3111 &self,
3112 len_ptrs: &CudaSlice<u64>,
3113 saved: &CudaSlice<i32>,
3114 acc: &CudaSlice<u32>,
3115 base: usize,
3116 n_layer: usize,
3117 ) -> Result<(), Box<dyn std::error::Error>> {
3118 let f = self.func("spec_rollback_kv");
3119 let (b, nl) = (base as i32, n_layer as i32);
3120 let cfg = LaunchConfig {
3121 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3122 block_dim: (64, 1, 1),
3123 shared_mem_bytes: 0,
3124 };
3125 let __s_bl = self.gpu.stream();
3126 let mut bl = __s_bl.launch_builder(&f);
3127 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3128 unsafe {
3129 bl.launch(cfg)?;
3130 }
3131 Ok(())
3132 }
3133
3134 pub fn spec_fork_valid(
3136 &self,
3137 acc: &CudaSlice<u32>,
3138 optimistic_pending: u32,
3139 valid: &mut CudaSlice<u32>,
3140 ) -> Result<(), Box<dyn std::error::Error>> {
3141 let f = self.func("spec_fork_valid");
3142 let cfg = LaunchConfig {
3143 grid_dim: (1, 1, 1),
3144 block_dim: (1, 1, 1),
3145 shared_mem_bytes: 0,
3146 };
3147 let __s_bl = self.gpu.stream();
3148 let mut bl = __s_bl.launch_builder(&f);
3149 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3150 unsafe {
3151 bl.launch(cfg)?;
3152 }
3153 Ok(())
3154 }
3155
3156 pub fn spec_fork_reconcile_kv(
3158 &self,
3159 len_ptrs: &CudaSlice<u64>,
3160 saved: &CudaSlice<i32>,
3161 acc: &CudaSlice<u32>,
3162 valid: &CudaSlice<u32>,
3163 base: usize,
3164 n_layer: usize,
3165 ) -> Result<(), Box<dyn std::error::Error>> {
3166 let f = self.func("spec_fork_reconcile_kv");
3167 let (b, nl) = (base as i32, n_layer as i32);
3168 let cfg = LaunchConfig {
3169 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3170 block_dim: (64, 1, 1),
3171 shared_mem_bytes: 0,
3172 };
3173 let __s_bl = self.gpu.stream();
3174 let mut bl = __s_bl.launch_builder(&f);
3175 bl.arg(len_ptrs)
3176 .arg(saved)
3177 .arg(acc)
3178 .arg(valid)
3179 .arg(&b)
3180 .arg(&nl);
3181 unsafe {
3182 bl.launch(cfg)?;
3183 }
3184 Ok(())
3185 }
3186
3187 pub fn spec_fork_restore_f32(
3189 &self,
3190 snapshot: &CudaSlice<f32>,
3191 state: &mut CudaSlice<f32>,
3192 valid: &CudaSlice<u32>,
3193 ) -> Result<(), Box<dyn std::error::Error>> {
3194 assert_eq!(
3195 snapshot.len(),
3196 state.len(),
3197 "fork recurrent snapshot shape mismatch"
3198 );
3199 let f = self.func("spec_fork_restore_f32");
3200 let n = state.len() as i32;
3201 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3202 let cfg = LaunchConfig {
3203 grid_dim: (blocks, 1, 1),
3204 block_dim: (256, 1, 1),
3205 shared_mem_bytes: 0,
3206 };
3207 let __s_bl = self.gpu.stream();
3208 let mut bl = __s_bl.launch_builder(&f);
3209 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3210 unsafe {
3211 bl.launch(cfg)?;
3212 }
3213 Ok(())
3214 }
3215
3216 pub fn spec_seed_gather(
3219 &self,
3220 vx: &CudaSlice<f32>,
3221 fill_prev: &CudaSlice<f32>,
3222 acc: &CudaSlice<u32>,
3223 h_seed: &mut CudaSlice<f32>,
3224 base: usize,
3225 n_embd: usize,
3226 ) -> Result<(), Box<dyn std::error::Error>> {
3227 let f = self.func("spec_seed_gather");
3228 let (b, ne) = (base as i32, n_embd as i32);
3229 let cfg = LaunchConfig {
3230 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3231 block_dim: (256, 1, 1),
3232 shared_mem_bytes: 0,
3233 };
3234 let __s_bl = self.gpu.stream();
3235 let mut bl = __s_bl.launch_builder(&f);
3236 bl.arg(vx)
3237 .arg(fill_prev)
3238 .arg(acc)
3239 .arg(h_seed)
3240 .arg(&b)
3241 .arg(&ne);
3242 unsafe {
3243 bl.launch(cfg)?;
3244 }
3245 Ok(())
3246 }
3247
3248 pub fn spec_accept_greedy(
3250 &self,
3251 preds: &CudaSlice<u32>,
3252 draft: &CudaSlice<u32>,
3253 last_pred: u32,
3254 base: usize,
3255 k_round: usize,
3256 out: &mut CudaSlice<u32>,
3257 ) -> Result<(), Box<dyn std::error::Error>> {
3258 let f = self.func("spec_accept_greedy");
3259 let (b, k) = (base as i32, k_round as i32);
3260 let cfg = LaunchConfig {
3261 grid_dim: (1, 1, 1),
3262 block_dim: (32, 1, 1),
3263 shared_mem_bytes: 0,
3264 };
3265 let __s_bl = self.gpu.stream();
3266 let mut bl = __s_bl.launch_builder(&f);
3267 bl.arg(preds)
3268 .arg(draft)
3269 .arg(&last_pred)
3270 .arg(&b)
3271 .arg(&k)
3272 .arg(out);
3273 unsafe {
3274 bl.launch(cfg)?;
3275 }
3276 Ok(())
3277 }
3278
3279 pub fn gumbel_perturb(
3286 &self,
3287 x: &CudaSlice<f32>,
3288 y: &mut CudaSlice<f32>,
3289 n: usize,
3290 seed: u64,
3291 stream_pos: u32,
3292 temp: f32,
3293 ) -> Result<(), Box<dyn std::error::Error>> {
3294 let f = self.func("gumbel_perturb_f32");
3295 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3296 let cfg = LaunchConfig {
3297 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3298 block_dim: (256, 1, 1),
3299 shared_mem_bytes: 0,
3300 };
3301 let __s_b = self.gpu.stream();
3302 let mut b = __s_b.launch_builder(&f);
3303 b.arg(x)
3304 .arg(&mut *y)
3305 .arg(&ni)
3306 .arg(&slo)
3307 .arg(&shi)
3308 .arg(&stream_pos)
3309 .arg(&temp);
3310 unsafe {
3311 b.launch(cfg)?;
3312 }
3313 Ok(())
3314 }
3315
3316 pub fn mask_logits_col(
3324 &self,
3325 logits: &mut CudaSlice<f32>,
3326 mask: &CudaSlice<u32>,
3327 col: usize,
3328 n: usize,
3329 mask_words: usize,
3330 ) -> Result<(), Box<dyn std::error::Error>> {
3331 let f = self.func("mask_logits_f32");
3332 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3333 let cfg = LaunchConfig {
3334 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3335 block_dim: (256, 1, 1),
3336 shared_mem_bytes: 0,
3337 };
3338 let __s_b = self.gpu.stream();
3339 let mut b = __s_b.launch_builder(&f);
3340 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3341 unsafe {
3342 b.launch(cfg)?;
3343 }
3344 Ok(())
3345 }
3346
3347 pub fn gumbel_perturb_col(
3354 &self,
3355 x: &CudaSlice<f32>,
3356 col: usize,
3357 y: &mut CudaSlice<f32>,
3358 n: usize,
3359 seed: u64,
3360 stream_pos: u32,
3361 temp: f32,
3362 ) -> Result<(), Box<dyn std::error::Error>> {
3363 let f = self.func("gumbel_perturb_f32");
3364 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3365 let col_view = x.slice(col * n..(col + 1) * n);
3366 let cfg = LaunchConfig {
3367 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3368 block_dim: (256, 1, 1),
3369 shared_mem_bytes: 0,
3370 };
3371 let __s_b = self.gpu.stream();
3372 let mut b = __s_b.launch_builder(&f);
3373 b.arg(&col_view)
3374 .arg(&mut *y)
3375 .arg(&ni)
3376 .arg(&slo)
3377 .arg(&shi)
3378 .arg(&stream_pos)
3379 .arg(&temp);
3380 unsafe {
3381 b.launch(cfg)?;
3382 }
3383 Ok(())
3384 }
3385
3386 #[allow(clippy::too_many_arguments)]
3392 pub fn gumbel_perturb_filtered_col(
3393 &self,
3394 x: &CudaSlice<f32>,
3395 col: usize,
3396 y: &mut CudaSlice<f32>,
3397 n: usize,
3398 seed: u64,
3399 stream_pos: u32,
3400 temp: f32,
3401 stat_max: &CudaSlice<f32>,
3402 stat_th: &CudaSlice<f32>,
3403 stat_idx: usize,
3404 ) -> Result<(), Box<dyn std::error::Error>> {
3405 let f = self.func("gumbel_perturb_filtered_col_f32");
3406 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3407 let (ci, si) = (col as i32, stat_idx as i32);
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(x)
3416 .arg(&ci)
3417 .arg(&mut *y)
3418 .arg(&ni)
3419 .arg(&slo)
3420 .arg(&shi)
3421 .arg(&stream_pos)
3422 .arg(&temp)
3423 .arg(stat_max)
3424 .arg(stat_th)
3425 .arg(&si);
3426 unsafe {
3427 b.launch(cfg)?;
3428 }
3429 Ok(())
3430 }
3431
3432 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3437 let f = self.func("memra_sctr_inc");
3438 let cfg = LaunchConfig {
3439 grid_dim: (1, 1, 1),
3440 block_dim: (1, 1, 1),
3441 shared_mem_bytes: 0,
3442 };
3443 let __s_b = self.gpu.stream();
3444 let mut b = __s_b.launch_builder(&f);
3445 b.arg(&mut *ctr);
3446 unsafe {
3447 b.launch(cfg)?;
3448 }
3449 Ok(())
3450 }
3451
3452 pub fn gumbel_perturb_ctr(
3457 &self,
3458 x: &CudaSlice<f32>,
3459 y: &mut CudaSlice<f32>,
3460 n: usize,
3461 seed: u64,
3462 ctr: &CudaSlice<u32>,
3463 temp: f32,
3464 ) -> Result<(), Box<dyn std::error::Error>> {
3465 let f = self.func("gumbel_perturb_ctr_f32");
3466 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3467 let cfg = LaunchConfig {
3468 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3469 block_dim: (256, 1, 1),
3470 shared_mem_bytes: 0,
3471 };
3472 let __s_b = self.gpu.stream();
3473 let mut b = __s_b.launch_builder(&f);
3474 b.arg(x)
3475 .arg(&mut *y)
3476 .arg(&ni)
3477 .arg(&slo)
3478 .arg(&shi)
3479 .arg(ctr)
3480 .arg(&temp);
3481 unsafe {
3482 b.launch(cfg)?;
3483 }
3484 Ok(())
3485 }
3486
3487 pub fn softmax_gather(
3491 &self,
3492 x: &CudaSlice<f32>,
3493 row_stride: usize,
3494 ids: &CudaSlice<u32>,
3495 rows: &CudaSlice<i32>,
3496 out: &mut CudaSlice<f32>,
3497 n: usize,
3498 npair: usize,
3499 temp: f32,
3500 ) -> Result<(), Box<dyn std::error::Error>> {
3501 let f = self.func("softmax_gather_f32");
3502 let (ni, rs) = (n as i32, row_stride as i64);
3503 let np = npair as i32;
3504 let cfg = LaunchConfig {
3505 grid_dim: (npair as u32, 1, 1),
3506 block_dim: (256, 1, 1),
3507 shared_mem_bytes: 0,
3508 };
3509 let __s_b = self.gpu.stream();
3510 let mut b = __s_b.launch_builder(&f);
3511 b.arg(x)
3512 .arg(&rs)
3513 .arg(ids)
3514 .arg(rows)
3515 .arg(&mut *out)
3516 .arg(&ni)
3517 .arg(&np)
3518 .arg(&temp);
3519 unsafe {
3520 b.launch(cfg)?;
3521 }
3522 Ok(())
3523 }
3524
3525 pub fn residual_sample(
3529 &self,
3530 p: &CudaSlice<f32>,
3531 q: Option<&CudaSlice<f32>>,
3532 n: usize,
3533 temp: f32,
3534 seed: u64,
3535 stream_pos: u32,
3536 out_tok: &mut CudaSlice<u32>,
3537 ) -> Result<(), Box<dyn std::error::Error>> {
3538 let f = self.func("residual_sample_f32");
3539 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3540 let nth = 1024u32;
3541 let cfg = LaunchConfig {
3542 grid_dim: (1, 1, 1),
3543 block_dim: (nth, 1, 1),
3544 shared_mem_bytes: 0,
3545 };
3546 let has_q: i32 = q.is_some() as i32;
3547 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
3549 let mut b = __s_b.launch_builder(&f);
3550 b.arg(p)
3551 .arg(qbuf)
3552 .arg(&has_q)
3553 .arg(&ni)
3554 .arg(&temp)
3555 .arg(&slo)
3556 .arg(&shi)
3557 .arg(&stream_pos)
3558 .arg(&mut *out_tok);
3559 unsafe {
3560 b.launch(cfg)?;
3561 }
3562 Ok(())
3563 }
3564
3565 pub fn with_moe_cache<R>(
3570 &self,
3571 max_block_bytes: usize,
3572 f: impl FnOnce(
3573 &mut crate::moe_cache::MoeSlotCache,
3574 &Engine,
3575 ) -> Result<R, Box<dyn std::error::Error>>,
3576 ) -> Result<R, Box<dyn std::error::Error>> {
3577 let mut guard = self.moe_cache.lock().unwrap();
3578 if guard.is_none() {
3579 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
3580 }
3581 let cache = guard.as_mut().unwrap();
3582 f(cache, self)
3583 }
3584
3585 pub fn freeze_moe_cache(&self) {
3588 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
3589 cache.freeze();
3590 }
3591 }
3592
3593 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
3596 self.moe_cache
3597 .lock()
3598 .unwrap()
3599 .as_ref()
3600 .map(crate::moe_cache::MoeSlotCache::export_residency)
3601 }
3602
3603 pub(crate) fn moe_cache_frozen(&self) -> bool {
3604 self.moe_cache
3605 .lock()
3606 .unwrap()
3607 .as_ref()
3608 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
3609 }
3610
3611 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
3618 crate::cpu_experts::configured()
3619 && self.moe_cache_frozen()
3620 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
3621 }
3622
3623 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
3625 assert!(
3626 self.moe_cache.lock().unwrap().is_none(),
3627 "MoE cache layout configured after cache construction"
3628 );
3629 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
3630 }
3631
3632 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
3633 self.moe_cache_layout.lock().unwrap().clone()
3634 }
3635
3636 pub fn moe_cache_enabled() -> bool {
3638 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
3639 }
3640
3641 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
3644 let guard = self.moe_cache.lock().unwrap();
3645 guard
3646 .as_ref()
3647 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
3648 }
3649
3650 pub fn cpu_expert_stats(
3654 &self,
3655 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
3656 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
3657 }
3658
3659 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
3662 crate::cpu_experts::predictor_stats()
3663 }
3664
3665 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
3666 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
3667 }
3668
3669 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
3672 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
3673 }
3674
3675 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
3678 let guard = self.moe_cache.lock().unwrap();
3679 guard
3680 .as_ref()
3681 .and_then(|cache| cache.pread_stats())
3682 .map(|stats| {
3683 (
3684 stats.reads,
3685 stats.bytes,
3686 stats.read_errors,
3687 stats.short_reads,
3688 stats.fallbacks,
3689 stats.buffer_waits,
3690 stats.ring_full,
3691 )
3692 })
3693 }
3694
3695 pub fn spill_config_fallbacks(&self) -> u64 {
3697 crate::spill_pread::config_fallbacks()
3698 }
3699
3700 pub fn moe_cache_reset_counters(&self) {
3702 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
3703 c.reset_counters();
3704 }
3705 }
3706
3707 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3708 Ok(self.gpu.stream().clone_htod(v)?)
3709 }
3710
3711 pub fn htod_bytes_padded(
3715 &self,
3716 v: &[u8],
3717 pad: usize,
3718 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3719 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
3720 {
3721 let mut view = d.slice_mut(0..v.len());
3722 self.gpu.stream().memcpy_htod(v, &mut view)?;
3723 }
3724 Ok(d)
3725 }
3726
3727 pub fn copy_into(
3729 &self,
3730 dst: &mut CudaSlice<f32>,
3731 off: usize,
3732 src: &CudaSlice<f32>,
3733 len: usize,
3734 ) -> Result<(), Box<dyn std::error::Error>> {
3735 let mut view = dst.slice_mut(off..off + len);
3736 self.gpu
3737 .stream()
3738 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3739 Ok(())
3740 }
3741
3742 pub fn copy_u8_into(
3745 &self,
3746 dst: &mut CudaSlice<u8>,
3747 off: usize,
3748 src: &CudaSlice<u8>,
3749 len: usize,
3750 ) -> Result<(), Box<dyn std::error::Error>> {
3751 let mut view = dst.slice_mut(off..off + len);
3752 self.gpu
3753 .stream()
3754 .memcpy_dtod(&src.slice(0..len), &mut view)?;
3755 Ok(())
3756 }
3757
3758 pub fn copy_u8_range_into(
3760 &self,
3761 dst: &mut CudaSlice<u8>,
3762 dst_off: usize,
3763 src: &CudaSlice<u8>,
3764 src_off: usize,
3765 len: usize,
3766 ) -> Result<(), Box<dyn std::error::Error>> {
3767 let mut dst_view = dst.slice_mut(dst_off..dst_off + len);
3768 self.gpu
3769 .stream()
3770 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
3771 Ok(())
3772 }
3773
3774 pub fn prepare_kv_append(
3778 &self,
3779 kv: &mut crate::cache::KvLayer,
3780 retain_from: usize,
3781 append_rows: usize,
3782 ) -> Result<usize, Box<dyn std::error::Error>> {
3783 let Some(plan) = kv
3784 .ring
3785 .as_ref()
3786 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
3787 .transpose()?
3788 else {
3789 return Ok(kv.len);
3790 };
3791 match plan {
3792 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
3793 crate::cache::KvRingAppend::Rebase {
3794 src_row,
3795 keep_rows,
3796 new_base,
3797 write_row,
3798 } => {
3799 if keep_rows > 0 {
3800 let k_len = keep_rows * kv.k_tok_bytes;
3801 let v_len = keep_rows * kv.v_tok_bytes;
3802 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
3803 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
3804 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
3805 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
3806 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
3807 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
3808 }
3809 kv.ring.as_mut().unwrap().apply_rebase(new_base);
3810 Ok(write_row)
3811 }
3812 }
3813 }
3814
3815 pub fn htod_u8_into(
3818 &self,
3819 dst: &mut CudaSlice<u8>,
3820 off: usize,
3821 src: &[u8],
3822 ) -> Result<(), Box<dyn std::error::Error>> {
3823 let mut view = dst.slice_mut(off..off + src.len());
3824 self.gpu.stream().memcpy_htod(src, &mut view)?;
3825 Ok(())
3826 }
3827
3828 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
3829 b.slice(0..len)
3830 }
3831
3832 pub fn view_u8_range<'a>(
3835 &self,
3836 b: &'a CudaSlice<u8>,
3837 start: usize,
3838 end: usize,
3839 ) -> cudarc::driver::CudaView<'a, u8> {
3840 b.slice(start..end)
3841 }
3842 pub fn view_u8<'a>(
3843 &self,
3844 b: &'a CudaSlice<u8>,
3845 len: usize,
3846 ) -> cudarc::driver::CudaView<'a, u8> {
3847 b.slice(0..len)
3848 }
3849
3850 pub fn append_kv_quantized(
3854 &self,
3855 k_row: &CudaSlice<f32>,
3856 v_row: &CudaSlice<f32>,
3857 kc: &mut CudaSlice<u8>,
3858 vc: &mut CudaSlice<u8>,
3859 t: usize,
3860 kv_dim_k: usize,
3861 kv_dim_v: usize,
3862 k_tok_bytes: usize,
3863 v_tok_bytes: usize,
3864 g: bool,
3865 ) -> Result<(), Box<dyn std::error::Error>> {
3866 let f = if g {
3867 self.func_g("append_quantize_kv_q8_0_q5_1")
3868 } else {
3869 self.func("append_quantize_kv_q8_0_q5_1")
3870 };
3871 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3872 let cfg = LaunchConfig {
3873 grid_dim: (nblk, 1, 1),
3874 block_dim: (32, 1, 1),
3875 shared_mem_bytes: 0,
3876 };
3877 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
3878 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3879 let __s_b = self.gpu.stream();
3880 let mut b = __s_b.launch_builder(&f);
3881 b.arg(k_row)
3882 .arg(v_row)
3883 .arg(kc)
3884 .arg(vc)
3885 .arg(&ti)
3886 .arg(&kdk)
3887 .arg(&kdv)
3888 .arg(&ktb)
3889 .arg(&vtb);
3890 unsafe {
3891 b.launch(cfg)?;
3892 }
3893 Ok(())
3894 }
3895
3896 pub fn append_kv_quantized_dc(
3900 &self,
3901 k_row: &CudaSlice<f32>,
3902 v_row: &CudaSlice<f32>,
3903 kc: &mut CudaSlice<u8>,
3904 vc: &mut CudaSlice<u8>,
3905 t_dev: &CudaSlice<i32>,
3906 kv_dim_k: usize,
3907 kv_dim_v: usize,
3908 k_tok_bytes: usize,
3909 v_tok_bytes: usize,
3910 g: bool,
3911 ) -> Result<(), Box<dyn std::error::Error>> {
3912 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3913 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3914 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3915 if Self::pdl_on() && Self::pdl_wb_on() {
3917 use cudarc::driver::{DevicePtr, DevicePtrMut};
3918 let s = &self.gpu.stream();
3919 let (pk, _g0) = k_row.device_ptr(s);
3920 let (pv, _g1) = v_row.device_ptr(s);
3921 let (pkc, _g2) = kc.device_ptr_mut(s);
3922 let (pvc, _g3) = vc.device_ptr_mut(s);
3923 let (pt, _g4) = t_dev.device_ptr(s);
3924 let mut ps = [
3925 &pk as *const _ as *mut std::ffi::c_void,
3926 &pv as *const _ as *mut _,
3927 &pkc as *const _ as *mut _,
3928 &pvc as *const _ as *mut _,
3929 &pt as *const _ as *mut _,
3930 &kdk as *const _ as *mut _,
3931 &kdv as *const _ as *mut _,
3932 &ktb as *const _ as *mut _,
3933 &vtb as *const _ as *mut _,
3934 ];
3935 unsafe {
3936 self.launch_pdl_flash(
3937 g,
3938 "append_quantize_kv_q8_0_q5_1_dc",
3939 (nblk, 1, 1),
3940 (32, 1, 1),
3941 0,
3942 &mut ps,
3943 )?;
3944 }
3945 return Ok(());
3946 }
3947 let f = if g {
3948 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
3949 } else {
3950 self.func("append_quantize_kv_q8_0_q5_1_dc")
3951 };
3952 let cfg = LaunchConfig {
3953 grid_dim: (nblk, 1, 1),
3954 block_dim: (32, 1, 1),
3955 shared_mem_bytes: 0,
3956 };
3957 let __s_b = self.gpu.stream();
3958 let mut b = __s_b.launch_builder(&f);
3959 b.arg(k_row)
3960 .arg(v_row)
3961 .arg(kc)
3962 .arg(vc)
3963 .arg(t_dev)
3964 .arg(&kdk)
3965 .arg(&kdv)
3966 .arg(&ktb)
3967 .arg(&vtb);
3968 unsafe {
3969 b.launch(cfg)?;
3970 }
3971 Ok(())
3972 }
3973
3974 #[allow(clippy::too_many_arguments)]
3981 pub fn append_kv_quantized_rows(
3982 &self,
3983 k_rows: &CudaSlice<f32>,
3984 v_rows: &CudaSlice<f32>,
3985 kc: &mut CudaSlice<u8>,
3986 vc: &mut CudaSlice<u8>,
3987 t0: usize,
3988 t: usize,
3989 kv_dim_k: usize,
3990 kv_dim_v: usize,
3991 k_tok_bytes: usize,
3992 v_tok_bytes: usize,
3993 g: bool,
3994 ) -> Result<(), Box<dyn std::error::Error>> {
3995 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
3996 for i in 0..t {
3997 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
3998 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
3999 self.append_kv_quantized_view(
4000 &k_row,
4001 &v_row,
4002 kc,
4003 vc,
4004 t0 + i,
4005 kv_dim_k,
4006 kv_dim_v,
4007 k_tok_bytes,
4008 v_tok_bytes,
4009 g,
4010 )?;
4011 }
4012 return Ok(());
4013 }
4014 let f = if g {
4015 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4016 } else {
4017 self.func("append_quantize_kv_q8_0_q5_1_rows")
4018 };
4019 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4020 let cfg = LaunchConfig {
4021 grid_dim: (nblk, t as u32, 1),
4022 block_dim: (32, 1, 1),
4023 shared_mem_bytes: 0,
4024 };
4025 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4026 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4027 let __s_b = self.gpu.stream();
4028 let mut b = __s_b.launch_builder(&f);
4029 b.arg(k_rows)
4030 .arg(v_rows)
4031 .arg(kc)
4032 .arg(vc)
4033 .arg(&t0i)
4034 .arg(&kdk)
4035 .arg(&kdv)
4036 .arg(&ktb)
4037 .arg(&vtb);
4038 unsafe {
4039 b.launch(cfg)?;
4040 }
4041 Ok(())
4042 }
4043
4044 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4048 let f = self.func("inc_i32");
4049 let cfg = LaunchConfig {
4050 grid_dim: (1, 1, 1),
4051 block_dim: (1, 1, 1),
4052 shared_mem_bytes: 0,
4053 };
4054 let __s_b = self.gpu.stream();
4055 let mut b = __s_b.launch_builder(&f);
4056 b.arg(p);
4057 unsafe {
4058 b.launch(cfg)?;
4059 }
4060 Ok(())
4061 }
4062
4063 pub fn append_kv_quantized_view(
4066 &self,
4067 k_row: &cudarc::driver::CudaView<f32>,
4068 v_row: &cudarc::driver::CudaView<f32>,
4069 kc: &mut CudaSlice<u8>,
4070 vc: &mut CudaSlice<u8>,
4071 t: usize,
4072 kv_dim_k: usize,
4073 kv_dim_v: usize,
4074 k_tok_bytes: usize,
4075 v_tok_bytes: usize,
4076 g: bool,
4077 ) -> Result<(), Box<dyn std::error::Error>> {
4078 let stream = self.gpu.stream();
4079 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4080 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4081 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4082 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4083 let f = if g {
4084 self.func_g("append_quantize_kv_q8_0_q5_1")
4085 } else {
4086 self.func("append_quantize_kv_q8_0_q5_1")
4087 };
4088 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4089 let cfg = LaunchConfig {
4090 grid_dim: (nblk, 1, 1),
4091 block_dim: (32, 1, 1),
4092 shared_mem_bytes: 0,
4093 };
4094 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4095 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4096 let mut b = stream.launch_builder(&f);
4097 b.arg(k_row)
4098 .arg(v_row)
4099 .arg(kc)
4100 .arg(vc)
4101 .arg(&ti)
4102 .arg(&kdk)
4103 .arg(&kdv)
4104 .arg(&ktb)
4105 .arg(&vtb);
4106 unsafe {
4107 b.launch(cfg)?;
4108 }
4109 Ok(())
4110 }
4111
4112 pub fn copy_view_into(
4115 &self,
4116 dst: &mut CudaSlice<f32>,
4117 off: usize,
4118 src: &cudarc::driver::CudaView<f32>,
4119 len: usize,
4120 ) -> Result<(), Box<dyn std::error::Error>> {
4121 let mut view = dst.slice_mut(off..off + len);
4122 self.gpu
4123 .stream()
4124 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4125 Ok(())
4126 }
4127
4128 pub fn clone_dtod(
4132 &self,
4133 src: &CudaSlice<f32>,
4134 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4135 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4136 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4137 Ok(dst)
4138 }
4139
4140 pub fn dtod_copy_view(
4143 &self,
4144 src: &cudarc::driver::CudaView<f32>,
4145 dst: &mut CudaSlice<f32>,
4146 ) -> Result<(), Box<dyn std::error::Error>> {
4147 self.gpu.stream().memcpy_dtod(src, dst)?;
4148 Ok(())
4149 }
4150
4151 pub fn dtod_copy_view_i8(
4153 &self,
4154 src: &cudarc::driver::CudaView<i8>,
4155 dst: &mut CudaSlice<i8>,
4156 ) -> Result<(), Box<dyn std::error::Error>> {
4157 self.gpu.stream().memcpy_dtod(src, dst)?;
4158 Ok(())
4159 }
4160
4161 pub fn dtod_copy_into(
4163 &self,
4164 src: &CudaSlice<f32>,
4165 dst: &mut CudaSlice<f32>,
4166 offset: usize,
4167 ) -> Result<(), Box<dyn std::error::Error>> {
4168 let n = src.len();
4169 let mut dv = dst.slice_mut(offset..offset + n);
4170 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4171 Ok(())
4172 }
4173
4174 pub fn copy_batch_uniform_f32(
4180 &self,
4181 table: &CudaSlice<u64>,
4182 n: usize,
4183 words: usize,
4184 ) -> Result<(), Box<dyn std::error::Error>> {
4185 if n == 0 || words == 0 {
4186 return Ok(());
4187 }
4188 debug_assert!(
4189 table.len() >= 2 * n,
4190 "pointer table must hold n srcs + n dsts"
4191 );
4192 let f = self.func("copy_batch_uniform_f32");
4193 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4196 let (ni, wi) = (n as i32, words as i32);
4197 let cfg = LaunchConfig {
4198 grid_dim: (chunks, n as u32, 1),
4199 block_dim: (256, 1, 1),
4200 shared_mem_bytes: 0,
4201 };
4202 let __s = self.gpu.stream();
4203 let mut b = __s.launch_builder(&f);
4204 b.arg(table).arg(&ni).arg(&wi);
4205 unsafe {
4206 b.launch(cfg)?;
4207 }
4208 Ok(())
4209 }
4210
4211 pub fn htod_u64_into(
4214 &self,
4215 v: &[u64],
4216 dst: &mut CudaSlice<u64>,
4217 ) -> Result<(), Box<dyn std::error::Error>> {
4218 let mut view = dst.slice_mut(0..v.len());
4219 self.gpu.stream().memcpy_htod(v, &mut view)?;
4220 Ok(())
4221 }
4222
4223 pub fn copy_indirect_src_f32(
4228 &self,
4229 src_entry: &cudarc::driver::CudaView<u64>,
4230 dst: &mut CudaSlice<f32>,
4231 dst_off: usize,
4232 words: usize,
4233 ) -> Result<(), Box<dyn std::error::Error>> {
4234 let f = self.func("copy_indirect_src_f32");
4235 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4236 let wi = words as i32;
4237 let cfg = LaunchConfig {
4238 grid_dim: (chunks, 1, 1),
4239 block_dim: (256, 1, 1),
4240 shared_mem_bytes: 0,
4241 };
4242 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4243 let __s = self.gpu.stream();
4244 let mut b = __s.launch_builder(&f);
4245 b.arg(src_entry).arg(&mut dv).arg(&wi);
4246 unsafe {
4247 b.launch(cfg)?;
4248 }
4249 Ok(())
4250 }
4251
4252 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4254 self.alloc_uninit::<i8>(n)
4255 }
4256
4257 pub fn qmatvec(
4259 &self,
4260 w: &CudaSlice<u8>,
4261 x: &CudaSlice<f32>,
4262 m: usize,
4263 in_f: usize,
4264 out_f: usize,
4265 qtype: i32,
4266 row_bytes: usize,
4267 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4268 let f = self.func("qmatvec_f32");
4269 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4271 grid_dim: (out_f as u32, m as u32, 1),
4272 block_dim: (256, 1, 1),
4273 shared_mem_bytes: 0,
4274 };
4275 let (inf, outf, mi, qt, rb) =
4276 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4277 let __s_b = self.gpu.stream();
4278 let mut b = __s_b.launch_builder(&f);
4279 b.arg(w)
4280 .arg(x)
4281 .arg(&mut y)
4282 .arg(&inf)
4283 .arg(&outf)
4284 .arg(&mi)
4285 .arg(&qt)
4286 .arg(&rb);
4287 unsafe {
4288 b.launch(cfg)?;
4289 }
4290 Ok(y)
4291 }
4292
4293 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4295 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4296 self.keep_if_capturing(&s);
4297 Ok(s)
4298 }
4299
4300 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4304 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4305 self.keep_if_capturing(&s);
4306 Ok(s)
4307 }
4308
4309 pub fn memset_zeros_view(
4312 &self,
4313 dst: &mut cudarc::driver::CudaViewMut<f32>,
4314 ) -> Result<(), Box<dyn std::error::Error>> {
4315 self.gpu.stream().memset_zeros(dst)?;
4316 Ok(())
4317 }
4318
4319 pub fn stage_expert(
4325 &self,
4326 host_bytes: &[u8],
4327 scratch: &mut CudaSlice<u8>,
4328 off: usize,
4329 ) -> Result<(), Box<dyn std::error::Error>> {
4330 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4333 }
4334
4335 pub fn moe_router_topk(
4341 &self,
4342 logits: &CudaSlice<f32>,
4343 t: usize,
4344 n_expert: usize,
4345 n_used: usize,
4346 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4347 let f = self.func("moe_router_topk_f32");
4348 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 {
4351 grid_dim: (t as u32, 1, 1),
4352 block_dim: (n_expert as u32, 1, 1),
4353 shared_mem_bytes: 0,
4354 };
4355 let (ne, nu) = (n_expert as i32, n_used as i32);
4356 let __s_b = self.gpu.stream();
4357 let mut b = __s_b.launch_builder(&f);
4358 b.arg(logits)
4359 .arg(&mut sel_idx)
4360 .arg(&mut sel_w)
4361 .arg(&ne)
4362 .arg(&nu);
4363 unsafe {
4364 b.launch(cfg)?;
4365 }
4366 Ok((sel_idx, sel_w))
4367 }
4368
4369 pub fn moe_router_topk_scaled(
4372 &self,
4373 logits: &CudaSlice<f32>,
4374 t: usize,
4375 n_expert: usize,
4376 n_used: usize,
4377 ex_scale: &CudaSlice<f32>,
4378 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4379 let f = self.func("moe_router_topk_scaled_f32");
4384 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4385 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4386 let cfg = LaunchConfig {
4387 grid_dim: (t as u32, 1, 1),
4388 block_dim: (n_expert as u32, 1, 1),
4389 shared_mem_bytes: 0,
4390 };
4391 let (ne, nu) = (n_expert as i32, n_used as i32);
4392 let __s_b = self.gpu.stream();
4393 let mut b = __s_b.launch_builder(&f);
4394 b.arg(logits)
4395 .arg(&mut sel_idx)
4396 .arg(&mut sel_w)
4397 .arg(&ne)
4398 .arg(&nu)
4399 .arg(ex_scale);
4400 unsafe {
4401 b.launch(cfg)?;
4402 }
4403 Ok((sel_idx, sel_w))
4404 }
4405
4406 pub fn moe_router_topk_host(
4414 &self,
4415 logits: &CudaSlice<f32>,
4416 t: usize,
4417 n_expert: usize,
4418 n_used: usize,
4419 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4420 let f = self.func("moe_router_topk_f32");
4421 let n = t * n_used;
4422 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
4423 let mut sel_w = self.alloc_uninit::<f32>(n)?;
4424 let cfg = LaunchConfig {
4425 grid_dim: (t as u32, 1, 1),
4426 block_dim: (n_expert as u32, 1, 1),
4427 shared_mem_bytes: 0,
4428 };
4429 let (ne, nu) = (n_expert as i32, n_used as i32);
4430 let __s_b = self.gpu.stream();
4431 let mut b = __s_b.launch_builder(&f);
4432 b.arg(logits)
4433 .arg(&mut sel_idx)
4434 .arg(&mut sel_w)
4435 .arg(&ne)
4436 .arg(&nu);
4437 unsafe {
4438 b.launch(cfg)?;
4439 }
4440 let bytes = n * 8;
4442 let mut guard = self.router_stage.lock().unwrap();
4443 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4444 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4445 }
4446 let stage = guard.as_mut().unwrap();
4447 let (si, sw) = unsafe {
4448 (
4449 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4450 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4451 )
4452 };
4453 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()))
4457 }
4458
4459 #[allow(clippy::too_many_arguments)]
4463 pub fn moe_router_sigmoid_topk(
4464 &self,
4465 logits: &CudaSlice<f32>,
4466 t: usize,
4467 n_expert: usize,
4468 n_used: usize,
4469 active_count: usize,
4470 correction_bias: &CudaSlice<f32>,
4471 active: &CudaSlice<u8>,
4472 scaling_factor: f32,
4473 route_norm: bool,
4474 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4475 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4476 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
4477 return Err(format!(
4478 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
4479 )
4480 .into());
4481 }
4482 if logits.len() < t * n_expert
4483 || correction_bias.len() != n_expert
4484 || active.len() != n_expert
4485 {
4486 return Err(format!(
4487 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
4488 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
4489 ).into());
4490 }
4491 let f = if crate::sig_expf_dev_on() && crate::topk_fast_on() {
4492 self.func("moe_router_sigmoid_topk_f32_dexp_fast")
4495 } else if crate::sig_expf_dev_on() {
4496 self.func("moe_router_sigmoid_topk_f32_dexp")
4497 } else if crate::topk_fast_on() {
4498 self.func("moe_router_sigmoid_topk_f32_fast")
4499 } else {
4500 self.func("moe_router_sigmoid_topk_f32")
4501 };
4502 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
4503 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
4504 let threads = n_expert.div_ceil(32) * 32;
4505 let cfg = LaunchConfig {
4506 grid_dim: (t as u32, 1, 1),
4507 block_dim: (threads as u32, 1, 1),
4508 shared_mem_bytes: 0,
4509 };
4510 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4511 let __s_b = self.gpu.stream();
4512 let mut b = __s_b.launch_builder(&f);
4513 b.arg(logits)
4514 .arg(correction_bias)
4515 .arg(active)
4516 .arg(&mut sel_idx)
4517 .arg(&mut sel_w)
4518 .arg(&ne)
4519 .arg(&nu)
4520 .arg(&scaling_factor)
4521 .arg(&rn);
4522 unsafe {
4523 b.launch(cfg)?;
4524 }
4525 Ok((sel_idx, sel_w))
4526 }
4527
4528 #[allow(clippy::too_many_arguments)]
4531 pub fn moe_router_sigmoid_topk_into(
4532 &self,
4533 logits: &CudaSlice<f32>,
4534 t: usize,
4535 n_expert: usize,
4536 n_used: usize,
4537 active_count: usize,
4538 correction_bias: &CudaSlice<f32>,
4539 active: &CudaSlice<u8>,
4540 scaling_factor: f32,
4541 route_norm: bool,
4542 sel_idx: &mut CudaSlice<i32>,
4543 sel_w: &mut CudaSlice<f32>,
4544 ) -> Result<(), Box<dyn std::error::Error>> {
4545 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
4546 if n_expert == 0
4547 || n_expert > 1024
4548 || n_used == 0
4549 || n_used > n_expert
4550 || logits.len() < t * n_expert
4551 || correction_bias.len() != n_expert
4552 || active.len() != n_expert
4553 || sel_idx.len() < t * n_used
4554 || sel_w.len() < t * n_used
4555 {
4556 return Err("sigmoid router _into geometry mismatch".into());
4557 }
4558 let f = if crate::sig_expf_dev_on() {
4559 self.func("moe_router_sigmoid_topk_f32_dexp")
4560 } else if crate::topk_fast_on() {
4561 self.func("moe_router_sigmoid_topk_f32_fast")
4562 } else {
4563 self.func("moe_router_sigmoid_topk_f32")
4564 };
4565 let threads = n_expert.div_ceil(32) * 32;
4566 let cfg = LaunchConfig {
4567 grid_dim: (t as u32, 1, 1),
4568 block_dim: (threads as u32, 1, 1),
4569 shared_mem_bytes: 0,
4570 };
4571 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
4572 let __s_b = self.gpu.stream();
4573 let mut b = __s_b.launch_builder(&f);
4574 b.arg(logits)
4575 .arg(correction_bias)
4576 .arg(active)
4577 .arg(&mut *sel_idx)
4578 .arg(&mut *sel_w)
4579 .arg(&ne)
4580 .arg(&nu)
4581 .arg(&scaling_factor)
4582 .arg(&rn);
4583 unsafe {
4584 b.launch(cfg)?;
4585 }
4586 Ok(())
4587 }
4588
4589 #[allow(clippy::too_many_arguments)]
4592 pub fn moe_router_sigmoid_topk_host(
4593 &self,
4594 logits: &CudaSlice<f32>,
4595 t: usize,
4596 n_expert: usize,
4597 n_used: usize,
4598 active_count: usize,
4599 correction_bias: &CudaSlice<f32>,
4600 active: &CudaSlice<u8>,
4601 scaling_factor: f32,
4602 route_norm: bool,
4603 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
4604 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
4605 logits,
4606 t,
4607 n_expert,
4608 n_used,
4609 active_count,
4610 correction_bias,
4611 active,
4612 scaling_factor,
4613 route_norm,
4614 )?;
4615 let n = t * n_used;
4616 let bytes = n * 8;
4617 let mut guard = self.router_stage.lock().unwrap();
4618 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
4619 *guard = Some(PinnedStage::new(bytes.max(4096))?);
4620 }
4621 let stage = guard.as_mut().unwrap();
4622 let (si, sw) = unsafe {
4623 (
4624 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
4625 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
4626 )
4627 };
4628 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
4629 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
4630 self.gpu.stream().synchronize()?;
4631 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
4632 }
4633
4634 pub fn stage_expert_async(
4638 &self,
4639 host_bytes: &[u8],
4640 scratch: &mut CudaSlice<u8>,
4641 off: usize,
4642 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
4643 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
4644 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
4645 Ok(self.copy_stream.record_event(None)?)
4646 }
4647
4648 pub fn compute_wait(
4650 &self,
4651 ev: &cudarc::driver::CudaEvent,
4652 ) -> Result<(), Box<dyn std::error::Error>> {
4653 self.gpu.stream().wait(ev)?;
4654 Ok(())
4655 }
4656
4657 pub fn qmatvec_view(
4662 &self,
4663 w: &CudaSlice<u8>,
4664 range: std::ops::Range<usize>,
4665 x: &cudarc::driver::CudaView<f32>,
4666 m: usize,
4667 in_f: usize,
4668 out_f: usize,
4669 qtype: i32,
4670 row_bytes: usize,
4671 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4672 let f = self.func("qmatvec_f32");
4673 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4676 grid_dim: (out_f as u32, m as u32, 1),
4677 block_dim: (256, 1, 1),
4678 shared_mem_bytes: 0,
4679 };
4680 let (inf, outf, mi, qt, rb) =
4681 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4682 let __s_b = self.gpu.stream();
4683 let mut b = __s_b.launch_builder(&f);
4684 b.arg(&wv)
4685 .arg(x)
4686 .arg(&mut y)
4687 .arg(&inf)
4688 .arg(&outf)
4689 .arg(&mi)
4690 .arg(&qt)
4691 .arg(&rb);
4692 unsafe {
4693 b.launch(cfg)?;
4694 }
4695 Ok(y)
4696 }
4697
4698 #[allow(clippy::too_many_arguments)]
4705 pub fn moe_gate_up_silu8_q8(
4709 &self,
4710 gp: WPtr8,
4711 up: WPtr8,
4712 aq: &CudaSlice<i8>,
4713 ad: &CudaSlice<f32>,
4714 in_f: usize,
4715 n_ff: usize,
4716 n_used: usize,
4717 qt_g: i32,
4718 qt_u: i32,
4719 rb_g: usize,
4720 rb_u: usize,
4721 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4722 let f = self.func("moe_gate_up_silu8_q8");
4723 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
4724 let cfg = LaunchConfig {
4725 grid_dim: (n_ff as u32, n_used as u32, 1),
4726 block_dim: (32, 1, 1),
4727 shared_mem_bytes: 0,
4728 };
4729 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4730 let __s_b = self.gpu.stream();
4731 let mut b = __s_b.launch_builder(&f);
4732 b.arg(&gp)
4733 .arg(&up)
4734 .arg(aq)
4735 .arg(ad)
4736 .arg(&mut act)
4737 .arg(&inf)
4738 .arg(&nff)
4739 .arg(&qt_g)
4740 .arg(&qt_u)
4741 .arg(&rbg)
4742 .arg(&rbu);
4743 unsafe {
4744 b.launch(cfg)?;
4745 }
4746 Ok(act)
4747 }
4748
4749 #[allow(clippy::too_many_arguments)]
4750 pub fn moe_down8_fma_q8(
4751 &self,
4752 dp: WPtr8,
4753 w: F32x8,
4754 aq2: &CudaSlice<i8>,
4755 ad2: &CudaSlice<f32>,
4756 dst: &mut cudarc::driver::CudaViewMut<f32>,
4757 in_f: usize,
4758 out_f: usize,
4759 n_used: usize,
4760 qt: i32,
4761 rb: usize,
4762 ) -> Result<(), Box<dyn std::error::Error>> {
4763 let f = self.func("moe_down8_fma_q8");
4764 let cfg = LaunchConfig {
4765 grid_dim: (out_f as u32, 1, 1),
4766 block_dim: (32, 1, 1),
4767 shared_mem_bytes: 0,
4768 };
4769 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4770 let __s_b = self.gpu.stream();
4771 let mut b = __s_b.launch_builder(&f);
4772 b.arg(&dp)
4773 .arg(&w)
4774 .arg(aq2)
4775 .arg(ad2)
4776 .arg(dst)
4777 .arg(&inf)
4778 .arg(&outf)
4779 .arg(&nu)
4780 .arg(&qt)
4781 .arg(&rbi);
4782 unsafe {
4783 b.launch(cfg)?;
4784 }
4785 Ok(())
4786 }
4787
4788 pub fn qmatvec_expert_q8(
4790 &self,
4791 w: &CudaSlice<u8>,
4792 range: std::ops::Range<usize>,
4793 aq: &CudaSlice<i8>,
4794 ad: &CudaSlice<f32>,
4795 m: usize,
4796 in_f: usize,
4797 out_f: usize,
4798 qtype: i32,
4799 row_bytes: usize,
4800 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4801 let f = self.func("qmatvec_expert_q8");
4802 let wv = w.slice(range);
4803 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
4804 const ROWS: u32 = 4; let cfg = LaunchConfig {
4806 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
4807 block_dim: (32, ROWS, 1),
4808 shared_mem_bytes: 0,
4809 };
4810 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
4811 let __s_b = self.gpu.stream();
4812 let mut b = __s_b.launch_builder(&f);
4813 b.arg(&wv)
4814 .arg(aq)
4815 .arg(ad)
4816 .arg(&mut y)
4817 .arg(&inf)
4818 .arg(&outf)
4819 .arg(&mi)
4820 .arg(&qtype)
4821 .arg(&rbi);
4822 unsafe {
4823 b.launch(cfg)?;
4824 }
4825 Ok(y)
4826 }
4827
4828 pub fn moe_gate_up_silu8(
4829 &self,
4830 gp: WPtr8,
4831 up: WPtr8,
4832 x: &cudarc::driver::CudaView<f32>,
4833 in_f: usize,
4834 n_ff: usize,
4835 n_used: usize,
4836 qt_g: i32,
4837 qt_u: i32,
4838 rb_g: usize,
4839 rb_u: usize,
4840 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4841 let f = self.func("moe_gate_up_silu8_f32");
4842 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
4844 grid_dim: (n_ff as u32, n_used as u32, 1),
4845 block_dim: (256, 1, 1),
4846 shared_mem_bytes: 0,
4847 };
4848 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
4849 let __s_b = self.gpu.stream();
4850 let mut b = __s_b.launch_builder(&f);
4851 b.arg(&gp)
4852 .arg(&up)
4853 .arg(x)
4854 .arg(&mut act)
4855 .arg(&inf)
4856 .arg(&nff)
4857 .arg(&qt_g)
4858 .arg(&qt_u)
4859 .arg(&rbg)
4860 .arg(&rbu);
4861 unsafe {
4862 b.launch(cfg)?;
4863 }
4864 Ok(act)
4865 }
4866
4867 #[allow(clippy::too_many_arguments)]
4873 pub fn moe_down8_fma_into(
4874 &self,
4875 dp: WPtr8,
4876 w: F32x8,
4877 act: &CudaSlice<f32>,
4878 dst: &mut cudarc::driver::CudaViewMut<f32>,
4879 in_f: usize,
4880 out_f: usize,
4881 n_used: usize,
4882 qt: i32,
4883 rb: usize,
4884 ) -> Result<(), Box<dyn std::error::Error>> {
4885 let f = self.func("moe_down8_fma_f32");
4886 let cfg = LaunchConfig {
4887 grid_dim: (out_f as u32, 1, 1),
4888 block_dim: (256, 1, 1),
4889 shared_mem_bytes: 0,
4890 };
4891 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
4892 let __s_b = self.gpu.stream();
4893 let mut b = __s_b.launch_builder(&f);
4894 b.arg(&dp)
4895 .arg(&w)
4896 .arg(act)
4897 .arg(dst)
4898 .arg(&inf)
4899 .arg(&outf)
4900 .arg(&nu)
4901 .arg(&qt)
4902 .arg(&rbv);
4903 unsafe {
4904 b.launch(cfg)?;
4905 }
4906 Ok(())
4907 }
4908
4909 #[allow(clippy::too_many_arguments)]
4914 #[allow(clippy::too_many_arguments)]
4929 #[allow(clippy::too_many_arguments)]
4931 pub fn moe_pairs_matvec_q8(
4932 &self,
4933 table: &CudaSlice<u64>,
4934 proj: i32,
4935 pair_tok: &CudaSlice<i32>,
4936 pair_ex: &CudaSlice<i32>,
4937 aq: &CudaSlice<i8>,
4938 ad: &CudaSlice<f32>,
4939 in_f: usize,
4940 out_f: usize,
4941 n_expert: usize,
4942 n_pairs: usize,
4943 qtype: i32,
4944 row_bytes: usize,
4945 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4946 let f = self.func("moe_pairs_matvec_q8");
4947 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
4948 const ROWS: u32 = 4;
4949 let cfg = LaunchConfig {
4950 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
4951 block_dim: (32, ROWS, 1),
4952 shared_mem_bytes: 0,
4953 };
4954 let (inf, outf, ne, np, rbi) = (
4955 in_f as i32,
4956 out_f as i32,
4957 n_expert as i32,
4958 n_pairs as i32,
4959 row_bytes as i64,
4960 );
4961 let __s_b = self.gpu.stream();
4962 let mut b = __s_b.launch_builder(&f);
4963 b.arg(table)
4964 .arg(&proj)
4965 .arg(pair_tok)
4966 .arg(pair_ex)
4967 .arg(aq)
4968 .arg(ad)
4969 .arg(&mut y)
4970 .arg(&inf)
4971 .arg(&outf)
4972 .arg(&ne)
4973 .arg(&np)
4974 .arg(&qtype)
4975 .arg(&rbi);
4976 unsafe {
4977 b.launch(cfg)?;
4978 }
4979 Ok(y)
4980 }
4981
4982 #[allow(clippy::too_many_arguments)]
4984 pub fn moe_pairs_matvec_q8_em(
4985 &self,
4986 table: &CudaSlice<u64>,
4987 proj: i32,
4988 ex_ids: &CudaSlice<i32>,
4989 ex_off: &CudaSlice<i32>,
4990 ex_pairs: &CudaSlice<i32>,
4991 pair_tok: &CudaSlice<i32>,
4992 aq: &CudaSlice<i8>,
4993 ad: &CudaSlice<f32>,
4994 in_f: usize,
4995 out_f: usize,
4996 n_expert: usize,
4997 n_active: usize,
4998 n_pairs: usize,
4999 qtype: i32,
5000 row_bytes: usize,
5001 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5002 let f = self.func("moe_pairs_matvec_q8_em");
5003 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5004 const ROWS: u32 = 4;
5005 let cfg = LaunchConfig {
5006 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5007 block_dim: (32, ROWS, 1),
5008 shared_mem_bytes: 0,
5009 };
5010 let (inf, outf, ne, na, rbi) = (
5011 in_f as i32,
5012 out_f as i32,
5013 n_expert as i32,
5014 n_active as i32,
5015 row_bytes as i64,
5016 );
5017 let __s_b = self.gpu.stream();
5018 let mut b = __s_b.launch_builder(&f);
5019 b.arg(table)
5020 .arg(&proj)
5021 .arg(ex_ids)
5022 .arg(ex_off)
5023 .arg(ex_pairs)
5024 .arg(pair_tok)
5025 .arg(aq)
5026 .arg(ad)
5027 .arg(&mut y)
5028 .arg(&inf)
5029 .arg(&outf)
5030 .arg(&ne)
5031 .arg(&na)
5032 .arg(&qtype)
5033 .arg(&rbi);
5034 unsafe {
5035 b.launch(cfg)?;
5036 }
5037 Ok(y)
5038 }
5039
5040 #[allow(clippy::too_many_arguments)]
5043 pub fn moe_pairs_matvec_q8_dec(
5044 &self,
5045 table: &CudaSlice<u64>,
5046 proj: i32,
5047 ex_ids: &CudaSlice<i32>,
5048 ex_off: &CudaSlice<i32>,
5049 ex_pairs: &CudaSlice<i32>,
5050 pair_tok: &CudaSlice<i32>,
5051 aq: &CudaSlice<i8>,
5052 ad: &CudaSlice<f32>,
5053 in_f: usize,
5054 out_f: usize,
5055 n_expert: usize,
5056 n_active: usize,
5057 n_pairs: usize,
5058 qtype: i32,
5059 row_bytes: usize,
5060 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5061 let f = self.func("moe_pairs_matvec_q8_dec");
5062 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5063 const ROWS: u32 = 4;
5064 let cfg = LaunchConfig {
5065 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5066 block_dim: (32, ROWS, 1),
5067 shared_mem_bytes: 0,
5068 };
5069 let (inf, outf, ne, na, rbi) = (
5070 in_f as i32,
5071 out_f as i32,
5072 n_expert as i32,
5073 n_active as i32,
5074 row_bytes as i64,
5075 );
5076 let __s_b = self.gpu.stream();
5077 let mut b = __s_b.launch_builder(&f);
5078 b.arg(table)
5079 .arg(&proj)
5080 .arg(ex_ids)
5081 .arg(ex_off)
5082 .arg(ex_pairs)
5083 .arg(pair_tok)
5084 .arg(aq)
5085 .arg(ad)
5086 .arg(&mut y)
5087 .arg(&inf)
5088 .arg(&outf)
5089 .arg(&ne)
5090 .arg(&na)
5091 .arg(&qtype)
5092 .arg(&rbi);
5093 unsafe {
5094 b.launch(cfg)?;
5095 }
5096 Ok(y)
5097 }
5098
5099 pub fn moe_pairs_gelu_mul(
5100 &self,
5101 gate: &CudaSlice<f32>,
5102 up: &CudaSlice<f32>,
5103 n: usize,
5104 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5105 let f = self.func("moe_pairs_gelu_mul");
5106 let mut act = self.alloc_uninit::<f32>(n)?;
5107 let cfg = LaunchConfig::for_num_elems(n as u32);
5108 let nl = n as i64;
5109 let __s_b = self.gpu.stream();
5110 let mut b = __s_b.launch_builder(&f);
5111 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5112 unsafe {
5113 b.launch(cfg)?;
5114 }
5115 Ok(act)
5116 }
5117
5118 pub fn moe_pairs_silu_mul(
5119 &self,
5120 gate: &CudaSlice<f32>,
5121 up: &CudaSlice<f32>,
5122 n: usize,
5123 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5124 let f = self.func("moe_pairs_silu_mul");
5125 let mut act = self.alloc_uninit::<f32>(n)?;
5126 let cfg = LaunchConfig::for_num_elems(n as u32);
5127 let nl = n as i64;
5128 let __s_b = self.gpu.stream();
5129 let mut b = __s_b.launch_builder(&f);
5130 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5131 unsafe {
5132 b.launch(cfg)?;
5133 }
5134 Ok(act)
5135 }
5136
5137 #[allow(clippy::too_many_arguments)]
5138 pub fn moe_pairs_scatter(
5139 &self,
5140 y_down: &CudaSlice<f32>,
5141 pair_w: &CudaSlice<f32>,
5142 tok_pair_off: &CudaSlice<i32>,
5143 tok_pair_ids: &CudaSlice<i32>,
5144 moe_out: &mut CudaSlice<f32>,
5145 t: usize,
5146 n_embd: usize,
5147 ) -> Result<(), Box<dyn std::error::Error>> {
5148 let f = self.func("moe_pairs_scatter");
5149 let cfg = LaunchConfig {
5150 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5151 block_dim: (256, 1, 1),
5152 shared_mem_bytes: 0,
5153 };
5154 let ne = n_embd as i32;
5155 let __s_b = self.gpu.stream();
5156 let mut b = __s_b.launch_builder(&f);
5157 b.arg(y_down)
5158 .arg(pair_w)
5159 .arg(tok_pair_off)
5160 .arg(tok_pair_ids)
5161 .arg(moe_out)
5162 .arg(&ne);
5163 unsafe {
5164 b.launch(cfg)?;
5165 }
5166 Ok(())
5167 }
5168
5169 #[allow(clippy::too_many_arguments)]
5173 pub fn moe_gate_up_gelu8_dev_q8(
5174 &self,
5175 table: &CudaSlice<u64>,
5176 sel: &cudarc::driver::CudaView<i32>,
5177 aq: &CudaSlice<i8>,
5178 ad: &CudaSlice<f32>,
5179 in_f: usize,
5180 n_ff: usize,
5181 n_used: usize,
5182 n_expert: usize,
5183 qt_g: i32,
5184 qt_u: i32,
5185 rb_g: usize,
5186 rb_u: usize,
5187 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5188 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5189 let (inf, nff, ne, rbg, rbu) = (
5190 in_f as i32,
5191 n_ff as i32,
5192 n_expert as i32,
5193 rb_g as i64,
5194 rb_u as i64,
5195 );
5196 let f = self.func("moe_gate_up_gelu8_dev_q8");
5197 let cfg = LaunchConfig {
5198 grid_dim: (n_ff as u32, n_used as u32, 1),
5199 block_dim: (32, 1, 1),
5200 shared_mem_bytes: 0,
5201 };
5202 let __s_b = self.gpu.stream();
5203 let mut b = __s_b.launch_builder(&f);
5204 b.arg(table)
5205 .arg(sel)
5206 .arg(aq)
5207 .arg(ad)
5208 .arg(&mut act)
5209 .arg(&inf)
5210 .arg(&nff)
5211 .arg(&ne)
5212 .arg(&qt_g)
5213 .arg(&qt_u)
5214 .arg(&rbg)
5215 .arg(&rbu);
5216 unsafe {
5217 b.launch(cfg)?;
5218 }
5219 Ok(act)
5220 }
5221
5222 #[allow(clippy::too_many_arguments)]
5224 pub fn moe_gate_up_gelu8_dev_q8_rows(
5225 &self,
5226 table: &CudaSlice<u64>,
5227 sel: &CudaSlice<i32>,
5228 aq: &CudaSlice<i8>,
5229 ad: &CudaSlice<f32>,
5230 t: usize,
5231 in_f: usize,
5232 n_ff: usize,
5233 n_used: usize,
5234 n_expert: usize,
5235 qt_g: i32,
5236 qt_u: i32,
5237 rb_g: usize,
5238 rb_u: usize,
5239 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5240 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5241 let (inf, nff, ne, rbg, rbu, nu) = (
5242 in_f as i32,
5243 n_ff as i32,
5244 n_expert as i32,
5245 rb_g as i64,
5246 rb_u as i64,
5247 n_used as i32,
5248 );
5249 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5250 let cfg = LaunchConfig {
5251 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5252 block_dim: (32, 1, 1),
5253 shared_mem_bytes: 0,
5254 };
5255 let __s_b = self.gpu.stream();
5256 let mut b = __s_b.launch_builder(&f);
5257 b.arg(table)
5258 .arg(sel)
5259 .arg(aq)
5260 .arg(ad)
5261 .arg(&mut act)
5262 .arg(&inf)
5263 .arg(&nff)
5264 .arg(&ne)
5265 .arg(&qt_g)
5266 .arg(&qt_u)
5267 .arg(&rbg)
5268 .arg(&rbu)
5269 .arg(&nu);
5270 unsafe {
5271 b.launch(cfg)?;
5272 }
5273 Ok(act)
5274 }
5275
5276 #[allow(clippy::too_many_arguments)]
5278 pub fn moe_gate_up_gelu8_dev_q8_csr(
5279 &self,
5280 table: &CudaSlice<u64>,
5281 sel: &CudaSlice<i32>,
5282 aq: &CudaSlice<i8>,
5283 ad: &CudaSlice<f32>,
5284 n_pairs: usize,
5285 in_f: usize,
5286 n_ff: usize,
5287 n_used: usize,
5288 n_expert: usize,
5289 qt_g: i32,
5290 qt_u: i32,
5291 rb_g: usize,
5292 rb_u: usize,
5293 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5294 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5295 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5296 in_f as i32,
5297 n_ff as i32,
5298 n_expert as i32,
5299 rb_g as i64,
5300 rb_u as i64,
5301 n_used as i32,
5302 n_pairs as i32,
5303 );
5304 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5305 let cfg = LaunchConfig {
5306 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5307 block_dim: (32, 1, 1),
5308 shared_mem_bytes: 0,
5309 };
5310 let __s_b = self.gpu.stream();
5311 let mut b = __s_b.launch_builder(&f);
5312 b.arg(table)
5313 .arg(sel)
5314 .arg(aq)
5315 .arg(ad)
5316 .arg(&mut act)
5317 .arg(&inf)
5318 .arg(&nff)
5319 .arg(&ne)
5320 .arg(&qt_g)
5321 .arg(&qt_u)
5322 .arg(&rbg)
5323 .arg(&rbu)
5324 .arg(&nu)
5325 .arg(&npi);
5326 unsafe {
5327 b.launch(cfg)?;
5328 }
5329 Ok(act)
5330 }
5331
5332 #[allow(clippy::too_many_arguments)]
5334 pub fn moe_down8_fma_dev_q8_rows_g(
5335 &self,
5336 table: &CudaSlice<u64>,
5337 sel: &CudaSlice<i32>,
5338 w: &CudaSlice<f32>,
5339 aq2: &CudaSlice<i8>,
5340 ad2: &CudaSlice<f32>,
5341 dst: &mut CudaSlice<f32>,
5342 t: usize,
5343 in_f: usize,
5344 out_f: usize,
5345 n_used: usize,
5346 n_expert: usize,
5347 qt: i32,
5348 rb: usize,
5349 ) -> Result<(), Box<dyn std::error::Error>> {
5350 let (inf, outf, nu, ne, rbi) = (
5351 in_f as i32,
5352 out_f as i32,
5353 n_used as i32,
5354 n_expert as i32,
5355 rb as i64,
5356 );
5357 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
5361 let f = self.func(if step_b1_w8 {
5362 "moe_down8_fma_dev_q8_rows_w8"
5363 } else {
5364 "moe_down8_fma_dev_q8_rows_g"
5365 });
5366 let cfg = LaunchConfig {
5367 grid_dim: (out_f as u32, 1, t as u32),
5368 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
5369 shared_mem_bytes: 0,
5370 };
5371 let __s_b = self.gpu.stream();
5372 let mut b = __s_b.launch_builder(&f);
5373 b.arg(table)
5374 .arg(sel)
5375 .arg(w)
5376 .arg(aq2)
5377 .arg(ad2)
5378 .arg(dst)
5379 .arg(&inf)
5380 .arg(&outf)
5381 .arg(&nu)
5382 .arg(&ne)
5383 .arg(&qt)
5384 .arg(&rbi);
5385 unsafe {
5386 b.launch(cfg)?;
5387 }
5388 Ok(())
5389 }
5390
5391 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
5395 let (out_f, in_f) = (2048usize, 2816usize);
5396 let nblk = in_f / 32;
5397 let mut seed = 0x9E3779B97F4A7C15u64;
5398 let mut rng = move || {
5399 seed = seed
5400 .wrapping_mul(6364136223846793005)
5401 .wrapping_add(1442695040888963407);
5402 (seed >> 33) as u8
5403 };
5404 let mut w = vec![0u8; out_f * nblk * 18];
5405 for b in w.iter_mut() {
5406 *b = rng();
5407 }
5408 for r in 0..out_f {
5409 for g in 0..nblk {
5410 let off = (r * nblk + g) * 18;
5411 w[off] = 0x00;
5412 w[off + 1] = 0x2C; }
5414 }
5415 let qplane = out_f * nblk * 16;
5416 let mut wrp = vec![0u8; w.len()];
5417 for r in 0..out_f {
5418 for g in 0..nblk {
5419 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
5420 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
5421 .copy_from_slice(&src[0..2]);
5422 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
5423 }
5424 }
5425 let w_d = self.htod_bytes(&w)?;
5426 let wrp_d = self.htod_bytes(&wrp)?;
5427 let mut aq = vec![0i8; m * in_f];
5428 for v in aq.iter_mut() {
5429 *v = rng() as i8;
5430 }
5431 let aq_d = self.htod_i8(&aq)?;
5432 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
5433 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5434 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5435 const RPB: u32 = 4;
5436 let cfg = LaunchConfig {
5437 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
5438 block_dim: (32, RPB, 1),
5439 shared_mem_bytes: 0,
5440 };
5441 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
5442 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
5443 let fb = self.func("qmatvec_q4_0_mmvq_b4");
5444 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
5445 {
5446 let __s_b = self.gpu.stream();
5447 let mut b = __s_b.launch_builder(&fb);
5448 b.arg(&w_d)
5449 .arg(&aq_d)
5450 .arg(&ad_d)
5451 .arg(&mut y0)
5452 .arg(&inf)
5453 .arg(&outf)
5454 .arg(&mi)
5455 .arg(&rb);
5456 unsafe {
5457 b.launch(cfg)?;
5458 }
5459 let __s_b = self.gpu.stream();
5460 let mut b = __s_b.launch_builder(&fr);
5461 b.arg(&wrp_d)
5462 .arg(&aq_d)
5463 .arg(&ad_d)
5464 .arg(&mut y1)
5465 .arg(&inf)
5466 .arg(&outf)
5467 .arg(&mi)
5468 .arg(&qp);
5469 unsafe {
5470 b.launch(cfg)?;
5471 }
5472 }
5473 self.gpu.stream().synchronize()?;
5474 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
5475 let nd = h0
5476 .iter()
5477 .zip(&h1)
5478 .filter(|(a, b)| a.to_bits() != b.to_bits())
5479 .count();
5480 if nd != 0 {
5481 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
5482 }
5483 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
5484 self.gpu.stream().synchronize()?;
5485 let t0 = std::time::Instant::now();
5486 for _ in 0..500 {
5487 if rp {
5488 let __s_b = self.gpu.stream();
5489 let mut b = __s_b.launch_builder(&fr);
5490 b.arg(&wrp_d)
5491 .arg(&aq_d)
5492 .arg(&ad_d)
5493 .arg(&mut y1)
5494 .arg(&inf)
5495 .arg(&outf)
5496 .arg(&mi)
5497 .arg(&qp);
5498 unsafe {
5499 b.launch(cfg)?;
5500 }
5501 } else {
5502 let __s_b = self.gpu.stream();
5503 let mut b = __s_b.launch_builder(&fb);
5504 b.arg(&w_d)
5505 .arg(&aq_d)
5506 .arg(&ad_d)
5507 .arg(&mut y0)
5508 .arg(&inf)
5509 .arg(&outf)
5510 .arg(&mi)
5511 .arg(&rb);
5512 unsafe {
5513 b.launch(cfg)?;
5514 }
5515 }
5516 }
5517 self.gpu.stream().synchronize()?;
5518 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
5519 };
5520 let _ = time(false)?;
5521 let _ = time(true)?; Ok((time(false)?, time(true)?))
5523 }
5524
5525 pub fn build_q4_rp4(
5530 &self,
5531 t: &mut crate::model::GpuTensor,
5532 ) -> Result<(), Box<dyn std::error::Error>> {
5533 use crate::model::GpuTensor;
5534 let GpuTensor::Quant {
5535 bytes,
5536 qtype,
5537 row_bytes,
5538 ne,
5539 rp4,
5540 ..
5541 } = t
5542 else {
5543 return Ok(());
5544 };
5545 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
5546 return Ok(());
5547 }
5548 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5549 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
5550 return Ok(());
5551 }
5552 let nblk = in_f / 32;
5553 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
5554 let f = self.func("q4_0_split_rp_build");
5555 let n = (out_f * nblk) as i32;
5556 let cfg = LaunchConfig {
5557 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5558 block_dim: (256, 1, 1),
5559 shared_mem_bytes: 0,
5560 };
5561 let (of, nb) = (out_f as i32, nblk as i32);
5562 let _ = n;
5563 let __s_b = self.gpu.stream();
5564 let mut b = __s_b.launch_builder(&f);
5565 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5566 unsafe {
5567 b.launch(cfg)?;
5568 }
5569 *rp4 = Some(dst);
5570 Ok(())
5571 }
5572
5573 pub fn build_q8_rp4(
5578 &self,
5579 t: &mut crate::model::GpuTensor,
5580 ) -> Result<(), Box<dyn std::error::Error>> {
5581 use crate::model::GpuTensor;
5582 let GpuTensor::Quant {
5583 bytes,
5584 qtype,
5585 row_bytes,
5586 ne,
5587 rp4,
5588 ..
5589 } = t
5590 else {
5591 return Ok(());
5592 };
5593 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
5594 return Ok(());
5595 }
5596 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5597 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
5598 return Ok(());
5599 }
5600 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
5601 Ok(())
5602 }
5603
5604 pub fn build_q8_rp4_raw(
5607 &self,
5608 bytes: &CudaSlice<u8>,
5609 in_f: usize,
5610 out_f: usize,
5611 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5612 assert!(in_f % 32 == 0);
5613 let nblk = in_f / 32;
5614 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
5615 let f = self.func("q8_0_split_rp_build");
5616 let cfg = LaunchConfig {
5617 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
5618 block_dim: (256, 1, 1),
5619 shared_mem_bytes: 0,
5620 };
5621 let (of, nb) = (out_f as i32, nblk as i32);
5622 let __s_b = self.gpu.stream();
5623 let mut b = __s_b.launch_builder(&f);
5624 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5625 unsafe {
5626 b.launch(cfg)?;
5627 }
5628 Ok(dst)
5629 }
5630
5631 pub fn build_q4k_rp4(
5639 &self,
5640 t: &mut crate::model::GpuTensor,
5641 ) -> Result<(), Box<dyn std::error::Error>> {
5642 use crate::model::GpuTensor;
5643 let GpuTensor::Quant {
5644 bytes,
5645 qtype,
5646 row_bytes,
5647 ne,
5648 rp4,
5649 ..
5650 } = t
5651 else {
5652 return Ok(());
5653 };
5654 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
5655 return Ok(());
5656 }
5657 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5658 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
5659 return Ok(());
5660 }
5661 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
5662 Ok(())
5663 }
5664
5665 pub fn build_q6k_rp4(
5666 &self,
5667 t: &mut crate::model::GpuTensor,
5668 ) -> Result<(), Box<dyn std::error::Error>> {
5669 use crate::model::GpuTensor;
5670 let GpuTensor::Quant {
5671 bytes,
5672 qtype,
5673 row_bytes,
5674 ne,
5675 rp4,
5676 ..
5677 } = t
5678 else {
5679 return Ok(());
5680 };
5681 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
5682 return Ok(());
5683 }
5684 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
5685 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
5686 return Ok(());
5687 }
5688 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
5689 Ok(())
5690 }
5691
5692 pub fn build_kq_rp4_raw(
5694 &self,
5695 bytes: &CudaSlice<u8>,
5696 in_f: usize,
5697 out_f: usize,
5698 qtype: i32,
5699 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
5700 assert!(in_f % 256 == 0);
5701 let nsbk = in_f / 256;
5702 let (sb_bytes, kname) = match qtype {
5703 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
5704 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
5705 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
5706 };
5707 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
5708 let f = self.func(kname);
5709 let cfg = LaunchConfig {
5710 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
5711 block_dim: (256, 1, 1),
5712 shared_mem_bytes: 0,
5713 };
5714 let (of, nb) = (out_f as i32, nsbk as i32);
5715 let __s_b = self.gpu.stream();
5716 let mut b = __s_b.launch_builder(&f);
5717 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
5718 unsafe {
5719 b.launch(cfg)?;
5720 }
5721 Ok(dst)
5722 }
5723
5724 pub fn kqrp_enabled() -> bool {
5728 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5729 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
5730 Ok("0") => false,
5731 Ok(_) => true,
5732 Err(_) => cfg!(memra_hopper_mma),
5733 })
5734 }
5735
5736 pub fn build_q4_rp_swap(
5742 &self,
5743 t: &mut crate::model::GpuTensor,
5744 ) -> Result<bool, Box<dyn std::error::Error>> {
5745 use crate::model::GpuTensor;
5746 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
5756 return Ok(false);
5757 }
5758 self.build_q4_rp4(t)?;
5759 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
5761 return Ok(false);
5762 };
5763 match rp4.take() {
5764 Some(split) => {
5765 *bytes = split; *rp = true;
5767 Ok(true)
5768 }
5769 None => Ok(false),
5770 }
5771 }
5772
5773 pub fn q4rp_enabled() -> bool {
5775 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5776 *ON.get_or_init(|| {
5777 std::env::var("MEMRA_Q4RP")
5778 .map(|v| v != "0")
5779 .unwrap_or(true)
5780 })
5781 }
5782
5783 pub fn copy_rows_strided(
5786 &self,
5787 src: &CudaSlice<f32>,
5788 dst: &mut CudaSlice<f32>,
5789 row_elems: usize,
5790 n_rows: usize,
5791 src_stride: usize,
5792 src_off: usize,
5793 ) -> Result<(), Box<dyn std::error::Error>> {
5794 let f = self.func("copy_rows_strided_f32");
5795 let cfg = LaunchConfig {
5796 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5797 block_dim: (256, 1, 1),
5798 shared_mem_bytes: 0,
5799 };
5800 let (re, nr) = (row_elems as i32, n_rows as i32);
5801 let (st, off) = (src_stride as i64, src_off as i64);
5802 let __s_b = self.gpu.stream();
5803 let mut b = __s_b.launch_builder(&f);
5804 b.arg(src)
5805 .arg(&mut *dst)
5806 .arg(&re)
5807 .arg(&nr)
5808 .arg(&st)
5809 .arg(&off);
5810 unsafe {
5811 b.launch(cfg)?;
5812 }
5813 Ok(())
5814 }
5815
5816 pub fn place_rows_strided(
5822 &self,
5823 src: &CudaSlice<f32>,
5824 dst: &mut CudaSlice<f32>,
5825 row_elems: usize,
5826 n_rows: usize,
5827 dst_stride: usize,
5828 dst_off: usize,
5829 ) -> Result<(), Box<dyn std::error::Error>> {
5830 if row_elems == 0 || n_rows == 0 {
5831 return Err("strided row placement requires nonzero rows and row width".into());
5832 }
5833 let src_len = n_rows
5834 .checked_mul(row_elems)
5835 .ok_or("strided row placement source size overflow")?;
5836 let dst_len = n_rows
5837 .checked_sub(1)
5838 .and_then(|rows| rows.checked_mul(dst_stride))
5839 .and_then(|base| base.checked_add(dst_off))
5840 .and_then(|base| base.checked_add(row_elems))
5841 .ok_or("strided row placement destination size overflow")?;
5842 let row_end = dst_off
5843 .checked_add(row_elems)
5844 .ok_or("strided row placement row size overflow")?;
5845 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
5846 return Err(format!(
5847 "strided row placement geometry mismatch: src={} need_src={src_len} \
5848 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
5849 dst_stride={dst_stride} dst_off={dst_off}",
5850 src.len(),
5851 dst.len(),
5852 )
5853 .into());
5854 }
5855 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
5856 return Err("strided row placement exceeds CUDA kernel geometry".into());
5857 }
5858 let f = self.func("place_rows_strided_f32");
5859 let cfg = LaunchConfig {
5860 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
5861 block_dim: (256, 1, 1),
5862 shared_mem_bytes: 0,
5863 };
5864 let (re, nr) = (row_elems as i32, n_rows as i32);
5865 let (st, off) = (dst_stride as i64, dst_off as i64);
5866 let __s_b = self.gpu.stream();
5867 let mut b = __s_b.launch_builder(&f);
5868 b.arg(src)
5869 .arg(&mut *dst)
5870 .arg(&re)
5871 .arg(&nr)
5872 .arg(&st)
5873 .arg(&off);
5874 unsafe {
5875 b.launch(cfg)?;
5876 }
5877 Ok(())
5878 }
5879
5880 pub fn u32_set_k(
5882 &self,
5883 dst: &mut CudaSlice<u32>,
5884 v: u32,
5885 idx: usize,
5886 ) -> Result<(), Box<dyn std::error::Error>> {
5887 let f = self.func("u32_set_k");
5888 let cfg = LaunchConfig {
5889 grid_dim: (1, 1, 1),
5890 block_dim: (1, 1, 1),
5891 shared_mem_bytes: 0,
5892 };
5893 let ii = idx as i32;
5894 let __s_b = self.gpu.stream();
5895 let mut b = __s_b.launch_builder(&f);
5896 b.arg(dst).arg(&v).arg(&ii);
5897 unsafe {
5898 b.launch(cfg)?;
5899 }
5900 Ok(())
5901 }
5902
5903 pub fn i32_add_k(
5905 &self,
5906 d: &mut CudaSlice<i32>,
5907 v: i32,
5908 ) -> Result<(), Box<dyn std::error::Error>> {
5909 let f = self.func("i32_add_k");
5910 let cfg = LaunchConfig {
5911 grid_dim: (1, 1, 1),
5912 block_dim: (32, 1, 1),
5913 shared_mem_bytes: 0,
5914 };
5915 let __s_b = self.gpu.stream();
5916 let mut b = __s_b.launch_builder(&f);
5917 b.arg(d).arg(&v);
5918 unsafe {
5919 b.launch(cfg)?;
5920 }
5921 Ok(())
5922 }
5923
5924 pub fn i32_iota_from(
5926 &self,
5927 ctr: &CudaSlice<i32>,
5928 dst: &mut CudaSlice<i32>,
5929 n: usize,
5930 ) -> Result<(), Box<dyn std::error::Error>> {
5931 let f = self.func("i32_iota_from");
5932 let cfg = LaunchConfig::for_num_elems(n as u32);
5933 let ni = n as i32;
5934 let __s_b = self.gpu.stream();
5935 let mut b = __s_b.launch_builder(&f);
5936 b.arg(ctr).arg(dst).arg(&ni);
5937 unsafe {
5938 b.launch(cfg)?;
5939 }
5940 Ok(())
5941 }
5942
5943 pub fn u32_map_k(
5945 &self,
5946 buf: &mut CudaSlice<u32>,
5947 map: &CudaSlice<u32>,
5948 idx: usize,
5949 ) -> Result<(), Box<dyn std::error::Error>> {
5950 let f = self.func("u32_map_k");
5951 let cfg = LaunchConfig {
5952 grid_dim: (1, 1, 1),
5953 block_dim: (1, 1, 1),
5954 shared_mem_bytes: 0,
5955 };
5956 let ii = idx as i32;
5957 let __s_b = self.gpu.stream();
5958 let mut b = __s_b.launch_builder(&f);
5959 b.arg(buf).arg(map).arg(&ii);
5960 unsafe {
5961 b.launch(cfg)?;
5962 }
5963 Ok(())
5964 }
5965
5966 #[allow(clippy::too_many_arguments)]
5968 pub fn u32_pack2(
5969 &self,
5970 a: &CudaSlice<u32>,
5971 off_a: usize,
5972 n1: usize,
5973 b_in: &CudaSlice<u32>,
5974 n2: usize,
5975 out: &mut CudaSlice<u32>,
5976 ) -> Result<(), Box<dyn std::error::Error>> {
5977 let f = self.func("u32_pack2");
5978 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
5979 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
5980 let __s_b = self.gpu.stream();
5981 let mut b = __s_b.launch_builder(&f);
5982 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
5983 unsafe {
5984 b.launch(cfg)?;
5985 }
5986 Ok(())
5987 }
5988
5989 pub fn moe_w_exscale(
5991 &self,
5992 w: &mut CudaSlice<f32>,
5993 sel: &CudaSlice<i32>,
5994 s: &CudaSlice<f32>,
5995 n: usize,
5996 ) -> Result<(), Box<dyn std::error::Error>> {
5997 let f = self.func("moe_w_exscale");
5998 let cfg = LaunchConfig::for_num_elems(n as u32);
5999 let ni = n as i32;
6000 let __s_b = self.gpu.stream();
6001 let mut b = __s_b.launch_builder(&f);
6002 b.arg(w).arg(sel).arg(s).arg(&ni);
6003 unsafe {
6004 b.launch(cfg)?;
6005 }
6006 Ok(())
6007 }
6008
6009 pub fn moe_w_scale_by_expert(
6012 &self,
6013 w: &mut CudaSlice<f32>,
6014 sel: &CudaSlice<i32>,
6015 macros: &CudaSlice<f32>,
6016 n_expert: usize,
6017 n: usize,
6018 ) -> Result<(), Box<dyn std::error::Error>> {
6019 let f = self.func("moe_w_scale_by_expert");
6020 let cfg = LaunchConfig {
6021 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6022 block_dim: (64, 1, 1),
6023 shared_mem_bytes: 0,
6024 };
6025 let (ne, nn) = (n_expert as i32, n as i32);
6026 let __s_b = self.gpu.stream();
6027 let mut b = __s_b.launch_builder(&f);
6028 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6029 unsafe {
6030 b.launch(cfg)?;
6031 }
6032 Ok(())
6033 }
6034
6035 pub fn moe_gate_up_silu8_dev_q8(
6036 &self,
6037 table: &CudaSlice<u64>,
6038 sel: &cudarc::driver::CudaView<i32>,
6039 aq: &CudaSlice<i8>,
6040 ad: &CudaSlice<f32>,
6041 in_f: usize,
6042 n_ff: usize,
6043 n_used: usize,
6044 n_expert: usize,
6045 qt_g: i32,
6046 qt_u: i32,
6047 rb_g: usize,
6048 rb_u: usize,
6049 macros: &CudaSlice<f32>,
6050 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6051 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6052 let (mode, wpb) = GU.get_or_init(|| {
6053 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6054 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6055 .ok()
6056 .and_then(|v| v.parse().ok())
6057 .unwrap_or(4u32)
6058 .clamp(1, 16);
6059 (mode, wpb)
6060 });
6061 let (mode, wpb) = (mode.as_str(), *wpb);
6062 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6063 let (inf, nff, ne, rbg, rbu) = (
6064 in_f as i32,
6065 n_ff as i32,
6066 n_expert as i32,
6067 rb_g as i64,
6068 rb_u as i64,
6069 );
6070 let (f, cfg) = match mode {
6071 "1" | "2" | "4" => {
6072 let rpw: u32 = mode.parse().unwrap();
6073 let f = self.func(match rpw {
6074 1 => "moe_gate_up_silu8_dev_q8_r1",
6075 2 => "moe_gate_up_silu8_dev_q8_r2",
6076 _ => "moe_gate_up_silu8_dev_q8_r4",
6077 });
6078 let rows_per_block = (rpw * wpb) as usize;
6079 let gx = n_ff.div_ceil(rows_per_block) as u32;
6080 (
6081 f,
6082 LaunchConfig {
6083 grid_dim: (gx, n_used as u32, 1),
6084 block_dim: (32, wpb, 1),
6085 shared_mem_bytes: 0,
6086 },
6087 )
6088 }
6089 "j8" if n_used <= 32 => (
6090 self.func("moe_gate_up_silu8_dev_q8_j8"),
6091 LaunchConfig {
6092 grid_dim: (n_ff as u32, 1, 1),
6093 block_dim: (32, n_used as u32, 1),
6094 shared_mem_bytes: 0,
6095 },
6096 ),
6097 "vsm2" => {
6099 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6100 let sh = (rb_g + rb_u) as u32;
6101 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6102 f.set_attribute(
6103 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6104 sh as i32,
6105 )?;
6106 (
6107 f,
6108 LaunchConfig {
6109 grid_dim: (n_ff as u32, n_used as u32, 1),
6110 block_dim: (32, 1, 1),
6111 shared_mem_bytes: sh,
6112 },
6113 )
6114 }
6115 "vsm" => {
6116 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6117 let sh = (rb_g + rb_u) as u32;
6118 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6119 f.set_attribute(
6120 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6121 sh as i32,
6122 )?;
6123 (
6124 f,
6125 LaunchConfig {
6126 grid_dim: (n_ff as u32, n_used as u32, 1),
6127 block_dim: (32, 1, 1),
6128 shared_mem_bytes: sh,
6129 },
6130 )
6131 }
6132 "sg" => (
6133 self.func("moe_gate_up_silu8_dev_q8_sg"),
6134 LaunchConfig {
6135 grid_dim: (n_ff as u32, n_used as u32, 1),
6136 block_dim: (32, 1, 1),
6137 shared_mem_bytes: 0,
6138 },
6139 ),
6140 "j8sg" if n_used <= 32 => (
6141 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6142 LaunchConfig {
6143 grid_dim: (n_ff as u32, 1, 1),
6144 block_dim: (32, n_used as u32, 1),
6145 shared_mem_bytes: 0,
6146 },
6147 ),
6148 "u64" if in_f == 2048 => (
6149 self.func("moe_gate_up_silu8_dev_q8_u64"),
6150 LaunchConfig {
6151 grid_dim: (n_ff as u32, n_used as u32, 1),
6152 block_dim: (32, 1, 1),
6153 shared_mem_bytes: 0,
6154 },
6155 ),
6156 "gs4" if in_f == 2048 => (
6157 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6158 LaunchConfig {
6159 grid_dim: (n_ff as u32, n_used as u32, 1),
6160 block_dim: (32, 4, 1),
6161 shared_mem_bytes: 0,
6162 },
6163 ),
6164 "v" | "" => (
6166 self.func("moe_gate_up_silu8_dev_q8_v"),
6167 LaunchConfig {
6168 grid_dim: (n_ff as u32, n_used as u32, 1),
6169 block_dim: (32, 1, 1),
6170 shared_mem_bytes: 0,
6171 },
6172 ),
6173 "s2" => (
6174 self.func("moe_gate_up_silu8_dev_q8_s2"),
6175 LaunchConfig {
6176 grid_dim: (n_ff as u32, n_used as u32, 1),
6177 block_dim: (32, 2, 1),
6178 shared_mem_bytes: 0,
6179 },
6180 ),
6181 "s2z" => {
6182 let rz = wpb.min(16); (
6184 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6185 LaunchConfig {
6186 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6187 block_dim: (32, 2, rz),
6188 shared_mem_bytes: 0,
6189 },
6190 )
6191 }
6192 _ => (
6193 self.func("moe_gate_up_silu8_dev_q8"),
6194 LaunchConfig {
6195 grid_dim: (n_ff as u32, n_used as u32, 1),
6196 block_dim: (32, 1, 1),
6197 shared_mem_bytes: 0,
6198 },
6199 ),
6200 };
6201 let __s_b = self.gpu.stream();
6202 let mut b = __s_b.launch_builder(&f);
6203 b.arg(table)
6204 .arg(sel)
6205 .arg(aq)
6206 .arg(ad)
6207 .arg(&mut act)
6208 .arg(&inf)
6209 .arg(&nff)
6210 .arg(&ne)
6211 .arg(&qt_g)
6212 .arg(&qt_u)
6213 .arg(&rbg)
6214 .arg(&rbu)
6215 .arg(macros);
6216 unsafe {
6217 b.launch(cfg)?;
6218 }
6219 Ok(act)
6220 }
6221
6222 #[allow(clippy::too_many_arguments)]
6223 pub fn moe_down8_fma_dev_q8(
6224 &self,
6225 table: &CudaSlice<u64>,
6226 sel: &cudarc::driver::CudaView<i32>,
6227 w: &cudarc::driver::CudaView<f32>,
6228 aq2: &CudaSlice<i8>,
6229 ad2: &CudaSlice<f32>,
6230 dst: &mut cudarc::driver::CudaViewMut<f32>,
6231 in_f: usize,
6232 out_f: usize,
6233 n_used: usize,
6234 n_expert: usize,
6235 qt: i32,
6236 rb: usize,
6237 ) -> Result<(), Box<dyn std::error::Error>> {
6238 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6239 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6240 let (inf, outf, nu, ne, rbi) = (
6241 in_f as i32,
6242 out_f as i32,
6243 n_used as i32,
6244 n_expert as i32,
6245 rb as i64,
6246 );
6247 let (f, cfg) = match mode.as_str() {
6250 m @ ("1" | "2" | "4") if n_used <= 8 => {
6251 let rpw: usize = m.parse().unwrap();
6252 let f = self.func(match rpw {
6253 1 => "moe_down8_fma_dev_q8_w8r1",
6254 2 => "moe_down8_fma_dev_q8_w8r2",
6255 _ => "moe_down8_fma_dev_q8_w8r4",
6256 });
6257 (
6258 f,
6259 LaunchConfig {
6260 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6261 block_dim: (32, n_used as u32, 1),
6262 shared_mem_bytes: 0,
6263 },
6264 )
6265 }
6266 "h2" if in_f == 512 => (
6267 self.func("moe_down8_fma_dev_q8_h2"),
6268 LaunchConfig {
6269 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6270 block_dim: (32, 1, 1),
6271 shared_mem_bytes: 0,
6272 },
6273 ),
6274 "" if in_f == 704 && n_used <= 8 => (
6277 self.func("moe_down8_fma_dev_q8_w8r2"),
6278 LaunchConfig {
6279 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6280 block_dim: (32, n_used as u32, 1),
6281 shared_mem_bytes: 0,
6282 },
6283 ),
6284 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6288 self.func("moe_down8_fma_dev_q8_w8h2v"),
6289 LaunchConfig {
6290 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6291 block_dim: (32, n_used as u32, 1),
6292 shared_mem_bytes: 0,
6293 },
6294 ),
6295 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6296 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6297 LaunchConfig {
6298 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6299 block_dim: (32, n_used as u32, 1),
6300 shared_mem_bytes: 0,
6301 },
6302 ),
6303 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6304 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6305 LaunchConfig {
6306 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6307 block_dim: (32, n_used as u32, 1),
6308 shared_mem_bytes: 0,
6309 },
6310 ),
6311 "w8h2" if in_f == 512 && n_used <= 8 => (
6312 self.func("moe_down8_fma_dev_q8_w8h2"),
6313 LaunchConfig {
6314 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6315 block_dim: (32, n_used as u32, 1),
6316 shared_mem_bytes: 0,
6317 },
6318 ),
6319 _ => (
6320 self.func("moe_down8_fma_dev_q8"),
6321 LaunchConfig {
6322 grid_dim: (out_f as u32, 1, 1),
6323 block_dim: (32, 1, 1),
6324 shared_mem_bytes: 0,
6325 },
6326 ),
6327 };
6328 let __s_b = self.gpu.stream();
6329 let mut b = __s_b.launch_builder(&f);
6330 b.arg(table)
6331 .arg(sel)
6332 .arg(w)
6333 .arg(aq2)
6334 .arg(ad2)
6335 .arg(dst)
6336 .arg(&inf)
6337 .arg(&outf)
6338 .arg(&nu)
6339 .arg(&ne)
6340 .arg(&qt)
6341 .arg(&rbi);
6342 unsafe {
6343 b.launch(cfg)?;
6344 }
6345 Ok(())
6346 }
6347
6348 #[allow(clippy::too_many_arguments)]
6355 pub fn moe_gate_up_silu8_dev_q8_rows(
6356 &self,
6357 table: &CudaSlice<u64>,
6358 sel: &CudaSlice<i32>,
6359 aq: &CudaSlice<i8>,
6360 ad: &CudaSlice<f32>,
6361 t: usize,
6362 in_f: usize,
6363 n_ff: usize,
6364 n_used: usize,
6365 n_expert: usize,
6366 qt_g: i32,
6367 qt_u: i32,
6368 rb_g: usize,
6369 rb_u: usize,
6370 macros: &CudaSlice<f32>,
6371 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6372 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
6373 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
6374 let cfg = LaunchConfig {
6375 grid_dim: (n_ff as u32, n_used as u32, t as u32),
6376 block_dim: (32, 1, 1),
6377 shared_mem_bytes: 0,
6378 };
6379 let (inf, nff, ne, nu, rbg, rbu) = (
6380 in_f as i32,
6381 n_ff as i32,
6382 n_expert as i32,
6383 n_used as i32,
6384 rb_g as i64,
6385 rb_u as i64,
6386 );
6387 let __s_b = self.gpu.stream();
6388 let mut b = __s_b.launch_builder(&f);
6389 b.arg(table)
6390 .arg(sel)
6391 .arg(aq)
6392 .arg(ad)
6393 .arg(&mut act)
6394 .arg(&inf)
6395 .arg(&nff)
6396 .arg(&ne)
6397 .arg(&qt_g)
6398 .arg(&qt_u)
6399 .arg(&rbg)
6400 .arg(&rbu)
6401 .arg(&nu)
6402 .arg(macros);
6403 unsafe {
6404 b.launch(cfg)?;
6405 }
6406 Ok(act)
6407 }
6408
6409 #[allow(clippy::too_many_arguments)]
6414 pub fn moe_down8_fma_dev_q8_rows(
6415 &self,
6416 table: &CudaSlice<u64>,
6417 sel: &CudaSlice<i32>,
6418 w: &CudaSlice<f32>,
6419 aq2: &CudaSlice<i8>,
6420 ad2: &CudaSlice<f32>,
6421 dst: &mut CudaSlice<f32>,
6422 t: usize,
6423 in_f: usize,
6424 out_f: usize,
6425 n_used: usize,
6426 n_expert: usize,
6427 qt: i32,
6428 rb: usize,
6429 ) -> Result<(), Box<dyn std::error::Error>> {
6430 assert!(
6431 in_f == 512 && n_used <= 8,
6432 "down rows twin is w8h2v shape-gated"
6433 );
6434 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
6435 let cfg = LaunchConfig {
6436 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
6437 block_dim: (32, n_used as u32, 1),
6438 shared_mem_bytes: 0,
6439 };
6440 let (inf, outf, nu, ne, rbi) = (
6441 in_f as i32,
6442 out_f as i32,
6443 n_used as i32,
6444 n_expert as i32,
6445 rb as i64,
6446 );
6447 let __s_b = self.gpu.stream();
6448 let mut b = __s_b.launch_builder(&f);
6449 b.arg(table)
6450 .arg(sel)
6451 .arg(w)
6452 .arg(aq2)
6453 .arg(ad2)
6454 .arg(dst)
6455 .arg(&inf)
6456 .arg(&outf)
6457 .arg(&nu)
6458 .arg(&ne)
6459 .arg(&qt)
6460 .arg(&rbi);
6461 unsafe {
6462 b.launch(cfg)?;
6463 }
6464 Ok(())
6465 }
6466
6467 #[allow(clippy::too_many_arguments)]
6471 pub fn moe_gate_up_silu8_dev_q8_csr(
6472 &self,
6473 table: &CudaSlice<u64>,
6474 sel: &CudaSlice<i32>,
6475 aq: &CudaSlice<i8>,
6476 ad: &CudaSlice<f32>,
6477 n_pairs: usize,
6478 in_f: usize,
6479 n_ff: usize,
6480 n_used: usize,
6481 n_expert: usize,
6482 qt_g: i32,
6483 qt_u: i32,
6484 rb_g: usize,
6485 rb_u: usize,
6486 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6487 let f = if qt_g == crate::QT_NVFP4 {
6490 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
6491 } else {
6492 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
6493 };
6494 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
6495 let cfg = LaunchConfig {
6496 grid_dim: (n_ff as u32, n_pairs as u32, 1),
6497 block_dim: (32, 1, 1),
6498 shared_mem_bytes: 0,
6499 };
6500 let (inf, nff, ne, nu, npi, rbg, rbu) = (
6501 in_f as i32,
6502 n_ff as i32,
6503 n_expert as i32,
6504 n_used as i32,
6505 n_pairs as i32,
6506 rb_g as i64,
6507 rb_u as i64,
6508 );
6509 let __s_b = self.gpu.stream();
6510 let mut b = __s_b.launch_builder(&f);
6511 b.arg(table)
6512 .arg(sel)
6513 .arg(aq)
6514 .arg(ad)
6515 .arg(&mut act)
6516 .arg(&inf)
6517 .arg(&nff)
6518 .arg(&ne)
6519 .arg(&qt_g)
6520 .arg(&qt_u)
6521 .arg(&rbg)
6522 .arg(&rbu)
6523 .arg(&nu)
6524 .arg(&npi);
6525 unsafe {
6526 b.launch(cfg)?;
6527 }
6528 Ok(act)
6529 }
6530
6531 #[allow(clippy::too_many_arguments)]
6535 pub fn moe_down8_fma_dev_q8_variant(
6536 &self,
6537 variant: &str,
6538 table: &CudaSlice<u64>,
6539 sel: &cudarc::driver::CudaView<i32>,
6540 w: &cudarc::driver::CudaView<f32>,
6541 aq2: &CudaSlice<i8>,
6542 ad2: &CudaSlice<f32>,
6543 dst: &mut cudarc::driver::CudaViewMut<f32>,
6544 in_f: usize,
6545 out_f: usize,
6546 n_used: usize,
6547 n_expert: usize,
6548 qt: i32,
6549 rb: usize,
6550 ) -> Result<(), Box<dyn std::error::Error>> {
6551 let (inf, outf, nu, ne, rbi) = (
6552 in_f as i32,
6553 out_f as i32,
6554 n_used as i32,
6555 n_expert as i32,
6556 rb as i64,
6557 );
6558 let (f, cfg) = match variant {
6559 "w8h2" | "w8h2v" => (
6560 self.func(if variant == "w8h2" {
6561 "moe_down8_fma_dev_q8_w8h2"
6562 } else {
6563 "moe_down8_fma_dev_q8_w8h2v"
6564 }),
6565 LaunchConfig {
6566 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6567 block_dim: (32, n_used as u32, 1),
6568 shared_mem_bytes: 0,
6569 },
6570 ),
6571 "w8h2r2" | "w8h2r2v" => (
6572 self.func(if variant == "w8h2r2" {
6573 "moe_down8_fma_dev_q8_w8h2r2"
6574 } else {
6575 "moe_down8_fma_dev_q8_w8h2r2v"
6576 }),
6577 LaunchConfig {
6578 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6579 block_dim: (32, n_used as u32, 1),
6580 shared_mem_bytes: 0,
6581 },
6582 ),
6583 _ => (
6584 self.func("moe_down8_fma_dev_q8"),
6585 LaunchConfig {
6586 grid_dim: (out_f as u32, 1, 1),
6587 block_dim: (32, 1, 1),
6588 shared_mem_bytes: 0,
6589 },
6590 ),
6591 };
6592 let __s_b = self.gpu.stream();
6593 let mut b = __s_b.launch_builder(&f);
6594 b.arg(table)
6595 .arg(sel)
6596 .arg(w)
6597 .arg(aq2)
6598 .arg(ad2)
6599 .arg(dst)
6600 .arg(&inf)
6601 .arg(&outf)
6602 .arg(&nu)
6603 .arg(&ne)
6604 .arg(&qt)
6605 .arg(&rbi);
6606 unsafe {
6607 b.launch(cfg)?;
6608 }
6609 Ok(())
6610 }
6611
6612 #[allow(clippy::too_many_arguments)]
6614 pub fn moe_gate_up_silu8_dev_q8_variant(
6615 &self,
6616 variant: &str,
6617 table: &CudaSlice<u64>,
6618 sel: &cudarc::driver::CudaView<i32>,
6619 aq: &CudaSlice<i8>,
6620 ad: &CudaSlice<f32>,
6621 in_f: usize,
6622 n_ff: usize,
6623 n_used: usize,
6624 n_expert: usize,
6625 qt_g: i32,
6626 qt_u: i32,
6627 rb_g: usize,
6628 rb_u: usize,
6629 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6630 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6631 let (inf, nff, ne, rbg, rbu) = (
6632 in_f as i32,
6633 n_ff as i32,
6634 n_expert as i32,
6635 rb_g as i64,
6636 rb_u as i64,
6637 );
6638 let f = self.func(if variant == "v" {
6639 "moe_gate_up_silu8_dev_q8_v"
6640 } else {
6641 "moe_gate_up_silu8_dev_q8"
6642 });
6643 let cfg = LaunchConfig {
6644 grid_dim: (n_ff as u32, n_used as u32, 1),
6645 block_dim: (32, 1, 1),
6646 shared_mem_bytes: 0,
6647 };
6648 let __s_b = self.gpu.stream();
6649 let mut b = __s_b.launch_builder(&f);
6650 b.arg(table)
6651 .arg(sel)
6652 .arg(aq)
6653 .arg(ad)
6654 .arg(&mut act)
6655 .arg(&inf)
6656 .arg(&nff)
6657 .arg(&ne)
6658 .arg(&qt_g)
6659 .arg(&qt_u)
6660 .arg(&rbg)
6661 .arg(&rbu);
6662 unsafe {
6663 b.launch(cfg)?;
6664 }
6665 Ok(act)
6666 }
6667
6668 pub fn moe_gate_up_silu8_dev(
6669 &self,
6670 table: &CudaSlice<u64>,
6671 sel: &cudarc::driver::CudaView<i32>,
6672 x: &cudarc::driver::CudaView<f32>,
6673 in_f: usize,
6674 n_ff: usize,
6675 n_used: usize,
6676 n_expert: usize,
6677 qt_g: i32,
6678 qt_u: i32,
6679 rb_g: usize,
6680 rb_u: usize,
6681 macros: &CudaSlice<f32>,
6682 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6683 let f = self.func("moe_gate_up_silu8_dev");
6684 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
6686 grid_dim: (n_ff as u32, n_used as u32, 1),
6687 block_dim: (256, 1, 1),
6688 shared_mem_bytes: 0,
6689 };
6690 let (inf, nff, ne, rbg, rbu) = (
6691 in_f as i32,
6692 n_ff as i32,
6693 n_expert as i32,
6694 rb_g as i64,
6695 rb_u as i64,
6696 );
6697 let __s_b = self.gpu.stream();
6698 let mut b = __s_b.launch_builder(&f);
6699 b.arg(table)
6700 .arg(sel)
6701 .arg(x)
6702 .arg(&mut act)
6703 .arg(&inf)
6704 .arg(&nff)
6705 .arg(&ne)
6706 .arg(&qt_g)
6707 .arg(&qt_u)
6708 .arg(&rbg)
6709 .arg(&rbu)
6710 .arg(macros);
6711 unsafe {
6712 b.launch(cfg)?;
6713 }
6714 Ok(act)
6715 }
6716
6717 #[allow(clippy::too_many_arguments)]
6720 pub fn moe_down8_fma_dev(
6721 &self,
6722 table: &CudaSlice<u64>,
6723 sel: &cudarc::driver::CudaView<i32>,
6724 w: &cudarc::driver::CudaView<f32>,
6725 act: &CudaSlice<f32>,
6726 dst: &mut cudarc::driver::CudaViewMut<f32>,
6727 in_f: usize,
6728 out_f: usize,
6729 n_used: usize,
6730 n_expert: usize,
6731 qt: i32,
6732 rb: usize,
6733 ) -> Result<(), Box<dyn std::error::Error>> {
6734 let f = self.func("moe_down8_fma_dev");
6735 let cfg = LaunchConfig {
6736 grid_dim: (out_f as u32, 1, 1),
6737 block_dim: (256, 1, 1),
6738 shared_mem_bytes: 0,
6739 };
6740 let (inf, outf, nu, ne, rbv) = (
6741 in_f as i32,
6742 out_f as i32,
6743 n_used as i32,
6744 n_expert as i32,
6745 rb as i64,
6746 );
6747 let __s_b = self.gpu.stream();
6748 let mut b = __s_b.launch_builder(&f);
6749 b.arg(table)
6750 .arg(sel)
6751 .arg(w)
6752 .arg(act)
6753 .arg(dst)
6754 .arg(&inf)
6755 .arg(&outf)
6756 .arg(&nu)
6757 .arg(&ne)
6758 .arg(&qt)
6759 .arg(&rbv);
6760 unsafe {
6761 b.launch(cfg)?;
6762 }
6763 Ok(())
6764 }
6765
6766 pub fn axpy_into(
6768 &self,
6769 src: &CudaSlice<f32>,
6770 alpha: f32,
6771 dst: &mut cudarc::driver::CudaViewMut<f32>,
6772 n: usize,
6773 ) -> Result<(), Box<dyn std::error::Error>> {
6774 let f = self.func("axpy_f32");
6775 let cfg = LaunchConfig::for_num_elems(n as u32);
6776 let (a, ni) = (alpha, n as i32);
6777 let __s_b = self.gpu.stream();
6778 let mut b = __s_b.launch_builder(&f);
6779 b.arg(src).arg(dst).arg(&a).arg(&ni);
6780 unsafe {
6781 b.launch(cfg)?;
6782 }
6783 Ok(())
6784 }
6785
6786 pub fn axpy_host_into(
6788 &self,
6789 src: &cudarc::driver::CudaView<'_, f32>,
6790 alpha: f32,
6791 dst: &mut cudarc::driver::CudaViewMut<f32>,
6792 n: usize,
6793 ) -> Result<(), Box<dyn std::error::Error>> {
6794 let f = self.func("axpy_host_f32");
6795 let cfg = LaunchConfig::for_num_elems(n as u32);
6796 let (a, ni) = (alpha, n as i32);
6797 let __s_b = self.gpu.stream();
6798 let mut b = __s_b.launch_builder(&f);
6799 b.arg(src).arg(dst).arg(&a).arg(&ni);
6800 unsafe {
6801 b.launch(cfg)?;
6802 }
6803 Ok(())
6804 }
6805
6806 pub fn add_scaled_rows(
6808 &self,
6809 src: &CudaSlice<f32>,
6810 scale: &CudaSlice<f32>,
6811 dst: &mut CudaSlice<f32>,
6812 ncols: usize,
6813 nrows: usize,
6814 ) -> Result<(), Box<dyn std::error::Error>> {
6815 let f = self.func("add_scaled_rows_f32");
6816 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
6817 let (nc, nr) = (ncols as i32, nrows as i32);
6818 let __s_b = self.gpu.stream();
6819 let mut b = __s_b.launch_builder(&f);
6820 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
6821 unsafe {
6822 b.launch(cfg)?;
6823 }
6824 Ok(())
6825 }
6826
6827 pub fn gather_rows(
6831 &self,
6832 src: &CudaSlice<f32>,
6833 idx: &CudaSlice<i32>,
6834 dst: &mut CudaSlice<f32>,
6835 ncols: usize,
6836 m_e: usize,
6837 ) -> Result<(), Box<dyn std::error::Error>> {
6838 let f = self.func("gather_rows_f32");
6839 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6840 let (nc, me) = (ncols as i32, m_e as i32);
6841 let __s_b = self.gpu.stream();
6842 let mut b = __s_b.launch_builder(&f);
6843 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
6844 unsafe {
6845 b.launch(cfg)?;
6846 }
6847 Ok(())
6848 }
6849
6850 pub fn scatter_slot(
6855 &self,
6856 src: &CudaSlice<f32>,
6857 tok_idx: &CudaSlice<i32>,
6858 slot_idx: &CudaSlice<i32>,
6859 weight: &CudaSlice<f32>,
6860 dst: &mut CudaSlice<f32>,
6861 wbuf: &mut CudaSlice<f32>,
6862 ncols: usize,
6863 n_used: usize,
6864 m_e: usize,
6865 ) -> Result<(), Box<dyn std::error::Error>> {
6866 let f = self.func("scatter_add_slot_f32");
6867 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
6868 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
6869 let __s_b = self.gpu.stream();
6870 let mut b = __s_b.launch_builder(&f);
6871 b.arg(src)
6872 .arg(tok_idx)
6873 .arg(slot_idx)
6874 .arg(weight)
6875 .arg(dst)
6876 .arg(wbuf)
6877 .arg(&nc)
6878 .arg(&nu)
6879 .arg(&me);
6880 unsafe {
6881 b.launch(cfg)?;
6882 }
6883 Ok(())
6884 }
6885
6886 pub fn reduce_slots(
6890 &self,
6891 slots: &CudaSlice<f32>,
6892 wbuf: &CudaSlice<f32>,
6893 dst: &mut CudaSlice<f32>,
6894 ncols: usize,
6895 n_used: usize,
6896 t: usize,
6897 ) -> Result<(), Box<dyn std::error::Error>> {
6898 let f = self.func("reduce_slots_f32");
6899 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6900 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6901 let __s_b = self.gpu.stream();
6902 let mut b = __s_b.launch_builder(&f);
6903 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6904 unsafe {
6905 b.launch(cfg)?;
6906 }
6907 Ok(())
6908 }
6909
6910 pub fn reduce_slots_host(
6915 &self,
6916 slots: &CudaSlice<f32>,
6917 wbuf: &CudaSlice<f32>,
6918 dst: &mut CudaSlice<f32>,
6919 ncols: usize,
6920 n_used: usize,
6921 t: usize,
6922 ) -> Result<(), Box<dyn std::error::Error>> {
6923 let f = self.func("reduce_slots_host_f32");
6924 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
6925 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
6926 let __s_b = self.gpu.stream();
6927 let mut b = __s_b.launch_builder(&f);
6928 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
6929 unsafe {
6930 b.launch(cfg)?;
6931 }
6932 Ok(())
6933 }
6934
6935 pub fn quantize_q8_1_view(
6942 &self,
6943 x: &cudarc::driver::CudaView<f32>,
6944 m: usize,
6945 in_f: usize,
6946 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6947 let f = self.func("quantize_q8_1");
6948 let nblk = in_f / 32;
6949 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
6950 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
6951 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
6952 let (inf, mi) = (in_f as i32, m as i32);
6953 let __s_b = self.gpu.stream();
6954 let mut b = __s_b.launch_builder(&f);
6955 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6956 unsafe {
6957 b.launch(cfg)?;
6958 }
6959 Ok((q, d))
6960 }
6961
6962 pub fn quantize_q8_1(
6963 &self,
6964 x: &CudaSlice<f32>,
6965 m: usize,
6966 in_f: usize,
6967 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6968 let nblk = in_f / 32;
6969 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);
6973 let (inf, mi) = (in_f as i32, m as i32);
6974 if Self::pdl_on() && Self::pdl_wb_on() {
6975 {
6976 use cudarc::driver::{DevicePtr, DevicePtrMut};
6977 let s = &self.gpu.stream();
6978 let (px, _g0) = x.device_ptr(s);
6979 let (pq, _g1) = q.device_ptr_mut(s);
6980 let (pd, _g2) = d.device_ptr_mut(s);
6981 let mut ps = [
6982 &px as *const _ as *mut std::ffi::c_void,
6983 &pq as *const _ as *mut _,
6984 &pd as *const _ as *mut _,
6985 &inf as *const _ as *mut _,
6986 &mi as *const _ as *mut _,
6987 ];
6988 unsafe {
6989 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
6990 }
6991 }
6992 return Ok((q, d));
6993 }
6994 let f = self.func("quantize_q8_1");
6995 let __s_b = self.gpu.stream();
6996 let mut b = __s_b.launch_builder(&f);
6997 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
6998 unsafe {
6999 b.launch(cfg)?;
7000 }
7001 Ok((q, d))
7002 }
7003
7004 pub fn quantize_fp4_act(
7008 &self,
7009 x: &CudaSlice<f32>,
7010 m: usize,
7011 in_f: usize,
7012 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7013 let f = self.func("quantize_fp4_act");
7014 let nb16 = in_f / 16;
7015 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);
7018 let (inf, mi) = (in_f as i32, m as i32);
7019 let __s_b = self.gpu.stream();
7020 let mut b = __s_b.launch_builder(&f);
7021 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7022 unsafe {
7023 b.launch(cfg)?;
7024 }
7025 Ok((aq4, ad4))
7026 }
7027
7028 pub fn qmatvec_gemm_nvfp4_fp4(
7033 &self,
7034 bytes: &CudaSlice<u8>,
7035 x: &CudaSlice<f32>,
7036 m: usize,
7037 in_f: usize,
7038 out_f: usize,
7039 row_bytes: usize,
7040 scale: f32,
7041 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7042 assert!(
7043 in_f % 64 == 0,
7044 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7045 );
7046 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7047 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7048 if scale != 1.0 {
7049 self.scale_inplace(&mut y, scale, m * out_f)?;
7050 }
7051 Ok(y)
7052 }
7053
7054 fn fp4_gemm_launch(
7057 &self,
7058 bytes: &CudaSlice<u8>,
7059 aq4: &CudaSlice<u32>,
7060 ad4: &CudaSlice<u8>,
7061 m: usize,
7062 in_f: usize,
7063 out_f: usize,
7064 row_bytes: usize,
7065 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7066 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7067 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7069 const BN: u32 = 256;
7070 let cfg = LaunchConfig {
7071 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7072 block_dim: (32, 4, 1),
7073 shared_mem_bytes: 0,
7074 };
7075 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7076 let __s_b = self.gpu.stream();
7077 let mut b = __s_b.launch_builder(&f);
7078 b.arg(bytes)
7079 .arg(aq4)
7080 .arg(ad4)
7081 .arg(&mut y)
7082 .arg(&inf)
7083 .arg(&outf)
7084 .arg(&mi)
7085 .arg(&rb);
7086 unsafe {
7087 b.launch(cfg)?;
7088 }
7089 Ok(y)
7090 }
7091
7092 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7094 &self,
7095 bytes: &CudaSlice<u8>,
7096 x: &CudaSlice<f32>,
7097 m: usize,
7098 in_f: usize,
7099 out_f: usize,
7100 row_bytes: usize,
7101 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7102 assert!(
7103 in_f % 64 == 0,
7104 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7105 );
7106 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7107 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7108 }
7109
7110 pub fn qmatvec_q8_0_fast(
7112 &self,
7113 w: &CudaSlice<u8>,
7114 x: &CudaSlice<f32>,
7115 m: usize,
7116 in_f: usize,
7117 out_f: usize,
7118 row_bytes: usize,
7119 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7120 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7121 let f = self.func("qmatvec_q8_0_dp4a");
7122 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7124 grid_dim: (out_f as u32, m as u32, 1),
7125 block_dim: (128, 1, 1),
7126 shared_mem_bytes: 0,
7127 };
7128 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7129 let __s_b = self.gpu.stream();
7130 let mut b = __s_b.launch_builder(&f);
7131 b.arg(w)
7132 .arg(&aq)
7133 .arg(&ad)
7134 .arg(&mut y)
7135 .arg(&inf)
7136 .arg(&outf)
7137 .arg(&mi)
7138 .arg(&rb);
7139 unsafe {
7140 b.launch(cfg)?;
7141 }
7142 Ok(y)
7143 }
7144
7145 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7148 &self,
7149 w: &CudaSlice<u8>,
7150 x: &CudaSlice<f32>,
7151 m: usize,
7152 in_f: usize,
7153 out_f: usize,
7154 row_bytes: usize,
7155 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7156 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7157 let f = self.func("qmatvec_q4_K_dp4a");
7158 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7160 grid_dim: (out_f as u32, m as u32, 1),
7161 block_dim: (128, 1, 1),
7162 shared_mem_bytes: 0,
7163 };
7164 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7165 let __s_b = self.gpu.stream();
7166 let mut b = __s_b.launch_builder(&f);
7167 b.arg(w)
7168 .arg(&aq)
7169 .arg(&ad)
7170 .arg(&mut y)
7171 .arg(&inf)
7172 .arg(&outf)
7173 .arg(&mi)
7174 .arg(&rb);
7175 unsafe {
7176 b.launch(cfg)?;
7177 }
7178 Ok(y)
7179 }
7180
7181 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7184 &self,
7185 w: &CudaSlice<u8>,
7186 x: &CudaSlice<f32>,
7187 m: usize,
7188 in_f: usize,
7189 out_f: usize,
7190 row_bytes: usize,
7191 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7192 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7193 let f = self.func("qmatvec_q6_K_dp4a");
7194 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7196 grid_dim: (out_f as u32, m as u32, 1),
7197 block_dim: (128, 1, 1),
7198 shared_mem_bytes: 0,
7199 };
7200 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7201 let __s_b = self.gpu.stream();
7202 let mut b = __s_b.launch_builder(&f);
7203 b.arg(w)
7204 .arg(&aq)
7205 .arg(&ad)
7206 .arg(&mut y)
7207 .arg(&inf)
7208 .arg(&outf)
7209 .arg(&mi)
7210 .arg(&rb);
7211 unsafe {
7212 b.launch(cfg)?;
7213 }
7214 Ok(y)
7215 }
7216
7217 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7220 &self,
7221 w: &CudaSlice<u8>,
7222 x: &CudaSlice<f32>,
7223 m: usize,
7224 in_f: usize,
7225 out_f: usize,
7226 row_bytes: usize,
7227 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7228 self.qmatvec_dp4a_named(
7229 "qmatvec_q5_K_dp4a",
7230 &w.slice(0..w.len()),
7231 x,
7232 m,
7233 in_f,
7234 out_f,
7235 row_bytes,
7236 )
7237 }
7238 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7241 &self,
7242 w: &CudaSlice<u8>,
7243 x: &CudaSlice<f32>,
7244 m: usize,
7245 in_f: usize,
7246 out_f: usize,
7247 row_bytes: usize,
7248 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7249 self.qmatvec_dp4a_named(
7250 "qmatvec_q3_K_dp4a",
7251 &w.slice(0..w.len()),
7252 x,
7253 m,
7254 in_f,
7255 out_f,
7256 row_bytes,
7257 )
7258 }
7259 pub fn qmatvec_nvfp4_fast_rp(
7261 &self,
7262 w: &CudaSlice<u8>,
7263 x: &CudaSlice<f32>,
7264 m: usize,
7265 in_f: usize,
7266 out_f: usize,
7267 row_bytes: usize,
7268 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7269 assert!(
7270 in_f % 64 == 0,
7271 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7272 );
7273 self.qmatvec_dp4a_named(
7274 "qmatvec_nvfp4_dp4a_rp",
7275 &w.slice(0..w.len()),
7276 x,
7277 m,
7278 in_f,
7279 out_f,
7280 row_bytes,
7281 )
7282 }
7283 pub fn qmatvec_nvfp4_fast(
7285 &self,
7286 w: &cudarc::driver::CudaView<'_, u8>,
7287 x: &CudaSlice<f32>,
7288 m: usize,
7289 in_f: usize,
7290 out_f: usize,
7291 row_bytes: usize,
7292 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7293 assert!(
7296 in_f % 64 == 0,
7297 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7298 );
7299 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
7300 }
7301 pub fn qmatvec_nvfp4_fast_v2(
7304 &self,
7305 w: &cudarc::driver::CudaView<'_, u8>,
7306 x: &CudaSlice<f32>,
7307 m: usize,
7308 in_f: usize,
7309 out_f: usize,
7310 row_bytes: usize,
7311 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7312 assert!(
7313 in_f % 64 == 0,
7314 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7315 );
7316 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
7317 }
7318 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
7321 &self,
7322 w: &CudaSlice<u8>,
7323 x: &CudaSlice<f32>,
7324 m: usize,
7325 in_f: usize,
7326 out_f: usize,
7327 row_bytes: usize,
7328 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7329 self.qmatvec_dp4a_named(
7330 "qmatvec_iq4_XS_dp4a",
7331 &w.slice(0..w.len()),
7332 x,
7333 m,
7334 in_f,
7335 out_f,
7336 row_bytes,
7337 )
7338 }
7339
7340 fn qmatvec_dp4a_named(
7342 &self,
7343 name: &str,
7344 w: &cudarc::driver::CudaView<'_, u8>,
7345 x: &CudaSlice<f32>,
7346 m: usize,
7347 in_f: usize,
7348 out_f: usize,
7349 row_bytes: usize,
7350 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7351 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7352 let f = self.func(name);
7353 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7355 grid_dim: (out_f as u32, m as u32, 1),
7356 block_dim: (128, 1, 1),
7357 shared_mem_bytes: 0,
7358 };
7359 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7360 let __s_b = self.gpu.stream();
7361 let mut b = __s_b.launch_builder(&f);
7362 b.arg(w)
7363 .arg(&aq)
7364 .arg(&ad)
7365 .arg(&mut y)
7366 .arg(&inf)
7367 .arg(&outf)
7368 .arg(&mi)
7369 .arg(&rb);
7370 unsafe {
7371 b.launch(cfg)?;
7372 }
7373 Ok(y)
7374 }
7375
7376 #[allow(clippy::too_many_arguments)]
7382 pub fn qmatvec_nvfp4_fast_prequant_into(
7383 &self,
7384 w: &CudaSlice<u8>,
7385 aq: &CudaSlice<i8>,
7386 ad: &CudaSlice<f32>,
7387 y: &mut CudaSlice<f32>,
7388 m: usize,
7389 in_f: usize,
7390 out_f: usize,
7391 row_bytes: usize,
7392 ) -> Result<(), Box<dyn std::error::Error>> {
7393 assert!(
7394 in_f % 64 == 0,
7395 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7396 );
7397 if y.len() < m * out_f {
7398 return Err(format!(
7399 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
7400 y.len()
7401 )
7402 .into());
7403 }
7404 let f = self.func("qmatvec_nvfp4_dp4a");
7405 let cfg = LaunchConfig {
7406 grid_dim: (out_f as u32, m as u32, 1),
7407 block_dim: (128, 1, 1),
7408 shared_mem_bytes: 0,
7409 };
7410 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7411 let __s_b = self.gpu.stream();
7412 let mut b = __s_b.launch_builder(&f);
7413 b.arg(w)
7414 .arg(aq)
7415 .arg(ad)
7416 .arg(y)
7417 .arg(&inf)
7418 .arg(&outf)
7419 .arg(&mi)
7420 .arg(&rb);
7421 unsafe {
7422 b.launch(cfg)?;
7423 }
7424 Ok(())
7425 }
7426
7427 #[allow(clippy::too_many_arguments)]
7430 pub fn matvec_f32_qkv_into(
7431 &self,
7432 wq: &CudaSlice<f32>,
7433 wk: &CudaSlice<f32>,
7434 wv: &CudaSlice<f32>,
7435 wg: &CudaSlice<f32>,
7436 x: &CudaSlice<f32>,
7437 yq: &mut CudaSlice<f32>,
7438 yk: &mut CudaSlice<f32>,
7439 yv: &mut CudaSlice<f32>,
7440 yg: &mut CudaSlice<f32>,
7441 in_f: usize,
7442 out_q: usize,
7443 out_kv: usize,
7444 out_g: usize,
7445 ) -> Result<(), Box<dyn std::error::Error>> {
7446 if in_f % 4 != 0
7447 || wq.len() != out_q * in_f
7448 || wk.len() != out_kv * in_f
7449 || wv.len() != out_kv * in_f
7450 || wg.len() < out_g * in_f
7451 || x.len() < in_f
7452 || yq.len() < out_q
7453 || yk.len() < out_kv
7454 || yv.len() < out_kv
7455 || (out_g > 0 && yg.len() < out_g)
7456 {
7457 return Err(format!(
7458 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
7459 wq={} wk={} wv={} wg={}",
7460 wq.len(),
7461 wk.len(),
7462 wv.len(),
7463 wg.len()
7464 )
7465 .into());
7466 }
7467 let f = self.func("matvec_f32_qkv");
7468 let cfg = LaunchConfig {
7469 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
7470 block_dim: (128, 1, 1),
7471 shared_mem_bytes: 0,
7472 };
7473 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
7474 let __s_b = self.gpu.stream();
7475 let mut b = __s_b.launch_builder(&f);
7476 b.arg(wq)
7477 .arg(wk)
7478 .arg(wv)
7479 .arg(wg)
7480 .arg(x)
7481 .arg(yq)
7482 .arg(yk)
7483 .arg(yv)
7484 .arg(yg)
7485 .arg(&inf)
7486 .arg(&oq)
7487 .arg(&okv)
7488 .arg(&og);
7489 unsafe {
7490 b.launch(cfg)?;
7491 }
7492 Ok(())
7493 }
7494
7495 #[allow(clippy::too_many_arguments)]
7501 #[allow(clippy::too_many_arguments)]
7504 pub fn qmatvec_nvfp4_sel_gu_into(
7505 &self,
7506 gate_bank: &CudaSlice<u8>,
7507 up_bank: &CudaSlice<u8>,
7508 sel: &CudaSlice<i32>,
7509 aq: &CudaSlice<i8>,
7510 ad: &CudaSlice<f32>,
7511 yg: &mut CudaSlice<f32>,
7512 yu: &mut CudaSlice<f32>,
7513 n_sel: usize,
7514 in_f: usize,
7515 out_f: usize,
7516 row_bytes: usize,
7517 expert_stride: usize,
7518 ) -> Result<(), Box<dyn std::error::Error>> {
7519 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
7520 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
7521 return Err("NVFP4 gu sel geometry".into());
7522 }
7523 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu");
7524 let cfg = LaunchConfig {
7525 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
7526 block_dim: (128, 1, 1),
7527 shared_mem_bytes: 0,
7528 };
7529 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7530 let (rb, es) = (row_bytes as i64, expert_stride as i64);
7531 let (ars, adrs) = (0i64, 0i64);
7532 let __s_b = self.gpu.stream();
7533 let mut b = __s_b.launch_builder(&f);
7534 b.arg(gate_bank)
7535 .arg(up_bank)
7536 .arg(sel)
7537 .arg(aq)
7538 .arg(ad)
7539 .arg(yg)
7540 .arg(yu)
7541 .arg(&inf)
7542 .arg(&outf)
7543 .arg(&ns)
7544 .arg(&rb)
7545 .arg(&es)
7546 .arg(&ars)
7547 .arg(&adrs);
7548 unsafe {
7549 b.launch(cfg)?;
7550 }
7551 Ok(())
7552 }
7553
7554 pub fn qmatvec_nvfp4_sel_into(
7555 &self,
7556 bank: &CudaSlice<u8>,
7557 sel: &CudaSlice<i32>,
7558 aq: &CudaSlice<i8>,
7559 ad: &CudaSlice<f32>,
7560 y: &mut CudaSlice<f32>,
7561 n_sel: usize,
7562 in_f: usize,
7563 out_f: usize,
7564 row_bytes: usize,
7565 expert_stride: usize,
7566 act_row_stride: usize,
7567 ad_row_stride: usize,
7568 ) -> Result<(), Box<dyn std::error::Error>> {
7569 assert!(
7570 in_f % 64 == 0,
7571 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
7572 );
7573 if y.len() < n_sel * out_f || sel.len() < n_sel {
7574 return Err(format!(
7575 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
7576 y.len(),
7577 sel.len()
7578 )
7579 .into());
7580 }
7581 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
7588 let mode = *MR.get_or_init(|| {
7589 if crate::tp::nvfp4_bank_v2_on() {
7590 3
7591 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
7592 2
7593 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
7594 1
7595 } else {
7596 0
7597 }
7598 });
7599 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
7600 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7604 let v2s = mode == 3
7605 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
7606 && row_bytes % 16 == 0
7607 && in_f <= 4096;
7608 let f = match (mode, v2s) {
7609 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
7610 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
7611 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
7612 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
7613 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
7614 };
7615 let nsb = in_f >> 5;
7620 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
7621 32
7622 } else if mode == 1 {
7623 512
7624 } else {
7625 128
7626 };
7627 let cfg = LaunchConfig {
7628 grid_dim: (
7629 if v2s {
7630 (out_f as u32).div_ceil(8)
7631 } else {
7632 match mode {
7633 2 => (out_f as u32).div_ceil(16),
7634 1 => (out_f as u32).div_ceil(4),
7635 _ => out_f as u32,
7636 }
7637 },
7638 n_sel as u32,
7639 1,
7640 ),
7641 block_dim: (fit_block, 1, 1),
7642 shared_mem_bytes: 0,
7643 };
7644 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
7645 let (rb, es, ars, adrs) = (
7646 row_bytes as i64,
7647 expert_stride as i64,
7648 act_row_stride as i64,
7649 ad_row_stride as i64,
7650 );
7651 let __s_b = self.gpu.stream();
7652 let mut b = __s_b.launch_builder(&f);
7653 b.arg(bank)
7654 .arg(sel)
7655 .arg(aq)
7656 .arg(ad)
7657 .arg(y)
7658 .arg(&inf)
7659 .arg(&outf)
7660 .arg(&ns)
7661 .arg(&rb)
7662 .arg(&es)
7663 .arg(&ars)
7664 .arg(&adrs);
7665 unsafe {
7666 b.launch(cfg)?;
7667 }
7668 Ok(())
7669 }
7670
7671 #[allow(clippy::too_many_arguments)]
7676 pub fn silu_mul_scaled_q8_1_sel_into(
7677 &self,
7678 gate: &CudaSlice<f32>,
7679 up: &CudaSlice<f32>,
7680 gmac: &CudaSlice<f32>,
7681 umac: &CudaSlice<f32>,
7682 sel: &CudaSlice<i32>,
7683 limit: Option<f32>,
7684 out_q: &mut CudaSlice<i8>,
7685 out_d: &mut CudaSlice<f32>,
7686 n_per: usize,
7687 n_sel: usize,
7688 ) -> Result<(), Box<dyn std::error::Error>> {
7689 let n = n_per * n_sel;
7690 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
7691 return Err(format!(
7692 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
7693 out_q.len(),
7694 out_d.len()
7695 )
7696 .into());
7697 }
7698 if let Some(limit) = limit {
7699 if limit <= 1e-6 {
7700 return Err(format!(
7701 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
7702 )
7703 .into());
7704 }
7705 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
7706 let cfg = LaunchConfig::for_num_elems(n as u32);
7707 let (np, ns) = (n_per as i32, n_sel as i32);
7708 let __s_b = self.gpu.stream();
7709 let mut b = __s_b.launch_builder(&f);
7710 b.arg(gate)
7711 .arg(up)
7712 .arg(gmac)
7713 .arg(umac)
7714 .arg(sel)
7715 .arg(&limit)
7716 .arg(out_q)
7717 .arg(out_d)
7718 .arg(&np)
7719 .arg(&ns);
7720 unsafe {
7721 b.launch(cfg)?;
7722 }
7723 return Ok(());
7724 }
7725 let f = self.func("silu_mul_scaled_q8_1_sel");
7726 let cfg = LaunchConfig::for_num_elems(n as u32);
7727 let (np, ns) = (n_per as i32, n_sel as i32);
7728 let __s_b = self.gpu.stream();
7729 let mut b = __s_b.launch_builder(&f);
7730 b.arg(gate)
7731 .arg(up)
7732 .arg(gmac)
7733 .arg(umac)
7734 .arg(sel)
7735 .arg(out_q)
7736 .arg(out_d)
7737 .arg(&np)
7738 .arg(&ns);
7739 unsafe {
7740 b.launch(cfg)?;
7741 }
7742 Ok(())
7743 }
7744
7745 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7746 Ok(self.gpu.stream().clone_htod(v)?)
7747 }
7748 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
7749 Ok(self.gpu.stream().clone_htod(v)?)
7750 }
7751 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
7753 Ok(self.gpu.stream().clone_htod(v)?)
7754 }
7755 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
7756 Ok(self.gpu.stream().clone_htod(v)?)
7757 }
7758 pub fn dtoh_view(
7760 &self,
7761 d: &cudarc::driver::CudaView<f32>,
7762 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7763 let v = self.gpu.stream().clone_dtoh(d)?;
7764 self.gpu.stream().synchronize()?;
7765 Ok(v)
7766 }
7767 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
7768 let v = self.gpu.stream().clone_dtoh(d)?;
7769 self.gpu.stream().synchronize()?;
7770 Ok(v)
7771 }
7772 pub fn dtoh_pair(
7776 &self,
7777 a: &CudaSlice<f32>,
7778 b: &CudaSlice<f32>,
7779 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7780 let av = self.gpu.stream().clone_dtoh(a)?;
7781 let bv = self.gpu.stream().clone_dtoh(b)?;
7782 self.gpu.stream().synchronize()?;
7783 Ok((av, bv))
7784 }
7785 pub fn dtoh_pair_views(
7788 &self,
7789 a: &cudarc::driver::CudaView<f32>,
7790 b: &cudarc::driver::CudaView<f32>,
7791 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
7792 let av = self.gpu.stream().clone_dtoh(a)?;
7793 let bv = self.gpu.stream().clone_dtoh(b)?;
7794 self.gpu.stream().synchronize()?;
7795 Ok((av, bv))
7796 }
7797 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
7799 let v = self.gpu.stream().clone_dtoh(d)?;
7800 self.gpu.stream().synchronize()?;
7801 Ok(v)
7802 }
7803 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7805 let v = self.gpu.stream().clone_dtoh(d)?;
7806 self.gpu.stream().synchronize()?;
7807 Ok(v)
7808 }
7809 pub fn dtoh_u8_view(
7810 &self,
7811 d: &cudarc::driver::CudaView<u8>,
7812 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
7813 let v = self.gpu.stream().clone_dtoh(d)?;
7814 self.gpu.stream().synchronize()?;
7815 Ok(v)
7816 }
7817 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7818 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
7819 self.keep_if_capturing(&s);
7820 Ok(s)
7821 }
7822
7823 pub fn prob_of_token_device(
7832 &self,
7833 logits: &CudaSlice<f32>,
7834 tok: &CudaSlice<u32>,
7835 n_vocab: usize,
7836 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7837 let nb = ARGMAX_NB;
7838 let mut part = self.alloc_uninit::<f32>(nb)?;
7839 let mut p = self.alloc_uninit::<f32>(1)?;
7840 let f1 = self.func("prob_of_token_partial_f32");
7841 let cfg1 = LaunchConfig {
7842 grid_dim: (nb as u32, 1, 1),
7843 block_dim: (256, 1, 1),
7844 shared_mem_bytes: 0,
7845 };
7846 let nv = n_vocab as i32;
7847 let __s_b1 = self.gpu.stream();
7848 let mut b1 = __s_b1.launch_builder(&f1);
7849 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7850 unsafe {
7851 b1.launch(cfg1)?;
7852 }
7853 let f2 = self.func("prob_of_token_final_f32");
7854 let cfg2 = LaunchConfig {
7855 grid_dim: (1, 1, 1),
7856 block_dim: (256, 1, 1),
7857 shared_mem_bytes: 0,
7858 };
7859 let nbi = nb as i32;
7860 let __s_b2 = self.gpu.stream();
7861 let mut b2 = __s_b2.launch_builder(&f2);
7862 b2.arg(&part).arg(&mut p).arg(&nbi);
7863 unsafe {
7864 b2.launch(cfg2)?;
7865 }
7866 Ok(p)
7867 }
7868
7869 pub fn prob_of_token_device_col(
7876 &self,
7877 logits: &CudaSlice<f32>,
7878 tok_all: &CudaSlice<u32>,
7879 tok_idx: usize,
7880 p_out: &mut CudaSlice<f32>,
7881 p_idx: usize,
7882 n_vocab: usize,
7883 ) -> Result<(), Box<dyn std::error::Error>> {
7884 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
7885 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
7886 let nb = ARGMAX_NB;
7887 let mut part = self.alloc_uninit::<f32>(nb)?;
7888 let f1 = self.func("prob_of_token_partial_f32");
7889 let cfg1 = LaunchConfig {
7890 grid_dim: (nb as u32, 1, 1),
7891 block_dim: (256, 1, 1),
7892 shared_mem_bytes: 0,
7893 };
7894 let nv = n_vocab as i32;
7895 let __s_b1 = self.gpu.stream();
7896 let mut b1 = __s_b1.launch_builder(&f1);
7897 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
7898 unsafe {
7899 b1.launch(cfg1)?;
7900 }
7901 let f2 = self.func("prob_of_token_final_f32");
7902 let cfg2 = LaunchConfig {
7903 grid_dim: (1, 1, 1),
7904 block_dim: (256, 1, 1),
7905 shared_mem_bytes: 0,
7906 };
7907 let nbi = nb as i32;
7908 let __s_b2 = self.gpu.stream();
7909 let mut b2 = __s_b2.launch_builder(&f2);
7910 b2.arg(&part).arg(&mut p_v).arg(&nbi);
7911 unsafe {
7912 b2.launch(cfg2)?;
7913 }
7914 Ok(())
7915 }
7916
7917 pub fn prob_of_token_device_into(
7918 &self,
7919 logits: &CudaSlice<f32>,
7920 tok: &CudaSlice<u32>,
7921 p_out: &mut CudaSlice<f32>,
7922 n_vocab: usize,
7923 ) -> Result<(), Box<dyn std::error::Error>> {
7924 let nb = ARGMAX_NB;
7925 let mut part = self.alloc_uninit::<f32>(nb)?;
7926 let f1 = self.func("prob_of_token_partial_f32");
7927 let cfg1 = LaunchConfig {
7928 grid_dim: (nb as u32, 1, 1),
7929 block_dim: (256, 1, 1),
7930 shared_mem_bytes: 0,
7931 };
7932 let nv = n_vocab as i32;
7933 let __s_b1 = self.gpu.stream();
7934 let mut b1 = __s_b1.launch_builder(&f1);
7935 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
7936 unsafe {
7937 b1.launch(cfg1)?;
7938 }
7939 let f2 = self.func("prob_of_token_final_f32");
7940 let cfg2 = LaunchConfig {
7941 grid_dim: (1, 1, 1),
7942 block_dim: (256, 1, 1),
7943 shared_mem_bytes: 0,
7944 };
7945 let nbi = nb as i32;
7946 let __s_b2 = self.gpu.stream();
7947 let mut b2 = __s_b2.launch_builder(&f2);
7948 b2.arg(&part).arg(p_out).arg(&nbi);
7949 unsafe {
7950 b2.launch(cfg2)?;
7951 }
7952 Ok(())
7953 }
7954
7955 pub fn u32_hist_append(
7958 &self,
7959 tok: &CudaSlice<u32>,
7960 hist: &mut CudaSlice<u32>,
7961 idx: &mut CudaSlice<i32>,
7962 ) -> Result<(), Box<dyn std::error::Error>> {
7963 let f = self.func("u32_hist_append");
7964 let cfg = LaunchConfig {
7965 grid_dim: (1, 1, 1),
7966 block_dim: (32, 1, 1),
7967 shared_mem_bytes: 0,
7968 };
7969 let __s_b = self.gpu.stream();
7970 let mut b = __s_b.launch_builder(&f);
7971 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
7972 unsafe {
7973 b.launch(cfg)?;
7974 }
7975 Ok(())
7976 }
7977
7978 pub fn argmax_token_device(
7979 &self,
7980 logits: &CudaSlice<f32>,
7981 n_vocab: usize,
7982 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
7983 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
7984 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
7985 Ok(tok)
7986 }
7987 pub fn argmax_token_device_into(
7994 &self,
7995 logits: &CudaSlice<f32>,
7996 tok: &mut CudaSlice<u32>,
7997 n_vocab: usize,
7998 ) -> Result<(), Box<dyn std::error::Error>> {
7999 let nb = ARGMAX_NB;
8000 let f1 = self.func("argmax_partial_f32");
8001 let f2 = self.func("argmax_final_f32");
8002 let mut guard = self.argmax_partials.lock().unwrap();
8003 if guard.is_none() {
8004 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8007 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8008 *guard = Some((pv, pi));
8009 }
8010 let (part_v, part_i) = guard.as_mut().unwrap();
8011 let nv = n_vocab as i32;
8012 let nbi = nb as i32;
8013 let cfg1 = LaunchConfig {
8015 grid_dim: (nb as u32, 1, 1),
8016 block_dim: (256, 1, 1),
8017 shared_mem_bytes: 0,
8018 };
8019 let __s_b1 = self.gpu.stream();
8020 let mut b1 = __s_b1.launch_builder(&f1);
8021 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
8022 unsafe {
8023 b1.launch(cfg1)?;
8024 }
8025 let cfg2 = LaunchConfig {
8027 grid_dim: (1, 1, 1),
8028 block_dim: (256, 1, 1),
8029 shared_mem_bytes: 0,
8030 };
8031 let __s_b2 = self.gpu.stream();
8032 let mut b2 = __s_b2.launch_builder(&f2);
8033 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
8034 unsafe {
8035 b2.launch(cfg2)?;
8036 }
8037 Ok(())
8038 }
8039 pub fn argmax_token_device_col(
8045 &self,
8046 logits: &CudaSlice<f32>,
8047 col: usize,
8048 n_vocab: usize,
8049 toks: &mut CudaSlice<u32>,
8050 out_idx: usize,
8051 ) -> Result<(), Box<dyn std::error::Error>> {
8052 let nb = ARGMAX_NB;
8053 let f1 = self.func("argmax_partial_f32");
8054 let f2 = self.func("argmax_final_f32");
8055 let mut guard = self.argmax_partials.lock().unwrap();
8056 if guard.is_none() {
8057 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
8058 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
8059 *guard = Some((pv, pi));
8060 }
8061 let (part_v, part_i) = guard.as_mut().unwrap();
8062 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
8063 let nv = n_vocab as i32;
8064 let nbi = nb as i32;
8065 let cfg1 = LaunchConfig {
8066 grid_dim: (nb as u32, 1, 1),
8067 block_dim: (256, 1, 1),
8068 shared_mem_bytes: 0,
8069 };
8070 let __s_b1 = self.gpu.stream();
8071 let mut b1 = __s_b1.launch_builder(&f1);
8072 b1.arg(&col_view)
8073 .arg(&mut *part_v)
8074 .arg(&mut *part_i)
8075 .arg(&nv);
8076 unsafe {
8077 b1.launch(cfg1)?;
8078 }
8079 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
8080 let cfg2 = LaunchConfig {
8081 grid_dim: (1, 1, 1),
8082 block_dim: (256, 1, 1),
8083 shared_mem_bytes: 0,
8084 };
8085 let __s_b2 = self.gpu.stream();
8086 let mut b2 = __s_b2.launch_builder(&f2);
8087 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
8088 unsafe {
8089 b2.launch(cfg2)?;
8090 }
8091 Ok(())
8092 }
8093 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8095 Ok(self.gpu.stream().clone_htod(v)?)
8096 }
8097 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
8098 let v = self.gpu.stream().clone_dtoh(d)?;
8099 self.gpu.stream().synchronize()?;
8100 Ok(v)
8101 }
8102 pub fn htod_u32_into(
8106 &self,
8107 dst: &mut CudaSlice<u32>,
8108 src: &[u32],
8109 ) -> Result<(), Box<dyn std::error::Error>> {
8110 let mut view = dst.slice_mut(0..src.len());
8111 self.gpu.stream().memcpy_htod(src, &mut view)?;
8112 Ok(())
8113 }
8114
8115 pub fn htod_i32_into(
8118 &self,
8119 dst: &mut CudaSlice<i32>,
8120 src: &[i32],
8121 ) -> Result<(), Box<dyn std::error::Error>> {
8122 let mut view = dst.slice_mut(0..src.len());
8123 self.gpu.stream().memcpy_htod(src, &mut view)?;
8124 Ok(())
8125 }
8126
8127 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
8128 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
8129 self.keep_if_capturing(&s);
8130 Ok(s)
8131 }
8132 pub fn embed_gather_device_into(
8135 &self,
8136 embd: &CudaSlice<u8>,
8137 token_d: &CudaSlice<u32>,
8138 x_out: &mut CudaSlice<f32>,
8139 n_embd: usize,
8140 qtype: i32,
8141 row_bytes: usize,
8142 ) -> Result<(), Box<dyn std::error::Error>> {
8143 let f = self.func("embed_gather_u32");
8144 let cfg = LaunchConfig {
8145 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8146 block_dim: (256, 1, 1),
8147 shared_mem_bytes: 0,
8148 };
8149 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8150 let __s_b = self.gpu.stream();
8151 let mut b = __s_b.launch_builder(&f);
8152 b.arg(embd)
8153 .arg(token_d)
8154 .arg(x_out)
8155 .arg(&ne)
8156 .arg(&qt)
8157 .arg(&rb);
8158 unsafe {
8159 b.launch(cfg)?;
8160 }
8161 Ok(())
8162 }
8163 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
8165 let v = self.gpu.stream().clone_dtoh(d)?;
8166 self.gpu.stream().synchronize()?;
8167 Ok(v[0])
8168 }
8169 pub fn i32_set_k(
8176 &self,
8177 dst: &mut CudaSlice<i32>,
8178 v: i32,
8179 ) -> Result<(), Box<dyn std::error::Error>> {
8180 let f = self.func("i32_set_k");
8181 let cfg = LaunchConfig {
8182 grid_dim: (1, 1, 1),
8183 block_dim: (1, 1, 1),
8184 shared_mem_bytes: 0,
8185 };
8186 let idx = 0i32;
8187 let __s_b = self.gpu.stream();
8188 let mut b = __s_b.launch_builder(&f);
8189 b.arg(dst).arg(&v).arg(&idx);
8190 unsafe {
8191 b.launch(cfg)?;
8192 }
8193 Ok(())
8194 }
8195
8196 pub fn set_i32_one(
8197 &self,
8198 d: &mut CudaSlice<i32>,
8199 v: i32,
8200 ) -> Result<(), Box<dyn std::error::Error>> {
8201 self.gpu.stream().memcpy_htod(&[v], d)?;
8202 Ok(())
8203 }
8204 pub fn set_u32_one(
8207 &self,
8208 d: &mut CudaSlice<u32>,
8209 v: u32,
8210 ) -> Result<(), Box<dyn std::error::Error>> {
8211 self.gpu.stream().memcpy_htod(&[v], d)?;
8212 Ok(())
8213 }
8214 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
8216 let v = self.gpu.stream().clone_dtoh(d)?;
8217 self.gpu.stream().synchronize()?;
8218 Ok(v[0])
8219 }
8220 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8222 Ok(self.gpu.stream().clone_htod(bytes)?)
8223 }
8224 pub fn embed_gather_device(
8228 &self,
8229 embd: &CudaSlice<u8>,
8230 token_d: &CudaSlice<u32>,
8231 n_embd: usize,
8232 qtype: i32,
8233 row_bytes: usize,
8234 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8235 let f = self.func("embed_gather_u32");
8236 let mut x = self.alloc_uninit::<f32>(n_embd)?;
8237 let cfg = LaunchConfig {
8238 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
8239 block_dim: (256, 1, 1),
8240 shared_mem_bytes: 0,
8241 };
8242 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
8243 let __s_b = self.gpu.stream();
8244 let mut b = __s_b.launch_builder(&f);
8245 b.arg(embd)
8246 .arg(token_d)
8247 .arg(&mut x)
8248 .arg(&ne)
8249 .arg(&qt)
8250 .arg(&rb);
8251 unsafe {
8252 b.launch(cfg)?;
8253 }
8254 Ok(x)
8255 }
8256
8257 pub fn embed_gather_device_t(
8261 &self,
8262 embd: &CudaSlice<u8>,
8263 tokens: &[u32],
8264 n_embd: usize,
8265 qtype: i32,
8266 row_bytes: usize,
8267 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8268 let t = tokens.len();
8269 let tok_d = self.gpu.stream().clone_htod(tokens)?;
8270 let f = self.func("embed_gather_u32_t");
8271 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8272 let cfg = LaunchConfig {
8273 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8274 block_dim: (256, 1, 1),
8275 shared_mem_bytes: 0,
8276 };
8277 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8278 let __s_b = self.gpu.stream();
8279 let mut b = __s_b.launch_builder(&f);
8280 b.arg(embd)
8281 .arg(&tok_d)
8282 .arg(&mut x)
8283 .arg(&ne)
8284 .arg(&qt)
8285 .arg(&rb)
8286 .arg(&ti);
8287 unsafe {
8288 b.launch(cfg)?;
8289 }
8290 Ok(x)
8291 }
8292
8293 pub fn embed_gather_device_tv(
8298 &self,
8299 embd: &CudaSlice<u8>,
8300 tok_v: &cudarc::driver::CudaView<u32>,
8301 t: usize,
8302 n_embd: usize,
8303 qtype: i32,
8304 row_bytes: usize,
8305 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8306 let f = self.func("embed_gather_u32_t");
8307 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8308 let cfg = LaunchConfig {
8309 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8310 block_dim: (256, 1, 1),
8311 shared_mem_bytes: 0,
8312 };
8313 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8314 let __s_b = self.gpu.stream();
8315 let mut b = __s_b.launch_builder(&f);
8316 b.arg(embd)
8317 .arg(tok_v)
8318 .arg(&mut x)
8319 .arg(&ne)
8320 .arg(&qt)
8321 .arg(&rb)
8322 .arg(&ti);
8323 unsafe {
8324 b.launch(cfg)?;
8325 }
8326 Ok(x)
8327 }
8328
8329 pub fn embed_gather_device_td(
8330 &self,
8331 embd: &CudaSlice<u8>,
8332 tok_d: &CudaSlice<u32>,
8333 t: usize,
8334 n_embd: usize,
8335 qtype: i32,
8336 row_bytes: usize,
8337 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8338 let f = self.func("embed_gather_u32_t");
8339 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
8340 let cfg = LaunchConfig {
8341 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
8342 block_dim: (256, 1, 1),
8343 shared_mem_bytes: 0,
8344 };
8345 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
8346 let __s_b = self.gpu.stream();
8347 let mut b = __s_b.launch_builder(&f);
8348 b.arg(embd)
8349 .arg(tok_d)
8350 .arg(&mut x)
8351 .arg(&ne)
8352 .arg(&qt)
8353 .arg(&rb)
8354 .arg(&ti);
8355 unsafe {
8356 b.launch(cfg)?;
8357 }
8358 Ok(x)
8359 }
8360
8361 #[inline]
8367 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
8369 if self
8370 .capture_keep_on
8371 .load(std::sync::atomic::Ordering::Relaxed)
8372 {
8373 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
8374 }
8375 }
8376
8377 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
8378 &self,
8379 n: usize,
8380 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
8381 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
8382 {
8386 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8387 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
8388 use cudarc::driver::DevicePtrMut;
8390 let n_bytes = s.len() * std::mem::size_of::<T>();
8391 let stream = self.gpu.stream();
8392 let (p_, _g) = s.device_ptr_mut(&stream);
8393 unsafe {
8394 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
8395 .result()?;
8396 }
8397 }
8398 }
8399 self.keep_if_capturing(&s);
8400 Ok(s)
8401 }
8402
8403 pub fn uninit_q8_pair(
8408 &self,
8409 n: usize,
8410 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8411 Ok((
8412 self.alloc_uninit::<i8>(n)?,
8413 self.alloc_uninit::<f32>(n / 32)?,
8414 ))
8415 }
8416
8417 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8418 self.alloc_uninit::<f32>(n)
8419 }
8420
8421 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8423 self.alloc_uninit::<i8>(n)
8424 }
8425
8426 #[allow(clippy::too_many_arguments)]
8430 pub fn rms_norm3(
8431 &self,
8432 x: &CudaSlice<f32>,
8433 w0: &CudaSlice<f32>,
8434 w1: &CudaSlice<f32>,
8435 w2: &CudaSlice<f32>,
8436 d0: &mut CudaSlice<f32>,
8437 d1: &mut CudaSlice<f32>,
8438 d2: &mut CudaSlice<f32>,
8439 ncols: usize,
8440 nrows: usize,
8441 eps: f32,
8442 ) -> Result<(), Box<dyn std::error::Error>> {
8443 let f = self.func("rms_norm3_f32");
8444 let cfg = LaunchConfig {
8445 grid_dim: (nrows as u32, 1, 1),
8446 block_dim: (rms_block(), 1, 1),
8447 shared_mem_bytes: 0,
8448 };
8449 let (nc, e) = (ncols as i32, eps);
8450 let __s_b = self.gpu.stream();
8451 let mut b = __s_b.launch_builder(&f);
8452 b.arg(x)
8453 .arg(w0)
8454 .arg(w1)
8455 .arg(w2)
8456 .arg(d0)
8457 .arg(d1)
8458 .arg(d2)
8459 .arg(&nc)
8460 .arg(&e);
8461 unsafe {
8462 b.launch(cfg)?;
8463 }
8464 Ok(())
8465 }
8466
8467 #[allow(clippy::too_many_arguments)]
8469 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
8472 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8473 *WARP_ON.get_or_init(|| {
8474 std::env::var("MEMRA_QKVNORM_W")
8475 .map(|v| v != "0")
8476 .unwrap_or(true)
8477 }) && ncols % 4 == 0
8478 && rows >= 64
8479 }
8480
8481 #[allow(clippy::too_many_arguments)]
8484 pub fn rms_norm_qkv_w4b(
8485 &self,
8486 q: &CudaSlice<f32>,
8487 k: &CudaSlice<f32>,
8488 v: &CudaSlice<f32>,
8489 wq: &CudaSlice<f32>,
8490 wk: &CudaSlice<f32>,
8491 wv: &CudaSlice<f32>,
8492 dq: &mut CudaSlice<f32>,
8493 dk: &mut CudaSlice<f32>,
8494 dv: &mut CudaSlice<f32>,
8495 dvb: &mut CudaSlice<u8>,
8496 ncols: usize,
8497 rq: usize,
8498 rk: usize,
8499 eps: f32,
8500 vf16: bool,
8501 ) -> Result<(), Box<dyn std::error::Error>> {
8502 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
8503 let f = self.func("rms_norm_qkv_w4b_f32");
8504 let rows = (rq + 2 * rk) as u32;
8505 let cfg = LaunchConfig {
8506 grid_dim: (rows.div_ceil(8), 1, 1),
8507 block_dim: (256, 1, 1),
8508 shared_mem_bytes: 0,
8509 };
8510 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8511 let vf = vf16 as i32;
8512 let __s_b = self.gpu.stream();
8513 let mut b = __s_b.launch_builder(&f);
8514 b.arg(q)
8515 .arg(k)
8516 .arg(v)
8517 .arg(wq)
8518 .arg(wk)
8519 .arg(wv)
8520 .arg(dq)
8521 .arg(dk)
8522 .arg(dv)
8523 .arg(&mut *dvb)
8524 .arg(&nc)
8525 .arg(&rqi)
8526 .arg(&rki)
8527 .arg(&rvi)
8528 .arg(&e)
8529 .arg(&vf);
8530 unsafe {
8531 b.launch(cfg)?;
8532 }
8533 Ok(())
8534 }
8535
8536 pub fn rms_norm_qkv(
8537 &self,
8538 q: &CudaSlice<f32>,
8539 k: &CudaSlice<f32>,
8540 v: &CudaSlice<f32>,
8541 wq: &CudaSlice<f32>,
8542 wk: &CudaSlice<f32>,
8543 wv: &CudaSlice<f32>,
8544 dq: &mut CudaSlice<f32>,
8545 dk: &mut CudaSlice<f32>,
8546 dv: &mut CudaSlice<f32>,
8547 ncols: usize,
8548 rq: usize,
8549 rk: usize,
8550 eps: f32,
8551 ) -> Result<(), Box<dyn std::error::Error>> {
8552 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8556 let warp_on = *WARP_ON.get_or_init(|| {
8557 std::env::var("MEMRA_QKVNORM_W")
8558 .map(|v| v != "0")
8559 .unwrap_or(true)
8560 });
8561 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
8564 let f = self.func("rms_norm_qkv_w4_f32");
8565 let rows = (rq + 2 * rk) as u32;
8566 let cfg = LaunchConfig {
8567 grid_dim: (rows.div_ceil(8), 1, 1),
8568 block_dim: (256, 1, 1),
8569 shared_mem_bytes: 0,
8570 };
8571 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
8572 let __s_b = self.gpu.stream();
8573 let mut b = __s_b.launch_builder(&f);
8574 b.arg(q)
8575 .arg(k)
8576 .arg(v)
8577 .arg(wq)
8578 .arg(wk)
8579 .arg(wv)
8580 .arg(dq)
8581 .arg(dk)
8582 .arg(dv)
8583 .arg(&nc)
8584 .arg(&rqi)
8585 .arg(&rki)
8586 .arg(&rvi)
8587 .arg(&e);
8588 unsafe {
8589 b.launch(cfg)?;
8590 }
8591 return Ok(());
8592 }
8593 let f = self.func("rms_norm_qkv_f32");
8594 let grid = (rq + 2 * rk) as u32;
8595 let cfg = LaunchConfig {
8596 grid_dim: (grid, 1, 1),
8597 block_dim: (rms_block(), 1, 1),
8598 shared_mem_bytes: 0,
8599 };
8600 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
8601 let __s_b = self.gpu.stream();
8602 let mut b = __s_b.launch_builder(&f);
8603 b.arg(q)
8604 .arg(k)
8605 .arg(v)
8606 .arg(wq)
8607 .arg(wk)
8608 .arg(wv)
8609 .arg(dq)
8610 .arg(dk)
8611 .arg(dv)
8612 .arg(&nc)
8613 .arg(&rqi)
8614 .arg(&rki)
8615 .arg(&e);
8616 unsafe {
8617 b.launch(cfg)?;
8618 }
8619 Ok(())
8620 }
8621
8622 #[allow(clippy::too_many_arguments)]
8624 pub fn rms_norm2x(
8625 &self,
8626 a: &CudaSlice<f32>,
8627 bb: &CudaSlice<f32>,
8628 wa: &CudaSlice<f32>,
8629 wb: &CudaSlice<f32>,
8630 da: &mut CudaSlice<f32>,
8631 db: &mut CudaSlice<f32>,
8632 ncols: usize,
8633 nrows: usize,
8634 eps: f32,
8635 ) -> Result<(), Box<dyn std::error::Error>> {
8636 let f = self.func("rms_norm2x_f32");
8637 let cfg = LaunchConfig {
8638 grid_dim: (2 * nrows as u32, 1, 1),
8639 block_dim: (rms_block(), 1, 1),
8640 shared_mem_bytes: 0,
8641 };
8642 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
8643 let __s_b = self.gpu.stream();
8644 let mut b = __s_b.launch_builder(&f);
8645 b.arg(a)
8646 .arg(bb)
8647 .arg(wa)
8648 .arg(wb)
8649 .arg(da)
8650 .arg(db)
8651 .arg(&nc)
8652 .arg(&nr)
8653 .arg(&e);
8654 unsafe {
8655 b.launch(cfg)?;
8656 }
8657 Ok(())
8658 }
8659
8660 pub fn softcap(
8662 &self,
8663 y: &mut CudaSlice<f32>,
8664 cap: f32,
8665 n: usize,
8666 ) -> Result<(), Box<dyn std::error::Error>> {
8667 let f = self.func("softcap_f32");
8668 let cfg = LaunchConfig::for_num_elems(n as u32);
8669 let ni = n as i32;
8670 let __s_b = self.gpu.stream();
8671 let mut b = __s_b.launch_builder(&f);
8672 b.arg(y).arg(&cap).arg(&ni);
8673 unsafe {
8674 b.launch(cfg)?;
8675 }
8676 Ok(())
8677 }
8678
8679 pub fn mask_ids_rows(
8682 &self,
8683 y: &mut CudaSlice<f32>,
8684 ids: &CudaSlice<i32>,
8685 n_ids: usize,
8686 n_vocab: usize,
8687 t: usize,
8688 ) -> Result<(), Box<dyn std::error::Error>> {
8689 let f = self.func("mask_ids_rows_f32");
8690 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
8691 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
8692 let __s_b = self.gpu.stream();
8693 let mut b = __s_b.launch_builder(&f);
8694 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
8695 unsafe {
8696 b.launch(cfg)?;
8697 }
8698 Ok(())
8699 }
8700
8701 #[allow(clippy::too_many_arguments)]
8703 pub fn add_scale_rms_norm(
8704 &self,
8705 a: &CudaSlice<f32>,
8706 b_in: &CudaSlice<f32>,
8707 c: f32,
8708 w: &CudaSlice<f32>,
8709 res: &mut CudaSlice<f32>,
8710 dst: &mut CudaSlice<f32>,
8711 ncols: usize,
8712 nrows: usize,
8713 eps: f32,
8714 ) -> Result<(), Box<dyn std::error::Error>> {
8715 let f = self.func("add_scale_rms_norm_f32");
8716 let cfg = LaunchConfig {
8717 grid_dim: (nrows as u32, 1, 1),
8718 block_dim: (rms_block(), 1, 1),
8719 shared_mem_bytes: 0,
8720 };
8721 let (nc, e2) = (ncols as i32, eps);
8722 let __s_b = self.gpu.stream();
8723 let mut b = __s_b.launch_builder(&f);
8724 b.arg(a)
8725 .arg(b_in)
8726 .arg(&c)
8727 .arg(w)
8728 .arg(res)
8729 .arg(dst)
8730 .arg(&nc)
8731 .arg(&e2);
8732 unsafe {
8733 b.launch(cfg)?;
8734 }
8735 Ok(())
8736 }
8737
8738 #[allow(clippy::too_many_arguments)]
8741 pub fn add_scale_rms_norm_q8_1(
8742 &self,
8743 a: &CudaSlice<f32>,
8744 b_in: &CudaSlice<f32>,
8745 c: f32,
8746 w: &CudaSlice<f32>,
8747 res: &mut CudaSlice<f32>,
8748 ncols: usize,
8749 nrows: usize,
8750 eps: f32,
8751 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8752 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8753 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8754 let (nc, e2) = (ncols as i32, eps);
8755 if Self::pdl_on() && Self::pdl_wb_on() {
8756 {
8757 use cudarc::driver::{DevicePtr, DevicePtrMut};
8758 let s = &self.gpu.stream();
8759 let (pa, _g0) = a.device_ptr(s);
8760 let (pb, _g1) = b_in.device_ptr(s);
8761 let (pw, _g2) = w.device_ptr(s);
8762 let (pr, _g3) = res.device_ptr_mut(s);
8763 let (pq, _g4) = out_q.device_ptr_mut(s);
8764 let (pd, _g5) = out_d.device_ptr_mut(s);
8765 let mut ps = [
8766 &pa as *const _ as *mut std::ffi::c_void,
8767 &pb as *const _ as *mut _,
8768 &c as *const _ as *mut _,
8769 &pw as *const _ as *mut _,
8770 &pr as *const _ as *mut _,
8771 &pq as *const _ as *mut _,
8772 &pd as *const _ as *mut _,
8773 &nc as *const _ as *mut _,
8774 &e2 as *const _ as *mut _,
8775 ];
8776 unsafe {
8777 self.launch_pdl(
8778 "add_scale_rms_norm_q8_1",
8779 (nrows as u32, 1, 1),
8780 (rms_block(), 1, 1),
8781 &mut ps,
8782 )?;
8783 }
8784 }
8785 return Ok((out_q, out_d));
8786 }
8787 let f = self.func("add_scale_rms_norm_q8_1");
8788 let cfg = LaunchConfig {
8789 grid_dim: (nrows as u32, 1, 1),
8790 block_dim: (rms_block(), 1, 1),
8791 shared_mem_bytes: 0,
8792 };
8793 let __s_b = self.gpu.stream();
8794 let mut b = __s_b.launch_builder(&f);
8795 b.arg(a)
8796 .arg(b_in)
8797 .arg(&c)
8798 .arg(w)
8799 .arg(res)
8800 .arg(&mut out_q)
8801 .arg(&mut out_d)
8802 .arg(&nc)
8803 .arg(&e2);
8804 unsafe {
8805 b.launch(cfg)?;
8806 }
8807 Ok((out_q, out_d))
8808 }
8809
8810 #[allow(clippy::too_many_arguments)]
8812 pub fn add_scale_rms_norm_q8_1_into(
8813 &self,
8814 a: &CudaSlice<f32>,
8815 b_in: &CudaSlice<f32>,
8816 c: f32,
8817 w: &CudaSlice<f32>,
8818 res: &mut CudaSlice<f32>,
8819 ncols: usize,
8820 nrows: usize,
8821 eps: f32,
8822 out_q: &mut CudaSlice<i8>,
8823 out_d: &mut CudaSlice<f32>,
8824 ) -> Result<(), Box<dyn std::error::Error>> {
8825 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
8826 let (nc, e2) = (ncols as i32, eps);
8827 if Self::pdl_on() && Self::pdl_wb_on() {
8828 use cudarc::driver::{DevicePtr, DevicePtrMut};
8829 let s = &self.gpu.stream();
8830 let (pa, _g0) = a.device_ptr(s);
8831 let (pb, _g1) = b_in.device_ptr(s);
8832 let (pw, _g2) = w.device_ptr(s);
8833 let (pr, _g3) = res.device_ptr_mut(s);
8834 let (pq, _g4) = out_q.device_ptr_mut(s);
8835 let (pd, _g5) = out_d.device_ptr_mut(s);
8836 let mut ps = [
8837 &pa as *const _ as *mut std::ffi::c_void,
8838 &pb as *const _ as *mut _,
8839 &c as *const _ as *mut _,
8840 &pw as *const _ as *mut _,
8841 &pr as *const _ as *mut _,
8842 &pq as *const _ as *mut _,
8843 &pd as *const _ as *mut _,
8844 &nc as *const _ as *mut _,
8845 &e2 as *const _ as *mut _,
8846 ];
8847 unsafe {
8848 self.launch_pdl(
8849 "add_scale_rms_norm_q8_1",
8850 (nrows as u32, 1, 1),
8851 (rms_block(), 1, 1),
8852 &mut ps,
8853 )?;
8854 }
8855 return Ok(());
8856 }
8857 let f = self.func("add_scale_rms_norm_q8_1");
8858 let cfg = LaunchConfig {
8859 grid_dim: (nrows as u32, 1, 1),
8860 block_dim: (rms_block(), 1, 1),
8861 shared_mem_bytes: 0,
8862 };
8863 let __s_b = self.gpu.stream();
8864 let mut b = __s_b.launch_builder(&f);
8865 b.arg(a)
8866 .arg(b_in)
8867 .arg(&c)
8868 .arg(w)
8869 .arg(res)
8870 .arg(&mut *out_q)
8871 .arg(&mut *out_d)
8872 .arg(&nc)
8873 .arg(&e2);
8874 unsafe {
8875 b.launch(cfg)?;
8876 }
8877 Ok(())
8878 }
8879
8880 #[allow(clippy::too_many_arguments)]
8883 pub fn rms_pre_add_scale_rms_norm_q8_1(
8884 &self,
8885 a: &CudaSlice<f32>,
8886 wa: &CudaSlice<f32>,
8887 b_in: &CudaSlice<f32>,
8888 c: f32,
8889 w: &CudaSlice<f32>,
8890 res: &mut CudaSlice<f32>,
8891 ncols: usize,
8892 nrows: usize,
8893 eps: f32,
8894 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8895 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8896 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8897 let (nc, e2) = (ncols as i32, eps);
8898 if Self::pdl_on() {
8899 {
8900 use cudarc::driver::{DevicePtr, DevicePtrMut};
8901 let s = &self.gpu.stream();
8902 let (pa, _g0) = a.device_ptr(s);
8903 let (pwa, _g1) = wa.device_ptr(s);
8904 let (pb, _g2) = b_in.device_ptr(s);
8905 let (pw, _g3) = w.device_ptr(s);
8906 let (pr, _g4) = res.device_ptr_mut(s);
8907 let (pq, _g5) = out_q.device_ptr_mut(s);
8908 let (pd, _g6) = out_d.device_ptr_mut(s);
8909 let mut ps = [
8910 &pa as *const _ as *mut std::ffi::c_void,
8911 &pwa as *const _ as *mut _,
8912 &pb as *const _ as *mut _,
8913 &c as *const _ as *mut _,
8914 &pw as *const _ as *mut _,
8915 &pr as *const _ as *mut _,
8916 &pq as *const _ as *mut _,
8917 &pd as *const _ as *mut _,
8918 &nc as *const _ as *mut _,
8919 &e2 as *const _ as *mut _,
8920 ];
8921 unsafe {
8922 self.launch_pdl(
8923 "rms_pre_add_scale_rms_norm_q8_1",
8924 (nrows as u32, 1, 1),
8925 (rms_block(), 1, 1),
8926 &mut ps,
8927 )?;
8928 }
8929 }
8930 return Ok((out_q, out_d));
8931 }
8932 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
8933 let cfg = LaunchConfig {
8934 grid_dim: (nrows as u32, 1, 1),
8935 block_dim: (rms_block(), 1, 1),
8936 shared_mem_bytes: 0,
8937 };
8938 let __s_b = self.gpu.stream();
8939 let mut b = __s_b.launch_builder(&f);
8940 b.arg(a)
8941 .arg(wa)
8942 .arg(b_in)
8943 .arg(&c)
8944 .arg(w)
8945 .arg(res)
8946 .arg(&mut out_q)
8947 .arg(&mut out_d)
8948 .arg(&nc)
8949 .arg(&e2);
8950 unsafe {
8951 b.launch(cfg)?;
8952 }
8953 Ok((out_q, out_d))
8954 }
8955
8956 pub fn gelu_tanh_mul_q8_1(
8959 &self,
8960 gate: &CudaSlice<f32>,
8961 up: &cudarc::driver::CudaView<f32>,
8962 act: &mut CudaSlice<f32>,
8963 ncols: usize,
8964 nrows: usize,
8965 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8966 debug_assert!(ncols % 128 == 0);
8967 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
8968 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
8969 let nc = ncols as i32;
8970 if Self::pdl_on() {
8971 {
8972 use cudarc::driver::{DevicePtr, DevicePtrMut};
8973 let s = &self.gpu.stream();
8974 let (pg, _g0) = gate.device_ptr(s);
8975 let (pu, _g1) = up.device_ptr(s);
8976 let (pact, _g2) = act.device_ptr_mut(s);
8977 let (pq, _g3) = out_q.device_ptr_mut(s);
8978 let (pd, _g4) = out_d.device_ptr_mut(s);
8979 let mut ps = [
8980 &pg as *const _ as *mut std::ffi::c_void,
8981 &pu as *const _ as *mut _,
8982 &pact as *const _ as *mut _,
8983 &pq as *const _ as *mut _,
8984 &pd as *const _ as *mut _,
8985 &nc as *const _ as *mut _,
8986 ];
8987 unsafe {
8988 self.launch_pdl(
8989 "gelu_tanh_mul_q8_1",
8990 (nrows as u32, 1, 1),
8991 (rms_block(), 1, 1),
8992 &mut ps,
8993 )?;
8994 }
8995 }
8996 return Ok((out_q, out_d));
8997 }
8998 let f = self.func("gelu_tanh_mul_q8_1");
8999 let cfg = LaunchConfig {
9000 grid_dim: (nrows as u32, 1, 1),
9001 block_dim: (rms_block(), 1, 1),
9002 shared_mem_bytes: 0,
9003 };
9004 let __s_b = self.gpu.stream();
9005 let mut b = __s_b.launch_builder(&f);
9006 b.arg(gate)
9007 .arg(up)
9008 .arg(act)
9009 .arg(&mut out_q)
9010 .arg(&mut out_d)
9011 .arg(&nc);
9012 unsafe {
9013 b.launch(cfg)?;
9014 }
9015 Ok((out_q, out_d))
9016 }
9017
9018 #[allow(clippy::too_many_arguments)]
9020 pub fn gelu_tanh_mul_q8_1_into(
9021 &self,
9022 gate: &CudaSlice<f32>,
9023 up: &cudarc::driver::CudaView<f32>,
9024 act: &mut CudaSlice<f32>,
9025 ncols: usize,
9026 nrows: usize,
9027 out_q: &mut CudaSlice<i8>,
9028 out_d: &mut CudaSlice<f32>,
9029 ) -> Result<(), Box<dyn std::error::Error>> {
9030 debug_assert!(ncols % 128 == 0);
9031 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9032 let nc = ncols as i32;
9033 if Self::pdl_on() {
9034 use cudarc::driver::{DevicePtr, DevicePtrMut};
9035 let s = &self.gpu.stream();
9036 let (pg, _g0) = gate.device_ptr(s);
9037 let (pu, _g1) = up.device_ptr(s);
9038 let (pact, _g2) = act.device_ptr_mut(s);
9039 let (pq, _g3) = out_q.device_ptr_mut(s);
9040 let (pd, _g4) = out_d.device_ptr_mut(s);
9041 let mut ps = [
9042 &pg as *const _ as *mut std::ffi::c_void,
9043 &pu as *const _ as *mut _,
9044 &pact as *const _ as *mut _,
9045 &pq as *const _ as *mut _,
9046 &pd as *const _ as *mut _,
9047 &nc as *const _ as *mut _,
9048 ];
9049 unsafe {
9050 self.launch_pdl(
9051 "gelu_tanh_mul_q8_1",
9052 (nrows as u32, 1, 1),
9053 (rms_block(), 1, 1),
9054 &mut ps,
9055 )?;
9056 }
9057 return Ok(());
9058 }
9059 let f = self.func("gelu_tanh_mul_q8_1");
9060 let cfg = LaunchConfig {
9061 grid_dim: (nrows as u32, 1, 1),
9062 block_dim: (rms_block(), 1, 1),
9063 shared_mem_bytes: 0,
9064 };
9065 let __s_b = self.gpu.stream();
9066 let mut b = __s_b.launch_builder(&f);
9067 b.arg(gate)
9068 .arg(up)
9069 .arg(&mut *act)
9070 .arg(&mut *out_q)
9071 .arg(&mut *out_d)
9072 .arg(&nc);
9073 unsafe {
9074 b.launch(cfg)?;
9075 }
9076 Ok(())
9077 }
9078
9079 #[allow(clippy::too_many_arguments)]
9081 pub fn add_rms_norm3_q8z(
9082 &self,
9083 a: &CudaSlice<f32>,
9084 b_in: &CudaSlice<f32>,
9085 w0: &CudaSlice<f32>,
9086 w1: &CudaSlice<f32>,
9087 w2: &CudaSlice<f32>,
9088 res: &mut CudaSlice<f32>,
9089 out1: &mut CudaSlice<f32>,
9090 ncols: usize,
9091 nrows: usize,
9092 eps: f32,
9093 ) -> Result<
9094 (
9095 (CudaSlice<i8>, CudaSlice<f32>),
9096 (CudaSlice<i8>, CudaSlice<f32>),
9097 ),
9098 Box<dyn std::error::Error>,
9099 > {
9100 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
9101 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9102 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
9103 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9104 let f = self.func("add_rms_norm3_q8z_f32");
9105 let cfg = LaunchConfig {
9106 grid_dim: (nrows as u32, 1, 1),
9107 block_dim: (rms_block(), 1, 1),
9108 shared_mem_bytes: 0,
9109 };
9110 let (nc, e2) = (ncols as i32, eps);
9111 let __s_b = self.gpu.stream();
9112 let mut b = __s_b.launch_builder(&f);
9113 b.arg(a)
9114 .arg(b_in)
9115 .arg(w0)
9116 .arg(w1)
9117 .arg(w2)
9118 .arg(res)
9119 .arg(&mut q0)
9120 .arg(&mut d0)
9121 .arg(out1)
9122 .arg(&mut q2)
9123 .arg(&mut d2)
9124 .arg(&nc)
9125 .arg(&e2);
9126 unsafe {
9127 b.launch(cfg)?;
9128 }
9129 Ok(((q0, d0), (q2, d2)))
9130 }
9131
9132 #[allow(clippy::too_many_arguments)]
9134 pub fn add_rms_norm3(
9135 &self,
9136 a: &CudaSlice<f32>,
9137 b_in: &CudaSlice<f32>,
9138 w0: &CudaSlice<f32>,
9139 w1: &CudaSlice<f32>,
9140 w2: &CudaSlice<f32>,
9141 res: &mut CudaSlice<f32>,
9142 d0: &mut CudaSlice<f32>,
9143 d1: &mut CudaSlice<f32>,
9144 d2: &mut CudaSlice<f32>,
9145 ncols: usize,
9146 nrows: usize,
9147 eps: f32,
9148 ) -> Result<(), Box<dyn std::error::Error>> {
9149 let f = self.func("add_rms_norm3_f32");
9150 let cfg = LaunchConfig {
9151 grid_dim: (nrows as u32, 1, 1),
9152 block_dim: (rms_block(), 1, 1),
9153 shared_mem_bytes: 0,
9154 };
9155 let (nc, e2) = (ncols as i32, eps);
9156 let __s_b = self.gpu.stream();
9157 let mut b = __s_b.launch_builder(&f);
9158 b.arg(a)
9159 .arg(b_in)
9160 .arg(w0)
9161 .arg(w1)
9162 .arg(w2)
9163 .arg(res)
9164 .arg(d0)
9165 .arg(d1)
9166 .arg(d2)
9167 .arg(&nc)
9168 .arg(&e2);
9169 unsafe {
9170 b.launch(cfg)?;
9171 }
9172 Ok(())
9173 }
9174
9175 pub fn add_scale(
9177 &self,
9178 a: &CudaSlice<f32>,
9179 b_in: &CudaSlice<f32>,
9180 c: f32,
9181 dst: &mut CudaSlice<f32>,
9182 n: usize,
9183 ) -> Result<(), Box<dyn std::error::Error>> {
9184 let f = self.func("add_scale_f32");
9185 let cfg = LaunchConfig::for_num_elems(n as u32);
9186 let ni = n as i32;
9187 let __s_b = self.gpu.stream();
9188 let mut b = __s_b.launch_builder(&f);
9189 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
9190 unsafe {
9191 b.launch(cfg)?;
9192 }
9193 Ok(())
9194 }
9195
9196 pub fn layer_norm_bias(
9198 &self,
9199 x: &CudaSlice<f32>,
9200 w: &CudaSlice<f32>,
9201 b: &CudaSlice<f32>,
9202 dst: &mut CudaSlice<f32>,
9203 ncols: usize,
9204 nrows: usize,
9205 eps: f32,
9206 ) -> Result<(), Box<dyn std::error::Error>> {
9207 let f = self.func("layer_norm_bias_f32");
9208 let (nc, e) = (ncols as i32, eps);
9209 let cfg = LaunchConfig {
9210 grid_dim: (nrows as u32, 1, 1),
9211 block_dim: (256, 1, 1),
9212 shared_mem_bytes: 0,
9213 };
9214 let __s_b = self.gpu.stream();
9215 let mut lb = __s_b.launch_builder(&f);
9216 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
9217 unsafe {
9218 lb.launch(cfg)?;
9219 }
9220 Ok(())
9221 }
9222
9223 pub fn gelu_tanh(
9225 &self,
9226 x: &CudaSlice<f32>,
9227 dst: &mut CudaSlice<f32>,
9228 n: usize,
9229 ) -> Result<(), Box<dyn std::error::Error>> {
9230 let f = self.func("gelu_tanh_f32");
9231 let ni = n as i64;
9232 let cfg = LaunchConfig {
9233 grid_dim: (n.div_ceil(256) as u32, 1, 1),
9234 block_dim: (256, 1, 1),
9235 shared_mem_bytes: 0,
9236 };
9237 let __s_b = self.gpu.stream();
9238 let mut lb = __s_b.launch_builder(&f);
9239 lb.arg(x).arg(&mut *dst).arg(&ni);
9240 unsafe {
9241 lb.launch(cfg)?;
9242 }
9243 Ok(())
9244 }
9245
9246 pub fn row_softmax(
9248 &self,
9249 x: &mut CudaSlice<f32>,
9250 ncols: usize,
9251 nrows: usize,
9252 ) -> Result<(), Box<dyn std::error::Error>> {
9253 let f = self.func("row_softmax_f32");
9254 let nc = ncols as i32;
9255 let cfg = LaunchConfig {
9256 grid_dim: (nrows as u32, 1, 1),
9257 block_dim: (256, 1, 1),
9258 shared_mem_bytes: 0,
9259 };
9260 let __s_b = self.gpu.stream();
9261 let mut lb = __s_b.launch_builder(&f);
9262 lb.arg(&mut *x).arg(&nc);
9263 unsafe {
9264 lb.launch(cfg)?;
9265 }
9266 Ok(())
9267 }
9268
9269 pub fn rms_norm(
9270 &self,
9271 x: &CudaSlice<f32>,
9272 w: &CudaSlice<f32>,
9273 dst: &mut CudaSlice<f32>,
9274 ncols: usize,
9275 nrows: usize,
9276 eps: f32,
9277 ) -> Result<(), Box<dyn std::error::Error>> {
9278 let (nc, e) = (ncols as i32, eps);
9279 let kname = if Self::norm_ilp_on() {
9280 "rms_norm_f32_v2"
9281 } else {
9282 "rms_norm_f32"
9283 };
9284 if Self::pdl_on() && Self::pdl_wb_on() {
9285 use cudarc::driver::{DevicePtr, DevicePtrMut};
9286 let s = &self.gpu.stream();
9287 let (px, _g0) = x.device_ptr(s);
9288 let (pw, _g1) = w.device_ptr(s);
9289 let (pd, _g2) = dst.device_ptr_mut(s);
9290 let mut ps = [
9291 &px as *const _ as *mut std::ffi::c_void,
9292 &pw as *const _ as *mut _,
9293 &pd as *const _ as *mut _,
9294 &nc as *const _ as *mut _,
9295 &e as *const _ as *mut _,
9296 ];
9297 unsafe {
9298 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9299 }
9300 return Ok(());
9301 }
9302 let f = self.func(kname);
9303 let cfg = LaunchConfig {
9304 grid_dim: (nrows as u32, 1, 1),
9305 block_dim: (rms_block(), 1, 1),
9306 shared_mem_bytes: 0,
9307 };
9308 let __s_b = self.gpu.stream();
9309 let mut b = __s_b.launch_builder(&f);
9310 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9311 unsafe {
9312 b.launch(cfg)?;
9313 }
9314 Ok(())
9315 }
9316
9317 pub fn rms_norm_decode(
9325 &self,
9326 x: &CudaSlice<f32>,
9327 w: &CudaSlice<f32>,
9328 dst: &mut CudaSlice<f32>,
9329 ncols: usize,
9330 nrows: usize,
9331 eps: f32,
9332 ) -> Result<(), Box<dyn std::error::Error>> {
9333 let f = self.func(if Self::norm_ilp_on() {
9334 "rms_norm_f32_v2"
9335 } else {
9336 "rms_norm_f32"
9337 });
9338 let cfg = LaunchConfig {
9339 grid_dim: (nrows as u32, 1, 1),
9340 block_dim: (1024, 1, 1),
9341 shared_mem_bytes: 0,
9342 };
9343 let (nc, e) = (ncols as i32, eps);
9344 let __s_b = self.gpu.stream();
9345 let mut b = __s_b.launch_builder(&f);
9346 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
9347 unsafe {
9348 b.launch(cfg)?;
9349 }
9350 Ok(())
9351 }
9352
9353 pub fn rms_norm_q8_1(
9357 &self,
9358 x: &CudaSlice<f32>,
9359 w: &CudaSlice<f32>,
9360 ncols: usize,
9361 nrows: usize,
9362 eps: f32,
9363 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9364 let nblk = ncols / 32;
9365 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9366 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9367 let (nc, e) = (ncols as i32, eps);
9368 if Self::pdl_on() {
9369 {
9370 use cudarc::driver::{DevicePtr, DevicePtrMut};
9371 let s = &self.gpu.stream();
9372 let (px, _g0) = x.device_ptr(s);
9373 let (pw, _g1) = w.device_ptr(s);
9374 let (pq, _g2) = q.device_ptr_mut(s);
9375 let (pd, _g3) = d.device_ptr_mut(s);
9376 let mut ps = [
9377 &px as *const _ as *mut std::ffi::c_void,
9378 &pw as *const _ as *mut _,
9379 &pq as *const _ as *mut _,
9380 &pd as *const _ as *mut _,
9381 &nc as *const _ as *mut _,
9382 &e as *const _ as *mut _,
9383 ];
9384 unsafe {
9385 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9386 }
9387 }
9388 return Ok((q, d));
9389 }
9390 let f = self.func("rms_norm_q8_1");
9391 let cfg = LaunchConfig {
9394 grid_dim: (nrows as u32, 1, 1),
9395 block_dim: (1024, 1, 1),
9396 shared_mem_bytes: 0,
9397 };
9398 let __s_b = self.gpu.stream();
9399 let mut b = __s_b.launch_builder(&f);
9400 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
9401 unsafe {
9402 b.launch(cfg)?;
9403 }
9404 Ok((q, d))
9405 }
9406
9407 pub fn rms_norm_q8_1_into(
9410 &self,
9411 x: &CudaSlice<f32>,
9412 w: &CudaSlice<f32>,
9413 ncols: usize,
9414 nrows: usize,
9415 eps: f32,
9416 q: &mut CudaSlice<i8>,
9417 d: &mut CudaSlice<f32>,
9418 ) -> Result<(), Box<dyn std::error::Error>> {
9419 let nblk = ncols / 32;
9420 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
9421 let (nc, e) = (ncols as i32, eps);
9422 if Self::pdl_on() {
9423 use cudarc::driver::{DevicePtr, DevicePtrMut};
9424 let s = &self.gpu.stream();
9425 let (px, _g0) = x.device_ptr(s);
9426 let (pw, _g1) = w.device_ptr(s);
9427 let (pq, _g2) = q.device_ptr_mut(s);
9428 let (pd, _g3) = d.device_ptr_mut(s);
9429 let mut ps = [
9430 &px as *const _ as *mut std::ffi::c_void,
9431 &pw as *const _ as *mut _,
9432 &pq as *const _ as *mut _,
9433 &pd as *const _ as *mut _,
9434 &nc as *const _ as *mut _,
9435 &e as *const _ as *mut _,
9436 ];
9437 unsafe {
9438 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
9439 }
9440 return Ok(());
9441 }
9442 let f = self.func("rms_norm_q8_1");
9443 let cfg = LaunchConfig {
9444 grid_dim: (nrows as u32, 1, 1),
9445 block_dim: (1024, 1, 1),
9446 shared_mem_bytes: 0,
9447 };
9448 let __s_b = self.gpu.stream();
9449 let mut b = __s_b.launch_builder(&f);
9450 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
9451 unsafe {
9452 b.launch(cfg)?;
9453 }
9454 Ok(())
9455 }
9456
9457 pub fn quantize_q8_1_into(
9459 &self,
9460 x: &CudaSlice<f32>,
9461 m: usize,
9462 in_f: usize,
9463 q: &mut CudaSlice<i8>,
9464 d: &mut CudaSlice<f32>,
9465 ) -> Result<(), Box<dyn std::error::Error>> {
9466 let nblk = in_f / 32;
9467 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
9468 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
9469 let (inf, mi) = (in_f as i32, m as i32);
9470 if Self::pdl_on() && Self::pdl_wb_on() {
9471 use cudarc::driver::{DevicePtr, DevicePtrMut};
9472 let s = &self.gpu.stream();
9473 let (px, _g0) = x.device_ptr(s);
9474 let (pq, _g1) = q.device_ptr_mut(s);
9475 let (pd, _g2) = d.device_ptr_mut(s);
9476 let mut ps = [
9477 &px as *const _ as *mut std::ffi::c_void,
9478 &pq as *const _ as *mut _,
9479 &pd as *const _ as *mut _,
9480 &inf as *const _ as *mut _,
9481 &mi as *const _ as *mut _,
9482 ];
9483 unsafe {
9484 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
9485 }
9486 return Ok(());
9487 }
9488 let f = self.func("quantize_q8_1");
9489 let __s_b = self.gpu.stream();
9490 let mut b = __s_b.launch_builder(&f);
9491 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
9492 unsafe {
9493 b.launch(cfg)?;
9494 }
9495 Ok(())
9496 }
9497
9498 pub fn add_rms_norm_q8_1(
9502 &self,
9503 a: &CudaSlice<f32>,
9504 b_in: &CudaSlice<f32>,
9505 w: &CudaSlice<f32>,
9506 res: &mut CudaSlice<f32>,
9507 ncols: usize,
9508 nrows: usize,
9509 eps: f32,
9510 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9511 let nblk = ncols / 32;
9512 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
9513 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
9514 let f = self.func("add_rms_norm_q8_1");
9515 let cfg = LaunchConfig {
9517 grid_dim: (nrows as u32, 1, 1),
9518 block_dim: (1024, 1, 1),
9519 shared_mem_bytes: 0,
9520 };
9521 let (nc, e) = (ncols as i32, eps);
9522 let __s_bld = self.gpu.stream();
9523 let mut bld = __s_bld.launch_builder(&f);
9524 bld.arg(a)
9525 .arg(b_in)
9526 .arg(w)
9527 .arg(res)
9528 .arg(&mut q)
9529 .arg(&mut d)
9530 .arg(&nc)
9531 .arg(&e);
9532 unsafe {
9533 bld.launch(cfg)?;
9534 }
9535 Ok((q, d))
9536 }
9537
9538 #[allow(clippy::too_many_arguments)]
9544 pub fn join_add_rms_norm_raw(
9545 &self,
9546 a0_raw: u64,
9547 a1_raw: u64,
9548 x: &CudaSlice<f32>,
9549 w: &CudaSlice<f32>,
9550 res: &mut CudaSlice<f32>,
9551 dst: &mut CudaSlice<f32>,
9552 ncols: usize,
9553 eps: f32,
9554 ) -> Result<(), Box<dyn std::error::Error>> {
9555 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
9556 return Err("join_add_rms_norm geometry".into());
9557 }
9558 let f = self.func("join_add_rms_norm_f32");
9559 let cfg = LaunchConfig {
9560 grid_dim: (1, 1, 1),
9561 block_dim: (rms_block(), 1, 1),
9562 shared_mem_bytes: 0,
9563 };
9564 let (nc, e) = (ncols as i32, eps);
9565 let __s_b = self.gpu.stream();
9566 let mut b = __s_b.launch_builder(&f);
9567 b.arg(&a0_raw)
9568 .arg(&a1_raw)
9569 .arg(x)
9570 .arg(w)
9571 .arg(&mut *res)
9572 .arg(&mut *dst)
9573 .arg(&nc)
9574 .arg(&e);
9575 unsafe {
9576 b.launch(cfg)?;
9577 }
9578 Ok(())
9579 }
9580
9581 pub fn add_rms_norm(
9582 &self,
9583 a: &CudaSlice<f32>,
9584 b: &CudaSlice<f32>,
9585 w: &CudaSlice<f32>,
9586 res: &mut CudaSlice<f32>,
9587 dst: &mut CudaSlice<f32>,
9588 ncols: usize,
9589 nrows: usize,
9590 eps: f32,
9591 ) -> Result<(), Box<dyn std::error::Error>> {
9592 let (nc, e) = (ncols as i32, eps);
9593 let kname = if Self::norm_ilp_on() {
9594 "add_rms_norm_f32_v2"
9595 } else {
9596 "add_rms_norm_f32"
9597 };
9598 if Self::pdl_on() && Self::pdl_wb_on() {
9599 use cudarc::driver::{DevicePtr, DevicePtrMut};
9600 let s = &self.gpu.stream();
9601 let (pa, _g0) = a.device_ptr(s);
9602 let (pb, _g1) = b.device_ptr(s);
9603 let (pw, _g2) = w.device_ptr(s);
9604 let (pr, _g3) = res.device_ptr_mut(s);
9605 let (pd, _g4) = dst.device_ptr_mut(s);
9606 let mut ps = [
9607 &pa as *const _ as *mut std::ffi::c_void,
9608 &pb as *const _ as *mut _,
9609 &pw as *const _ as *mut _,
9610 &pr as *const _ as *mut _,
9611 &pd as *const _ as *mut _,
9612 &nc as *const _ as *mut _,
9613 &e as *const _ as *mut _,
9614 ];
9615 unsafe {
9616 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
9617 }
9618 return Ok(());
9619 }
9620 let f = self.func(kname);
9621 let cfg = LaunchConfig {
9622 grid_dim: (nrows as u32, 1, 1),
9623 block_dim: (rms_block(), 1, 1),
9624 shared_mem_bytes: 0,
9625 };
9626 let __s_b2 = self.gpu.stream();
9627 let mut b2 = __s_b2.launch_builder(&f);
9628 b2.arg(a)
9629 .arg(b)
9630 .arg(w)
9631 .arg(&mut *res)
9632 .arg(&mut *dst)
9633 .arg(&nc)
9634 .arg(&e);
9635 unsafe {
9636 b2.launch(cfg)?;
9637 }
9638 Ok(())
9639 }
9640
9641 #[allow(clippy::too_many_arguments)]
9644 pub fn rms_pre_add_rms_norm(
9645 &self,
9646 a: &CudaSlice<f32>,
9647 wa: &CudaSlice<f32>,
9648 b: &CudaSlice<f32>,
9649 w: &CudaSlice<f32>,
9650 res: &mut CudaSlice<f32>,
9651 dst: &mut CudaSlice<f32>,
9652 ncols: usize,
9653 nrows: usize,
9654 eps: f32,
9655 ) -> Result<(), Box<dyn std::error::Error>> {
9656 let f = self.func("rms_pre_add_rms_norm_f32");
9657 let cfg = LaunchConfig {
9658 grid_dim: (nrows as u32, 1, 1),
9659 block_dim: (rms_block(), 1, 1),
9660 shared_mem_bytes: 0,
9661 };
9662 let (nc, e) = (ncols as i32, eps);
9663 let __s_b2 = self.gpu.stream();
9664 let mut b2 = __s_b2.launch_builder(&f);
9665 b2.arg(a)
9666 .arg(wa)
9667 .arg(b)
9668 .arg(w)
9669 .arg(&mut *res)
9670 .arg(&mut *dst)
9671 .arg(&nc)
9672 .arg(&e);
9673 unsafe {
9674 b2.launch(cfg)?;
9675 }
9676 Ok(())
9677 }
9678
9679 #[allow(clippy::too_many_arguments)]
9681 pub fn rms_pre_add_rms_norm_q8z(
9682 &self,
9683 a: &CudaSlice<f32>,
9684 wa: &CudaSlice<f32>,
9685 b: &CudaSlice<f32>,
9686 w: &CudaSlice<f32>,
9687 res: &mut CudaSlice<f32>,
9688 dst: &mut CudaSlice<f32>,
9689 ncols: usize,
9690 nrows: usize,
9691 eps: f32,
9692 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9693 debug_assert!(ncols % 128 == 0);
9694 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9695 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9696 let (nc, e) = (ncols as i32, eps);
9697 if Self::pdl_on() {
9698 {
9699 use cudarc::driver::{DevicePtr, DevicePtrMut};
9700 let s = &self.gpu.stream();
9701 let (pa, _g0) = a.device_ptr(s);
9702 let (pwa, _g1) = wa.device_ptr(s);
9703 let (pb, _g2) = b.device_ptr(s);
9704 let (pw, _g3) = w.device_ptr(s);
9705 let (pr, _g4) = res.device_ptr_mut(s);
9706 let (pdst, _g5) = dst.device_ptr_mut(s);
9707 let (pq, _g6) = out_q.device_ptr_mut(s);
9708 let (pd, _g7) = out_d.device_ptr_mut(s);
9709 let mut ps = [
9710 &pa as *const _ as *mut std::ffi::c_void,
9711 &pwa as *const _ as *mut _,
9712 &pb as *const _ as *mut _,
9713 &pw as *const _ as *mut _,
9714 &pr as *const _ as *mut _,
9715 &pdst as *const _ as *mut _,
9716 &pq as *const _ as *mut _,
9717 &pd as *const _ as *mut _,
9718 &nc as *const _ as *mut _,
9719 &e as *const _ as *mut _,
9720 ];
9721 unsafe {
9722 self.launch_pdl(
9723 "rms_pre_add_rms_norm_q8z_f32",
9724 (nrows as u32, 1, 1),
9725 (rms_block(), 1, 1),
9726 &mut ps,
9727 )?;
9728 }
9729 }
9730 return Ok((out_q, out_d));
9731 }
9732 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9733 let cfg = LaunchConfig {
9734 grid_dim: (nrows as u32, 1, 1),
9735 block_dim: (rms_block(), 1, 1),
9736 shared_mem_bytes: 0,
9737 };
9738 let __s_b2 = self.gpu.stream();
9739 let mut b2 = __s_b2.launch_builder(&f);
9740 b2.arg(a)
9741 .arg(wa)
9742 .arg(b)
9743 .arg(w)
9744 .arg(&mut *res)
9745 .arg(&mut *dst)
9746 .arg(&mut out_q)
9747 .arg(&mut out_d)
9748 .arg(&nc)
9749 .arg(&e);
9750 unsafe {
9751 b2.launch(cfg)?;
9752 }
9753 Ok((out_q, out_d))
9754 }
9755
9756 #[allow(clippy::too_many_arguments)]
9760 pub fn rms_pre_add_rms_norm_q8z_into(
9761 &self,
9762 a: &CudaSlice<f32>,
9763 wa: &CudaSlice<f32>,
9764 b: &CudaSlice<f32>,
9765 w: &CudaSlice<f32>,
9766 res: &mut CudaSlice<f32>,
9767 dst: &mut CudaSlice<f32>,
9768 ncols: usize,
9769 nrows: usize,
9770 eps: f32,
9771 out_q: &mut CudaSlice<i8>,
9772 out_d: &mut CudaSlice<f32>,
9773 ) -> Result<(), Box<dyn std::error::Error>> {
9774 debug_assert!(ncols % 128 == 0);
9775 let (nc, e) = (ncols as i32, eps);
9776 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
9777 let cfg = LaunchConfig {
9778 grid_dim: (nrows as u32, 1, 1),
9779 block_dim: (rms_block(), 1, 1),
9780 shared_mem_bytes: 0,
9781 };
9782 let __s_b = self.gpu.stream();
9783 let mut b2 = __s_b.launch_builder(&f);
9784 b2.arg(a)
9785 .arg(wa)
9786 .arg(b)
9787 .arg(w)
9788 .arg(&mut *res)
9789 .arg(&mut *dst)
9790 .arg(&mut *out_q)
9791 .arg(&mut *out_d)
9792 .arg(&nc)
9793 .arg(&e);
9794 unsafe {
9795 b2.launch(cfg)?;
9796 }
9797 Ok(())
9798 }
9799
9800 #[allow(clippy::too_many_arguments)]
9803 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
9804 &self,
9805 a: &CudaSlice<f32>,
9806 wa: &CudaSlice<f32>,
9807 b_in: &CudaSlice<f32>,
9808 c: f32,
9809 w: &CudaSlice<f32>,
9810 res: &mut CudaSlice<f32>,
9811 ncols: usize,
9812 nrows: usize,
9813 eps: f32,
9814 out_q: &mut CudaSlice<i8>,
9815 out_d: &mut CudaSlice<f32>,
9816 ) -> Result<(), Box<dyn std::error::Error>> {
9817 debug_assert!(ncols % 128 == 0);
9818 let (nc, e2) = (ncols as i32, eps);
9819 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
9820 let cfg = LaunchConfig {
9821 grid_dim: (nrows as u32, 1, 1),
9822 block_dim: (rms_block(), 1, 1),
9823 shared_mem_bytes: 0,
9824 };
9825 let __s_b = self.gpu.stream();
9826 let mut b2 = __s_b.launch_builder(&f);
9827 b2.arg(a)
9828 .arg(wa)
9829 .arg(b_in)
9830 .arg(&c)
9831 .arg(w)
9832 .arg(&mut *res)
9833 .arg(&mut *out_q)
9834 .arg(&mut *out_d)
9835 .arg(&nc)
9836 .arg(&e2);
9837 unsafe {
9838 b2.launch(cfg)?;
9839 }
9840 Ok(())
9841 }
9842
9843 pub fn g4_pnfold_on() -> bool {
9851 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9852 *ON.get_or_init(|| {
9853 std::env::var("MEMRA_G4_PNFOLD")
9854 .map(|v| v != "0")
9855 .unwrap_or(true)
9856 })
9857 }
9858
9859 pub fn build_q4_out_concat3(
9863 &self,
9864 w0: &crate::model::GpuTensor,
9865 w1: &crate::model::GpuTensor,
9866 w2: &crate::model::GpuTensor,
9867 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
9868 use crate::model::GpuTensor;
9869 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
9870 match w {
9871 GpuTensor::Quant {
9872 qtype,
9873 row_bytes,
9874 rp,
9875 ..
9876 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
9877 _ => None,
9878 }
9879 };
9880 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
9881 else {
9882 return Ok(None);
9883 };
9884 if rb0 != rb1
9885 || rb0 != rb2
9886 || w0.in_features() != w1.in_features()
9887 || w0.in_features() != w2.in_features()
9888 {
9889 return Ok(None);
9890 }
9891 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
9892 match w {
9893 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
9894 _ => unreachable!(),
9895 }
9896 }
9897 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
9898 let total = rb0 * (o0 + o1 + o2);
9899 let mut cat = self.alloc_u8(total)?;
9900 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
9901 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
9902 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
9903 Ok(Some(GpuTensor::Quant {
9904 bytes: cat,
9905 qtype: QT_Q4_0,
9906 row_bytes: rb0,
9907 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
9908 scale: 1.0,
9909 rp: false,
9910 #[cfg(memra_cutlass)]
9911 cutlass: None,
9912 fp8: None,
9913 blk: None,
9914 rp4: None,
9915 f16: None,
9916 }))
9917 }
9918
9919 fn full_width_rope_only(
9937 kernel: &str,
9938 n_rot: usize,
9939 head_dim: usize,
9940 ) -> Result<(), Box<dyn std::error::Error>> {
9941 if n_rot == head_dim {
9942 return Ok(());
9943 }
9944 Err(format!(
9945 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
9946 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
9947 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
9948 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
9949 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
9950 )
9951 .into())
9952 }
9953
9954 #[allow(clippy::too_many_arguments)]
9958 pub fn rms_norm_qkv_rope_cat(
9959 &self,
9960 qkv: &CudaSlice<f32>,
9961 wq: &CudaSlice<f32>,
9962 wk: &CudaSlice<f32>,
9963 wv: &CudaSlice<f32>,
9964 q: &mut CudaSlice<f32>,
9965 k: &mut CudaSlice<f32>,
9966 v: &mut CudaSlice<f32>,
9967 head_dim: usize,
9968 n_rot: usize,
9969 rq: usize,
9970 rk: usize,
9971 pos: &CudaSlice<i32>,
9972 nh_q: usize,
9973 nh_k: usize,
9974 base: f32,
9975 freq_scale: f32,
9976 ff: Option<&CudaSlice<f32>>,
9977 eps: f32,
9978 ) -> Result<(), Box<dyn std::error::Error>> {
9979 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
9980 let rows = rq + rk + rk;
9981 let theta_scale = base.powf(-2.0 / head_dim as f32);
9982 let (nc, rqi, rki, nhq, nhk) = (
9983 head_dim as i32,
9984 rq as i32,
9985 rk as i32,
9986 nh_q as i32,
9987 nh_k as i32,
9988 );
9989 if Self::pdl_on() {
9990 use cudarc::driver::{DevicePtr, DevicePtrMut};
9991 let s = &self.gpu.stream();
9992 let (pqkv, _g0) = qkv.device_ptr(s);
9993 let (pwq, _g1) = wq.device_ptr(s);
9994 let (pwk, _g2) = wk.device_ptr(s);
9995 let (pwv, _g3) = wv.device_ptr(s);
9996 let (pq, _g4) = q.device_ptr_mut(s);
9997 let (pk, _g5) = k.device_ptr_mut(s);
9998 let (pv, _g6) = v.device_ptr_mut(s);
9999 let (ppos, _g7) = pos.device_ptr(s);
10000 let (pff, _g8) = match ff {
10001 Some(t) => {
10002 let (p, g) = t.device_ptr(s);
10003 (p, Some(g))
10004 }
10005 None => (0, None),
10006 };
10007 let mut ps = [
10008 &pqkv as *const _ as *mut std::ffi::c_void,
10009 &pwq as *const _ as *mut _,
10010 &pwk as *const _ as *mut _,
10011 &pwv as *const _ as *mut _,
10012 &pq as *const _ as *mut _,
10013 &pk as *const _ as *mut _,
10014 &pv as *const _ as *mut _,
10015 &nc as *const _ as *mut _,
10016 &rqi as *const _ as *mut _,
10017 &rki as *const _ as *mut _,
10018 &ppos as *const _ as *mut _,
10019 &nhq as *const _ as *mut _,
10020 &nhk as *const _ as *mut _,
10021 &theta_scale as *const _ as *mut _,
10022 &freq_scale as *const _ as *mut _,
10023 &pff as *const _ as *mut _,
10024 &eps as *const _ as *mut _,
10025 ];
10026 unsafe {
10027 self.launch_pdl(
10028 "rms_norm_qkv_rope_cat_f32",
10029 (rows as u32, 1, 1),
10030 (rms_block(), 1, 1),
10031 &mut ps,
10032 )?;
10033 }
10034 return Ok(());
10035 }
10036 let f = self.func("rms_norm_qkv_rope_cat_f32");
10037 let cfg = LaunchConfig {
10038 grid_dim: (rows as u32, 1, 1),
10039 block_dim: (rms_block(), 1, 1),
10040 shared_mem_bytes: 0,
10041 };
10042 let __s_b = self.gpu.stream();
10043 let mut b = __s_b.launch_builder(&f);
10044 match ff {
10045 Some(t) => {
10046 b.arg(qkv)
10047 .arg(wq)
10048 .arg(wk)
10049 .arg(wv)
10050 .arg(&mut *q)
10051 .arg(&mut *k)
10052 .arg(&mut *v)
10053 .arg(&nc)
10054 .arg(&rqi)
10055 .arg(&rki)
10056 .arg(pos)
10057 .arg(&nhq)
10058 .arg(&nhk)
10059 .arg(&theta_scale)
10060 .arg(&freq_scale)
10061 .arg(t)
10062 .arg(&eps);
10063 unsafe {
10064 b.launch(cfg)?;
10065 }
10066 }
10067 None => {
10068 let null: u64 = 0;
10069 b.arg(qkv)
10070 .arg(wq)
10071 .arg(wk)
10072 .arg(wv)
10073 .arg(&mut *q)
10074 .arg(&mut *k)
10075 .arg(&mut *v)
10076 .arg(&nc)
10077 .arg(&rqi)
10078 .arg(&rki)
10079 .arg(pos)
10080 .arg(&nhq)
10081 .arg(&nhk)
10082 .arg(&theta_scale)
10083 .arg(&freq_scale)
10084 .arg(&null)
10085 .arg(&eps);
10086 unsafe {
10087 b.launch(cfg)?;
10088 }
10089 }
10090 }
10091 Ok(())
10092 }
10093
10094 #[allow(clippy::too_many_arguments)]
10098 pub fn rms_norm_qkv_rope(
10099 &self,
10100 q0: &CudaSlice<f32>,
10101 k0: &CudaSlice<f32>,
10102 v0: &CudaSlice<f32>,
10103 wq: &CudaSlice<f32>,
10104 wk: &CudaSlice<f32>,
10105 wv: &CudaSlice<f32>,
10106 q: &mut CudaSlice<f32>,
10107 k: &mut CudaSlice<f32>,
10108 v: &mut CudaSlice<f32>,
10109 head_dim: usize,
10110 n_rot: usize,
10111 rq: usize,
10112 rk: usize,
10113 pos: &CudaSlice<i32>,
10114 nh_q: usize,
10115 nh_k: usize,
10116 base: f32,
10117 freq_scale: f32,
10118 ff: Option<&CudaSlice<f32>>,
10119 eps: f32,
10120 ) -> Result<(), Box<dyn std::error::Error>> {
10121 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
10122 let f = self.func("rms_norm_qkv_rope_f32");
10123 let rows = rq + rk + rk; let cfg = LaunchConfig {
10125 grid_dim: (rows as u32, 1, 1),
10126 block_dim: (rms_block(), 1, 1),
10127 shared_mem_bytes: 0,
10128 };
10129 let theta_scale = base.powf(-2.0 / head_dim as f32);
10130 let (nc, rqi, rki, nhq, nhk) = (
10131 head_dim as i32,
10132 rq as i32,
10133 rk as i32,
10134 nh_q as i32,
10135 nh_k as i32,
10136 );
10137 let __s_b = self.gpu.stream();
10138 let mut b = __s_b.launch_builder(&f);
10139 match ff {
10140 Some(t) => {
10141 b.arg(q0)
10142 .arg(k0)
10143 .arg(v0)
10144 .arg(wq)
10145 .arg(wk)
10146 .arg(wv)
10147 .arg(&mut *q)
10148 .arg(&mut *k)
10149 .arg(&mut *v)
10150 .arg(&nc)
10151 .arg(&rqi)
10152 .arg(&rki)
10153 .arg(pos)
10154 .arg(&nhq)
10155 .arg(&nhk)
10156 .arg(&theta_scale)
10157 .arg(&freq_scale)
10158 .arg(t)
10159 .arg(&eps);
10160 unsafe {
10161 b.launch(cfg)?;
10162 }
10163 }
10164 None => {
10165 let null: u64 = 0;
10166 b.arg(q0)
10167 .arg(k0)
10168 .arg(v0)
10169 .arg(wq)
10170 .arg(wk)
10171 .arg(wv)
10172 .arg(&mut *q)
10173 .arg(&mut *k)
10174 .arg(&mut *v)
10175 .arg(&nc)
10176 .arg(&rqi)
10177 .arg(&rki)
10178 .arg(pos)
10179 .arg(&nhq)
10180 .arg(&nhk)
10181 .arg(&theta_scale)
10182 .arg(&freq_scale)
10183 .arg(&null)
10184 .arg(&eps);
10185 unsafe {
10186 b.launch(cfg)?;
10187 }
10188 }
10189 }
10190 Ok(())
10191 }
10192
10193 #[allow(clippy::too_many_arguments)]
10199 pub fn rms_norm_qkv_rope_append_dc(
10200 &self,
10201 q0: &CudaSlice<f32>,
10202 k0: &CudaSlice<f32>,
10203 v0: &CudaSlice<f32>,
10204 wq: &CudaSlice<f32>,
10205 wk: &CudaSlice<f32>,
10206 wv: &CudaSlice<f32>,
10207 q: &mut CudaSlice<f32>,
10208 k: &mut CudaSlice<f32>,
10209 v: &mut CudaSlice<f32>,
10210 head_dim: usize,
10211 n_rot: usize,
10212 rq: usize,
10213 rk: usize,
10214 pos: &CudaSlice<i32>,
10215 nh_q: usize,
10216 nh_k: usize,
10217 base: f32,
10218 freq_scale: f32,
10219 ff: Option<&CudaSlice<f32>>,
10220 eps: f32,
10221 kc: &mut CudaSlice<u8>,
10222 vc: &mut CudaSlice<u8>,
10223 t_dev: &CudaSlice<i32>,
10224 k_tok_bytes: usize,
10225 v_tok_bytes: usize,
10226 g: bool,
10227 ) -> Result<(), Box<dyn std::error::Error>> {
10228 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
10229 let rows = rq + rk + rk;
10230 let theta_scale = base.powf(-2.0 / head_dim as f32);
10231 let (nc, rqi, rki, nhq, nhk) = (
10232 head_dim as i32,
10233 rq as i32,
10234 rk as i32,
10235 nh_q as i32,
10236 nh_k as i32,
10237 );
10238 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10239 if Self::pdl_on() && Self::pdl_wb_on() {
10240 use cudarc::driver::{DevicePtr, DevicePtrMut};
10241 let s = &self.gpu.stream();
10242 let (p0, _a0) = q0.device_ptr(s);
10243 let (p1, _a1) = k0.device_ptr(s);
10244 let (p2, _a2) = v0.device_ptr(s);
10245 let (pwq, _a3) = wq.device_ptr(s);
10246 let (pwk, _a4) = wk.device_ptr(s);
10247 let (pwv, _a5) = wv.device_ptr(s);
10248 let (pq, _a6) = q.device_ptr_mut(s);
10249 let (pk, _a7) = k.device_ptr_mut(s);
10250 let (pv, _a8) = v.device_ptr_mut(s);
10251 let (pp, _a9) = pos.device_ptr(s);
10252 let pff: u64 = match ff {
10253 Some(t) => {
10254 let (p, _gg) = t.device_ptr(s);
10255 p as u64
10256 }
10257 None => 0,
10258 };
10259 let (pkc, _a10) = kc.device_ptr_mut(s);
10260 let (pvc, _a11) = vc.device_ptr_mut(s);
10261 let (pt, _a12) = t_dev.device_ptr(s);
10262 let mut ps = [
10263 &p0 as *const _ as *mut std::ffi::c_void,
10264 &p1 as *const _ as *mut _,
10265 &p2 as *const _ as *mut _,
10266 &pwq as *const _ as *mut _,
10267 &pwk as *const _ as *mut _,
10268 &pwv as *const _ as *mut _,
10269 &pq as *const _ as *mut _,
10270 &pk as *const _ as *mut _,
10271 &pv as *const _ as *mut _,
10272 &nc as *const _ as *mut _,
10273 &rqi as *const _ as *mut _,
10274 &rki as *const _ as *mut _,
10275 &pp as *const _ as *mut _,
10276 &nhq as *const _ as *mut _,
10277 &nhk as *const _ as *mut _,
10278 &theta_scale as *const _ as *mut _,
10279 &freq_scale as *const _ as *mut _,
10280 &pff as *const _ as *mut _,
10281 &eps as *const _ as *mut _,
10282 &pkc as *const _ as *mut _,
10283 &pvc as *const _ as *mut _,
10284 &pt as *const _ as *mut _,
10285 &ktb as *const _ as *mut _,
10286 &vtb as *const _ as *mut _,
10287 ];
10288 unsafe {
10289 self.launch_pdl_flash(
10290 g,
10291 "rms_norm_qkv_rope_append_dc_f32",
10292 (rows as u32, 1, 1),
10293 (rms_block(), 1, 1),
10294 0,
10295 &mut ps,
10296 )?;
10297 }
10298 return Ok(());
10299 }
10300 let f = if g {
10301 self.func_g("rms_norm_qkv_rope_append_dc_f32")
10302 } else {
10303 self.func("rms_norm_qkv_rope_append_dc_f32")
10304 };
10305 let cfg = LaunchConfig {
10306 grid_dim: (rows as u32, 1, 1),
10307 block_dim: (rms_block(), 1, 1),
10308 shared_mem_bytes: 0,
10309 };
10310 let __s_b = self.gpu.stream();
10311 let mut b = __s_b.launch_builder(&f);
10312 match ff {
10313 Some(t) => {
10314 b.arg(q0)
10315 .arg(k0)
10316 .arg(v0)
10317 .arg(wq)
10318 .arg(wk)
10319 .arg(wv)
10320 .arg(&mut *q)
10321 .arg(&mut *k)
10322 .arg(&mut *v)
10323 .arg(&nc)
10324 .arg(&rqi)
10325 .arg(&rki)
10326 .arg(pos)
10327 .arg(&nhq)
10328 .arg(&nhk)
10329 .arg(&theta_scale)
10330 .arg(&freq_scale)
10331 .arg(t)
10332 .arg(&eps)
10333 .arg(&mut *kc)
10334 .arg(&mut *vc)
10335 .arg(t_dev)
10336 .arg(&ktb)
10337 .arg(&vtb);
10338 unsafe {
10339 b.launch(cfg)?;
10340 }
10341 }
10342 None => {
10343 let null: u64 = 0;
10344 b.arg(q0)
10345 .arg(k0)
10346 .arg(v0)
10347 .arg(wq)
10348 .arg(wk)
10349 .arg(wv)
10350 .arg(&mut *q)
10351 .arg(&mut *k)
10352 .arg(&mut *v)
10353 .arg(&nc)
10354 .arg(&rqi)
10355 .arg(&rki)
10356 .arg(pos)
10357 .arg(&nhq)
10358 .arg(&nhk)
10359 .arg(&theta_scale)
10360 .arg(&freq_scale)
10361 .arg(&null)
10362 .arg(&eps)
10363 .arg(&mut *kc)
10364 .arg(&mut *vc)
10365 .arg(t_dev)
10366 .arg(&ktb)
10367 .arg(&vtb);
10368 unsafe {
10369 b.launch(cfg)?;
10370 }
10371 }
10372 }
10373 Ok(())
10374 }
10375
10376 #[allow(clippy::too_many_arguments)]
10384 pub fn rms_norm_qkv_rope_append(
10385 &self,
10386 q0: &CudaSlice<f32>,
10387 k0: &CudaSlice<f32>,
10388 v0: &CudaSlice<f32>,
10389 wq: &CudaSlice<f32>,
10390 wk: &CudaSlice<f32>,
10391 wv: &CudaSlice<f32>,
10392 q: &mut CudaSlice<f32>,
10393 k: &mut CudaSlice<f32>,
10394 v: &mut CudaSlice<f32>,
10395 head_dim: usize,
10396 n_rot: usize,
10397 rq: usize,
10398 rk: usize,
10399 pos: &CudaSlice<i32>,
10400 nh_q: usize,
10401 nh_k: usize,
10402 base: f32,
10403 freq_scale: f32,
10404 ff: Option<&CudaSlice<f32>>,
10405 eps: f32,
10406 kc: &mut CudaSlice<u8>,
10407 vc: &mut CudaSlice<u8>,
10408 t: usize,
10409 k_tok_bytes: usize,
10410 v_tok_bytes: usize,
10411 g: bool,
10412 ) -> Result<(), Box<dyn std::error::Error>> {
10413 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
10414 let rows = rq + rk + rk;
10415 let theta_scale = base.powf(-2.0 / head_dim as f32);
10416 let (nc, rqi, rki, nhq, nhk) = (
10417 head_dim as i32,
10418 rq as i32,
10419 rk as i32,
10420 nh_q as i32,
10421 nh_k as i32,
10422 );
10423 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10424 let ti = t as i32;
10425 if Self::pdl_on() && Self::pdl_wb_on() {
10426 use cudarc::driver::{DevicePtr, DevicePtrMut};
10427 let s = &self.gpu.stream();
10428 let (p0, _a0) = q0.device_ptr(s);
10429 let (p1, _a1) = k0.device_ptr(s);
10430 let (p2, _a2) = v0.device_ptr(s);
10431 let (pwq, _a3) = wq.device_ptr(s);
10432 let (pwk, _a4) = wk.device_ptr(s);
10433 let (pwv, _a5) = wv.device_ptr(s);
10434 let (pq, _a6) = q.device_ptr_mut(s);
10435 let (pk, _a7) = k.device_ptr_mut(s);
10436 let (pv, _a8) = v.device_ptr_mut(s);
10437 let (pp, _a9) = pos.device_ptr(s);
10438 let pff: u64 = match ff {
10439 Some(t) => {
10440 let (p, _gg) = t.device_ptr(s);
10441 p as u64
10442 }
10443 None => 0,
10444 };
10445 let (pkc, _a10) = kc.device_ptr_mut(s);
10446 let (pvc, _a11) = vc.device_ptr_mut(s);
10447 let mut ps = [
10448 &p0 as *const _ as *mut std::ffi::c_void,
10449 &p1 as *const _ as *mut _,
10450 &p2 as *const _ as *mut _,
10451 &pwq as *const _ as *mut _,
10452 &pwk as *const _ as *mut _,
10453 &pwv as *const _ as *mut _,
10454 &pq as *const _ as *mut _,
10455 &pk as *const _ as *mut _,
10456 &pv as *const _ as *mut _,
10457 &nc as *const _ as *mut _,
10458 &rqi as *const _ as *mut _,
10459 &rki as *const _ as *mut _,
10460 &pp as *const _ as *mut _,
10461 &nhq as *const _ as *mut _,
10462 &nhk as *const _ as *mut _,
10463 &theta_scale as *const _ as *mut _,
10464 &freq_scale as *const _ as *mut _,
10465 &pff as *const _ as *mut _,
10466 &eps as *const _ as *mut _,
10467 &pkc as *const _ as *mut _,
10468 &pvc as *const _ as *mut _,
10469 &ti as *const _ as *mut _,
10470 &ktb as *const _ as *mut _,
10471 &vtb as *const _ as *mut _,
10472 ];
10473 unsafe {
10474 self.launch_pdl_flash(
10475 g,
10476 "rms_norm_qkv_rope_append_f32",
10477 (rows as u32, 1, 1),
10478 (rms_block(), 1, 1),
10479 0,
10480 &mut ps,
10481 )?;
10482 }
10483 return Ok(());
10484 }
10485 let f = if g {
10486 self.func_g("rms_norm_qkv_rope_append_f32")
10487 } else {
10488 self.func("rms_norm_qkv_rope_append_f32")
10489 };
10490 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 __s_b = self.gpu.stream();
10496 let mut b = __s_b.launch_builder(&f);
10497 let null: u64 = 0;
10498 b.arg(q0)
10499 .arg(k0)
10500 .arg(v0)
10501 .arg(wq)
10502 .arg(wk)
10503 .arg(wv)
10504 .arg(&mut *q)
10505 .arg(&mut *k)
10506 .arg(&mut *v)
10507 .arg(&nc)
10508 .arg(&rqi)
10509 .arg(&rki)
10510 .arg(pos)
10511 .arg(&nhq)
10512 .arg(&nhk)
10513 .arg(&theta_scale)
10514 .arg(&freq_scale);
10515 match ff {
10516 Some(t) => {
10517 b.arg(t);
10518 }
10519 None => {
10520 b.arg(&null);
10521 }
10522 }
10523 b.arg(&eps)
10524 .arg(&mut *kc)
10525 .arg(&mut *vc)
10526 .arg(&ti)
10527 .arg(&ktb)
10528 .arg(&vtb);
10529 unsafe {
10530 b.launch(cfg)?;
10531 }
10532 Ok(())
10533 }
10534
10535 pub fn add_q8_1(
10536 &self,
10537 a: &CudaSlice<f32>,
10538 b: &CudaSlice<f32>,
10539 res: &mut CudaSlice<f32>,
10540 ncols: usize,
10541 nrows: usize,
10542 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10543 debug_assert!(ncols % 128 == 0);
10544 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10545 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10546 let f = self.func("add_q8_1_f32");
10547 let cfg = LaunchConfig {
10548 grid_dim: (nrows as u32, 1, 1),
10549 block_dim: (rms_block(), 1, 1),
10550 shared_mem_bytes: 0,
10551 };
10552 let nc = ncols as i32;
10553 let __s_b2 = self.gpu.stream();
10554 let mut b2 = __s_b2.launch_builder(&f);
10555 b2.arg(a)
10556 .arg(b)
10557 .arg(&mut *res)
10558 .arg(&mut out_q)
10559 .arg(&mut out_d)
10560 .arg(&nc);
10561 unsafe {
10562 b2.launch(cfg)?;
10563 }
10564 Ok((out_q, out_d))
10565 }
10566
10567 pub fn rms_pre_add_q8_1(
10571 &self,
10572 a: &CudaSlice<f32>,
10573 wa: &CudaSlice<f32>,
10574 b: &CudaSlice<f32>,
10575 res: &mut CudaSlice<f32>,
10576 ncols: usize,
10577 nrows: usize,
10578 eps: f32,
10579 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10580 debug_assert!(ncols % 128 == 0);
10581 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10582 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10583 let f = self.func("rms_pre_add_q8_1_f32");
10584 let cfg = LaunchConfig {
10585 grid_dim: (nrows as u32, 1, 1),
10586 block_dim: (rms_block(), 1, 1),
10587 shared_mem_bytes: 0,
10588 };
10589 let (nc, ep) = (ncols as i32, eps);
10590 let __s_b2 = self.gpu.stream();
10591 let mut b2 = __s_b2.launch_builder(&f);
10592 b2.arg(a)
10593 .arg(wa)
10594 .arg(b)
10595 .arg(&mut *res)
10596 .arg(&mut out_q)
10597 .arg(&mut out_d)
10598 .arg(&nc)
10599 .arg(&ep);
10600 unsafe {
10601 b2.launch(cfg)?;
10602 }
10603 Ok((out_q, out_d))
10604 }
10605
10606 pub fn l2_v2_on(ncols: usize) -> bool {
10610 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
10611 }
10612
10613 pub fn l2_norm_pp(
10614 &self,
10615 x: &CudaSlice<f32>,
10616 dst: &mut CudaSlice<f32>,
10617 dst16: Option<&mut CudaSlice<u8>>,
10618 ncols: usize,
10619 nrows: usize,
10620 eps: f32,
10621 ) -> Result<(), Box<dyn std::error::Error>> {
10622 if Self::l2_v2_on(ncols) {
10623 let f = self.func("l2_norm_pp_v2_f32");
10624 let rows_per_block = 8u32; let cfg = LaunchConfig {
10626 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
10627 block_dim: (256, 1, 1),
10628 shared_mem_bytes: 0,
10629 };
10630 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
10631 let d16: u64 = match dst16 {
10633 Some(d) => self.addr_u8(d),
10634 None => 0,
10635 };
10636 let __s_b = self.gpu.stream();
10637 let mut b = __s_b.launch_builder(&f);
10638 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
10639 unsafe {
10640 b.launch(cfg)?;
10641 }
10642 return Ok(());
10643 }
10644 self.l2_norm(x, dst, ncols, nrows, eps)
10645 }
10646
10647 pub fn l2_norm(
10648 &self,
10649 x: &CudaSlice<f32>,
10650 dst: &mut CudaSlice<f32>,
10651 ncols: usize,
10652 nrows: usize,
10653 eps: f32,
10654 ) -> Result<(), Box<dyn std::error::Error>> {
10655 let f = self.func("l2_norm_f32");
10656 let cfg = LaunchConfig {
10657 grid_dim: (nrows as u32, 1, 1),
10658 block_dim: (256, 1, 1),
10659 shared_mem_bytes: 0,
10660 };
10661 let (nc, e) = (ncols as i32, eps);
10662 let __s_b = self.gpu.stream();
10663 let mut b = __s_b.launch_builder(&f);
10664 b.arg(x).arg(dst).arg(&nc).arg(&e);
10665 unsafe {
10666 b.launch(cfg)?;
10667 }
10668 Ok(())
10669 }
10670
10671 pub fn l2_norm_decode(
10677 &self,
10678 x: &CudaSlice<f32>,
10679 dst: &mut CudaSlice<f32>,
10680 ncols: usize,
10681 nrows: usize,
10682 eps: f32,
10683 ) -> Result<(), Box<dyn std::error::Error>> {
10684 let f = self.func("l2_norm_f32");
10685 let cfg = LaunchConfig {
10686 grid_dim: (nrows as u32, 1, 1),
10687 block_dim: (32, 1, 1),
10688 shared_mem_bytes: 0,
10689 };
10690 let (nc, e) = (ncols as i32, eps);
10691 let __s_b = self.gpu.stream();
10692 let mut b = __s_b.launch_builder(&f);
10693 b.arg(x).arg(dst).arg(&nc).arg(&e);
10694 unsafe {
10695 b.launch(cfg)?;
10696 }
10697 Ok(())
10698 }
10699
10700 pub fn rope_neox(
10702 &self,
10703 x: &mut CudaSlice<f32>,
10704 pos: &CudaSlice<i32>,
10705 head_dim: usize,
10706 n_dims: usize,
10707 n_heads: usize,
10708 n_tokens: usize,
10709 freq_base: f32,
10710 freq_scale: f32,
10711 ) -> Result<(), Box<dyn std::error::Error>> {
10712 let f = self.func("rope_neox_f32");
10713 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10714 let grid = (n_heads * n_tokens) as u32;
10715 let cfg = LaunchConfig {
10716 grid_dim: (grid, 1, 1),
10717 block_dim: ((head_dim / 2) as u32, 1, 1),
10718 shared_mem_bytes: 0,
10719 };
10720 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10721 let __s_b = self.gpu.stream();
10722 let mut b = __s_b.launch_builder(&f);
10723 b.arg(x)
10724 .arg(pos)
10725 .arg(&hd)
10726 .arg(&nd)
10727 .arg(&nh)
10728 .arg(&theta_scale)
10729 .arg(&freq_scale);
10730 unsafe {
10731 b.launch(cfg)?;
10732 }
10733 Ok(())
10734 }
10735
10736 pub fn rope_neox_ff(
10738 &self,
10739 x: &mut CudaSlice<f32>,
10740 pos: &CudaSlice<i32>,
10741 head_dim: usize,
10742 n_dims: usize,
10743 n_heads: usize,
10744 n_tokens: usize,
10745 freq_base: f32,
10746 freq_scale: f32,
10747 ff: &CudaSlice<f32>,
10748 ) -> Result<(), Box<dyn std::error::Error>> {
10749 let f = self.func("rope_neox_ff_f32");
10750 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10751 let grid = (n_heads * n_tokens) as u32;
10752 let cfg = LaunchConfig {
10753 grid_dim: (grid, 1, 1),
10754 block_dim: ((head_dim / 2) as u32, 1, 1),
10755 shared_mem_bytes: 0,
10756 };
10757 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
10758 let __s_b = self.gpu.stream();
10759 let mut b = __s_b.launch_builder(&f);
10760 b.arg(x)
10761 .arg(pos)
10762 .arg(&hd)
10763 .arg(&nd)
10764 .arg(&nh)
10765 .arg(&theta_scale)
10766 .arg(&freq_scale)
10767 .arg(ff);
10768 unsafe {
10769 b.launch(cfg)?;
10770 }
10771 Ok(())
10772 }
10773
10774 #[allow(clippy::too_many_arguments)]
10776 pub fn rope_neox2(
10777 &self,
10778 q: &mut CudaSlice<f32>,
10779 k: &mut CudaSlice<f32>,
10780 pos: &CudaSlice<i32>,
10781 head_dim: usize,
10782 n_dims: usize,
10783 nh_q: usize,
10784 nh_k: usize,
10785 n_tokens: usize,
10786 freq_base: f32,
10787 freq_scale: f32,
10788 ff: Option<&CudaSlice<f32>>,
10789 ) -> Result<(), Box<dyn std::error::Error>> {
10790 let f = self.func("rope_neox2_f32");
10791 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
10792 let grid = ((nh_q + nh_k) * n_tokens) as u32;
10793 let cfg = LaunchConfig {
10794 grid_dim: (grid, 1, 1),
10795 block_dim: ((head_dim / 2) as u32, 1, 1),
10796 shared_mem_bytes: 0,
10797 };
10798 let (hd, nd, nq, nk, nt) = (
10799 head_dim as i32,
10800 n_dims as i32,
10801 nh_q as i32,
10802 nh_k as i32,
10803 n_tokens as i32,
10804 );
10805 let __s_b = self.gpu.stream();
10806 let mut b = __s_b.launch_builder(&f);
10807 b.arg(q)
10808 .arg(k)
10809 .arg(pos)
10810 .arg(&hd)
10811 .arg(&nd)
10812 .arg(&nq)
10813 .arg(&nk)
10814 .arg(&nt)
10815 .arg(&theta_scale)
10816 .arg(&freq_scale);
10817 match ff {
10818 Some(ffv) => {
10819 b.arg(ffv);
10820 unsafe {
10821 b.launch(cfg)?;
10822 }
10823 }
10824 None => {
10825 let null: u64 = 0;
10826 b.arg(&null);
10827 unsafe {
10828 b.launch(cfg)?;
10829 }
10830 }
10831 }
10832 Ok(())
10833 }
10834
10835 pub fn gelu_tanh_mul(
10837 &self,
10838 gate: &CudaSlice<f32>,
10839 up: &CudaSlice<f32>,
10840 dst: &mut CudaSlice<f32>,
10841 n: usize,
10842 ) -> Result<(), Box<dyn std::error::Error>> {
10843 let f = self.func("gelu_tanh_mul_f32");
10844 let cfg = LaunchConfig::for_num_elems(n as u32);
10845 let ni = n as i32;
10846 let __s_b = self.gpu.stream();
10847 let mut b = __s_b.launch_builder(&f);
10848 b.arg(gate).arg(up).arg(dst).arg(&ni);
10849 unsafe {
10850 b.launch(cfg)?;
10851 }
10852 Ok(())
10853 }
10854
10855 pub fn silu_mul(
10856 &self,
10857 gate: &CudaSlice<f32>,
10858 up: &CudaSlice<f32>,
10859 dst: &mut CudaSlice<f32>,
10860 n: usize,
10861 ) -> Result<(), Box<dyn std::error::Error>> {
10862 let f = self.func("silu_mul_f32");
10863 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10865 let ni = n as i32;
10866 let __s_b = self.gpu.stream();
10867 let mut b = __s_b.launch_builder(&f);
10868 b.arg(gate).arg(up).arg(dst).arg(&ni);
10869 unsafe {
10870 b.launch(cfg)?;
10871 }
10872 Ok(())
10873 }
10874
10875 pub fn silu_mul_host_expf(
10877 &self,
10878 gate: &CudaSlice<f32>,
10879 up: &CudaSlice<f32>,
10880 dst: &mut CudaSlice<f32>,
10881 n: usize,
10882 ) -> Result<(), Box<dyn std::error::Error>> {
10883 let f = self.func("silu_mul_host_expf_f32");
10884 let cfg = LaunchConfig::for_num_elems(n as u32);
10885 let ni = n as i32;
10886 let __s_b = self.gpu.stream();
10887 let mut b = __s_b.launch_builder(&f);
10888 b.arg(gate).arg(up).arg(dst).arg(&ni);
10889 unsafe {
10890 b.launch(cfg)?;
10891 }
10892 Ok(())
10893 }
10894
10895 pub fn silu_clamped_mul_host_expf(
10897 &self,
10898 gate: &CudaSlice<f32>,
10899 up: &CudaSlice<f32>,
10900 limit: f32,
10901 dst: &mut CudaSlice<f32>,
10902 n: usize,
10903 ) -> Result<(), Box<dyn std::error::Error>> {
10904 if !limit.is_finite() || limit <= 0.0 {
10905 return Err(
10906 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
10907 );
10908 }
10909 let f = self.func("silu_clamped_mul_host_expf_f32");
10910 let cfg = LaunchConfig::for_num_elems(n as u32);
10911 let ni = n as i32;
10912 let __s_b = self.gpu.stream();
10913 let mut b = __s_b.launch_builder(&f);
10914 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
10915 unsafe {
10916 b.launch(cfg)?;
10917 }
10918 Ok(())
10919 }
10920
10921 pub fn silu_mul_f16out(
10924 &self,
10925 gate: &CudaSlice<f32>,
10926 up: &CudaSlice<f32>,
10927 dst: &mut CudaSlice<f32>,
10928 dst16: &mut CudaSlice<u8>,
10929 n: usize,
10930 ) -> Result<(), Box<dyn std::error::Error>> {
10931 let f = self.func("silu_mul_f16out_f32");
10932 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10933 let ni = n as i32;
10934 let __s_b = self.gpu.stream();
10935 let mut b = __s_b.launch_builder(&f);
10936 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
10937 unsafe {
10938 b.launch(cfg)?;
10939 }
10940 Ok(())
10941 }
10942
10943 pub fn silu_mul_scaled(
10950 &self,
10951 gate: &CudaSlice<f32>,
10952 up: &CudaSlice<f32>,
10953 gs: f32,
10954 us: f32,
10955 dst: &mut CudaSlice<f32>,
10956 n: usize,
10957 ) -> Result<(), Box<dyn std::error::Error>> {
10958 let f = self.func("silu_mul_scaled_f32");
10959 let cfg = LaunchConfig::for_num_elems(n as u32);
10960 let ni = n as i32;
10961 let (gsf, usf) = (gs, us);
10962 let __s_b = self.gpu.stream();
10963 let mut b = __s_b.launch_builder(&f);
10964 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
10965 unsafe {
10966 b.launch(cfg)?;
10967 }
10968 Ok(())
10969 }
10970
10971 #[allow(clippy::too_many_arguments)]
10975 pub fn swigluoai_mul_scaled(
10976 &self,
10977 gate: &CudaSlice<f32>,
10978 up: &CudaSlice<f32>,
10979 gs: f32,
10980 us: f32,
10981 alpha: f32,
10982 limit: f32,
10983 dst: &mut CudaSlice<f32>,
10984 n: usize,
10985 ) -> Result<(), Box<dyn std::error::Error>> {
10986 let f = self.func("swigluoai_mul_scaled_f32");
10987 let cfg = LaunchConfig::for_num_elems(n as u32);
10988 let ni = n as i32;
10989 let __s_b = self.gpu.stream();
10990 let mut b = __s_b.launch_builder(&f);
10991 b.arg(gate)
10992 .arg(up)
10993 .arg(&gs)
10994 .arg(&us)
10995 .arg(&alpha)
10996 .arg(&limit)
10997 .arg(dst)
10998 .arg(&ni);
10999 unsafe {
11000 b.launch(cfg)?;
11001 }
11002 Ok(())
11003 }
11004
11005 pub fn silu_mul_scaled_q8_1(
11013 &self,
11014 gate: &CudaSlice<f32>,
11015 up: &CudaSlice<f32>,
11016 gs: f32,
11017 us: f32,
11018 n: usize,
11019 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11020 let f = self.func("silu_mul_scaled_q8_1");
11021 let nblk = n / 32;
11022 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);
11026 let (gsf, usf, ni) = (gs, us, n as i32);
11027 let __s_b = self.gpu.stream();
11028 let mut b = __s_b.launch_builder(&f);
11029 b.arg(gate)
11030 .arg(up)
11031 .arg(&gsf)
11032 .arg(&usf)
11033 .arg(&mut aq)
11034 .arg(&mut ad)
11035 .arg(&ni);
11036 unsafe {
11037 b.launch(cfg)?;
11038 }
11039 Ok((aq, ad))
11040 }
11041
11042 pub fn add(
11043 &self,
11044 a: &CudaSlice<f32>,
11045 b_in: &CudaSlice<f32>,
11046 dst: &mut CudaSlice<f32>,
11047 n: usize,
11048 ) -> Result<(), Box<dyn std::error::Error>> {
11049 let f = self.func("add_f32");
11050 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11052 let ni = n as i32;
11053 let __s_bld = self.gpu.stream();
11054 let mut bld = __s_bld.launch_builder(&f);
11055 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11056 unsafe {
11057 bld.launch(cfg)?;
11058 }
11059 Ok(())
11060 }
11061
11062 pub fn mul(
11063 &self,
11064 a: &CudaSlice<f32>,
11065 b_in: &CudaSlice<f32>,
11066 dst: &mut CudaSlice<f32>,
11067 n: usize,
11068 ) -> Result<(), Box<dyn std::error::Error>> {
11069 let f = self.func("mul_f32");
11070 let cfg = LaunchConfig::for_num_elems(n as u32);
11071 let ni = n as i32;
11072 let __s_bld = self.gpu.stream();
11073 let mut bld = __s_bld.launch_builder(&f);
11074 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
11075 unsafe {
11076 bld.launch(cfg)?;
11077 }
11078 Ok(())
11079 }
11080
11081 pub fn matmul(
11084 &self,
11085 w: &crate::model::GpuTensor,
11086 x: &CudaSlice<f32>,
11087 m: usize,
11088 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11089 use crate::model::GpuTensor;
11090 let in_f = w.in_features();
11091 let out_f = w.out_features();
11092 #[allow(non_snake_case)]
11100 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
11103 usize::MAX
11104 } else {
11105 16usize
11106 };
11107
11108 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
11133 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
11134 return Ok(y);
11135 }
11136 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
11143 return Ok(y);
11144 }
11145 if let Some(y) = self.try_f16_gemm(w, x, m)? {
11148 return Ok(y);
11149 }
11150 }
11151 if let GpuTensor::Quant { qtype, .. } = w {
11166 if *qtype == QT_F8_E4M3_BLK {
11167 if m >= GEMM_M_THRESHOLD {
11168 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
11169 return Ok(y);
11170 }
11171 }
11172 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11173 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11174 return Ok(y);
11175 }
11176 }
11177 }
11178 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
11179 return self.qmatvec_mmq(w, x, m);
11180 }
11181 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
11182 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11183 return self.qmatvec_gemm(w, &aq, &ad, m);
11184 }
11185 if m >= GEMM_M_THRESHOLD {
11188 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
11189 return Ok(y);
11190 }
11191 }
11192 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
11196 if m == 1 && fast {
11201 if let GpuTensor::Quant {
11202 bytes,
11203 qtype,
11204 row_bytes,
11205 rp,
11206 rp4,
11207 scale,
11208 ..
11209 } = w
11210 {
11211 if self.mmvq_supports(*qtype) {
11212 let (bytes, rp) = match rp4 {
11216 Some(m4) => (m4, true),
11217 None => (bytes, *rp),
11218 };
11219 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11220 return self.qmatvec_mmvq(
11221 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
11222 );
11223 }
11224 }
11225 }
11226 if (2..=16).contains(&m)
11242 && fast
11243 && std::env::var("MEMRA_NO_BATCHED").is_err()
11244 && (m <= 4 || Self::b8_enabled())
11245 {
11246 let m_ok = m <= 8
11256 || matches!(w, GpuTensor::Quant { qtype, .. }
11257 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
11258 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
11259 if m_ok {
11260 if let GpuTensor::Quant {
11261 bytes,
11262 qtype,
11263 row_bytes,
11264 rp,
11265 rp4,
11266 ..
11267 } = w
11268 {
11269 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
11270 let (bytes, rp) = match rp4 {
11271 Some(m4) => (m4, true),
11272 None => (bytes, *rp),
11273 };
11274 let mcols = Self::batched_mcols(m);
11275 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11276 let mut y = self.qmatvec_mmvq_batched(
11277 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
11278 )?;
11279 if let GpuTensor::Quant { scale, .. } = w {
11280 if *scale != 1.0 {
11281 self.scale_inplace(&mut y, *scale, m * out_f)?;
11282 }
11283 }
11284 return Ok(y);
11285 }
11286 }
11287 }
11288 }
11289 if fast {
11295 if let GpuTensor::Quant {
11296 bytes,
11297 qtype,
11298 row_bytes,
11299 scale,
11300 ..
11301 } = w
11302 {
11303 if *qtype == QT_F8_E4M3 {
11304 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11305 return self.qmatvec_mmvq(
11306 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
11307 );
11308 }
11309 }
11310 }
11311 let mut y = match w {
11312 GpuTensor::Quant {
11313 bytes,
11314 qtype,
11315 row_bytes,
11316 ..
11317 } if fast && *qtype == QT_Q8_0 => {
11318 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11319 }
11320 GpuTensor::Quant {
11321 bytes,
11322 qtype,
11323 row_bytes,
11324 ..
11325 } if fast && *qtype == QT_Q4_K => {
11326 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11327 }
11328 GpuTensor::Quant {
11329 bytes,
11330 qtype,
11331 row_bytes,
11332 ..
11333 } if fast && *qtype == QT_Q6_K => {
11334 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11335 }
11336 GpuTensor::Quant {
11337 bytes,
11338 qtype,
11339 row_bytes,
11340 ..
11341 } if fast && *qtype == QT_Q5_K => {
11342 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11343 }
11344 GpuTensor::Quant {
11345 bytes,
11346 qtype,
11347 row_bytes,
11348 ..
11349 } if fast && *qtype == QT_Q3_K => {
11350 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11351 }
11352 GpuTensor::Quant {
11353 bytes,
11354 qtype,
11355 row_bytes,
11356 rp,
11357 ..
11358 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
11359 if *rp {
11360 "qmatvec_nvfp4_dp4a_rp"
11361 } else {
11362 "qmatvec_nvfp4_dp4a"
11363 },
11364 &bytes.slice(0..bytes.len()),
11365 x,
11366 m,
11367 in_f,
11368 out_f,
11369 *row_bytes,
11370 )?,
11371 GpuTensor::Quant {
11375 bytes,
11376 qtype,
11377 row_bytes,
11378 ..
11379 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
11380 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
11381 }
11382 GpuTensor::Quant {
11387 bytes,
11388 qtype,
11389 row_bytes,
11390 rp,
11391 ..
11392 } =>
11393 {
11396 self.qmatvec(
11397 bytes,
11398 x,
11399 m,
11400 in_f,
11401 out_f,
11402 if *rp && *qtype == QT_NVFP4 {
11403 QT_NVFP4_RP
11404 } else {
11405 *qtype
11406 },
11407 *row_bytes,
11408 )?
11409 }
11410 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
11411 GpuTensor::FloatBf16 { data, .. } => {
11414 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
11415 }
11416 };
11417 if let GpuTensor::Quant { scale, .. } = w {
11419 if *scale != 1.0 {
11420 self.scale_inplace(&mut y, *scale, m * out_f)?;
11421 }
11422 }
11423 Ok(y)
11424 }
11425
11426 pub fn stage_a_raw_needed() -> bool {
11436 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11437 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
11438 }
11439
11440 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
11443 use crate::model::GpuTensor;
11444 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
11445 return false;
11446 }
11447 match w {
11448 GpuTensor::Quant { qtype, .. } => {
11455 matches!(
11456 *qtype,
11457 QT_Q8_0
11458 | QT_Q4_K
11459 | QT_Q6_K
11460 | QT_Q5_K
11461 | QT_Q3_K
11462 | QT_NVFP4
11463 | QT_F8_E4M3
11464 | QT_F8_E4M3_BLK
11465 | QT_Q4_0
11466 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
11467 }
11468 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
11469 }
11470 }
11471
11472 pub fn matmul_pre(
11477 &self,
11478 w: &crate::model::GpuTensor,
11479 aq: &CudaSlice<i8>,
11480 ad: &CudaSlice<f32>,
11481 x_fallback: &CudaSlice<f32>,
11482 m: usize,
11483 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11484 use crate::model::GpuTensor;
11485 let x_raw_ok = x_fallback.len() >= m * w.in_features();
11491 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11494 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
11495 return Ok(y);
11496 }
11497 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
11500 return Ok(y);
11501 }
11502 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
11504 return Ok(y);
11505 }
11506 }
11507 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11513 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
11514 return Ok(y);
11515 }
11516 }
11517 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11518 return Ok(y);
11519 }
11520 if m >= 16
11525 && w.out_features() >= 128
11526 && self.mmq_supports(w)
11527 && !self.verify_exact_on()
11528 && x_raw_ok
11529 {
11530 return self.qmatvec_mmq(w, x_fallback, m);
11531 }
11532 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
11535 if let Some(y) =
11536 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
11537 {
11538 return Ok(y);
11539 }
11540 }
11541 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
11544 return self.qmatvec_gemm(w, aq, ad, m);
11545 }
11546 if !self.uses_q8_1_fast(w) {
11565 if !x_raw_ok {
11566 return Err(format!(
11567 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
11568 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
11569 dropped the f32, so there is nothing to fall back to — pass the real f32 \
11570 activation (see Engine::rms_norm_decode, which is bit-identical to \
11571 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
11572 x_fallback.len(),
11573 m,
11574 w.in_features(),
11575 m * w.in_features()
11576 )
11577 .into());
11578 }
11579 return self.matmul(w, x_fallback, m);
11580 }
11581 let in_f = w.in_features();
11582 let out_f = w.out_features();
11583 let (bytes, qtype, row_bytes, scale, rp) = match w {
11584 GpuTensor::Quant {
11585 bytes,
11586 qtype,
11587 row_bytes,
11588 scale,
11589 rp,
11590 ..
11591 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11592 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
11593 };
11594 let (mbytes, mrp) = match w {
11597 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11598 _ => (bytes, rp),
11599 };
11600 if m == 1 && self.mmvq_supports(qtype) {
11604 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
11605 }
11606 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11619 && std::env::var("MEMRA_NO_BATCHED").is_err()
11620 && (m <= 4 || Self::b8_enabled())
11621 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
11625 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
11626 {
11627 let mcols = Self::batched_mcols(m);
11628 return self.qmatvec_mmvq_batched(
11629 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
11630 );
11631 }
11632 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
11638 let (b2, r2) = if qtype == QT_Q4_0 {
11639 (mbytes, mrp)
11640 } else {
11641 (bytes, rp)
11642 };
11643 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
11644 }
11645 let name = match qtype {
11646 QT_Q8_0 => "qmatvec_q8_0_dp4a",
11647 QT_Q4_K => "qmatvec_q4_K_dp4a",
11648 QT_Q6_K => "qmatvec_q6_K_dp4a",
11649 QT_Q5_K => "qmatvec_q5_K_dp4a",
11650 QT_Q3_K => "qmatvec_q3_K_dp4a",
11651 QT_NVFP4 => {
11652 if rp {
11653 "qmatvec_nvfp4_dp4a_rp"
11654 } else {
11655 "qmatvec_nvfp4_dp4a"
11656 }
11657 }
11658 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
11659 _ => unreachable!(),
11660 };
11661 let f = self.func(name);
11662 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
11664 grid_dim: (out_f as u32, m as u32, 1),
11665 block_dim: (128, 1, 1),
11666 shared_mem_bytes: 0,
11667 };
11668 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
11669 let __s_b = self.gpu.stream();
11670 let mut b = __s_b.launch_builder(&f);
11671 b.arg(bytes)
11672 .arg(aq)
11673 .arg(ad)
11674 .arg(&mut y)
11675 .arg(&inf)
11676 .arg(&outf)
11677 .arg(&mi)
11678 .arg(&rb);
11679 unsafe {
11680 b.launch(cfg)?;
11681 }
11682 if scale != 1.0 {
11683 self.scale_inplace(&mut y, scale, m * out_f)?;
11684 }
11685 Ok(y)
11686 }
11687
11688 pub fn matmul_decode_exact(
11696 &self,
11697 w: &crate::model::GpuTensor,
11698 x: &CudaSlice<f32>,
11699 m: usize,
11700 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11701 use crate::model::GpuTensor;
11702 if let GpuTensor::Float { data, .. } = w {
11710 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
11711 }
11712 if let GpuTensor::FloatBf16 { data, .. } = w {
11715 let (in_f, out_f) = (w.in_features(), w.out_features());
11716 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
11717 }
11718 if !self.uses_q8_1_fast(w) {
11719 return self.matmul(w, x, m);
11720 }
11721 let in_f = w.in_features();
11722 let out_f = w.out_features();
11723 let (bytes, qtype, row_bytes, scale, rp) = match w {
11724 GpuTensor::Quant {
11725 bytes,
11726 qtype,
11727 row_bytes,
11728 scale,
11729 rp,
11730 ..
11731 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11732 _ => return self.matmul(w, x, m),
11733 };
11734 let (bytes, rp) = match w {
11737 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11738 _ => (bytes, rp),
11739 };
11740 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
11741 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
11745 return Ok(y);
11746 }
11747 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
11756 && std::env::var("MEMRA_NO_BATCHED").is_err()
11757 && (m <= 4 || Self::b8_enabled())
11758 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
11761 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
11762 {
11763 let mcols = Self::batched_mcols(m);
11764 return self.qmatvec_mmvq_batched(
11765 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11766 );
11767 }
11768 if self.mmvq_supports(qtype) {
11769 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11772 }
11773 self.matmul_pre(w, &aq, &ad, x, m)
11776 }
11777
11778 pub fn matmul_decode_exact_pre(
11788 &self,
11789 w: &crate::model::GpuTensor,
11790 aq: &CudaSlice<i8>,
11791 ad: &CudaSlice<f32>,
11792 m: usize,
11793 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11794 use crate::model::GpuTensor;
11795 debug_assert!(
11796 self.uses_q8_1_fast(w),
11797 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
11798 );
11799 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
11801 return Ok(y);
11802 }
11803 let in_f = w.in_features();
11804 let out_f = w.out_features();
11805 let (bytes, qtype, row_bytes, scale, rp) = match w {
11806 GpuTensor::Quant {
11807 bytes,
11808 qtype,
11809 row_bytes,
11810 scale,
11811 rp,
11812 ..
11813 } => (bytes, *qtype, *row_bytes, *scale, *rp),
11814 _ => {
11815 return Err(
11816 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
11817 );
11818 }
11819 };
11820 let (bytes, rp) = match w {
11822 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
11823 _ => (bytes, rp),
11824 };
11825 if (2..=16).contains(&m)
11827 && self.batched_supports(qtype)
11828 && self.mmvq_supports(qtype)
11829 && std::env::var("MEMRA_NO_BATCHED").is_err()
11830 && (m <= 4 || Self::b8_enabled())
11831 && (m <= 8
11832 || qtype == QT_Q4_0
11833 || qtype == QT_Q6_K
11834 || qtype == QT_F8_E4M3
11835 || qtype == QT_NVFP4
11836 || qtype == QT_Q4_K
11837 || qtype == QT_Q5_K
11838 || qtype == QT_Q8_0)
11839 {
11840 let mcols = Self::batched_mcols(m);
11841 return self.qmatvec_mmvq_batched(
11842 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
11843 );
11844 }
11845 if self.mmvq_supports(qtype) {
11846 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
11847 }
11848 let x0 = self.zeros(0)?;
11851 self.matmul_pre(w, aq, ad, &x0, m)
11852 }
11853
11854 pub fn matmul_decode_exact_dual_pre(
11863 &self,
11864 w0: &crate::model::GpuTensor,
11865 w1: &crate::model::GpuTensor,
11866 aq: &CudaSlice<i8>,
11867 ad: &CudaSlice<f32>,
11868 m: usize,
11869 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
11870 {
11871 use crate::model::GpuTensor;
11872 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11873 let on = *ON.get_or_init(|| {
11874 std::env::var("MEMRA_SPEC_DUAL_T")
11875 .map(|v| v != "0")
11876 .unwrap_or(true)
11877 });
11878 if !on
11879 || !(2..=7).contains(&m)
11880 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11881 || !self.uses_q8_1_fast(w0)
11882 || !self.uses_q8_1_fast(w1)
11883 {
11884 return Ok(None);
11885 }
11886 if !self.mmvq_supports(QT_NVFP4) {
11891 return Ok(None);
11892 }
11893 let (in_f, out_f) = (w0.in_features(), w0.out_features());
11894 if w1.in_features() != in_f || w1.out_features() != out_f {
11895 return Ok(None);
11896 }
11897 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
11898 (
11899 GpuTensor::Quant {
11900 bytes: b0,
11901 qtype: q0,
11902 row_bytes: rb0,
11903 scale: s0,
11904 rp: rp0,
11905 rp4: None,
11906 ..
11907 },
11908 GpuTensor::Quant {
11909 bytes: b1,
11910 qtype: q1,
11911 row_bytes: rb1,
11912 scale: s1,
11913 rp: rp1,
11914 rp4: None,
11915 ..
11916 },
11917 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
11918 (b0, b1, *rb0, *s0, *s1, *rp0)
11919 }
11920 _ => return Ok(None),
11921 };
11922 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
11925 {
11926 return Ok(None);
11927 }
11928 let (y0, y1) =
11929 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
11930 Ok(Some(((y0, s0), (y1, s1))))
11931 }
11932
11933 pub fn matmul_decode_exact_group4_pre(
11945 &self,
11946 ws: [&crate::model::GpuTensor; 4],
11947 aq: &CudaSlice<i8>,
11948 ad: &CudaSlice<f32>,
11949 m: usize,
11950 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11951 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11952 let on = *ON.get_or_init(|| {
11953 std::env::var("MEMRA_TK_GDN_GROUP")
11954 .map(|v| v != "0")
11955 .unwrap_or(true)
11956 });
11957 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
11958 }
11959
11960 pub fn matmul_decode_exact_group3_pre(
11965 &self,
11966 ws: [&crate::model::GpuTensor; 3],
11967 aq: &CudaSlice<i8>,
11968 ad: &CudaSlice<f32>,
11969 m: usize,
11970 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11971 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11972 let on = *ON.get_or_init(|| {
11973 std::env::var("MEMRA_TK_FA_GROUP")
11974 .map(|v| v != "0")
11975 .unwrap_or(true)
11976 });
11977 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
11978 }
11979
11980 fn matmul_decode_exact_group_pre(
11984 &self,
11985 ws: &[&crate::model::GpuTensor],
11986 aq: &CudaSlice<i8>,
11987 ad: &CudaSlice<f32>,
11988 m: usize,
11989 on: bool,
11990 tag: &'static str,
11991 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
11992 use crate::model::GpuTensor;
11993 if !on
11994 || !(2..=16).contains(&m)
11995 || std::env::var("MEMRA_NO_BATCHED").is_ok()
11996 || (m > 4 && !Self::b8_enabled())
11997 || !self.mmvq_supports(QT_NVFP4)
11998 || !self.batched_supports(QT_NVFP4)
11999 {
12000 return Ok(None);
12001 }
12002 let in_f = ws[0].in_features();
12003 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
12004 for w in ws {
12005 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
12006 return Ok(None);
12007 }
12008 match w {
12009 GpuTensor::Quant {
12010 bytes,
12011 qtype,
12012 scale,
12013 rp: true,
12014 rp4: None,
12015 ..
12016 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
12017 parts.push((bytes, w.out_features(), *scale));
12018 }
12019 _ => return Ok(None),
12020 }
12021 }
12022 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12024 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
12025 let mcols = if (5..=7).contains(&m) && b567 {
12026 m
12027 } else {
12028 Self::batched_mcols(m)
12029 };
12030 let kname: &'static str = match mcols {
12031 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
12032 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
12033 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
12034 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
12035 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
12036 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
12037 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
12038 _ => return Ok(None),
12039 };
12040 if std::env::var("MEMRA_DEBUG").is_ok() {
12043 use std::sync::Mutex;
12044 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
12045 let mut seen = SEEN.lock().unwrap();
12046 if !seen.contains(&tag) {
12047 seen.push(tag);
12048 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
12049 }
12050 }
12051 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
12054 let three = parts.len() == 3;
12055 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
12056 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
12057 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
12058 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
12061 let cfg = LaunchConfig {
12062 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
12063 block_dim: (32, ROWS_PER_BLOCK, 1),
12064 shared_mem_bytes: 0,
12065 };
12066 let (inf, mi) = (in_f as i32, m as i32);
12067 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
12068 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
12069 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
12070 let s3 = if three { 1.0f32 } else { parts[3].2 };
12071 let w3 = if three { parts[0].0 } else { parts[3].0 };
12072 let f = self.func(kname);
12073 let __s_b = self.gpu.stream();
12074 let mut b = __s_b.launch_builder(&f);
12075 b.arg(parts[0].0)
12076 .arg(parts[1].0)
12077 .arg(parts[2].0)
12078 .arg(w3)
12079 .arg(aq)
12080 .arg(ad)
12081 .arg(&mut y0)
12082 .arg(&mut y1)
12083 .arg(&mut y2)
12084 .arg(&mut y3)
12085 .arg(&inf)
12086 .arg(&n0)
12087 .arg(&n1)
12088 .arg(&n2)
12089 .arg(&n3)
12090 .arg(&mi)
12091 .arg(&s0)
12092 .arg(&s1)
12093 .arg(&s2)
12094 .arg(&s3);
12095 unsafe {
12096 b.launch(cfg)?;
12097 }
12098 Ok(Some(if three {
12099 vec![y0, y1, y2]
12100 } else {
12101 vec![y0, y1, y2, y3]
12102 }))
12103 }
12104
12105 pub fn matmul_decode_exact_dual(
12121 &self,
12122 w0: &crate::model::GpuTensor,
12123 w1: &crate::model::GpuTensor,
12124 x: &CudaSlice<f32>,
12125 m: usize,
12126 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12127 use crate::model::GpuTensor;
12128 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12129 let on = *ON.get_or_init(|| {
12130 std::env::var("MEMRA_SPEC_DUAL_T")
12131 .map(|v| v != "0")
12132 .unwrap_or(true)
12133 });
12134 if !on
12135 || !(2..=4).contains(&m)
12136 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12137 || !self.uses_q8_1_fast(w0)
12138 || !self.uses_q8_1_fast(w1)
12139 {
12140 return Ok(None);
12141 }
12142 if !self.mmvq_supports(QT_NVFP4) {
12147 return Ok(None);
12148 }
12149 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12150 if w1.in_features() != in_f || w1.out_features() != out_f {
12151 return Ok(None);
12152 }
12153 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
12154 (
12155 GpuTensor::Quant {
12156 bytes: b0,
12157 qtype: q0,
12158 row_bytes: rb0,
12159 scale: s0,
12160 rp: rp0,
12161 rp4: None,
12162 ..
12163 },
12164 GpuTensor::Quant {
12165 bytes: b1,
12166 qtype: q1,
12167 row_bytes: rb1,
12168 scale: s1,
12169 rp: rp1,
12170 rp4: None,
12171 ..
12172 },
12173 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
12174 (b0, b1, *rb0, *s0, *s1, *rp0)
12175 }
12176 _ => return Ok(None),
12177 };
12178 if std::env::var("MEMRA_DEBUG").is_ok() {
12181 static ONCE: std::sync::Once = std::sync::Once::new();
12182 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
12183 }
12184 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12185 let (y0, y1) =
12186 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
12187 let mut y0 = y0;
12188 let mut y1 = y1;
12189 if s0 != 1.0 {
12190 self.scale_inplace(&mut y0, s0, m * out_f)?;
12191 }
12192 if s1 != 1.0 {
12193 self.scale_inplace(&mut y1, s1, m * out_f)?;
12194 }
12195 Ok(Some((y0, y1)))
12196 }
12197
12198 #[allow(clippy::too_many_arguments)]
12203 pub fn qmatvec_batched_dual_raw(
12204 &self,
12205 b0: &CudaSlice<u8>,
12206 b1: &CudaSlice<u8>,
12207 aq: &CudaSlice<i8>,
12208 ad: &CudaSlice<f32>,
12209 m: usize,
12210 in_f: usize,
12211 out_f: usize,
12212 row_bytes: usize,
12213 rp: bool,
12214 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12215 const ROWS_PER_BLOCK: u32 = 4;
12216 let mcols = Self::batched_mcols(m);
12217 let tiny_rp1 = rp
12220 && mcols == 4
12221 && out_f <= 128
12222 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
12223 let (name, rows_per_block) = if tiny_rp1 {
12224 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
12225 } else {
12226 match (mcols, rp, m) {
12227 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
12228 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
12229 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
12230 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
12231 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
12232 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
12233 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
12234 _ => {
12235 return Err(
12236 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
12237 );
12238 }
12239 }
12240 };
12241 let f = self.func(name);
12242 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
12243 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
12244 let cfg = LaunchConfig {
12245 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12246 block_dim: (32, ROWS_PER_BLOCK, 1),
12247 shared_mem_bytes: 0,
12248 };
12249 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12250 let __s_b = self.gpu.stream();
12251 let mut b = __s_b.launch_builder(&f);
12252 b.arg(b0)
12253 .arg(b1)
12254 .arg(aq)
12255 .arg(ad)
12256 .arg(&mut y0)
12257 .arg(&mut y1)
12258 .arg(&inf)
12259 .arg(&outf)
12260 .arg(&mi)
12261 .arg(&rb);
12262 unsafe {
12263 b.launch(cfg)?;
12264 }
12265 Ok((y0, y1))
12266 }
12267
12268 pub fn matmul_pre_dual_noscale(
12280 &self,
12281 w0: &crate::model::GpuTensor,
12282 w1: &crate::model::GpuTensor,
12283 aq: &CudaSlice<i8>,
12284 ad: &CudaSlice<f32>,
12285 m: usize,
12286 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12287 {
12288 use crate::model::GpuTensor;
12289 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
12290 return Ok(None);
12291 }
12292 if !self.mmvq_supports(QT_NVFP4) {
12302 return Ok(None);
12303 }
12304 let (in_f, out_f) = (w0.in_features(), w0.out_features());
12305 if w1.in_features() != in_f || w1.out_features() != out_f {
12306 return Ok(None);
12307 }
12308 let no_mirror =
12321 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
12322 if self.q8_ffn_fuse2_on()
12323 && no_mirror(w0)
12324 && no_mirror(w1)
12325 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
12326 {
12327 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
12328 return Ok(Some(((y0, 1.0), (y1, 1.0))));
12329 }
12330 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12340 let (y0, y1) =
12341 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
12342 return Ok(Some(((y0, p0.3), (y1, p1.3))));
12343 }
12344 let (b0, q0, rb0, s0, rp0) = match w0 {
12345 GpuTensor::Quant {
12346 bytes,
12347 qtype,
12348 row_bytes,
12349 scale,
12350 rp,
12351 ..
12352 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12353 _ => return Ok(None),
12354 };
12355 let (b1, q1, rb1, s1, rp1) = match w1 {
12356 GpuTensor::Quant {
12357 bytes,
12358 qtype,
12359 row_bytes,
12360 scale,
12361 rp,
12362 ..
12363 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12364 _ => return Ok(None),
12365 };
12366 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
12367 return Ok(None);
12368 }
12369 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12371 let rows_per_block = ROWS_PER_BLOCK * RPW;
12372 let f = self.func(if rp0 {
12373 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
12374 } else {
12375 "qmatvec_nvfp4_mmvq_dual_mr2"
12376 });
12377 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
12378 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
12379 let cfg = LaunchConfig {
12380 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
12381 block_dim: (32, ROWS_PER_BLOCK, 1),
12382 shared_mem_bytes: 0,
12383 };
12384 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
12385 let one = 1.0f32;
12388 let __s_b = self.gpu.stream();
12389 let mut b = __s_b.launch_builder(&f);
12390 b.arg(b0)
12391 .arg(b1)
12392 .arg(aq)
12393 .arg(ad)
12394 .arg(&mut y0)
12395 .arg(&mut y1)
12396 .arg(&inf)
12397 .arg(&outf)
12398 .arg(&mi)
12399 .arg(&rb)
12400 .arg(&one)
12401 .arg(&one);
12402 unsafe {
12403 b.launch(cfg)?;
12404 }
12405 Ok(Some(((y0, s0), (y1, s1))))
12406 }
12407
12408 #[allow(clippy::too_many_arguments)]
12416 pub fn matmul_nvfp4_fused3(
12417 &self,
12418 w0: &crate::model::GpuTensor,
12419 w1: &crate::model::GpuTensor,
12420 w2: &crate::model::GpuTensor,
12421 aq: &CudaSlice<i8>,
12422 ad: &CudaSlice<f32>,
12423 m: usize,
12424 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
12425 {
12426 use crate::model::GpuTensor;
12427 if !self.mmvq_supports(QT_NVFP4)
12434 || !self.uses_q8_1_fast(w0)
12435 || !self.uses_q8_1_fast(w1)
12436 || !self.uses_q8_1_fast(w2)
12437 {
12438 return Ok(None);
12439 }
12440 if (9..=16).contains(&m) {
12443 return Ok(
12444 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
12445 Some(mut ys) => {
12446 let y2 = ys.pop().unwrap();
12447 let y1 = ys.pop().unwrap();
12448 let y0 = ys.pop().unwrap();
12449 Some((y0, y1, y2))
12450 }
12451 None => None,
12452 },
12453 );
12454 }
12455 if !(1..=8).contains(&m) {
12456 return Ok(None);
12457 }
12458 if m > 1 {
12459 let in_f = w0.in_features();
12460 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
12461 || !self.batched_supports(QT_NVFP4)
12462 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12463 || (m > 4 && !Self::b8_enabled())
12464 || in_f % 512 != 0
12465 || in_f / 64 > 272
12466 {
12467 return Ok(None);
12468 }
12469 }
12470 let unpack = |w: &crate::model::GpuTensor| match w {
12471 GpuTensor::Quant {
12472 bytes,
12473 qtype,
12474 scale,
12475 rp,
12476 ..
12477 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12478 _ => None,
12479 };
12480 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
12481 return Ok(None);
12482 };
12483 let in_f = w0.in_features();
12484 if w1.in_features() != in_f || w2.in_features() != in_f {
12485 return Ok(None);
12486 }
12487 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
12488 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12490 let rows_pb = ROWS_PER_BLOCK * RPW;
12491 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12492 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12493 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12494 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12495 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
12496 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
12499 if m > 1 {
12500 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
12502 return Ok(None);
12503 }
12504 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
12505 let cfg = LaunchConfig {
12506 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
12507 block_dim: (32, ROWS_PER_BLOCK, 1),
12508 shared_mem_bytes: 0,
12509 };
12510 let __s_b = self.gpu.stream();
12511 let mut b = __s_b.launch_builder(&f);
12512 b.arg(b0)
12513 .arg(b1)
12514 .arg(b2)
12515 .arg(aq)
12516 .arg(ad)
12517 .arg(&mut y0)
12518 .arg(&mut y1)
12519 .arg(&mut y2)
12520 .arg(&inf)
12521 .arg(&oi0)
12522 .arg(&oi1)
12523 .arg(&oi2)
12524 .arg(&mi);
12525 unsafe {
12526 b.launch(cfg)?;
12527 }
12528 return Ok(Some((y0, y1, y2)));
12529 }
12530 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
12531 let cfg = LaunchConfig {
12532 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
12533 block_dim: (32, ROWS_PER_BLOCK, 1),
12534 shared_mem_bytes: 0,
12535 };
12536 let __s_b = self.gpu.stream();
12537 let mut b = __s_b.launch_builder(&f);
12538 b.arg(b0)
12539 .arg(b1)
12540 .arg(b2)
12541 .arg(aq)
12542 .arg(ad)
12543 .arg(&mut y0)
12544 .arg(&mut y1)
12545 .arg(&mut y2)
12546 .arg(&inf)
12547 .arg(&oi0)
12548 .arg(&oi1)
12549 .arg(&oi2)
12550 .arg(&mi)
12551 .arg(&p0.1)
12552 .arg(&p1.1)
12553 .arg(&p2.1);
12554 unsafe {
12555 b.launch(cfg)?;
12556 }
12557 Ok(Some((y0, y1, y2)))
12558 }
12559
12560 pub fn matmul_nvfp4_fused2(
12569 &self,
12570 w0: &crate::model::GpuTensor,
12571 w1: &crate::model::GpuTensor,
12572 aq: &CudaSlice<i8>,
12573 ad: &CudaSlice<f32>,
12574 m: usize,
12575 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12576 use crate::model::GpuTensor;
12577 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12578 let off =
12579 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12580 if off
12583 || m != 1
12584 || !self.mmvq_supports(QT_NVFP4)
12585 || !self.uses_q8_1_fast(w0)
12586 || !self.uses_q8_1_fast(w1)
12587 {
12588 return Ok(None);
12589 }
12590 let unpack = |w: &crate::model::GpuTensor| match w {
12591 GpuTensor::Quant {
12592 bytes,
12593 qtype,
12594 scale,
12595 rp,
12596 ..
12597 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12598 _ => None,
12599 };
12600 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12601 return Ok(None);
12602 };
12603 let in_f = w0.in_features();
12604 if w1.in_features() != in_f {
12605 return Ok(None);
12606 }
12607 let (o0, o1) = (w0.out_features(), w1.out_features());
12608 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12610 let rows_pb = ROWS_PER_BLOCK * RPW;
12611 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12612 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12613 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12614 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12615 let cfg = LaunchConfig {
12616 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
12617 block_dim: (32, ROWS_PER_BLOCK, 1),
12618 shared_mem_bytes: 0,
12619 };
12620 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
12621 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12624 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
12627 {
12628 use cudarc::driver::{DevicePtr, DevicePtrMut};
12629 let s = &self.gpu.stream();
12630 let (pw0, _g0) = b0.device_ptr(s);
12631 let (pw1, _g1) = b1.device_ptr(s);
12632 let (paq, _g2) = aq.device_ptr(s);
12633 let (pad, _g3) = ad.device_ptr(s);
12634 let (py0, _g4) = y0.device_ptr_mut(s);
12635 let (py1, _g5) = y1.device_ptr_mut(s);
12636 let (s0, s1) = (p0.1, p1.1);
12637 let mut ps = [
12638 &pw0 as *const _ as *mut std::ffi::c_void,
12639 &pw1 as *const _ as *mut _,
12640 &paq as *const _ as *mut _,
12641 &pad as *const _ as *mut _,
12642 &py0 as *const _ as *mut _,
12643 &py1 as *const _ as *mut _,
12644 &inf as *const _ as *mut _,
12645 &oi0 as *const _ as *mut _,
12646 &oi1 as *const _ as *mut _,
12647 &mi as *const _ as *mut _,
12648 &s0 as *const _ as *mut _,
12649 &s1 as *const _ as *mut _,
12650 ];
12651 unsafe {
12652 self.launch_pdl(
12653 "qmatvec_nvfp4_mmvq_fused2_rp",
12654 cfg.grid_dim,
12655 cfg.block_dim,
12656 &mut ps,
12657 )?;
12658 }
12659 }
12660 return Ok(Some((y0, y1)));
12661 }
12662 let __s_b = self.gpu.stream();
12663 let mut b = __s_b.launch_builder(&f);
12664 b.arg(b0)
12665 .arg(b1)
12666 .arg(aq)
12667 .arg(ad)
12668 .arg(&mut y0)
12669 .arg(&mut y1)
12670 .arg(&inf)
12671 .arg(&oi0)
12672 .arg(&oi1)
12673 .arg(&mi)
12674 .arg(&p0.1)
12675 .arg(&p1.1);
12676 unsafe {
12677 b.launch(cfg)?;
12678 }
12679 Ok(Some((y0, y1)))
12680 }
12681
12682 pub fn matmul_nvfp4_fused2_into(
12687 &self,
12688 w0: &crate::model::GpuTensor,
12689 w1: &crate::model::GpuTensor,
12690 aq: &CudaSlice<i8>,
12691 ad: &CudaSlice<f32>,
12692 y0: &mut CudaSlice<f32>,
12693 y1: &mut CudaSlice<f32>,
12694 ) -> Result<bool, Box<dyn std::error::Error>> {
12695 use crate::model::GpuTensor;
12696 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12697 let off =
12698 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
12699 if off
12700 || !self.mmvq_supports(QT_NVFP4)
12701 || !self.uses_q8_1_fast(w0)
12702 || !self.uses_q8_1_fast(w1)
12703 {
12704 return Ok(false);
12705 }
12706 let unpack = |w: &crate::model::GpuTensor| match w {
12707 GpuTensor::Quant {
12708 bytes,
12709 qtype,
12710 scale,
12711 rp,
12712 ..
12713 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12714 _ => None,
12715 };
12716 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
12717 return Ok(false);
12718 };
12719 let in_f = w0.in_features();
12720 if w1.in_features() != in_f {
12721 return Ok(false);
12722 }
12723 let (o0, o1) = (w0.out_features(), w1.out_features());
12724 if y0.len() < o0 || y1.len() < o1 {
12725 return Ok(false);
12726 }
12727 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12729 let rows_pb = ROWS_PER_BLOCK * RPW;
12730 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12731 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
12732 let cfg = LaunchConfig {
12733 grid_dim: (nb(o0) + nb(o1), 1, 1),
12734 block_dim: (32, ROWS_PER_BLOCK, 1),
12735 shared_mem_bytes: 0,
12736 };
12737 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
12738 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
12741 let __s_b = self.gpu.stream();
12742 let mut b = __s_b.launch_builder(&f);
12743 b.arg(b0)
12744 .arg(b1)
12745 .arg(aq)
12746 .arg(ad)
12747 .arg(&mut *y0)
12748 .arg(&mut *y1)
12749 .arg(&inf)
12750 .arg(&oi0)
12751 .arg(&oi1)
12752 .arg(&mi)
12753 .arg(&p0.1)
12754 .arg(&p1.1);
12755 unsafe {
12756 b.launch(cfg)?;
12757 }
12758 Ok(true)
12759 }
12760
12761 #[allow(clippy::type_complexity)]
12766 pub fn matmul_nvfp4_fused4(
12767 &self,
12768 w0: &crate::model::GpuTensor,
12769 w1: &crate::model::GpuTensor,
12770 w2: &crate::model::GpuTensor,
12771 w3: &crate::model::GpuTensor,
12772 aq: &CudaSlice<i8>,
12773 ad: &CudaSlice<f32>,
12774 m: usize,
12775 ) -> Result<
12776 Option<(
12777 CudaSlice<f32>,
12778 CudaSlice<f32>,
12779 CudaSlice<f32>,
12780 CudaSlice<f32>,
12781 )>,
12782 Box<dyn std::error::Error>,
12783 > {
12784 use crate::model::GpuTensor;
12785 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
12792 || !self.mmvq_supports(QT_NVFP4)
12793 || !self.uses_q8_1_fast(w0)
12794 || !self.uses_q8_1_fast(w1)
12795 || !self.uses_q8_1_fast(w2)
12796 || !self.uses_q8_1_fast(w3)
12797 {
12798 return Ok(None);
12799 }
12800 if (9..=16).contains(&m) {
12805 return Ok(
12806 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
12807 Some(mut ys) => {
12808 let y3 = ys.pop().unwrap();
12809 let y2 = ys.pop().unwrap();
12810 let y1 = ys.pop().unwrap();
12811 let y0 = ys.pop().unwrap();
12812 Some((y0, y1, y2, y3))
12813 }
12814 None => None,
12815 },
12816 );
12817 }
12818 if !(1..=8).contains(&m) {
12819 return Ok(None);
12820 }
12821 if m > 1 {
12822 let in_f = w0.in_features();
12825 if !self.batched_supports(QT_NVFP4)
12826 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12827 || (m > 4 && !Self::b8_enabled())
12828 || in_f % 512 != 0
12829 || in_f / 64 > 272
12830 {
12831 return Ok(None);
12832 }
12833 }
12834 let unpack = |w: &crate::model::GpuTensor| match w {
12835 GpuTensor::Quant {
12836 bytes,
12837 qtype,
12838 scale,
12839 rp,
12840 ..
12841 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
12842 _ => None,
12843 };
12844 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
12845 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
12846 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 || w3.in_features() != in_f {
12851 return Ok(None);
12852 }
12853 let (o0, o1, o2, o3) = (
12854 w0.out_features(),
12855 w1.out_features(),
12856 w2.out_features(),
12857 w3.out_features(),
12858 );
12859 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
12861 let rows_pb = ROWS_PER_BLOCK * RPW;
12862 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
12863 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
12864 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
12865 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
12866 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
12867 let (inf, oi0, oi1, oi2, oi3, mi) = (
12868 in_f as i32,
12869 o0 as i32,
12870 o1 as i32,
12871 o2 as i32,
12872 o3 as i32,
12873 m as i32,
12874 );
12875 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
12878 if m > 1 {
12879 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
12882 return Ok(None);
12883 }
12884 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
12885 let cfg = LaunchConfig {
12886 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
12887 block_dim: (32, ROWS_PER_BLOCK, 1),
12888 shared_mem_bytes: 0,
12889 };
12890 let __s_b = self.gpu.stream();
12891 let mut b = __s_b.launch_builder(&f);
12892 b.arg(b0)
12893 .arg(b1)
12894 .arg(b2)
12895 .arg(b3)
12896 .arg(aq)
12897 .arg(ad)
12898 .arg(&mut y0)
12899 .arg(&mut y1)
12900 .arg(&mut y2)
12901 .arg(&mut y3)
12902 .arg(&inf)
12903 .arg(&oi0)
12904 .arg(&oi1)
12905 .arg(&oi2)
12906 .arg(&oi3)
12907 .arg(&mi);
12908 unsafe {
12909 b.launch(cfg)?;
12910 }
12911 return Ok(Some((y0, y1, y2, y3)));
12912 }
12913 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
12914 let cfg = LaunchConfig {
12915 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
12916 block_dim: (32, ROWS_PER_BLOCK, 1),
12917 shared_mem_bytes: 0,
12918 };
12919 let __s_b = self.gpu.stream();
12920 let mut b = __s_b.launch_builder(&f);
12921 b.arg(b0)
12922 .arg(b1)
12923 .arg(b2)
12924 .arg(b3)
12925 .arg(aq)
12926 .arg(ad)
12927 .arg(&mut y0)
12928 .arg(&mut y1)
12929 .arg(&mut y2)
12930 .arg(&mut y3)
12931 .arg(&inf)
12932 .arg(&oi0)
12933 .arg(&oi1)
12934 .arg(&oi2)
12935 .arg(&oi3)
12936 .arg(&mi)
12937 .arg(&p0.1)
12938 .arg(&p1.1)
12939 .arg(&p2.1)
12940 .arg(&p3.1);
12941 unsafe {
12942 b.launch(cfg)?;
12943 }
12944 Ok(Some((y0, y1, y2, y3)))
12945 }
12946
12947 pub fn matmul_q8_fused2(
12955 &self,
12956 w0: &crate::model::GpuTensor,
12957 w1: &crate::model::GpuTensor,
12958 aq: &CudaSlice<i8>,
12959 ad: &CudaSlice<f32>,
12960 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
12961 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
12967 return Ok(Some(self.e4m3_fused2_core(
12968 p0.0,
12969 p1.0,
12970 aq,
12971 ad,
12972 w0.in_features(),
12973 p0.1,
12974 p1.1,
12975 p0.2,
12976 p0.3,
12977 p1.3,
12978 )?));
12979 }
12980 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
12981 return Ok(None);
12982 };
12983 Ok(Some(self.q8_fused2_core(
12984 p0.0,
12985 p1.0,
12986 aq,
12987 ad,
12988 w0.in_features(),
12989 p0.1,
12990 p1.1,
12991 p0.2,
12992 )?))
12993 }
12994
12995 #[allow(clippy::too_many_arguments)]
12996 fn q8_fused2_core(
12997 &self,
12998 b0: &CudaSlice<u8>,
12999 b1: &CudaSlice<u8>,
13000 aq: &CudaSlice<i8>,
13001 ad: &CudaSlice<f32>,
13002 in_f: usize,
13003 out0: usize,
13004 out1: usize,
13005 row_bytes: usize,
13006 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13007 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13009 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13010 let f = self.func("qmatvec_q8_0_mmvq_fused2");
13011 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13012 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13013 let cfg = LaunchConfig {
13014 grid_dim: (nb0 + nb1, 1, 1),
13015 block_dim: (32, ROWS_PER_BLOCK, 1),
13016 shared_mem_bytes: 0,
13017 };
13018 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
13019 let __s_b = self.gpu.stream();
13020 let mut b = __s_b.launch_builder(&f);
13021 b.arg(b0)
13022 .arg(b1)
13023 .arg(aq)
13024 .arg(ad)
13025 .arg(&mut y0)
13026 .arg(&mut y1)
13027 .arg(&inf)
13028 .arg(&o0)
13029 .arg(&o1)
13030 .arg(&rbl);
13031 unsafe {
13032 b.launch(cfg)?;
13033 }
13034 Ok((y0, y1))
13035 }
13036
13037 pub fn matmul_q8_fused2_x(
13043 &self,
13044 w0: &crate::model::GpuTensor,
13045 w1: &crate::model::GpuTensor,
13046 x: &CudaSlice<f32>,
13047 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13048 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13049 return Ok(None);
13050 }
13051 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13052 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13053 return Ok(Some(self.e4m3_fused2_core(
13054 p0.0,
13055 p1.0,
13056 &aq,
13057 &ad,
13058 w0.in_features(),
13059 p0.1,
13060 p1.1,
13061 p0.2,
13062 p0.3,
13063 p1.3,
13064 )?));
13065 }
13066 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13067 return Ok(None);
13068 };
13069 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
13070 Ok(Some(self.q8_fused2_core(
13071 p0.0,
13072 p1.0,
13073 &aq,
13074 &ad,
13075 w0.in_features(),
13076 p0.1,
13077 p1.1,
13078 p0.2,
13079 )?))
13080 }
13081
13082 #[allow(clippy::too_many_arguments)]
13085 pub fn qmatvec_q8_fused2_raw(
13086 &self,
13087 b0: &CudaSlice<u8>,
13088 b1: &CudaSlice<u8>,
13089 x: &CudaSlice<f32>,
13090 in_f: usize,
13091 out0: usize,
13092 out1: usize,
13093 row_bytes: usize,
13094 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13095 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13096 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
13097 }
13098
13099 pub fn matmul_q4_fused3(
13105 &self,
13106 w0: &crate::model::GpuTensor,
13107 w1: &crate::model::GpuTensor,
13108 w2: &crate::model::GpuTensor,
13109 aq: &CudaSlice<i8>,
13110 ad: &CudaSlice<f32>,
13111 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13112 {
13113 use crate::model::GpuTensor;
13114 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13115 match w {
13116 GpuTensor::Quant {
13117 qtype, row_bytes, ..
13118 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13119 _ => None,
13120 }
13121 };
13122 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13123 return Ok(None);
13124 };
13125 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13126 return Ok(None);
13127 }
13128 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13132 match w {
13133 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13134 Some(m) => (m, true),
13135 None => (bytes, *rp),
13136 },
13137 _ => unreachable!(),
13138 }
13139 }
13140 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13141 if rp0 != rp1 || rp1 != rp2 {
13142 return Ok(None);
13143 }
13144 let rp = rp0;
13145 let rpb: u32 = 4;
13146 let mr1 = rp && Self::q40_mr1_on();
13150 let nb = |o: usize| {
13151 if mr1 {
13152 (o as u32).div_ceil(rpb)
13153 } else {
13154 (o as u32).div_ceil(2).div_ceil(rpb)
13155 }
13156 };
13157 let grid = nb(o0) + nb(o1) + nb(o2);
13158 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13159 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13160 let mut y2 = self.alloc_uninit::<f32>(o2)?;
13161 let f = self.func(if mr1 {
13162 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13163 } else if rp {
13164 "qmatvec_q4_0_mmvq_fused3_rp"
13165 } else {
13166 "qmatvec_q4_0_mmvq_fused3"
13167 });
13168 let cfg = LaunchConfig {
13169 grid_dim: (grid, 1, 1),
13170 block_dim: (32, rpb, 1),
13171 shared_mem_bytes: 0,
13172 };
13173 let inf = w0.in_features() as i32;
13174 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13175 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13176 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13179 {
13180 use cudarc::driver::{DevicePtr, DevicePtrMut};
13181 let s = &self.gpu.stream();
13182 let (p0, _g0) = b0.device_ptr(s);
13183 let (p1, _g1) = b1.device_ptr(s);
13184 let (p2, _g2) = b2.device_ptr(s);
13185 let (paq, _g3) = aq.device_ptr(s);
13186 let (pad, _g4) = ad.device_ptr(s);
13187 let (py0, _g5) = y0.device_ptr_mut(s);
13188 let (py1, _g6) = y1.device_ptr_mut(s);
13189 let (py2, _g7) = y2.device_ptr_mut(s);
13190 let mut ps = [
13191 &p0 as *const _ as *mut std::ffi::c_void,
13192 &p1 as *const _ as *mut _,
13193 &p2 as *const _ as *mut _,
13194 &paq as *const _ as *mut _,
13195 &pad as *const _ as *mut _,
13196 &py0 as *const _ as *mut _,
13197 &py1 as *const _ as *mut _,
13198 &py2 as *const _ as *mut _,
13199 &inf as *const _ as *mut _,
13200 &oo0 as *const _ as *mut _,
13201 &oo1 as *const _ as *mut _,
13202 &oo2 as *const _ as *mut _,
13203 &r0 as *const _ as *mut _,
13204 &r1 as *const _ as *mut _,
13205 &r2 as *const _ as *mut _,
13206 ];
13207 unsafe {
13208 self.launch_pdl(
13209 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13210 (grid, 1, 1),
13211 (32, rpb, 1),
13212 &mut ps,
13213 )?;
13214 }
13215 }
13216 return Ok(Some((y0, y1, y2)));
13217 }
13218 let __s_b = self.gpu.stream();
13219 let mut b = __s_b.launch_builder(&f);
13220 b.arg(b0)
13221 .arg(b1)
13222 .arg(b2)
13223 .arg(aq)
13224 .arg(ad)
13225 .arg(&mut y0)
13226 .arg(&mut y1)
13227 .arg(&mut y2)
13228 .arg(&inf)
13229 .arg(&oo0)
13230 .arg(&oo1)
13231 .arg(&oo2)
13232 .arg(&r0)
13233 .arg(&r1)
13234 .arg(&r2);
13235 unsafe {
13236 b.launch(cfg)?;
13237 }
13238 Ok(Some((y0, y1, y2)))
13239 }
13240
13241 #[allow(clippy::too_many_arguments)]
13244 pub fn matmul_q4_fused3_into(
13245 &self,
13246 w0: &crate::model::GpuTensor,
13247 w1: &crate::model::GpuTensor,
13248 w2: &crate::model::GpuTensor,
13249 aq: &CudaSlice<i8>,
13250 ad: &CudaSlice<f32>,
13251 y0: &mut CudaSlice<f32>,
13252 y1: &mut CudaSlice<f32>,
13253 y2: &mut CudaSlice<f32>,
13254 ) -> Result<bool, Box<dyn std::error::Error>> {
13255 use crate::model::GpuTensor;
13256 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13257 match w {
13258 GpuTensor::Quant {
13259 qtype, row_bytes, ..
13260 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13261 _ => None,
13262 }
13263 };
13264 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
13265 return Ok(false);
13266 };
13267 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13268 return Ok(false);
13269 }
13270 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13271 match w {
13272 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13273 Some(m) => (m, true),
13274 None => (bytes, *rp),
13275 },
13276 _ => unreachable!(),
13277 }
13278 }
13279 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13280 if rp0 != rp1 || rp1 != rp2 {
13281 return Ok(false);
13282 }
13283 let rp = rp0;
13284 let rpb: u32 = 4;
13285 let mr1 = rp && Self::q40_mr1_on();
13286 let nb = |o: usize| {
13287 if mr1 {
13288 (o as u32).div_ceil(rpb)
13289 } else {
13290 (o as u32).div_ceil(2).div_ceil(rpb)
13291 }
13292 };
13293 let grid = nb(o0) + nb(o1) + nb(o2);
13294 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
13295 let f = self.func(if mr1 {
13296 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
13297 } else if rp {
13298 "qmatvec_q4_0_mmvq_fused3_rp"
13299 } else {
13300 "qmatvec_q4_0_mmvq_fused3"
13301 });
13302 let cfg = LaunchConfig {
13303 grid_dim: (grid, 1, 1),
13304 block_dim: (32, rpb, 1),
13305 shared_mem_bytes: 0,
13306 };
13307 let inf = w0.in_features() as i32;
13308 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
13309 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
13310 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13312 use cudarc::driver::{DevicePtr, DevicePtrMut};
13313 let s = &self.gpu.stream();
13314 let (p0, _g0) = b0.device_ptr(s);
13315 let (p1, _g1) = b1.device_ptr(s);
13316 let (p2, _g2) = b2.device_ptr(s);
13317 let (paq, _g3) = aq.device_ptr(s);
13318 let (pad, _g4) = ad.device_ptr(s);
13319 let (py0, _g5) = y0.device_ptr_mut(s);
13320 let (py1, _g6) = y1.device_ptr_mut(s);
13321 let (py2, _g7) = y2.device_ptr_mut(s);
13322 let mut ps = [
13323 &p0 as *const _ as *mut std::ffi::c_void,
13324 &p1 as *const _ as *mut _,
13325 &p2 as *const _ as *mut _,
13326 &paq as *const _ as *mut _,
13327 &pad as *const _ as *mut _,
13328 &py0 as *const _ as *mut _,
13329 &py1 as *const _ as *mut _,
13330 &py2 as *const _ as *mut _,
13331 &inf as *const _ as *mut _,
13332 &oo0 as *const _ as *mut _,
13333 &oo1 as *const _ as *mut _,
13334 &oo2 as *const _ as *mut _,
13335 &r0 as *const _ as *mut _,
13336 &r1 as *const _ as *mut _,
13337 &r2 as *const _ as *mut _,
13338 ];
13339 unsafe {
13340 self.launch_pdl(
13341 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
13342 (grid, 1, 1),
13343 (32, rpb, 1),
13344 &mut ps,
13345 )?;
13346 }
13347 return Ok(true);
13348 }
13349 let __s_b = self.gpu.stream();
13350 let mut b = __s_b.launch_builder(&f);
13351 b.arg(b0)
13352 .arg(b1)
13353 .arg(b2)
13354 .arg(aq)
13355 .arg(ad)
13356 .arg(&mut *y0)
13357 .arg(&mut *y1)
13358 .arg(&mut *y2)
13359 .arg(&inf)
13360 .arg(&oo0)
13361 .arg(&oo1)
13362 .arg(&oo2)
13363 .arg(&r0)
13364 .arg(&r1)
13365 .arg(&r2);
13366 unsafe {
13367 b.launch(cfg)?;
13368 }
13369 Ok(true)
13370 }
13371
13372 pub fn matmul_q4_fused2(
13374 &self,
13375 w0: &crate::model::GpuTensor,
13376 w1: &crate::model::GpuTensor,
13377 aq: &CudaSlice<i8>,
13378 ad: &CudaSlice<f32>,
13379 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13380 use crate::model::GpuTensor;
13381 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13382 match w {
13383 GpuTensor::Quant {
13384 qtype, row_bytes, ..
13385 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13386 _ => None,
13387 }
13388 };
13389 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13390 return Ok(None);
13391 };
13392 if w0.in_features() != w1.in_features() {
13393 return Ok(None);
13394 }
13395 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13397 match w {
13398 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13399 Some(m) => (m, true),
13400 None => (bytes, *rp),
13401 },
13402 _ => unreachable!(),
13403 }
13404 }
13405 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13406 if rp0 != rp1 {
13407 return Ok(None);
13408 }
13409 let rp = rp0;
13410 let rpb: u32 = 4;
13411 let mr1 = rp && Self::q40_mr1_on();
13413 let nb = |o: usize| {
13414 if mr1 {
13415 (o as u32).div_ceil(rpb)
13416 } else {
13417 (o as u32).div_ceil(2).div_ceil(rpb)
13418 }
13419 };
13420 let grid = nb(o0) + nb(o1);
13421 let mut y0 = self.alloc_uninit::<f32>(o0)?;
13422 let mut y1 = self.alloc_uninit::<f32>(o1)?;
13423 let f = self.func(if mr1 {
13424 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13425 } else if rp {
13426 "qmatvec_q4_0_mmvq_fused2_rp"
13427 } else {
13428 "qmatvec_q4_0_mmvq_fused2"
13429 });
13430 let cfg = LaunchConfig {
13431 grid_dim: (grid, 1, 1),
13432 block_dim: (32, rpb, 1),
13433 shared_mem_bytes: 0,
13434 };
13435 let inf = w0.in_features() as i32;
13436 let (oo0, oo1) = (o0 as i32, o1 as i32);
13437 let (r0, r1) = (rb0 as i64, rb1 as i64);
13438 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13440 {
13441 use cudarc::driver::{DevicePtr, DevicePtrMut};
13442 let s = &self.gpu.stream();
13443 let (p0, _g0) = b0.device_ptr(s);
13444 let (p1, _g1) = b1.device_ptr(s);
13445 let (paq, _g2) = aq.device_ptr(s);
13446 let (pad, _g3) = ad.device_ptr(s);
13447 let (py0, _g4) = y0.device_ptr_mut(s);
13448 let (py1, _g5) = y1.device_ptr_mut(s);
13449 let mut ps = [
13450 &p0 as *const _ as *mut std::ffi::c_void,
13451 &p1 as *const _ as *mut _,
13452 &paq as *const _ as *mut _,
13453 &pad as *const _ as *mut _,
13454 &py0 as *const _ as *mut _,
13455 &py1 as *const _ as *mut _,
13456 &inf as *const _ as *mut _,
13457 &oo0 as *const _ as *mut _,
13458 &oo1 as *const _ as *mut _,
13459 &r0 as *const _ as *mut _,
13460 &r1 as *const _ as *mut _,
13461 ];
13462 unsafe {
13463 self.launch_pdl(
13464 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13465 (grid, 1, 1),
13466 (32, rpb, 1),
13467 &mut ps,
13468 )?;
13469 }
13470 }
13471 return Ok(Some((y0, y1)));
13472 }
13473 let __s_b = self.gpu.stream();
13474 let mut b = __s_b.launch_builder(&f);
13475 b.arg(b0)
13476 .arg(b1)
13477 .arg(aq)
13478 .arg(ad)
13479 .arg(&mut y0)
13480 .arg(&mut y1)
13481 .arg(&inf)
13482 .arg(&oo0)
13483 .arg(&oo1)
13484 .arg(&r0)
13485 .arg(&r1);
13486 unsafe {
13487 b.launch(cfg)?;
13488 }
13489 Ok(Some((y0, y1)))
13490 }
13491
13492 pub fn matmul_q4_fused2_into(
13494 &self,
13495 w0: &crate::model::GpuTensor,
13496 w1: &crate::model::GpuTensor,
13497 aq: &CudaSlice<i8>,
13498 ad: &CudaSlice<f32>,
13499 y0: &mut CudaSlice<f32>,
13500 y1: &mut CudaSlice<f32>,
13501 ) -> Result<bool, Box<dyn std::error::Error>> {
13502 use crate::model::GpuTensor;
13503 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13504 match w {
13505 GpuTensor::Quant {
13506 qtype, row_bytes, ..
13507 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13508 _ => None,
13509 }
13510 };
13511 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
13512 return Ok(false);
13513 };
13514 if w0.in_features() != w1.in_features() {
13515 return Ok(false);
13516 }
13517 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13518 match w {
13519 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13520 Some(m) => (m, true),
13521 None => (bytes, *rp),
13522 },
13523 _ => unreachable!(),
13524 }
13525 }
13526 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13527 if rp0 != rp1 {
13528 return Ok(false);
13529 }
13530 let rp = rp0;
13531 let rpb: u32 = 4;
13532 let mr1 = rp && Self::q40_mr1_on();
13533 let nb = |o: usize| {
13534 if mr1 {
13535 (o as u32).div_ceil(rpb)
13536 } else {
13537 (o as u32).div_ceil(2).div_ceil(rpb)
13538 }
13539 };
13540 let grid = nb(o0) + nb(o1);
13541 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
13542 let f = self.func(if mr1 {
13543 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
13544 } else if rp {
13545 "qmatvec_q4_0_mmvq_fused2_rp"
13546 } else {
13547 "qmatvec_q4_0_mmvq_fused2"
13548 });
13549 let cfg = LaunchConfig {
13550 grid_dim: (grid, 1, 1),
13551 block_dim: (32, rpb, 1),
13552 shared_mem_bytes: 0,
13553 };
13554 let inf = w0.in_features() as i32;
13555 let (oo0, oo1) = (o0 as i32, o1 as i32);
13556 let (r0, r1) = (rb0 as i64, rb1 as i64);
13557 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
13559 use cudarc::driver::{DevicePtr, DevicePtrMut};
13560 let s = &self.gpu.stream();
13561 let (p0, _g0) = b0.device_ptr(s);
13562 let (p1, _g1) = b1.device_ptr(s);
13563 let (paq, _g2) = aq.device_ptr(s);
13564 let (pad, _g3) = ad.device_ptr(s);
13565 let (py0, _g4) = y0.device_ptr_mut(s);
13566 let (py1, _g5) = y1.device_ptr_mut(s);
13567 let mut ps = [
13568 &p0 as *const _ as *mut std::ffi::c_void,
13569 &p1 as *const _ as *mut _,
13570 &paq as *const _ as *mut _,
13571 &pad as *const _ as *mut _,
13572 &py0 as *const _ as *mut _,
13573 &py1 as *const _ as *mut _,
13574 &inf as *const _ as *mut _,
13575 &oo0 as *const _ as *mut _,
13576 &oo1 as *const _ as *mut _,
13577 &r0 as *const _ as *mut _,
13578 &r1 as *const _ as *mut _,
13579 ];
13580 unsafe {
13581 self.launch_pdl(
13582 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
13583 (grid, 1, 1),
13584 (32, rpb, 1),
13585 &mut ps,
13586 )?;
13587 }
13588 return Ok(true);
13589 }
13590 let __s_b = self.gpu.stream();
13591 let mut b = __s_b.launch_builder(&f);
13592 b.arg(b0)
13593 .arg(b1)
13594 .arg(aq)
13595 .arg(ad)
13596 .arg(&mut *y0)
13597 .arg(&mut *y1)
13598 .arg(&inf)
13599 .arg(&oo0)
13600 .arg(&oo1)
13601 .arg(&r0)
13602 .arg(&r1);
13603 unsafe {
13604 b.launch(cfg)?;
13605 }
13606 Ok(true)
13607 }
13608
13609 pub fn matmul_q4_fused2_batched(
13614 &self,
13615 w0: &crate::model::GpuTensor,
13616 w1: &crate::model::GpuTensor,
13617 aq: &CudaSlice<i8>,
13618 ad: &CudaSlice<f32>,
13619 m: usize,
13620 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13621 use crate::model::GpuTensor;
13622 if m < 2 || m > 8 {
13623 return Ok(None);
13624 }
13625 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
13626 match w {
13627 GpuTensor::Quant {
13628 qtype, row_bytes, ..
13629 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
13630 _ => None,
13631 }
13632 };
13633 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
13634 return Ok(None);
13635 };
13636 if w0.in_features() != w1.in_features() {
13637 return Ok(None);
13638 }
13639 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13640 match w {
13641 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13642 Some(mr) => (mr, true),
13643 None => (bytes, *rp),
13644 },
13645 _ => unreachable!(),
13646 }
13647 }
13648 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
13649 if !rp0 || !rp1 {
13650 return Ok(None);
13651 }
13652 let mcols = Self::batched_mcols(m);
13653 let rpb: u32 = 4;
13654 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13655 let grid = nb(o0) + nb(o1);
13656 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13657 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13658 let f = self.func(match mcols {
13659 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
13660 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
13661 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
13662 });
13663 let cfg = LaunchConfig {
13664 grid_dim: (grid, 1, 1),
13665 block_dim: (32, rpb, 1),
13666 shared_mem_bytes: 0,
13667 };
13668 let inf = w0.in_features() as i32;
13669 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
13670 let rb = rb0 as i64;
13671 let __s_b = self.gpu.stream();
13672 let mut b = __s_b.launch_builder(&f);
13673 b.arg(b0)
13674 .arg(b1)
13675 .arg(aq)
13676 .arg(ad)
13677 .arg(&mut y0)
13678 .arg(&mut y1)
13679 .arg(&inf)
13680 .arg(&oo0)
13681 .arg(&oo1)
13682 .arg(&mi)
13683 .arg(&rb);
13684 unsafe {
13685 b.launch(cfg)?;
13686 }
13687 Ok(Some((y0, y1)))
13688 }
13689
13690 #[allow(clippy::too_many_arguments)]
13693 pub fn matmul_q4_fused3_batched(
13694 &self,
13695 w0: &crate::model::GpuTensor,
13696 w1: &crate::model::GpuTensor,
13697 w2: &crate::model::GpuTensor,
13698 aq: &CudaSlice<i8>,
13699 ad: &CudaSlice<f32>,
13700 m: usize,
13701 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13702 {
13703 use crate::model::GpuTensor;
13704 if m < 2 || m > 8 {
13705 return Ok(None);
13706 }
13707 let q4 = |w: &GpuTensor| -> Option<usize> {
13708 match w {
13709 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
13710 _ => None,
13711 }
13712 };
13713 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
13714 return Ok(None);
13715 };
13716 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
13717 return Ok(None);
13718 }
13719 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
13720 match w {
13721 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
13722 Some(mr) => (mr, true),
13723 None => (bytes, *rp),
13724 },
13725 _ => unreachable!(),
13726 }
13727 }
13728 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
13729 if !rp0 || !rp1 || !rp2 {
13730 return Ok(None);
13731 }
13732 let mcols = Self::batched_mcols(m);
13733 let rpb: u32 = 4;
13734 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
13735 let grid = nb(o0) + nb(o1) + nb(o2);
13736 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13737 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13738 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13739 let f = self.func(match mcols {
13740 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
13741 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
13742 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
13743 });
13744 let cfg = LaunchConfig {
13745 grid_dim: (grid, 1, 1),
13746 block_dim: (32, rpb, 1),
13747 shared_mem_bytes: 0,
13748 };
13749 let inf = w0.in_features() as i32;
13750 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
13751 let rb = 0i64;
13752 let __s_b = self.gpu.stream();
13753 let mut b = __s_b.launch_builder(&f);
13754 b.arg(b0)
13755 .arg(b1)
13756 .arg(b2)
13757 .arg(aq)
13758 .arg(ad)
13759 .arg(&mut y0)
13760 .arg(&mut y1)
13761 .arg(&mut y2)
13762 .arg(&inf)
13763 .arg(&oo0)
13764 .arg(&oo1)
13765 .arg(&oo2)
13766 .arg(&mi)
13767 .arg(&rb);
13768 unsafe {
13769 b.launch(cfg)?;
13770 }
13771 Ok(Some((y0, y1, y2)))
13772 }
13773
13774 pub fn matmul_q8_fused3(
13775 &self,
13776 w0: &crate::model::GpuTensor,
13777 w1: &crate::model::GpuTensor,
13778 w2: &crate::model::GpuTensor,
13779 aq: &CudaSlice<i8>,
13780 ad: &CudaSlice<f32>,
13781 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13782 {
13783 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
13786 return Ok(Some(self.e4m3_fused3_core(
13787 p0.0,
13788 p1.0,
13789 p2.0,
13790 aq,
13791 ad,
13792 w0.in_features(),
13793 p0.1,
13794 p1.1,
13795 p2.1,
13796 p0.2,
13797 p0.3,
13798 p1.3,
13799 p2.3,
13800 )?));
13801 }
13802 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
13803 return Ok(None);
13804 };
13805 Ok(Some(self.q8_fused3_core(
13806 p0.0,
13807 p1.0,
13808 p2.0,
13809 aq,
13810 ad,
13811 w0.in_features(),
13812 p0.1,
13813 p1.1,
13814 p2.1,
13815 p0.2,
13816 )?))
13817 }
13818
13819 #[allow(clippy::too_many_arguments)]
13820 fn q8_fused3_core(
13821 &self,
13822 b0: &CudaSlice<u8>,
13823 b1: &CudaSlice<u8>,
13824 b2: &CudaSlice<u8>,
13825 aq: &CudaSlice<i8>,
13826 ad: &CudaSlice<f32>,
13827 in_f: usize,
13828 out0: usize,
13829 out1: usize,
13830 out2: usize,
13831 row_bytes: usize,
13832 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13833 const ROWS_PER_BLOCK: u32 = 4;
13834 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13835 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13836 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
13837 let f = self.func("qmatvec_q8_0_mmvq_fused3");
13838 let mut y0 = self.alloc_uninit::<f32>(out0)?;
13839 let mut y1 = self.alloc_uninit::<f32>(out1)?;
13840 let mut y2 = self.alloc_uninit::<f32>(out2)?;
13841 let cfg = LaunchConfig {
13842 grid_dim: (nb0 + nb1 + nb2, 1, 1),
13843 block_dim: (32, ROWS_PER_BLOCK, 1),
13844 shared_mem_bytes: 0,
13845 };
13846 let (inf, o0, o1, o2, rbl) = (
13847 in_f as i32,
13848 out0 as i32,
13849 out1 as i32,
13850 out2 as i32,
13851 row_bytes as i64,
13852 );
13853 let __s_b = self.gpu.stream();
13854 let mut b = __s_b.launch_builder(&f);
13855 b.arg(b0)
13856 .arg(b1)
13857 .arg(b2)
13858 .arg(aq)
13859 .arg(ad)
13860 .arg(&mut y0)
13861 .arg(&mut y1)
13862 .arg(&mut y2)
13863 .arg(&inf)
13864 .arg(&o0)
13865 .arg(&o1)
13866 .arg(&o2)
13867 .arg(&rbl);
13868 unsafe {
13869 b.launch(cfg)?;
13870 }
13871 Ok((y0, y1, y2))
13872 }
13873
13874 #[allow(clippy::too_many_arguments)]
13876 pub fn qmatvec_q8_fused3_raw(
13877 &self,
13878 b0: &CudaSlice<u8>,
13879 b1: &CudaSlice<u8>,
13880 b2: &CudaSlice<u8>,
13881 x: &CudaSlice<f32>,
13882 in_f: usize,
13883 out0: usize,
13884 out1: usize,
13885 out2: usize,
13886 row_bytes: usize,
13887 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13888 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
13889 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
13890 }
13891
13892 pub fn matmul_q8_fused2_t(
13903 &self,
13904 w0: &crate::model::GpuTensor,
13905 w1: &crate::model::GpuTensor,
13906 aq: &CudaSlice<i8>,
13907 ad: &CudaSlice<f32>,
13908 m: usize,
13909 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13910 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
13914 return Ok(None);
13915 }
13916 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13919 if m > 4 && !Self::b8_enabled() {
13920 return Ok(None);
13921 }
13922 return Ok(Some(self.e4m3_fused2_t_core(
13923 p0.0,
13924 p1.0,
13925 aq,
13926 ad,
13927 m,
13928 w0.in_features(),
13929 p0.1,
13930 p1.1,
13931 p0.2,
13932 p0.3,
13933 p1.3,
13934 )?));
13935 }
13936 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
13937 return Ok(None);
13938 };
13939 Ok(Some(self.q8_fused2_t_core(
13940 p0.0,
13941 p1.0,
13942 aq,
13943 ad,
13944 m,
13945 w0.in_features(),
13946 p0.1,
13947 p1.1,
13948 p0.2,
13949 )?))
13950 }
13951
13952 #[allow(clippy::too_many_arguments)]
13953 fn q8_fused2_t_core(
13954 &self,
13955 b0: &CudaSlice<u8>,
13956 b1: &CudaSlice<u8>,
13957 aq: &CudaSlice<i8>,
13958 ad: &CudaSlice<f32>,
13959 m: usize,
13960 in_f: usize,
13961 out0: usize,
13962 out1: usize,
13963 row_bytes: usize,
13964 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13965 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
13967 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
13968 let f = self.func(match Self::batched_mcols(m) {
13969 2 => "qmatvec_q8_0_mmvq_fused2_b2",
13970 4 => "qmatvec_q8_0_mmvq_fused2_b4",
13971 _ => "qmatvec_q8_0_mmvq_fused2_b8",
13973 });
13974 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
13975 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
13976 let cfg = LaunchConfig {
13977 grid_dim: (nb0 + nb1, 1, 1),
13978 block_dim: (32, ROWS_PER_BLOCK, 1),
13979 shared_mem_bytes: 0,
13980 };
13981 let (inf, o0, o1, mi, rbl) = (
13982 in_f as i32,
13983 out0 as i32,
13984 out1 as i32,
13985 m as i32,
13986 row_bytes as i64,
13987 );
13988 let __s_b = self.gpu.stream();
13989 let mut b = __s_b.launch_builder(&f);
13990 b.arg(b0)
13991 .arg(b1)
13992 .arg(aq)
13993 .arg(ad)
13994 .arg(&mut y0)
13995 .arg(&mut y1)
13996 .arg(&inf)
13997 .arg(&o0)
13998 .arg(&o1)
13999 .arg(&mi)
14000 .arg(&rbl);
14001 unsafe {
14002 b.launch(cfg)?;
14003 }
14004 Ok((y0, y1))
14005 }
14006
14007 #[allow(clippy::too_many_arguments)]
14010 pub fn qmatvec_q8_fused2_t_raw(
14011 &self,
14012 b0: &CudaSlice<u8>,
14013 b1: &CudaSlice<u8>,
14014 x: &CudaSlice<f32>,
14015 m: usize,
14016 in_f: usize,
14017 out0: usize,
14018 out1: usize,
14019 row_bytes: usize,
14020 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14021 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14022 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
14023 }
14024
14025 #[allow(clippy::too_many_arguments)]
14028 pub fn matmul_q8_fused3_t(
14029 &self,
14030 w0: &crate::model::GpuTensor,
14031 w1: &crate::model::GpuTensor,
14032 w2: &crate::model::GpuTensor,
14033 aq: &CudaSlice<i8>,
14034 ad: &CudaSlice<f32>,
14035 m: usize,
14036 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14037 {
14038 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
14039 return Ok(None);
14040 }
14041 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14042 return Ok(Some(self.e4m3_fused3_t_core(
14043 p0.0,
14044 p1.0,
14045 p2.0,
14046 aq,
14047 ad,
14048 m,
14049 w0.in_features(),
14050 p0.1,
14051 p1.1,
14052 p2.1,
14053 p0.2,
14054 p0.3,
14055 p1.3,
14056 p2.3,
14057 )?));
14058 }
14059 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14060 return Ok(None);
14061 };
14062 Ok(Some(self.q8_fused3_t_core(
14063 p0.0,
14064 p1.0,
14065 p2.0,
14066 aq,
14067 ad,
14068 m,
14069 w0.in_features(),
14070 p0.1,
14071 p1.1,
14072 p2.1,
14073 p0.2,
14074 )?))
14075 }
14076
14077 #[allow(clippy::too_many_arguments)]
14078 fn q8_fused3_t_core(
14079 &self,
14080 b0: &CudaSlice<u8>,
14081 b1: &CudaSlice<u8>,
14082 b2: &CudaSlice<u8>,
14083 aq: &CudaSlice<i8>,
14084 ad: &CudaSlice<f32>,
14085 m: usize,
14086 in_f: usize,
14087 out0: usize,
14088 out1: usize,
14089 out2: usize,
14090 row_bytes: usize,
14091 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14092 const ROWS_PER_BLOCK: u32 = 4;
14093 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14094 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14095 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14096 let f = self.func(if Self::batched_mcols(m) == 2 {
14097 "qmatvec_q8_0_mmvq_fused3_b2"
14098 } else {
14099 "qmatvec_q8_0_mmvq_fused3_b4"
14100 });
14101 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14102 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14103 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14104 let cfg = LaunchConfig {
14105 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14106 block_dim: (32, ROWS_PER_BLOCK, 1),
14107 shared_mem_bytes: 0,
14108 };
14109 let (inf, o0, o1, o2, mi, rbl) = (
14110 in_f as i32,
14111 out0 as i32,
14112 out1 as i32,
14113 out2 as i32,
14114 m as i32,
14115 row_bytes as i64,
14116 );
14117 let __s_b = self.gpu.stream();
14118 let mut b = __s_b.launch_builder(&f);
14119 b.arg(b0)
14120 .arg(b1)
14121 .arg(b2)
14122 .arg(aq)
14123 .arg(ad)
14124 .arg(&mut y0)
14125 .arg(&mut y1)
14126 .arg(&mut y2)
14127 .arg(&inf)
14128 .arg(&o0)
14129 .arg(&o1)
14130 .arg(&o2)
14131 .arg(&mi)
14132 .arg(&rbl);
14133 unsafe {
14134 b.launch(cfg)?;
14135 }
14136 Ok((y0, y1, y2))
14137 }
14138
14139 #[allow(clippy::too_many_arguments)]
14141 pub fn qmatvec_q8_fused3_t_raw(
14142 &self,
14143 b0: &CudaSlice<u8>,
14144 b1: &CudaSlice<u8>,
14145 b2: &CudaSlice<u8>,
14146 x: &CudaSlice<f32>,
14147 m: usize,
14148 in_f: usize,
14149 out0: usize,
14150 out1: usize,
14151 out2: usize,
14152 row_bytes: usize,
14153 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14154 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14155 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
14156 }
14157
14158 pub fn q8_ffn_fuse2_on(&self) -> bool {
14162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14163 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
14164 }
14165
14166 #[allow(clippy::type_complexity)]
14172 fn q8_fused_params<'w, const N: usize>(
14173 &self,
14174 ws: &[&'w crate::model::GpuTensor; N],
14175 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
14176 use crate::model::GpuTensor;
14177 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
14178 return None;
14179 }
14180 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
14181 return None;
14182 }
14183 let in_f = ws[0].in_features();
14184 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
14185 for (i, w) in ws.iter().enumerate() {
14186 match w {
14187 GpuTensor::Quant {
14188 bytes,
14189 qtype,
14190 row_bytes,
14191 scale,
14192 ..
14193 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
14194 out[i] = Some((bytes, w.out_features(), *row_bytes))
14195 }
14196 _ => return None,
14197 }
14198 }
14199 Some(out.map(|o| o.unwrap()))
14200 }
14201
14202 pub fn e4m3_dual_on(&self) -> bool {
14205 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
14206 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
14207 }
14208
14209 #[allow(clippy::type_complexity)]
14221 fn e4m3_fused_params<'w, const N: usize>(
14222 &self,
14223 ws: &[&'w crate::model::GpuTensor; N],
14224 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
14225 use crate::model::GpuTensor;
14226 if !self.e4m3_dual_on() {
14227 return None;
14228 }
14229 let in_f = ws[0].in_features();
14230 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
14231 for (i, w) in ws.iter().enumerate() {
14232 match w {
14233 GpuTensor::Quant {
14234 bytes,
14235 qtype,
14236 row_bytes,
14237 scale,
14238 rp,
14239 rp4,
14240 ..
14241 } if *qtype == QT_F8_E4M3
14242 && w.in_features() == in_f
14243 && *row_bytes == in_f
14244 && !*rp
14245 && rp4.is_none() =>
14246 {
14247 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
14248 }
14249 _ => return None,
14250 }
14251 }
14252 Some(out.map(|o| o.unwrap()))
14253 }
14254
14255 #[allow(clippy::too_many_arguments)]
14259 fn e4m3_fused2_core(
14260 &self,
14261 b0: &CudaSlice<u8>,
14262 b1: &CudaSlice<u8>,
14263 aq: &CudaSlice<i8>,
14264 ad: &CudaSlice<f32>,
14265 in_f: usize,
14266 out0: usize,
14267 out1: usize,
14268 row_bytes: usize,
14269 ws0: f32,
14270 ws1: f32,
14271 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14272 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14274 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14275 let f = self.func("qmatvec_e4m3_mmvq_fused2");
14276 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14277 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14278 let cfg = LaunchConfig {
14279 grid_dim: (nb0 + nb1, 1, 1),
14280 block_dim: (32, ROWS_PER_BLOCK, 1),
14281 shared_mem_bytes: 0,
14282 };
14283 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14284 let __s_b = self.gpu.stream();
14285 let mut b = __s_b.launch_builder(&f);
14286 b.arg(b0)
14287 .arg(b1)
14288 .arg(aq)
14289 .arg(ad)
14290 .arg(&mut y0)
14291 .arg(&mut y1)
14292 .arg(&inf)
14293 .arg(&o0)
14294 .arg(&o1)
14295 .arg(&rbl)
14296 .arg(&ws0)
14297 .arg(&ws1);
14298 unsafe {
14299 b.launch(cfg)?;
14300 }
14301 Ok((y0, y1))
14302 }
14303
14304 #[allow(clippy::too_many_arguments)]
14306 fn e4m3_fused3_core(
14307 &self,
14308 b0: &CudaSlice<u8>,
14309 b1: &CudaSlice<u8>,
14310 b2: &CudaSlice<u8>,
14311 aq: &CudaSlice<i8>,
14312 ad: &CudaSlice<f32>,
14313 in_f: usize,
14314 out0: usize,
14315 out1: usize,
14316 out2: usize,
14317 row_bytes: usize,
14318 ws0: f32,
14319 ws1: f32,
14320 ws2: f32,
14321 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14322 const ROWS_PER_BLOCK: u32 = 4;
14323 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14324 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14325 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14326 let f = self.func("qmatvec_e4m3_mmvq_fused3");
14327 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14328 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14329 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14330 let cfg = LaunchConfig {
14331 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14332 block_dim: (32, ROWS_PER_BLOCK, 1),
14333 shared_mem_bytes: 0,
14334 };
14335 let (inf, o0, o1, o2, rbl) = (
14336 in_f as i32,
14337 out0 as i32,
14338 out1 as i32,
14339 out2 as i32,
14340 row_bytes as i64,
14341 );
14342 let __s_b = self.gpu.stream();
14343 let mut b = __s_b.launch_builder(&f);
14344 b.arg(b0)
14345 .arg(b1)
14346 .arg(b2)
14347 .arg(aq)
14348 .arg(ad)
14349 .arg(&mut y0)
14350 .arg(&mut y1)
14351 .arg(&mut y2)
14352 .arg(&inf)
14353 .arg(&o0)
14354 .arg(&o1)
14355 .arg(&o2)
14356 .arg(&rbl)
14357 .arg(&ws0)
14358 .arg(&ws1)
14359 .arg(&ws2);
14360 unsafe {
14361 b.launch(cfg)?;
14362 }
14363 Ok((y0, y1, y2))
14364 }
14365
14366 #[allow(clippy::too_many_arguments)]
14370 fn e4m3_fused2_t_core(
14371 &self,
14372 b0: &CudaSlice<u8>,
14373 b1: &CudaSlice<u8>,
14374 aq: &CudaSlice<i8>,
14375 ad: &CudaSlice<f32>,
14376 m: usize,
14377 in_f: usize,
14378 out0: usize,
14379 out1: usize,
14380 row_bytes: usize,
14381 ws0: f32,
14382 ws1: f32,
14383 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14384 const ROWS_PER_BLOCK: u32 = 4;
14385 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14386 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14387 let f = self.func(match Self::batched_mcols(m) {
14388 2 => "qmatvec_e4m3_mmvq_fused2_b2",
14389 4 => "qmatvec_e4m3_mmvq_fused2_b4",
14390 _ => "qmatvec_e4m3_mmvq_fused2_b8",
14391 });
14392 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14393 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14394 let cfg = LaunchConfig {
14395 grid_dim: (nb0 + nb1, 1, 1),
14396 block_dim: (32, ROWS_PER_BLOCK, 1),
14397 shared_mem_bytes: 0,
14398 };
14399 let (inf, o0, o1, mi, rbl) = (
14400 in_f as i32,
14401 out0 as i32,
14402 out1 as i32,
14403 m as i32,
14404 row_bytes as i64,
14405 );
14406 let __s_b = self.gpu.stream();
14407 let mut b = __s_b.launch_builder(&f);
14408 b.arg(b0)
14409 .arg(b1)
14410 .arg(aq)
14411 .arg(ad)
14412 .arg(&mut y0)
14413 .arg(&mut y1)
14414 .arg(&inf)
14415 .arg(&o0)
14416 .arg(&o1)
14417 .arg(&mi)
14418 .arg(&rbl);
14419 unsafe {
14420 b.launch(cfg)?;
14421 }
14422 if ws0 != 1.0 {
14423 self.scale_inplace(&mut y0, ws0, m * out0)?;
14424 }
14425 if ws1 != 1.0 {
14426 self.scale_inplace(&mut y1, ws1, m * out1)?;
14427 }
14428 Ok((y0, y1))
14429 }
14430
14431 #[allow(clippy::too_many_arguments)]
14433 fn e4m3_fused3_t_core(
14434 &self,
14435 b0: &CudaSlice<u8>,
14436 b1: &CudaSlice<u8>,
14437 b2: &CudaSlice<u8>,
14438 aq: &CudaSlice<i8>,
14439 ad: &CudaSlice<f32>,
14440 m: usize,
14441 in_f: usize,
14442 out0: usize,
14443 out1: usize,
14444 out2: usize,
14445 row_bytes: usize,
14446 ws0: f32,
14447 ws1: f32,
14448 ws2: f32,
14449 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14450 const ROWS_PER_BLOCK: u32 = 4;
14451 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14452 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14453 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14454 let f = self.func(if Self::batched_mcols(m) == 2 {
14455 "qmatvec_e4m3_mmvq_fused3_b2"
14456 } else {
14457 "qmatvec_e4m3_mmvq_fused3_b4"
14458 });
14459 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
14460 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
14461 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
14462 let cfg = LaunchConfig {
14463 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14464 block_dim: (32, ROWS_PER_BLOCK, 1),
14465 shared_mem_bytes: 0,
14466 };
14467 let (inf, o0, o1, o2, mi, rbl) = (
14468 in_f as i32,
14469 out0 as i32,
14470 out1 as i32,
14471 out2 as i32,
14472 m as i32,
14473 row_bytes as i64,
14474 );
14475 let __s_b = self.gpu.stream();
14476 let mut b = __s_b.launch_builder(&f);
14477 b.arg(b0)
14478 .arg(b1)
14479 .arg(b2)
14480 .arg(aq)
14481 .arg(ad)
14482 .arg(&mut y0)
14483 .arg(&mut y1)
14484 .arg(&mut y2)
14485 .arg(&inf)
14486 .arg(&o0)
14487 .arg(&o1)
14488 .arg(&o2)
14489 .arg(&mi)
14490 .arg(&rbl);
14491 unsafe {
14492 b.launch(cfg)?;
14493 }
14494 if ws0 != 1.0 {
14495 self.scale_inplace(&mut y0, ws0, m * out0)?;
14496 }
14497 if ws1 != 1.0 {
14498 self.scale_inplace(&mut y1, ws1, m * out1)?;
14499 }
14500 if ws2 != 1.0 {
14501 self.scale_inplace(&mut y2, ws2, m * out2)?;
14502 }
14503 Ok((y0, y1, y2))
14504 }
14505
14506 pub fn qmatvec_e4m3_blk_mmvq(
14516 &self,
14517 bytes: &CudaSlice<u8>,
14518 aq: &CudaSlice<i8>,
14519 ad: &CudaSlice<f32>,
14520 scales: &CudaSlice<f32>,
14521 m: usize,
14522 in_f: usize,
14523 out_f: usize,
14524 row_bytes: usize,
14525 scale_cols: usize,
14526 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14527 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
14529 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
14530 )?;
14531 Ok(y)
14532 }
14533
14534 #[allow(clippy::too_many_arguments)]
14536 pub fn qmatvec_e4m3_blk_mmvq_into(
14537 &self,
14538 bytes: &CudaSlice<u8>,
14539 aq: &CudaSlice<i8>,
14540 ad: &CudaSlice<f32>,
14541 scales: &CudaSlice<f32>,
14542 m: usize,
14543 in_f: usize,
14544 out_f: usize,
14545 row_bytes: usize,
14546 scale_cols: usize,
14547 y: &mut CudaSlice<f32>,
14548 ) -> Result<(), Box<dyn std::error::Error>> {
14549 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
14551 let cfg = LaunchConfig {
14552 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
14553 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
14556 let (inf, outf, mi, rb, sc) = (
14557 in_f as i32,
14558 out_f as i32,
14559 m as i32,
14560 row_bytes as i64,
14561 scale_cols as i32,
14562 );
14563 let __s_b = self.gpu.stream();
14564 let mut b = __s_b.launch_builder(&f);
14565 b.arg(bytes)
14566 .arg(aq)
14567 .arg(ad)
14568 .arg(scales)
14569 .arg(&mut *y)
14570 .arg(&inf)
14571 .arg(&outf)
14572 .arg(&mi)
14573 .arg(&rb)
14574 .arg(&sc);
14575 unsafe {
14576 b.launch(cfg)?;
14577 }
14578 Ok(())
14579 }
14580
14581 #[allow(clippy::too_many_arguments)]
14587 pub fn qmatvec_e4m3_blk_mmvq_batched(
14588 &self,
14589 bytes: &CudaSlice<u8>,
14590 aq: &CudaSlice<i8>,
14591 ad: &CudaSlice<f32>,
14592 scales: &CudaSlice<f32>,
14593 m: usize,
14594 in_f: usize,
14595 out_f: usize,
14596 row_bytes: usize,
14597 scale_cols: usize,
14598 mcols: usize,
14599 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14600 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
14602 let name = match mcols {
14603 2 => "qmatvec_e4m3_blk_mmvq_b2",
14604 4 => "qmatvec_e4m3_blk_mmvq_b4",
14605 8 => "qmatvec_e4m3_blk_mmvq_b8",
14606 16 => "qmatvec_e4m3_blk_mmvq_b16",
14607 _ => {
14608 return Err(
14609 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
14610 );
14611 }
14612 };
14613 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14614 let f = self.func(name);
14615 let cfg = LaunchConfig {
14616 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
14617 block_dim: (32, ROWS_PER_BLOCK, 1),
14618 shared_mem_bytes: 0,
14619 };
14620 let (inf, outf, mi, rb, sc) = (
14621 in_f as i32,
14622 out_f as i32,
14623 m as i32,
14624 row_bytes as i64,
14625 scale_cols as i32,
14626 );
14627 let __s_b = self.gpu.stream();
14628 let mut b = __s_b.launch_builder(&f);
14629 b.arg(bytes)
14630 .arg(aq)
14631 .arg(ad)
14632 .arg(scales)
14633 .arg(&mut y)
14634 .arg(&inf)
14635 .arg(&outf)
14636 .arg(&mi)
14637 .arg(&rb)
14638 .arg(&sc);
14639 unsafe {
14640 b.launch(cfg)?;
14641 }
14642 Ok(y)
14643 }
14644
14645 #[allow(clippy::too_many_arguments)]
14648 pub fn qmatvec_e4m3_blk_batched_raw(
14649 &self,
14650 bytes: &CudaSlice<u8>,
14651 x: &CudaSlice<f32>,
14652 scales: &CudaSlice<f32>,
14653 m: usize,
14654 in_f: usize,
14655 out_f: usize,
14656 row_bytes: usize,
14657 scale_cols: usize,
14658 mcols: usize,
14659 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14660 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14661 self.qmatvec_e4m3_blk_mmvq_batched(
14662 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
14663 )
14664 }
14665
14666 #[allow(clippy::too_many_arguments)]
14669 pub fn qmatvec_e4m3_blk_mmvq_raw(
14670 &self,
14671 bytes: &CudaSlice<u8>,
14672 x: &CudaSlice<f32>,
14673 scales: &CudaSlice<f32>,
14674 m: usize,
14675 in_f: usize,
14676 out_f: usize,
14677 row_bytes: usize,
14678 scale_cols: usize,
14679 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
14680 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14681 self.qmatvec_e4m3_blk_mmvq(
14682 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
14683 )
14684 }
14685
14686 #[allow(clippy::too_many_arguments)]
14689 pub fn qmatvec_e4m3_fused2_raw(
14690 &self,
14691 b0: &CudaSlice<u8>,
14692 b1: &CudaSlice<u8>,
14693 x: &CudaSlice<f32>,
14694 in_f: usize,
14695 out0: usize,
14696 out1: usize,
14697 row_bytes: usize,
14698 ws0: f32,
14699 ws1: f32,
14700 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14701 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14702 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
14703 }
14704
14705 #[allow(clippy::too_many_arguments)]
14706 pub fn qmatvec_e4m3_fused3_raw(
14707 &self,
14708 b0: &CudaSlice<u8>,
14709 b1: &CudaSlice<u8>,
14710 b2: &CudaSlice<u8>,
14711 x: &CudaSlice<f32>,
14712 in_f: usize,
14713 out0: usize,
14714 out1: usize,
14715 out2: usize,
14716 row_bytes: usize,
14717 ws0: f32,
14718 ws1: f32,
14719 ws2: f32,
14720 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14721 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14722 self.e4m3_fused3_core(
14723 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14724 )
14725 }
14726
14727 #[allow(clippy::too_many_arguments)]
14728 pub fn qmatvec_e4m3_fused2_t_raw(
14729 &self,
14730 b0: &CudaSlice<u8>,
14731 b1: &CudaSlice<u8>,
14732 x: &CudaSlice<f32>,
14733 m: usize,
14734 in_f: usize,
14735 out0: usize,
14736 out1: usize,
14737 row_bytes: usize,
14738 ws0: f32,
14739 ws1: f32,
14740 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14741 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14742 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
14743 }
14744
14745 #[allow(clippy::too_many_arguments)]
14746 pub fn qmatvec_e4m3_fused3_t_raw(
14747 &self,
14748 b0: &CudaSlice<u8>,
14749 b1: &CudaSlice<u8>,
14750 b2: &CudaSlice<u8>,
14751 x: &CudaSlice<f32>,
14752 m: usize,
14753 in_f: usize,
14754 out0: usize,
14755 out1: usize,
14756 out2: usize,
14757 row_bytes: usize,
14758 ws0: f32,
14759 ws1: f32,
14760 ws2: f32,
14761 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14762 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
14763 self.e4m3_fused3_t_core(
14764 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
14765 )
14766 }
14767
14768 fn try_e4m3_blk_pre(
14779 &self,
14780 w: &crate::model::GpuTensor,
14781 aq: &CudaSlice<i8>,
14782 ad: &CudaSlice<f32>,
14783 m: usize,
14784 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14785 use crate::model::GpuTensor;
14786 if let GpuTensor::Quant {
14787 bytes,
14788 qtype,
14789 row_bytes,
14790 blk: Some(g),
14791 ..
14792 } = w
14793 {
14794 if *qtype == QT_F8_E4M3_BLK {
14795 if (2..=16).contains(&m)
14801 && std::env::var("MEMRA_NO_BATCHED").is_err()
14802 && (m <= 4 || Self::b8_enabled())
14803 {
14804 let mcols = Self::batched_mcols(m);
14805 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
14806 bytes,
14807 aq,
14808 ad,
14809 &g.scales,
14810 m,
14811 w.in_features(),
14812 w.out_features(),
14813 *row_bytes,
14814 g.cols,
14815 mcols,
14816 )?));
14817 }
14818 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
14819 bytes,
14820 aq,
14821 ad,
14822 &g.scales,
14823 m,
14824 w.in_features(),
14825 w.out_features(),
14826 *row_bytes,
14827 g.cols,
14828 )?));
14829 }
14830 }
14831 Ok(None)
14832 }
14833
14834 fn try_e4m3_blk_prefill(
14881 &self,
14882 w: &crate::model::GpuTensor,
14883 x: &CudaSlice<f32>,
14884 m: usize,
14885 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
14886 use crate::model::GpuTensor;
14887 let GpuTensor::Quant {
14888 bytes,
14889 qtype,
14890 blk: Some(g),
14891 ..
14892 } = w
14893 else {
14894 return Ok(None);
14895 };
14896 if *qtype != QT_F8_E4M3_BLK {
14897 return Ok(None);
14898 }
14899 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
14904 return Ok(Some(y));
14905 }
14906 let (in_f, out_f) = (w.in_features(), w.out_features());
14907 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
14908 let tmp = GpuTensor::Quant {
14909 bytes: slab,
14910 qtype: QT_Q8_0,
14911 row_bytes: in_f / 32 * 34,
14912 ne: vec![in_f as u64, out_f as u64],
14913 scale: 1.0,
14914 rp: false,
14915 #[cfg(memra_cutlass)]
14916 cutlass: None,
14917 fp8: None,
14918 blk: None,
14919 f16: None,
14920 rp4: None,
14921 };
14922 Ok(Some(self.matmul(&tmp, x, m)?))
14924 }
14925
14926 pub fn matmul_pre_noscale(
14927 &self,
14928 w: &crate::model::GpuTensor,
14929 aq: &CudaSlice<i8>,
14930 ad: &CudaSlice<f32>,
14931 m: usize,
14932 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
14933 use crate::model::GpuTensor;
14934 if m == 1 {
14938 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
14939 return Ok(Some((y, 1.0)));
14940 }
14941 }
14942 if m != 1 || !self.uses_q8_1_fast(w) {
14944 return Ok(None);
14945 }
14946 let in_f = w.in_features();
14947 let out_f = w.out_features();
14948 let (bytes, qtype, row_bytes, scale, rp) = match w {
14949 GpuTensor::Quant {
14950 bytes,
14951 qtype,
14952 row_bytes,
14953 scale,
14954 rp,
14955 ..
14956 } => (bytes, *qtype, *row_bytes, *scale, *rp),
14957 _ => return Ok(None),
14958 };
14959 if self.mmvq_supports(qtype) {
14961 let (mbytes, mrp) = match w {
14963 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
14964 _ => (bytes, rp),
14965 };
14966 let y = self.qmatvec_mmvq(
14967 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
14968 )?;
14969 return Ok(Some((y, scale)));
14970 }
14971 let name = match qtype {
14973 QT_Q8_0 => "qmatvec_q8_0_dp4a",
14974 QT_Q4_K => "qmatvec_q4_K_dp4a",
14975 QT_Q6_K => "qmatvec_q6_K_dp4a",
14976 QT_Q5_K => "qmatvec_q5_K_dp4a",
14977 QT_Q3_K => "qmatvec_q3_K_dp4a",
14978 QT_NVFP4 => {
14979 if rp {
14980 "qmatvec_nvfp4_dp4a_rp"
14981 } else {
14982 "qmatvec_nvfp4_dp4a"
14983 }
14984 }
14985 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
14986 _ => return Ok(None),
14987 };
14988 let f = self.func(name);
14989 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
14990 let cfg = LaunchConfig {
14991 grid_dim: (out_f as u32, m as u32, 1),
14992 block_dim: (128, 1, 1),
14993 shared_mem_bytes: 0,
14994 };
14995 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
14996 let __s_b = self.gpu.stream();
14997 let mut b = __s_b.launch_builder(&f);
14998 b.arg(bytes)
14999 .arg(aq)
15000 .arg(ad)
15001 .arg(&mut y)
15002 .arg(&inf)
15003 .arg(&outf)
15004 .arg(&mi)
15005 .arg(&rb);
15006 unsafe {
15007 b.launch(cfg)?;
15008 }
15009 Ok(Some((y, scale)))
15010 }
15011
15012 pub fn mmvq_supports(&self, qtype: i32) -> bool {
15015 if qtype == QT_F8_E4M3 {
15020 return true;
15021 }
15022 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15023 return false;
15024 }
15025 matches!(
15026 qtype,
15027 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
15028 )
15029 }
15030
15031 pub fn qmatvec_mmvq(
15036 &self,
15037 bytes: &CudaSlice<u8>,
15038 aq: &CudaSlice<i8>,
15039 ad: &CudaSlice<f32>,
15040 m: usize,
15041 in_f: usize,
15042 out_f: usize,
15043 qtype: i32,
15044 row_bytes: usize,
15045 scale: f32,
15046 rp: bool,
15047 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15048 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
15050 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
15051 )?;
15052 Ok(y)
15053 }
15054
15055 #[allow(clippy::too_many_arguments)]
15057 pub fn qmatvec_mmvq_into(
15058 &self,
15059 bytes: &CudaSlice<u8>,
15060 aq: &CudaSlice<i8>,
15061 ad: &CudaSlice<f32>,
15062 m: usize,
15063 in_f: usize,
15064 out_f: usize,
15065 qtype: i32,
15066 row_bytes: usize,
15067 scale: f32,
15068 rp: bool,
15069 y: &mut CudaSlice<f32>,
15070 ) -> Result<(), Box<dyn std::error::Error>> {
15071 debug_assert!(y.len() >= m * out_f);
15072 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
15078 && rp
15079 && m == 1
15080 && out_f >= 64
15081 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
15082 && {
15083 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15084 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
15085 }
15086 {
15087 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
15088 let cfg = LaunchConfig {
15089 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
15090 block_dim: (32, 2, 1),
15091 shared_mem_bytes: 0,
15092 };
15093 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
15094 let __s_b = self.gpu.stream();
15095 let mut b = __s_b.launch_builder(&f);
15096 b.arg(bytes)
15097 .arg(aq)
15098 .arg(ad)
15099 .arg(&mut *y)
15100 .arg(&inf)
15101 .arg(&outf)
15102 .arg(&mi)
15103 .arg(&rb);
15104 unsafe {
15105 b.launch(cfg)?;
15106 }
15107 if scale != 1.0 {
15108 self.scale_inplace(y, scale, out_f)?;
15109 }
15110 return Ok(());
15111 }
15112 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
15121 2
15122 } else {
15123 1
15124 };
15125 if m == 1 && qtype == QT_Q4_0 {
15130 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15131 mr = *Q40MR.get_or_init(|| {
15134 std::env::var("MEMRA_Q40_MR")
15135 .ok()
15136 .and_then(|v| v.parse().ok())
15137 .unwrap_or(1)
15138 });
15139 }
15140 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
15151 let q5_force = q5_mode.as_deref() == Some("2");
15152 let q5_il = qtype == QT_Q5_K
15155 && m == 1
15156 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
15157 if q5_il && !q5_force && out_f > 65536 {
15158 mr = 1;
15159 }
15160 if qtype == QT_Q4_0 && rp && mr != 1 {
15163 mr = 2;
15164 }
15165 if qtype == QT_Q8_0 && rp {
15169 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
15170 mr = *Q80MR.get_or_init(|| {
15171 std::env::var("MEMRA_Q80_MR")
15172 .ok()
15173 .and_then(|v| v.parse().ok())
15174 .unwrap_or(1)
15175 });
15176 }
15177 let name = match (qtype, mr, rp) {
15178 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
15179 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
15180 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
15181 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
15182 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
15183 (QT_Q5_K, 2, _) => {
15184 if q5_il {
15185 "qmatvec_q5_K_mmvq_mr2_il"
15186 } else {
15187 "qmatvec_q5_K_mmvq_mr2"
15188 }
15189 }
15190 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
15191 (QT_Q8_0, _, true)
15196 if in_f % 1024 == 0 && {
15197 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15198 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
15199 } =>
15200 {
15201 "qmatvec_q8_0_mmvq_rpca"
15202 }
15203 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
15204 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
15205 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
15209 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
15210 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
15211 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
15212 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
15213 (QT_Q5_K, _, _) => {
15214 if q5_il {
15215 "qmatvec_q5_K_mmvq_il"
15216 } else {
15217 "qmatvec_q5_K_mmvq"
15218 }
15219 }
15220 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
15221 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
15222 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
15223 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
15224 };
15225 let f = self.func(name);
15226 let rows_per_block = ROWS_PER_BLOCK * mr;
15228 let cfg = LaunchConfig {
15229 grid_dim: (
15230 (out_f as u32 + rows_per_block - 1) / rows_per_block,
15231 m as u32,
15232 1,
15233 ),
15234 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15237 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15238 let __s_b = self.gpu.stream();
15239 let mut b = __s_b.launch_builder(&f);
15240 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
15245 if Self::pdl_on()
15248 && Self::pdl_mmvq_on()
15249 && Self::pdl_nvfp4q8_on()
15250 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
15251 {
15252 use cudarc::driver::{DevicePtr, DevicePtrMut};
15253 let s = &self.gpu.stream();
15254 let (pw, _g0) = bytes.device_ptr(s);
15255 let (paq, _g1) = aq.device_ptr(s);
15256 let (pad, _g2) = ad.device_ptr(s);
15257 let (py, _g3) = y.device_ptr_mut(s);
15258 let mut ps = [
15259 &pw as *const _ as *mut std::ffi::c_void,
15260 &paq as *const _ as *mut _,
15261 &pad as *const _ as *mut _,
15262 &py as *const _ as *mut _,
15263 &inf as *const _ as *mut _,
15264 &outf as *const _ as *mut _,
15265 &mi as *const _ as *mut _,
15266 &rb as *const _ as *mut _,
15267 &scale as *const _ as *mut _,
15268 ];
15269 unsafe {
15270 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15271 }
15272 return Ok(());
15273 }
15274 b.arg(bytes)
15275 .arg(aq)
15276 .arg(ad)
15277 .arg(&mut *y)
15278 .arg(&inf)
15279 .arg(&outf)
15280 .arg(&mi)
15281 .arg(&rb)
15282 .arg(&scale);
15283 unsafe {
15284 b.launch(cfg)?;
15285 }
15286 } else if Self::pdl_on()
15287 && Self::pdl_mmvq_on()
15288 && (matches!(
15289 name,
15290 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
15291 ) || (Self::pdl_nvfp4q8_on()
15292 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
15293 {
15294 {
15298 use cudarc::driver::{DevicePtr, DevicePtrMut};
15299 let s = &self.gpu.stream();
15300 let (pw, _g0) = bytes.device_ptr(s);
15301 let (paq, _g1) = aq.device_ptr(s);
15302 let (pad, _g2) = ad.device_ptr(s);
15303 let (py, _g3) = y.device_ptr_mut(s);
15304 let mut ps = [
15305 &pw as *const _ as *mut std::ffi::c_void,
15306 &paq as *const _ as *mut _,
15307 &pad as *const _ as *mut _,
15308 &py as *const _ as *mut _,
15309 &inf as *const _ as *mut _,
15310 &outf as *const _ as *mut _,
15311 &mi as *const _ as *mut _,
15312 &rb as *const _ as *mut _,
15313 ];
15314 unsafe {
15315 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
15316 }
15317 }
15318 if scale != 1.0 {
15319 self.scale_inplace(y, scale, m * out_f)?;
15320 }
15321 } else {
15322 b.arg(bytes)
15323 .arg(aq)
15324 .arg(ad)
15325 .arg(&mut *y)
15326 .arg(&inf)
15327 .arg(&outf)
15328 .arg(&mi)
15329 .arg(&rb);
15330 unsafe {
15331 b.launch(cfg)?;
15332 }
15333 if scale != 1.0 {
15334 self.scale_inplace(y, scale, m * out_f)?;
15335 }
15336 }
15337 Ok(())
15338 }
15339
15340 pub fn qmatvec_mmvq_raw(
15344 &self,
15345 bytes: &CudaSlice<u8>,
15346 x: &CudaSlice<f32>,
15347 m: usize,
15348 in_f: usize,
15349 out_f: usize,
15350 qtype: i32,
15351 row_bytes: usize,
15352 rp: bool,
15353 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15354 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15355 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
15356 }
15357
15358 pub fn batched_supports(&self, qtype: i32) -> bool {
15362 matches!(
15363 qtype,
15364 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
15365 )
15366 }
15367
15368 pub fn iq_fast_enabled() -> bool {
15376 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15377 *ON.get_or_init(|| {
15378 std::env::var("MEMRA_IQ_FAST")
15379 .map(|v| v != "0")
15380 .unwrap_or(true)
15381 })
15382 }
15383
15384 pub fn b8_enabled() -> bool {
15387 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15388 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
15389 }
15390
15391 pub fn batched_mcols(m: usize) -> usize {
15393 if m == 2 {
15394 2
15395 } else if m <= 4 {
15396 4
15397 } else if m <= 8 {
15398 8
15399 } else {
15400 16
15401 }
15402 }
15403
15404 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
15409 Some(match (qtype, mcols) {
15410 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
15411 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
15412 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
15413 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
15419 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
15420 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
15421 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
15422 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
15425 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
15426 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
15427 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
15428 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
15431 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
15432 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
15433 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
15434 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
15435 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
15436 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
15437 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
15438 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
15442 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
15443 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
15444 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
15445 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
15449 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
15450 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
15451 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
15452 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
15453 _ => return None,
15454 })
15455 }
15456
15457 pub fn sm_count(&self) -> i32 {
15492 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15493 *SMS.get_or_init(|| {
15494 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15495 self.gpu
15496 .ctx
15497 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15498 .unwrap_or(82)
15499 })
15500 }
15501
15502 pub fn batched_variant(
15503 &self,
15504 _m: usize,
15505 in_f: usize,
15506 out_f: usize,
15507 qtype: i32,
15508 row_bytes: usize,
15509 mcols: usize,
15510 rp: bool,
15511 ) -> &'static str {
15512 if qtype == QT_Q8_0 {
15517 return if rp { "rp" } else { "base" };
15518 }
15519 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15520 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
15521 Ok("base") => "base",
15522 Ok("pf") => "pf",
15523 Ok("r2") => "r2",
15524 Ok("r2w8") => "r2w8",
15525 Ok("pfr2") => "pfr2",
15526 Ok("ca") => "ca",
15527 Ok("car2") => "car2",
15528 Ok("rp") => "rp",
15531 Ok("rpr2") => "rpr2",
15532 Ok("rpr2w8") => "rpr2w8",
15533 Ok("rpca") => "rpca",
15536 Ok("rpcar2") => "rpcar2",
15537 Ok("rpsc") => "rpsc",
15544 Ok("rpms") => "rpms",
15545 Ok("rpmsc") => "rpmsc",
15546 Ok("rpks") => "rpks",
15547 Ok("rpksc") => "rpksc",
15548 _ => "auto",
15549 });
15550 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
15554 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15559 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
15560 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
15561 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
15562 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
15563 let sms = *SMS.get_or_init(|| {
15564 use cudarc::driver::sys::CUdevice_attribute_enum as A;
15565 self.gpu
15566 .ctx
15567 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
15568 .unwrap_or(82)
15569 });
15570 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
15590 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15593 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
15594 Ok("base") => "base",
15595 Ok("r2") => "r2",
15596 Ok("r2w8") => "r2w8",
15597 _ => "auto",
15598 });
15599 let variant: &'static str = if qtype == QT_Q4_0 {
15600 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
15604 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
15605 Ok("base") => "base",
15611 Ok("r2") => "r2",
15612 Ok("ms") => "ms",
15613 Ok("sm") => "sm",
15614 Ok("la") => "la",
15615 _ => "auto",
15616 });
15617 let v = if q40 != "auto" {
15618 q40
15619 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
15620 "r2"
15621 } else {
15622 "base"
15623 };
15624 if rp {
15629 match v {
15630 "ms" => "r2ms_rp",
15631 "sm" => "r2sm_rp",
15632 "la" => "r2la_rp",
15633 "r2" => "r2_rp",
15634 _ => "rp",
15635 }
15636 } else if matches!(v, "ms" | "sm" | "la") {
15637 "r2"
15638 } else {
15639 v
15640 }
15641 } else if qtype != QT_NVFP4 && !kq_r2 {
15642 "base"
15643 } else if kq_r2 && rp {
15644 "rp"
15648 } else if kq_r2 {
15649 if kq_bv != "auto" {
15652 if kq_bv == "r2w8" && mcols != 4 {
15653 "r2"
15654 } else {
15655 kq_bv
15656 }
15657 } else if bv != "auto" {
15658 match bv {
15659 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
15660 "r2w8" | "rpr2w8" => {
15661 if mcols != 4 {
15662 "r2"
15663 } else {
15664 "r2w8"
15665 }
15666 }
15667 _ => "base", }
15669 } else {
15670 let blocks = (out_f + 7) / 8;
15671 let waves = blocks as f64 / (7 * sms as usize) as f64;
15672 let filled = blocks >= 4 * sms as usize;
15673 let use_r2 = if qtype == QT_Q4_K {
15674 filled
15675 } else {
15676 waves >= 2.0
15677 };
15678 if use_r2 { "r2" } else { "base" }
15679 }
15680 } else if bv != "auto" {
15681 let v = if bv == "r2w8" && mcols == 2 {
15686 "r2"
15687 } else if bv == "ca" && (!ca_ok || mcols == 8) {
15688 "pf"
15689 } else if bv == "car2" && (!ca_ok || mcols == 8) {
15690 "r2"
15691 } else if bv == "pfr2" && mcols == 8 {
15692 "r2"
15693 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
15694 "rpr2"
15695 }
15696 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
15698 if mcols == 8 { "rpr2w8" } else { "rpr2" }
15699 } else if bv == "rpcar2" && mcols == 2 {
15700 "rpca"
15701 }
15702 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
15705 "rpr2"
15706 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
15707 "rpr2"
15708 } else {
15709 bv
15710 };
15711 if rp {
15712 match v {
15713 "base" | "pf" | "ca" | "rp" => "rp",
15714 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
15715 "r2w8" | "rpr2w8" => {
15716 if mcols == 2 {
15717 "rpr2"
15718 } else {
15719 "rpr2w8"
15720 }
15721 }
15722 other => other, }
15724 } else {
15725 v
15726 }
15727 } else if mcols == 8 {
15728 if rp {
15739 if sc_ok { "rpsc" } else { "rpr2w8" }
15740 } else {
15741 "r2w8"
15742 }
15743 } else if mcols >= 4 {
15744 let blocks = (out_f + 7) / 8;
15748 let r7 = 7 * sms as usize;
15749 let r8 = 8 * sms as usize;
15750 let waves = blocks as f64 / r7 as f64;
15751 let filled = blocks >= 4 * sms as usize;
15752 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
15756 if rp { "rpr2w8" } else { "r2w8" }
15760 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
15761 if rp { "rpr2" } else { "r2" }
15764 } else {
15765 if rp { "rp" } else { "pf" }
15769 }
15770 } else if in_f >= 6144 {
15771 if rp { "rpr2" } else { "r2" }
15775 } else if rp {
15776 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
15781 if sc_ok && waves >= 0.9 && waves <= 1.1 {
15782 "rpsc"
15783 } else {
15784 "rp"
15785 }
15786 } else {
15787 "base"
15788 };
15789 variant
15790 }
15791
15792 pub fn qmatvec_mmvq_batched(
15793 &self,
15794 bytes: &CudaSlice<u8>,
15795 aq: &CudaSlice<i8>,
15796 ad: &CudaSlice<f32>,
15797 m: usize,
15798 in_f: usize,
15799 out_f: usize,
15800 qtype: i32,
15801 row_bytes: usize,
15802 mcols: usize,
15803 scale: f32,
15804 rp: bool,
15805 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15806 const ROWS_PER_BLOCK: u32 = 4;
15807 let forced: Option<&'static str> = {
15812 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
15813 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
15814 .as_deref()
15815 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
15816 };
15817 let variant = match forced {
15818 Some(v) if !rp || v.contains("rp") => v,
15819 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
15820 };
15821 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
15822 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
15823 })?;
15824 let variant = if mcols == 16 {
15828 if rp { "rp" } else { "base" }
15829 } else {
15830 variant
15831 };
15832 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15839 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
15840 if b567
15841 && qtype == QT_NVFP4
15842 && rp
15843 && mcols == 8
15844 && (5..=7).contains(&m)
15845 && matches!(variant, "rpsc" | "rpr2w8")
15846 {
15847 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
15848 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15850 let cfg = LaunchConfig {
15851 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15852 block_dim: (32, ROWS_PER_BLOCK, 1),
15853 shared_mem_bytes: 0,
15854 };
15855 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15856 let __s_b = self.gpu.stream();
15857 let mut b = __s_b.launch_builder(&f);
15858 b.arg(bytes)
15859 .arg(aq)
15860 .arg(ad)
15861 .arg(&mut y)
15862 .arg(&inf)
15863 .arg(&outf)
15864 .arg(&mi)
15865 .arg(&rb);
15866 unsafe {
15867 b.launch(cfg)?;
15868 }
15869 if scale != 1.0 {
15870 self.scale_inplace(&mut y, scale, m * out_f)?;
15871 }
15872 return Ok(y);
15873 }
15874 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
15875 "base" => (base_name.into(), ROWS_PER_BLOCK),
15876 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
15877 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
15878 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
15879 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
15883 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
15884 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
15885 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
15886 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
15887 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
15888 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
15889 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
15891 debug_assert!(
15892 !rp || name.contains("_rp"),
15893 "rp weight dispatched to a GGUF-layout kernel"
15894 );
15895 let f = self.func(&name);
15896 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15897 let smem = if name.contains("_r2sm_rp") {
15899 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
15900 } else {
15901 0
15902 };
15903 let cfg = LaunchConfig {
15904 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
15905 block_dim: (32, ROWS_PER_BLOCK, 1),
15906 shared_mem_bytes: smem,
15907 };
15908 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
15909 let __s_b = self.gpu.stream();
15910 let mut b = __s_b.launch_builder(&f);
15911 b.arg(bytes)
15912 .arg(aq)
15913 .arg(ad)
15914 .arg(&mut y)
15915 .arg(&inf)
15916 .arg(&outf)
15917 .arg(&mi)
15918 .arg(&rb);
15919 unsafe {
15920 b.launch(cfg)?;
15921 }
15922 if scale != 1.0 {
15923 self.scale_inplace(&mut y, scale, m * out_f)?;
15924 }
15925 Ok(y)
15926 }
15927
15928 pub fn qmatvec_batched_raw(
15932 &self,
15933 bytes: &CudaSlice<u8>,
15934 x: &CudaSlice<f32>,
15935 m: usize,
15936 in_f: usize,
15937 out_f: usize,
15938 qtype: i32,
15939 row_bytes: usize,
15940 mcols: usize,
15941 rp: bool,
15942 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15943 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15944 self.qmatvec_mmvq_batched(
15945 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
15946 )
15947 }
15948
15949 pub fn qmatvec_nvfp4_batched_raw(
15951 &self,
15952 bytes: &CudaSlice<u8>,
15953 x: &CudaSlice<f32>,
15954 m: usize,
15955 in_f: usize,
15956 out_f: usize,
15957 row_bytes: usize,
15958 mcols: usize,
15959 rp: bool,
15960 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15961 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
15962 }
15963
15964 fn try_fp4_gemm(
15968 &self,
15969 w: &crate::model::GpuTensor,
15970 x: &CudaSlice<f32>,
15971 m: usize,
15972 in_f: usize,
15973 out_f: usize,
15974 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15975 use crate::model::GpuTensor;
15976 if cfg!(memra_portable_cuda) {
15977 return Ok(None);
15978 }
15979 if std::env::var("MEMRA_FP4").is_err() {
15980 return Ok(None);
15981 }
15982 #[cfg(memra_cutlass)]
15991 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
15992 if let GpuTensor::Quant {
15993 bytes,
15994 qtype,
15995 scale,
15996 row_bytes,
15997 cutlass,
15998 ..
15999 } = w
16000 {
16001 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
16002 if let Some(cw) = cutlass {
16003 let y = self.cutlass_fp4_gemm(
16005 &cw.b_packed,
16006 &cw.sfb_swizzled,
16007 x,
16008 *scale,
16009 m,
16010 out_f,
16011 in_f,
16012 )?;
16013 return Ok(Some(y));
16014 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
16015 let (b_packed, sfb_sw) =
16020 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
16021 let y =
16022 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
16023 return Ok(Some(y));
16024 }
16025 }
16026 }
16027 }
16028 if let GpuTensor::Quant {
16029 bytes,
16030 qtype,
16031 row_bytes,
16032 scale,
16033 rp,
16034 ..
16035 } = w
16036 {
16037 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
16040 let y =
16041 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
16042 return Ok(Some(y));
16043 }
16044 }
16045 Ok(None)
16046 }
16047
16048 pub fn rms_norm_f16out(
16052 &self,
16053 x: &CudaSlice<f32>,
16054 w: &CudaSlice<f32>,
16055 dst: &mut CudaSlice<f32>,
16056 dst16: &mut CudaSlice<u8>,
16057 ncols: usize,
16058 nrows: usize,
16059 eps: f32,
16060 ) -> Result<(), Box<dyn std::error::Error>> {
16061 let f = self.func("rms_norm_f16out_f32");
16062 let cfg = LaunchConfig {
16063 grid_dim: (nrows as u32, 1, 1),
16064 block_dim: (rms_block(), 1, 1),
16065 shared_mem_bytes: 0,
16066 };
16067 let (nc, e) = (ncols as i32, eps);
16068 let __s_b = self.gpu.stream();
16069 let mut b = __s_b.launch_builder(&f);
16070 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
16071 unsafe {
16072 b.launch(cfg)?;
16073 }
16074 Ok(())
16075 }
16076
16077 #[allow(clippy::too_many_arguments)]
16080 pub fn add_rms_norm_f16out(
16081 &self,
16082 a: &CudaSlice<f32>,
16083 b: &CudaSlice<f32>,
16084 w: &CudaSlice<f32>,
16085 res: &mut CudaSlice<f32>,
16086 dst: &mut CudaSlice<f32>,
16087 dst16: &mut CudaSlice<u8>,
16088 ncols: usize,
16089 nrows: usize,
16090 eps: f32,
16091 ) -> Result<(), Box<dyn std::error::Error>> {
16092 let f = self.func("add_rms_norm_f16out_f32");
16093 let cfg = LaunchConfig {
16094 grid_dim: (nrows as u32, 1, 1),
16095 block_dim: (rms_block(), 1, 1),
16096 shared_mem_bytes: 0,
16097 };
16098 let (nc, e) = (ncols as i32, eps);
16099 let __s_lb = self.gpu.stream();
16100 let mut lb = __s_lb.launch_builder(&f);
16101 lb.arg(a)
16102 .arg(b)
16103 .arg(w)
16104 .arg(res)
16105 .arg(dst)
16106 .arg(dst16)
16107 .arg(&nc)
16108 .arg(&e);
16109 unsafe {
16110 lb.launch(cfg)?;
16111 }
16112 Ok(())
16113 }
16114
16115 pub fn matmul_group_xh(
16118 &self,
16119 ws: &[&crate::model::GpuTensor],
16120 x: &CudaSlice<f32>,
16121 xh: &CudaSlice<u8>,
16122 m: usize,
16123 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16124 let mut out = Vec::with_capacity(ws.len());
16125 let in_f = ws[0].in_features();
16126 for w in ws {
16127 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
16128 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
16129 out.push(y);
16130 continue;
16131 }
16132 }
16133 out.push(self.matmul(w, x, m)?);
16134 }
16135 Ok(out)
16136 }
16137
16138 pub fn gdn_pad_mask(
16141 &self,
16142 beta: &mut CudaSlice<f32>,
16143 g_log: &mut CudaSlice<f32>,
16144 len_d: &CudaSlice<i32>,
16145 h: usize,
16146 t: usize,
16147 ) -> Result<(), Box<dyn std::error::Error>> {
16148 let f = self.func("gdn_pad_mask_f32");
16149 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
16150 let (hi, ti) = (h as i32, t as i32);
16151 let __s_b = self.gpu.stream();
16152 let mut b = __s_b.launch_builder(&f);
16153 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
16154 unsafe {
16155 b.launch(cfg)?;
16156 }
16157 Ok(())
16158 }
16159
16160 pub fn row_gather_dev(
16163 &self,
16164 src: &CudaSlice<f32>,
16165 dst: &mut CudaSlice<f32>,
16166 len_d: &CudaSlice<i32>,
16167 ncols: usize,
16168 ) -> Result<(), Box<dyn std::error::Error>> {
16169 let f = self.func("row_gather_dev_f32");
16170 let cfg = LaunchConfig::for_num_elems(ncols as u32);
16171 let nc = ncols as i32;
16172 let __s_b = self.gpu.stream();
16173 let mut b = __s_b.launch_builder(&f);
16174 b.arg(src).arg(dst).arg(len_d).arg(&nc);
16175 unsafe {
16176 b.launch(cfg)?;
16177 }
16178 Ok(())
16179 }
16180
16181 pub fn matmul_group(
16188 &self,
16189 ws: &[&crate::model::GpuTensor],
16190 x: &CudaSlice<f32>,
16191 m: usize,
16192 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
16193 use crate::model::GpuTensor;
16194 let mut out = Vec::with_capacity(ws.len());
16195 let any_mirror = ws
16196 .iter()
16197 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
16198 if m >= 16 && any_mirror && !self.verify_exact_on() {
16199 let in_f = ws[0].in_features();
16200 let xh = self.f16_act(x, m * in_f, in_f)?;
16201 for w in ws {
16202 if w.in_features() == in_f {
16203 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
16204 out.push(y);
16205 continue;
16206 }
16207 }
16208 out.push(self.matmul(w, x, m)?);
16209 }
16210 return Ok(out);
16211 }
16212 for w in ws {
16213 out.push(self.matmul(w, x, m)?);
16214 }
16215 Ok(out)
16216 }
16217
16218 pub fn matmul_group_multi(
16225 &self,
16226 ws: &[&crate::model::GpuTensor],
16227 xs: &[&CudaSlice<f32>],
16228 ms: &[usize],
16229 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
16230 assert_eq!(xs.len(), ms.len());
16231 let in_f = ws[0].in_features();
16232 let total: usize = ms.iter().sum();
16233 let mut xcat = self.uninit(total * in_f)?;
16234 let mut off = 0usize;
16235 for (x, &m) in xs.iter().zip(ms) {
16236 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
16237 off += m;
16238 }
16239 let ys = self.matmul_group(ws, &xcat, total)?;
16240 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
16241 for (w, y) in ws.iter().zip(ys) {
16242 let out_f = w.out_features();
16243 let mut off = 0usize;
16244 for (s, &m) in ms.iter().enumerate() {
16245 let mut ys_s = self.uninit(m * out_f)?;
16246 let src = y.slice(off * out_f..(off + m) * out_f);
16247 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
16248 out[s].push(ys_s);
16249 off += m;
16250 }
16251 }
16252 Ok(out)
16253 }
16254
16255 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
16265 use crate::model::GpuTensor;
16266 if !legacy_quant_gemm_allowed(
16267 cfg!(memra_portable_cuda),
16268 cfg!(memra_hopper_mma),
16269 std::env::var_os("MEMRA_NO_GEMM").is_some(),
16270 ) {
16271 return false;
16272 }
16273 match w {
16274 GpuTensor::Quant { qtype, .. } => {
16275 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
16276 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
16277 }
16278 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
16279 }
16280 }
16281
16282 pub fn qmatvec_gemm(
16289 &self,
16290 w: &crate::model::GpuTensor,
16291 aq: &CudaSlice<i8>,
16292 ad: &CudaSlice<f32>,
16293 m: usize,
16294 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16295 use crate::model::GpuTensor;
16296 let in_f = w.in_features();
16297 let out_f = w.out_features();
16298 let (bytes, qtype, row_bytes, scale, rp) = match w {
16299 GpuTensor::Quant {
16300 bytes,
16301 qtype,
16302 row_bytes,
16303 scale,
16304 rp,
16305 ..
16306 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16307 _ => unreachable!("gemm_supports guaranteed Quant"),
16308 };
16309 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
16315 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
16316 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
16317 if scale != 1.0 {
16318 self.scale_inplace(&mut y, scale, m * out_f)?;
16319 }
16320 return Ok(y);
16321 }
16322 }
16323 let name = match qtype {
16324 QT_Q8_0 => "qmatvec_gemm_q8_0",
16325 QT_Q4_K => "qmatvec_gemm_q4_K",
16326 QT_Q4_0 => {
16327 if rp {
16328 "qmatvec_gemm_q4_0_rp"
16329 } else {
16330 "qmatvec_gemm_q4_0"
16331 }
16332 }
16333 QT_Q5_K => "qmatvec_gemm_q5_K",
16334 QT_Q6_K => "qmatvec_gemm_q6_K",
16335 QT_NVFP4 => {
16336 if rp {
16337 "qmatvec_gemm_nvfp4_rp"
16338 } else {
16339 "qmatvec_gemm_nvfp4"
16340 }
16341 }
16342 _ => unreachable!(),
16343 };
16344 let f = self.func(name);
16345 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);
16350 let k1_tile = if is_k1 {
16352 k1_launch_override().unwrap_or((128, 128, 8))
16353 } else {
16354 (128, 128, 8)
16355 };
16356 let (bm, bn): (u32, u32) = if is_k1 {
16357 (k1_tile.0, k1_tile.1)
16358 } else {
16359 (64, 256)
16360 };
16361 let warps: u32 = if is_k1 {
16362 k1_tile.2
16363 } else {
16364 match qtype {
16365 QT_NVFP4 => 8,
16366 _ => 4,
16367 }
16368 };
16369 let cfg = LaunchConfig {
16370 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16371 block_dim: (32, warps, 1),
16372 shared_mem_bytes: 0,
16373 };
16374 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16375 let __s_b = self.gpu.stream();
16376 let mut b = __s_b.launch_builder(&f);
16377 b.arg(bytes)
16378 .arg(aq)
16379 .arg(ad)
16380 .arg(&mut y)
16381 .arg(&inf)
16382 .arg(&outf)
16383 .arg(&mi)
16384 .arg(&rb);
16385 unsafe {
16386 b.launch(cfg)?;
16387 }
16388 if scale != 1.0 {
16389 self.scale_inplace(&mut y, scale, m * out_f)?;
16390 }
16391 Ok(y)
16392 }
16393
16394 pub fn qmatvec_gemm_raw(
16399 &self,
16400 bytes: &CudaSlice<u8>,
16401 x: &CudaSlice<f32>,
16402 m: usize,
16403 in_f: usize,
16404 out_f: usize,
16405 qtype: i32,
16406 row_bytes: usize,
16407 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16408 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16409 let name = match qtype {
16410 QT_Q8_0 => "qmatvec_gemm_q8_0",
16411 QT_Q4_K => "qmatvec_gemm_q4_K",
16412 QT_Q4_0 => "qmatvec_gemm_q4_0",
16413 QT_Q5_K => "qmatvec_gemm_q5_K",
16414 QT_Q6_K => "qmatvec_gemm_q6_K",
16415 QT_NVFP4 => "qmatvec_gemm_nvfp4",
16416 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
16417 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
16418 };
16419 let f = self.func(name);
16420 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);
16424 let k1_tile = if is_k1 {
16426 k1_launch_override().unwrap_or((128, 128, 8))
16427 } else {
16428 (128, 128, 8)
16429 };
16430 let (bm, bn): (u32, u32) = if is_k1 {
16431 (k1_tile.0, k1_tile.1)
16432 } else {
16433 (64, 256)
16434 };
16435 let warps: u32 = if is_k1 {
16436 k1_tile.2
16437 } else {
16438 match qtype {
16439 QT_NVFP4 | QT_NVFP4_RP => 8,
16440 _ => 4,
16441 }
16442 };
16443 let cfg = LaunchConfig {
16444 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
16445 block_dim: (32, warps, 1),
16446 shared_mem_bytes: 0,
16447 };
16448 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16449 let __s_b = self.gpu.stream();
16450 let mut b = __s_b.launch_builder(&f);
16451 b.arg(bytes)
16452 .arg(&aq)
16453 .arg(&ad)
16454 .arg(&mut y)
16455 .arg(&inf)
16456 .arg(&outf)
16457 .arg(&mi)
16458 .arg(&rb);
16459 unsafe {
16460 b.launch(cfg)?;
16461 }
16462 Ok(y)
16463 }
16464
16465 pub fn qmatvec_gemm_q8_0_wgmma_raw(
16472 &self,
16473 rp4: &CudaSlice<u8>,
16474 aq: &CudaSlice<i8>,
16475 ad: &CudaSlice<f32>,
16476 m: usize,
16477 in_f: usize,
16478 out_f: usize,
16479 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16480 assert!(
16481 out_f % 64 == 0 && in_f % 32 == 0,
16482 "wgmma GEMM needs out_f%64==0, in_f%32==0"
16483 );
16484 let f = self.func("qmatvec_gemm_q8_0_wgmma");
16485 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
16487 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
16488 block_dim: (128, 1, 1),
16489 shared_mem_bytes: 0,
16490 };
16491 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
16492 let __s_b = self.gpu.stream();
16493 let mut b = __s_b.launch_builder(&f);
16494 b.arg(rp4)
16495 .arg(aq)
16496 .arg(ad)
16497 .arg(&mut y)
16498 .arg(&inf)
16499 .arg(&outf)
16500 .arg(&mi);
16501 unsafe {
16502 b.launch(cfg)?;
16503 }
16504 Ok(y)
16505 }
16506
16507 pub fn scale_inplace(
16509 &self,
16510 y: &mut CudaSlice<f32>,
16511 s: f32,
16512 n: usize,
16513 ) -> Result<(), Box<dyn std::error::Error>> {
16514 let f = self.func("scale_f32");
16515 let cfg = LaunchConfig::for_num_elems(n as u32);
16516 let (sf, ni) = (s, n as i32);
16517 let __s_b = self.gpu.stream();
16518 let mut b = __s_b.launch_builder(&f);
16519 b.arg(y).arg(&sf).arg(&ni);
16520 unsafe {
16521 b.launch(cfg)?;
16522 }
16523 Ok(())
16524 }
16525
16526 pub fn bf16_to_f32(
16531 &self,
16532 data: &cudarc::driver::CudaView<'_, u8>,
16533 n: usize,
16534 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16535 let mut out = self.alloc_uninit::<f32>(n)?;
16536 let f = self.func("bf16_to_f32");
16537 let cfg = LaunchConfig::for_num_elems(n as u32);
16538 let ni = n as i32;
16539 let __s_b = self.gpu.stream();
16540 let mut b = __s_b.launch_builder(&f);
16541 b.arg(data).arg(&mut out).arg(&ni);
16542 unsafe {
16543 b.launch(cfg)?;
16544 }
16545 Ok(out)
16546 }
16547
16548 fn linear_bf16_chunked(
16555 &self,
16556 x: &CudaSlice<f32>,
16557 data: &CudaSlice<u8>,
16558 m: usize,
16559 in_f: usize,
16560 out_f: usize,
16561 exact: bool,
16562 canonical_chunk_rows: Option<usize>,
16563 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16564 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16568 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16569 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
16570 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
16571 let started = timing.then(std::time::Instant::now);
16572 let result =
16573 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
16574 if let Some(started) = started {
16575 use std::sync::atomic::Ordering;
16576 self.stream().synchronize()?;
16577 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
16578 + started.elapsed().as_nanos() as u64;
16579 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
16580 + (in_f * out_f * 2) as u64;
16581 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
16582 if calls % 1024 == 0 {
16583 eprintln!(
16584 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
16585 weight_gb={:.2}",
16586 ns as f64 / 1.0e6,
16587 ns as f64 / calls as f64 / 1.0e3,
16588 wb as f64 / 1.0e9,
16589 );
16590 }
16591 }
16592 result
16593 }
16594
16595 pub(crate) fn bf16_mmv_on() -> bool {
16600 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16601 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
16602 }
16603
16604 fn matvec_bf16(
16607 &self,
16608 data: &CudaSlice<u8>,
16609 x: &CudaSlice<f32>,
16610 in_f: usize,
16611 out_f: usize,
16612 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16613 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
16614 return Err(format!(
16615 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
16616 data.len(),
16617 x.len()
16618 )
16619 .into());
16620 }
16621 let mut y = self.alloc_uninit::<f32>(out_f)?;
16622 let f = self.func("matvec_bf16_f32acc");
16623 let cfg = LaunchConfig {
16624 grid_dim: (out_f as u32, 1, 1),
16625 block_dim: (mmv_block(), 1, 1),
16626 shared_mem_bytes: 0,
16627 };
16628 let ini = in_f as i32;
16629 let __s_bld = self.gpu.stream();
16630 let mut bld = __s_bld.launch_builder(&f);
16631 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
16632 unsafe {
16633 bld.launch(cfg)?;
16634 }
16635 Ok(y)
16636 }
16637
16638 #[allow(clippy::too_many_arguments)]
16642 #[allow(clippy::too_many_arguments)]
16647 pub fn qk_norm_rope_append_inc_dcw(
16648 &self,
16649 q_raw: &CudaSlice<f32>,
16650 k_raw: &CudaSlice<f32>,
16651 v_raw: &CudaSlice<f32>,
16652 qw: &CudaSlice<f32>,
16653 kw: &CudaSlice<f32>,
16654 q_out: &mut CudaSlice<f32>,
16655 k_out: &mut CudaSlice<f32>,
16656 pos: &CudaSlice<i32>,
16657 k_plane: &mut CudaSlice<u8>,
16658 v_plane: &mut CudaSlice<u8>,
16659 len_dev: &CudaSlice<i32>,
16662 base_dev: Option<&CudaSlice<i32>>,
16663 done_ctr: &mut CudaSlice<u32>,
16664 kv_dim_k: usize,
16665 kv_dim_v: usize,
16666 k_tok_bytes: usize,
16667 v_tok_bytes: usize,
16668 head_dim: usize,
16669 n_dims: usize,
16670 nh_q: usize,
16671 nh_k: usize,
16672 eps: f32,
16673 freq_base: f32,
16674 freq_scale: f32,
16675 ff: Option<&CudaSlice<f32>>,
16676 ) -> Result<(), Box<dyn std::error::Error>> {
16677 if head_dim != 128
16678 || kv_dim_v != kv_dim_k
16679 || kv_dim_k != nh_k * head_dim
16680 || q_raw.len() < nh_q * head_dim
16681 || k_raw.len() < nh_k * head_dim
16682 || v_raw.len() < kv_dim_v
16683 || q_out.len() < nh_q * head_dim
16684 || k_out.len() < nh_k * head_dim
16685 || pos.is_empty()
16686 || done_ctr.is_empty()
16687 {
16688 return Err(format!(
16689 "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}"
16690 )
16691 .into());
16692 }
16693 let f = self.func("qk_norm_rope_append_inc_dcw");
16694 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16695 let cfg = LaunchConfig {
16696 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16697 block_dim: (128, 1, 1),
16698 shared_mem_bytes: 0,
16699 };
16700 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
16701 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
16702 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16703 let null: u64 = 0;
16704 let __s_b = self.gpu.stream();
16705 let mut b = __s_b.launch_builder(&f);
16706 b.arg(q_raw)
16707 .arg(k_raw)
16708 .arg(v_raw)
16709 .arg(qw)
16710 .arg(kw)
16711 .arg(q_out)
16712 .arg(k_out)
16713 .arg(pos)
16714 .arg(&mut *k_plane)
16715 .arg(&mut *v_plane)
16716 .arg(len_dev);
16717 match base_dev {
16718 Some(base) => {
16719 b.arg(base);
16720 }
16721 None => {
16722 b.arg(&null);
16723 }
16724 }
16725 b.arg(&mut *done_ctr)
16726 .arg(&kvk)
16727 .arg(&kvv)
16728 .arg(&ktb)
16729 .arg(&vtb)
16730 .arg(&hd)
16731 .arg(&nd)
16732 .arg(&nq)
16733 .arg(&eps)
16734 .arg(&theta_scale)
16735 .arg(&freq_scale);
16736 match ff {
16737 Some(freqs) => {
16738 b.arg(freqs);
16739 }
16740 None => {
16741 b.arg(&null);
16742 }
16743 }
16744 unsafe {
16745 b.launch(cfg)?;
16746 }
16747 Ok(())
16748 }
16749
16750 pub fn qk_norm_rope_into(
16751 &self,
16752 q_raw: &CudaSlice<f32>,
16753 k_raw: &CudaSlice<f32>,
16754 qw: &CudaSlice<f32>,
16755 kw: &CudaSlice<f32>,
16756 q_out: &mut CudaSlice<f32>,
16757 k_out: &mut CudaSlice<f32>,
16758 pos: &CudaSlice<i32>,
16759 head_dim: usize,
16760 n_dims: usize,
16761 nh_q: usize,
16762 nh_k: usize,
16763 eps: f32,
16764 freq_base: f32,
16765 freq_scale: f32,
16766 ff: Option<&CudaSlice<f32>>,
16767 ) -> Result<(), Box<dyn std::error::Error>> {
16768 if head_dim > 512
16769 || q_raw.len() < nh_q * head_dim
16770 || k_raw.len() < nh_k * head_dim
16771 || q_out.len() < nh_q * head_dim
16772 || k_out.len() < nh_k * head_dim
16773 || qw.len() < head_dim
16774 || kw.len() < head_dim
16775 || pos.is_empty()
16776 {
16777 return Err(format!(
16778 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
16779 )
16780 .into());
16781 }
16782 let f = self.func("qk_norm_rope_f32");
16783 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
16784 let cfg = LaunchConfig {
16785 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
16786 block_dim: (128, 1, 1),
16787 shared_mem_bytes: 0,
16788 };
16789 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
16790 let __s_b = self.gpu.stream();
16791 let mut b = __s_b.launch_builder(&f);
16792 b.arg(q_raw)
16793 .arg(k_raw)
16794 .arg(qw)
16795 .arg(kw)
16796 .arg(q_out)
16797 .arg(k_out)
16798 .arg(pos)
16799 .arg(&hd)
16800 .arg(&nd)
16801 .arg(&nq)
16802 .arg(&eps)
16803 .arg(&theta_scale)
16804 .arg(&freq_scale);
16805 match ff {
16806 Some(ffv) => {
16807 b.arg(ffv);
16808 unsafe {
16809 b.launch(cfg)?;
16810 }
16811 }
16812 None => {
16813 let null: u64 = 0;
16814 b.arg(&null);
16815 unsafe {
16816 b.launch(cfg)?;
16817 }
16818 }
16819 }
16820 Ok(())
16821 }
16822
16823 #[allow(clippy::too_many_arguments)]
16826 pub fn matvec_f32_b4_into(
16827 &self,
16828 w: [&CudaSlice<f32>; 4],
16829 x: &CudaSlice<f32>,
16830 y: &mut CudaSlice<f32>,
16831 block_cols: usize,
16832 out_f: usize,
16833 ) -> Result<(), Box<dyn std::error::Error>> {
16834 if block_cols % 4 != 0
16835 || x.len() < 4 * block_cols
16836 || y.len() < out_f
16837 || w.iter().any(|w| w.len() != out_f * block_cols)
16838 {
16839 return Err(format!(
16840 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
16841 x.len()
16842 )
16843 .into());
16844 }
16845 let f = self.func("matvec_f32_b4");
16846 let cfg = LaunchConfig {
16847 grid_dim: (out_f as u32, 1, 1),
16848 block_dim: (128, 1, 1),
16849 shared_mem_bytes: 0,
16850 };
16851 let (bc, of) = (block_cols as i32, out_f as i32);
16852 let __s_b = self.gpu.stream();
16853 let mut b = __s_b.launch_builder(&f);
16854 b.arg(w[0])
16855 .arg(w[1])
16856 .arg(w[2])
16857 .arg(w[3])
16858 .arg(x)
16859 .arg(y)
16860 .arg(&bc)
16861 .arg(&of);
16862 unsafe {
16863 b.launch(cfg)?;
16864 }
16865 Ok(())
16866 }
16867
16868 pub fn axpy_rows_seq_into(
16871 &self,
16872 x: &CudaSlice<f32>,
16873 w: &CudaSlice<f32>,
16874 y: &mut CudaSlice<f32>,
16875 width: usize,
16876 n_rows: usize,
16877 ) -> Result<(), Box<dyn std::error::Error>> {
16878 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
16879 return Err(format!(
16880 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
16881 x.len(),
16882 w.len(),
16883 y.len()
16884 )
16885 .into());
16886 }
16887 let f = self.func("axpy_rows_seq_f32");
16888 let cfg = LaunchConfig::for_num_elems(width as u32);
16889 let (wi, nr) = (width as i32, n_rows as i32);
16890 let __s_b = self.gpu.stream();
16891 let mut b = __s_b.launch_builder(&f);
16892 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
16893 unsafe {
16894 b.launch(cfg)?;
16895 }
16896 Ok(())
16897 }
16898
16899 #[allow(clippy::too_many_arguments)]
16903 pub fn axpy_rows_seq_md_off_into(
16904 &self,
16905 x: &CudaSlice<f32>,
16906 w_route: &CudaSlice<f32>,
16907 md: &CudaSlice<f32>,
16908 sel: &CudaSlice<i32>,
16909 y: &mut CudaSlice<f32>,
16910 width: usize,
16911 n_rows: usize,
16912 row0: usize,
16913 ) -> Result<(), Box<dyn std::error::Error>> {
16914 if x.len() < (row0 + n_rows) * width
16915 || w_route.len() < row0 + n_rows
16916 || sel.len() < row0 + n_rows
16917 || y.len() < width
16918 {
16919 return Err(format!(
16920 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
16921 rows={n_rows} row0={row0}",
16922 x.len(),
16923 w_route.len(),
16924 sel.len(),
16925 y.len()
16926 )
16927 .into());
16928 }
16929 let f = self.func("axpy_rows_seq_md_off_f32");
16930 let cfg = LaunchConfig::for_num_elems(width as u32);
16931 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
16932 let __s_b = self.gpu.stream();
16933 let mut b = __s_b.launch_builder(&f);
16934 b.arg(x)
16935 .arg(w_route)
16936 .arg(md)
16937 .arg(sel)
16938 .arg(y)
16939 .arg(&wi)
16940 .arg(&nr)
16941 .arg(&r0);
16942 unsafe {
16943 b.launch(cfg)?;
16944 }
16945 Ok(())
16946 }
16947
16948 #[allow(clippy::too_many_arguments)]
16953 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
16954 &self,
16955 gate_bank: &CudaSlice<u8>,
16956 up_bank: &CudaSlice<u8>,
16957 sel: &CudaSlice<i32>,
16958 aq: &CudaSlice<i8>,
16959 ad: &CudaSlice<f32>,
16960 yg: &mut CudaSlice<f32>,
16961 yu: &mut CudaSlice<f32>,
16962 n_sel: usize,
16963 n_sel_col: usize,
16964 in_f: usize,
16965 out_f: usize,
16966 row_bytes: usize,
16967 expert_stride: usize,
16968 act_row_stride: usize,
16969 ad_row_stride: usize,
16970 ) -> Result<(), Box<dyn std::error::Error>> {
16971 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
16972 if yg.len() < n_sel * out_f
16973 || yu.len() < n_sel * out_f
16974 || sel.len() < n_sel
16975 || n_sel_col == 0
16976 || n_sel % n_sel_col != 0
16977 {
16978 return Err("NVFP4 gu tcol geometry".into());
16979 }
16980 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
16981 let cfg = LaunchConfig {
16982 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
16983 block_dim: (128, 1, 1),
16984 shared_mem_bytes: 0,
16985 };
16986 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
16987 let (rb, es) = (row_bytes as i64, expert_stride as i64);
16988 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
16989 let __s_b = self.gpu.stream();
16990 let mut b = __s_b.launch_builder(&f);
16991 b.arg(gate_bank)
16992 .arg(up_bank)
16993 .arg(sel)
16994 .arg(aq)
16995 .arg(ad)
16996 .arg(yg)
16997 .arg(yu)
16998 .arg(&inf)
16999 .arg(&outf)
17000 .arg(&ns)
17001 .arg(&rb)
17002 .arg(&es)
17003 .arg(&ars)
17004 .arg(&adrs)
17005 .arg(&nsc);
17006 unsafe {
17007 b.launch(cfg)?;
17008 }
17009 Ok(())
17010 }
17011
17012 #[allow(clippy::too_many_arguments)]
17015 pub fn axpy_rows_seq_md_into(
17016 &self,
17017 x: &CudaSlice<f32>,
17018 w_route: &CudaSlice<f32>,
17019 md: &CudaSlice<f32>,
17020 sel: &CudaSlice<i32>,
17021 y: &mut CudaSlice<f32>,
17022 width: usize,
17023 n_rows: usize,
17024 ) -> Result<(), Box<dyn std::error::Error>> {
17025 if x.len() < n_rows * width
17026 || w_route.len() < n_rows
17027 || sel.len() < n_rows
17028 || y.len() < width
17029 {
17030 return Err(format!(
17031 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
17032 x.len(),
17033 w_route.len(),
17034 sel.len(),
17035 y.len()
17036 )
17037 .into());
17038 }
17039 let f = self.func("axpy_rows_seq_md_f32");
17040 let cfg = LaunchConfig::for_num_elems(width as u32);
17041 let (wi, nr) = (width as i32, n_rows as i32);
17042 let __s_b = self.gpu.stream();
17043 let mut b = __s_b.launch_builder(&f);
17044 b.arg(x)
17045 .arg(w_route)
17046 .arg(md)
17047 .arg(sel)
17048 .arg(y)
17049 .arg(&wi)
17050 .arg(&nr);
17051 unsafe {
17052 b.launch(cfg)?;
17053 }
17054 Ok(())
17055 }
17056
17057 #[allow(clippy::too_many_arguments)]
17059 #[allow(clippy::too_many_arguments)]
17063 pub fn matvec_bf16_qkvg_tcol_into(
17064 &self,
17065 wq: &CudaSlice<u8>,
17066 wk: &CudaSlice<u8>,
17067 wv: &CudaSlice<u8>,
17068 wg: &CudaSlice<u8>,
17069 x_t: &CudaSlice<f32>,
17070 yq: &mut CudaSlice<f32>,
17071 yk: &mut CudaSlice<f32>,
17072 yv: &mut CudaSlice<f32>,
17073 yg: &mut CudaSlice<f32>,
17074 in_f: usize,
17075 out_q: usize,
17076 out_kv: usize,
17077 out_g: usize,
17078 t: usize,
17079 ) -> Result<(), Box<dyn std::error::Error>> {
17080 if t == 0
17081 || t > 8
17082 || in_f % 8 != 0
17083 || x_t.len() < t * in_f
17084 || yq.len() < t * out_q
17085 || yk.len() < t * out_kv
17086 || yv.len() < t * out_kv
17087 || (out_g > 0 && yg.len() < t * out_g)
17088 {
17089 return Err("matvec_bf16_qkvg_tcol geometry".into());
17090 }
17091 let f = self.func("matvec_bf16_qkvg_tcol");
17092 let grid = out_q + 2 * out_kv + out_g;
17093 let cfg = LaunchConfig {
17094 grid_dim: (grid as u32, 1, 1),
17095 block_dim: (mmv_block(), 1, 1),
17096 shared_mem_bytes: 0,
17097 };
17098 let (ini, oq, okv, og, ti) = (
17099 in_f as i32,
17100 out_q as i32,
17101 out_kv as i32,
17102 out_g as i32,
17103 t as i32,
17104 );
17105 let __s_b = self.gpu.stream();
17106 let mut b = __s_b.launch_builder(&f);
17107 b.arg(wq)
17108 .arg(wk)
17109 .arg(wv)
17110 .arg(wg)
17111 .arg(x_t)
17112 .arg(yq)
17113 .arg(yk)
17114 .arg(yv)
17115 .arg(yg)
17116 .arg(&ini)
17117 .arg(&oq)
17118 .arg(&okv)
17119 .arg(&og)
17120 .arg(&ti);
17121 unsafe {
17122 b.launch(cfg)?;
17123 }
17124 Ok(())
17125 }
17126
17127 pub fn matvec_bf16_qkvg_into(
17128 &self,
17129 wq: &CudaSlice<u8>,
17130 wk: &CudaSlice<u8>,
17131 wv: &CudaSlice<u8>,
17132 wg: &CudaSlice<u8>,
17133 x: &CudaSlice<f32>,
17134 yq: &mut CudaSlice<f32>,
17135 yk: &mut CudaSlice<f32>,
17136 yv: &mut CudaSlice<f32>,
17137 yg: &mut CudaSlice<f32>,
17138 in_f: usize,
17139 out_q: usize,
17140 out_kv: usize,
17141 out_g: usize,
17142 ) -> Result<(), Box<dyn std::error::Error>> {
17143 if in_f % 8 != 0
17144 || wq.len() != out_q * in_f * 2
17145 || wk.len() != out_kv * in_f * 2
17146 || wv.len() != out_kv * in_f * 2
17147 || wg.len() < out_g * in_f * 2
17148 || x.len() < in_f
17149 || yq.len() < out_q
17150 || yk.len() < out_kv
17151 || yv.len() < out_kv
17152 || (out_g > 0 && yg.len() < out_g)
17153 {
17154 return Err(format!(
17155 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
17156 )
17157 .into());
17158 }
17159 let f = self.func("matvec_bf16_qkvg");
17160 let cfg = LaunchConfig {
17161 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
17162 block_dim: (mmv_block(), 1, 1),
17163 shared_mem_bytes: 0,
17164 };
17165 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
17166 let __s_b = self.gpu.stream();
17167 let mut b = __s_b.launch_builder(&f);
17168 b.arg(wq)
17169 .arg(wk)
17170 .arg(wv)
17171 .arg(wg)
17172 .arg(x)
17173 .arg(yq)
17174 .arg(yk)
17175 .arg(yv)
17176 .arg(yg)
17177 .arg(&inf)
17178 .arg(&oq)
17179 .arg(&okv)
17180 .arg(&og);
17181 unsafe {
17182 b.launch(cfg)?;
17183 }
17184 Ok(())
17185 }
17186
17187 pub fn matvec_bf16_b4_into(
17189 &self,
17190 w: [&CudaSlice<u8>; 4],
17191 x: &CudaSlice<f32>,
17192 y: &mut CudaSlice<f32>,
17193 block_cols: usize,
17194 out_f: usize,
17195 ) -> Result<(), Box<dyn std::error::Error>> {
17196 if block_cols % 8 != 0
17197 || x.len() < 4 * block_cols
17198 || y.len() < out_f
17199 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17200 {
17201 return Err(format!(
17202 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
17203 x.len()
17204 )
17205 .into());
17206 }
17207 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17210 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
17211 let f = self.func(if x2 {
17212 "matvec_bf16_b4_x2"
17213 } else {
17214 "matvec_bf16_b4"
17215 });
17216 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
17217 let cfg = LaunchConfig {
17218 grid_dim: (grid as u32, 1, 1),
17219 block_dim: (mmv_block(), 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 matvec_bf16_b4_tcol_into(
17245 &self,
17246 w: [&CudaSlice<u8>; 4],
17247 x_t: &CudaSlice<f32>,
17248 y_t: &mut CudaSlice<f32>,
17249 block_cols: usize,
17250 out_f: usize,
17251 t: usize,
17252 ) -> Result<(), Box<dyn std::error::Error>> {
17253 if block_cols % 8 != 0
17254 || t == 0
17255 || t > 8
17256 || x_t.len() < t * 4 * block_cols
17257 || y_t.len() < t * out_f
17258 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
17259 {
17260 return Err(format!(
17261 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
17262 x_t.len()
17263 )
17264 .into());
17265 }
17266 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
17267 return Err(
17268 "b4 tcol verify is qualified against the plain b4 kernel only \
17269 (MEMRA_B4_X2=1 is a different t=1 program)"
17270 .into(),
17271 );
17272 }
17273 let f = self.func("matvec_bf16_b4_tcol");
17274 let cfg = LaunchConfig {
17275 grid_dim: (out_f as u32, 1, 1),
17276 block_dim: (mmv_block(), 1, 1),
17277 shared_mem_bytes: 0,
17278 };
17279 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
17280 let __s_b = self.gpu.stream();
17281 let mut b = __s_b.launch_builder(&f);
17282 b.arg(w[0])
17283 .arg(w[1])
17284 .arg(w[2])
17285 .arg(w[3])
17286 .arg(x_t)
17287 .arg(y_t)
17288 .arg(&bc)
17289 .arg(&of)
17290 .arg(&ti);
17291 unsafe {
17292 b.launch(cfg)?;
17293 }
17294 Ok(())
17295 }
17296
17297 pub fn matvec_bf16_into(
17299 &self,
17300 data: &CudaSlice<u8>,
17301 x: &CudaSlice<f32>,
17302 y: &mut CudaSlice<f32>,
17303 in_f: usize,
17304 out_f: usize,
17305 ) -> Result<(), Box<dyn std::error::Error>> {
17306 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17307 return Err(format!(
17308 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17309 data.len(),
17310 x.len(),
17311 y.len()
17312 )
17313 .into());
17314 }
17315 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17319 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
17320 && in_f <= 2048;
17321 if x4 {
17322 let f = self.func("matvec_bf16_f32acc_x4");
17323 let cfg = LaunchConfig {
17324 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
17325 block_dim: (mmv_block(), 1, 1),
17326 shared_mem_bytes: 0,
17327 };
17328 let (ini, outi) = (in_f as i32, out_f as i32);
17329 let __s_b = self.gpu.stream();
17330 let mut b = __s_b.launch_builder(&f);
17331 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
17332 unsafe {
17333 b.launch(cfg)?;
17334 }
17335 return Ok(());
17336 }
17337 let f = self.func("matvec_bf16_f32acc");
17338 let cfg = LaunchConfig {
17339 grid_dim: (out_f as u32, 1, 1),
17340 block_dim: (mmv_block(), 1, 1),
17341 shared_mem_bytes: 0,
17342 };
17343 let ini = in_f as i32;
17344 let __s_b = self.gpu.stream();
17345 let mut b = __s_b.launch_builder(&f);
17346 b.arg(data).arg(x).arg(y).arg(&ini);
17347 unsafe {
17348 b.launch(cfg)?;
17349 }
17350 Ok(())
17351 }
17352
17353 pub fn matvec_bf16_view_into(
17356 &self,
17357 data: &cudarc::driver::CudaView<'_, u8>,
17358 x: &CudaSlice<f32>,
17359 y: &mut CudaSlice<f32>,
17360 in_f: usize,
17361 out_f: usize,
17362 ) -> Result<(), Box<dyn std::error::Error>> {
17363 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
17364 return Err(format!(
17365 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
17366 data.len(),
17367 x.len(),
17368 y.len()
17369 )
17370 .into());
17371 }
17372 let f = self.func("matvec_bf16_f32acc");
17373 let cfg = LaunchConfig {
17374 grid_dim: (out_f as u32, 1, 1),
17375 block_dim: (mmv_block(), 1, 1),
17376 shared_mem_bytes: 0,
17377 };
17378 let ini = in_f as i32;
17379 let __s_b = self.gpu.stream();
17380 let mut b = __s_b.launch_builder(&f);
17381 b.arg(data).arg(x).arg(y).arg(&ini);
17382 unsafe {
17383 b.launch(cfg)?;
17384 }
17385 Ok(())
17386 }
17387
17388 pub fn matvec_bf16_raw_out(
17391 &self,
17392 w: &CudaSlice<u8>,
17393 x: &CudaSlice<f32>,
17394 y_raw: u64,
17395 in_f: usize,
17396 out_f: usize,
17397 ) -> Result<(), Box<dyn std::error::Error>> {
17398 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
17399 return Err("matvec_bf16_raw_out geometry".into());
17400 }
17401 let f = self.func("matvec_bf16_f32acc");
17402 let cfg = LaunchConfig {
17403 grid_dim: (out_f as u32, 1, 1),
17404 block_dim: (mmv_block(), 1, 1),
17405 shared_mem_bytes: 0,
17406 };
17407 let ini = in_f as i32;
17408 let __s_b = self.gpu.stream();
17409 let mut b = __s_b.launch_builder(&f);
17410 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
17411 unsafe {
17412 b.launch(cfg)?;
17413 }
17414 Ok(())
17415 }
17416
17417 pub fn add3_raw(
17421 &self,
17422 a: &CudaSlice<f32>,
17423 b: &CudaSlice<f32>,
17424 sh_raw: u64,
17425 scale_raw: u64,
17426 dst: &mut CudaSlice<f32>,
17427 n: usize,
17428 ) -> Result<(), Box<dyn std::error::Error>> {
17429 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
17430 return Err("add3_raw geometry".into());
17431 }
17432 let f = self.func("add3_f32");
17433 let cfg = LaunchConfig {
17434 grid_dim: ((n as u32).div_ceil(256), 1, 1),
17435 block_dim: (256, 1, 1),
17436 shared_mem_bytes: 0,
17437 };
17438 let ni = n as i32;
17439 let __s_b = self.gpu.stream();
17440 let mut bld = __s_b.launch_builder(&f);
17441 bld.arg(a)
17442 .arg(b)
17443 .arg(&sh_raw)
17444 .arg(&scale_raw)
17445 .arg(dst)
17446 .arg(&ni);
17447 unsafe {
17448 bld.launch(cfg)?;
17449 }
17450 Ok(())
17451 }
17452
17453 pub fn matvec_bf16_down_addscale_into(
17456 &self,
17457 w: &CudaSlice<u8>,
17458 x: &CudaSlice<f32>,
17459 scale: &CudaSlice<f32>,
17460 dst: &mut CudaSlice<f32>,
17461 in_f: usize,
17462 out_f: usize,
17463 ) -> Result<(), Box<dyn std::error::Error>> {
17464 if w.len() != in_f * out_f * 2
17465 || x.len() < in_f
17466 || in_f % 8 != 0
17467 || dst.len() < out_f
17468 || scale.is_empty()
17469 {
17470 return Err("matvec_bf16_down_addscale geometry".into());
17471 }
17472 let f = self.func("matvec_bf16_down_addscale");
17473 let cfg = LaunchConfig {
17474 grid_dim: (out_f as u32, 1, 1),
17475 block_dim: (mmv_block(), 1, 1),
17476 shared_mem_bytes: 0,
17477 };
17478 let ini = in_f as i32;
17479 let __s_b = self.gpu.stream();
17480 let mut b = __s_b.launch_builder(&f);
17481 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
17482 unsafe {
17483 b.launch(cfg)?;
17484 }
17485 Ok(())
17486 }
17487
17488 pub fn matvec_bf16_dual_silu_into(
17491 &self,
17492 wg: &CudaSlice<u8>,
17493 wu: &CudaSlice<u8>,
17494 x: &CudaSlice<f32>,
17495 act: &mut CudaSlice<f32>,
17496 in_f: usize,
17497 out_f: usize,
17498 limit: Option<f32>,
17499 ) -> Result<(), Box<dyn std::error::Error>> {
17500 if wg.len() != in_f * out_f * 2
17501 || wu.len() != in_f * out_f * 2
17502 || x.len() < in_f
17503 || in_f % 8 != 0
17504 || act.len() < out_f
17505 {
17506 return Err("matvec_bf16_dual_silu geometry".into());
17507 }
17508 let f = self.func("matvec_bf16_dual_silu");
17509 let cfg = LaunchConfig {
17510 grid_dim: (out_f as u32, 1, 1),
17511 block_dim: (mmv_block(), 1, 1),
17512 shared_mem_bytes: 0,
17513 };
17514 let (ini, outi) = (in_f as i32, out_f as i32);
17515 let lim = limit.unwrap_or(0.0);
17516 let __s_b = self.gpu.stream();
17517 let mut b = __s_b.launch_builder(&f);
17518 b.arg(wg)
17519 .arg(wu)
17520 .arg(x)
17521 .arg(act)
17522 .arg(&ini)
17523 .arg(&outi)
17524 .arg(&lim);
17525 unsafe {
17526 b.launch(cfg)?;
17527 }
17528 Ok(())
17529 }
17530
17531 #[allow(clippy::too_many_arguments)]
17534 pub fn matvec_bf16_dual_view_into(
17535 &self,
17536 wg: &cudarc::driver::CudaView<'_, u8>,
17537 wu: &cudarc::driver::CudaView<'_, u8>,
17538 x: &CudaSlice<f32>,
17539 yg: &mut CudaSlice<f32>,
17540 yu: &mut CudaSlice<f32>,
17541 in_f: usize,
17542 out_f: usize,
17543 ) -> Result<(), Box<dyn std::error::Error>> {
17544 if wg.len() != in_f * out_f * 2
17545 || wu.len() != in_f * out_f * 2
17546 || x.len() < in_f
17547 || in_f % 8 != 0
17548 || yg.len() < out_f
17549 || yu.len() < out_f
17550 {
17551 return Err(format!(
17552 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17553 wg.len(),
17554 wu.len(),
17555 x.len()
17556 )
17557 .into());
17558 }
17559 let f = self.func("matvec_bf16_dual");
17560 let cfg = LaunchConfig {
17561 grid_dim: ((2 * out_f) as u32, 1, 1),
17562 block_dim: (mmv_block(), 1, 1),
17563 shared_mem_bytes: 0,
17564 };
17565 let (ini, outi) = (in_f as i32, out_f as i32);
17566 let __s_b = self.gpu.stream();
17567 let mut b = __s_b.launch_builder(&f);
17568 b.arg(wg)
17569 .arg(wu)
17570 .arg(x)
17571 .arg(yg)
17572 .arg(yu)
17573 .arg(&ini)
17574 .arg(&outi);
17575 unsafe {
17576 b.launch(cfg)?;
17577 }
17578 Ok(())
17579 }
17580
17581 #[allow(clippy::too_many_arguments)]
17583 pub fn matvec_bf16_dual_into(
17584 &self,
17585 wg: &CudaSlice<u8>,
17586 wu: &CudaSlice<u8>,
17587 x: &CudaSlice<f32>,
17588 yg: &mut CudaSlice<f32>,
17589 yu: &mut CudaSlice<f32>,
17590 in_f: usize,
17591 out_f: usize,
17592 ) -> Result<(), Box<dyn std::error::Error>> {
17593 if wg.len() != in_f * out_f * 2
17594 || wu.len() != in_f * out_f * 2
17595 || x.len() < in_f
17596 || in_f % 8 != 0
17597 || yg.len() < out_f
17598 || yu.len() < out_f
17599 {
17600 return Err(format!(
17601 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
17602 wg.len(),
17603 wu.len(),
17604 x.len()
17605 )
17606 .into());
17607 }
17608 let f = self.func("matvec_bf16_dual");
17609 let cfg = LaunchConfig {
17610 grid_dim: ((2 * out_f) as u32, 1, 1),
17611 block_dim: (mmv_block(), 1, 1),
17612 shared_mem_bytes: 0,
17613 };
17614 let (ini, outi) = (in_f as i32, out_f as i32);
17615 let __s_b = self.gpu.stream();
17616 let mut b = __s_b.launch_builder(&f);
17617 b.arg(wg)
17618 .arg(wu)
17619 .arg(x)
17620 .arg(yg)
17621 .arg(yu)
17622 .arg(&ini)
17623 .arg(&outi);
17624 unsafe {
17625 b.launch(cfg)?;
17626 }
17627 Ok(())
17628 }
17629
17630 pub(crate) fn matvec_bf16_dual(
17633 &self,
17634 wg: &CudaSlice<u8>,
17635 wu: &CudaSlice<u8>,
17636 x: &CudaSlice<f32>,
17637 in_f: usize,
17638 out_f: usize,
17639 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
17640 if wg.len() != in_f * out_f * 2
17641 || wu.len() != in_f * out_f * 2
17642 || x.len() < in_f
17643 || in_f % 8 != 0
17644 {
17645 return Err(format!(
17646 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
17647 wg.len(),
17648 wu.len(),
17649 x.len()
17650 )
17651 .into());
17652 }
17653 let mut yg = self.alloc_uninit::<f32>(out_f)?;
17654 let mut yu = self.alloc_uninit::<f32>(out_f)?;
17655 let f = self.func("matvec_bf16_dual");
17656 let cfg = LaunchConfig {
17657 grid_dim: ((2 * out_f) as u32, 1, 1),
17658 block_dim: (mmv_block(), 1, 1),
17659 shared_mem_bytes: 0,
17660 };
17661 let (ini, outi) = (in_f as i32, out_f as i32);
17662 let __s_b = self.gpu.stream();
17663 let mut b = __s_b.launch_builder(&f);
17664 b.arg(wg)
17665 .arg(wu)
17666 .arg(x)
17667 .arg(&mut yg)
17668 .arg(&mut yu)
17669 .arg(&ini)
17670 .arg(&outi);
17671 unsafe {
17672 b.launch(cfg)?;
17673 }
17674 Ok((yg, yu))
17675 }
17676
17677 #[allow(clippy::too_many_arguments)]
17678 fn linear_bf16_chunked_inner(
17679 &self,
17680 x: &CudaSlice<f32>,
17681 data: &CudaSlice<u8>,
17682 m: usize,
17683 in_f: usize,
17684 out_f: usize,
17685 exact: bool,
17686 canonical_chunk_rows: Option<usize>,
17687 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17688 const CHUNK_BYTES: usize = 256 << 20;
17689 if m == 1
17692 && !exact
17693 && canonical_chunk_rows.is_none()
17694 && in_f % 8 == 0
17695 && Self::bf16_mmv_on()
17696 {
17697 return self.matvec_bf16(data, x, in_f, out_f);
17698 }
17699 let row_bytes = in_f
17700 .checked_mul(std::mem::size_of::<f32>())
17701 .ok_or("BF16 chunk row byte count overflow")?;
17702 if row_bytes == 0 || out_f == 0 {
17703 return Err("BF16 chunk dimensions must be nonzero".into());
17704 }
17705 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
17706 let chunk_rows = match canonical_chunk_rows {
17707 Some(rows) if rows == 0 => {
17708 return Err("canonical BF16 chunk rows must be nonzero".into());
17709 }
17710 Some(rows) if rows > max_chunk_rows => {
17711 return Err(format!(
17712 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
17713 )
17714 .into());
17715 }
17716 Some(rows) if out_f % rows != 0 => {
17717 return Err(format!(
17718 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
17719 )
17720 .into());
17721 }
17722 Some(rows) => rows,
17723 None => max_chunk_rows,
17724 };
17725 if chunk_rows >= out_f {
17726 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
17727 return if exact {
17728 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
17729 } else {
17730 self.linear(x, &wf32, m, in_f, out_f)
17731 };
17732 }
17733 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17734 let mut r0 = 0usize;
17735 while r0 < out_f {
17736 let rows = chunk_rows.min(out_f - r0);
17737 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
17738 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
17739 let yc = if exact {
17740 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
17741 } else {
17742 self.linear(x, &wf32, m, in_f, rows)?
17743 };
17744 for mi in 0..m {
17746 let src = yc.slice(mi * rows..(mi + 1) * rows);
17747 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
17748 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
17749 }
17750 r0 += rows;
17751 }
17752 Ok(y)
17753 }
17754
17755 pub fn linear_bf16_resident(
17759 &self,
17760 x: &CudaSlice<f32>,
17761 data: &CudaSlice<u8>,
17762 m: usize,
17763 in_f: usize,
17764 out_f: usize,
17765 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17766 if data.len() != in_f * out_f * 2 {
17767 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17768 }
17769 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
17770 }
17771
17772 pub fn linear_bf16_resident_canonical_rows(
17778 &self,
17779 x: &CudaSlice<f32>,
17780 data: &CudaSlice<u8>,
17781 m: usize,
17782 in_f: usize,
17783 out_f: usize,
17784 canonical_chunk_rows: usize,
17785 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17786 if data.len() != in_f * out_f * 2 {
17787 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
17788 }
17789 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
17790 }
17791
17792 pub fn linear_f32_resident_canonical_rows(
17797 &self,
17798 x: &CudaSlice<f32>,
17799 data: &CudaSlice<f32>,
17800 m: usize,
17801 in_f: usize,
17802 out_f: usize,
17803 canonical_chunk_rows: usize,
17804 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17805 self.linear_f32_resident_canonical_rows_inner(
17806 x,
17807 data,
17808 m,
17809 in_f,
17810 out_f,
17811 canonical_chunk_rows,
17812 false,
17813 )
17814 }
17815
17816 pub fn linear_f32_resident_canonical_rows_strided(
17822 &self,
17823 x: &CudaSlice<f32>,
17824 data: &CudaSlice<f32>,
17825 m: usize,
17826 in_f: usize,
17827 out_f: usize,
17828 canonical_chunk_rows: usize,
17829 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17830 self.linear_f32_resident_canonical_rows_inner(
17831 x,
17832 data,
17833 m,
17834 in_f,
17835 out_f,
17836 canonical_chunk_rows,
17837 true,
17838 )
17839 }
17840
17841 fn linear_f32_resident_canonical_rows_inner(
17842 &self,
17843 x: &CudaSlice<f32>,
17844 data: &CudaSlice<f32>,
17845 m: usize,
17846 in_f: usize,
17847 out_f: usize,
17848 canonical_chunk_rows: usize,
17849 strided_output: bool,
17850 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17851 if data.len() != in_f * out_f {
17852 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
17853 }
17854 if canonical_chunk_rows == 0
17855 || canonical_chunk_rows > out_f
17856 || out_f % canonical_chunk_rows != 0
17857 {
17858 return Err(format!(
17859 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
17860 )
17861 .into());
17862 }
17863 if canonical_chunk_rows == out_f {
17864 return self.linear(x, data, m, in_f, out_f);
17865 }
17866
17867 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17868 let input = x.slice(0..x.len());
17869 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
17870 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
17871 if m == 1 {
17872 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
17873 self.linear_device_into(
17874 &input,
17875 &weights,
17876 &mut destination,
17877 1,
17878 in_f,
17879 canonical_chunk_rows,
17880 )?;
17881 continue;
17882 }
17883 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
17884 if strided_output {
17885 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
17886 } else {
17887 for token in 0..m {
17888 let source = chunk
17889 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
17890 let mut destination =
17891 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
17892 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
17893 }
17894 }
17895 }
17896 Ok(y)
17897 }
17898
17899 pub fn linear_f32_resident_canonical_rows_t1_into(
17904 &self,
17905 x: &CudaSlice<f32>,
17906 data: &CudaSlice<f32>,
17907 y: &mut CudaSlice<f32>,
17908 in_f: usize,
17909 out_f: usize,
17910 canonical_chunk_rows: usize,
17911 ) -> Result<(), Box<dyn std::error::Error>> {
17912 if data.len() != in_f * out_f {
17913 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
17914 }
17915 if y.len() != out_f || x.len() != in_f {
17916 return Err(format!(
17917 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
17918 x.len(),
17919 y.len()
17920 )
17921 .into());
17922 }
17923 if canonical_chunk_rows == 0
17924 || canonical_chunk_rows > out_f
17925 || out_f % canonical_chunk_rows != 0
17926 {
17927 return Err(format!(
17928 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
17929 )
17930 .into());
17931 }
17932 let input = x.slice(0..x.len());
17933 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
17934 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
17935 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
17936 self.linear_device_into(
17937 &input,
17938 &weights,
17939 &mut destination,
17940 1,
17941 in_f,
17942 canonical_chunk_rows,
17943 )?;
17944 }
17945 Ok(())
17946 }
17947
17948 pub fn linear_t1_into(
17951 &self,
17952 x: &cudarc::driver::CudaView<'_, f32>,
17953 w: &cudarc::driver::CudaView<'_, f32>,
17954 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
17955 in_f: usize,
17956 out_f: usize,
17957 ) -> Result<(), Box<dyn std::error::Error>> {
17958 self.linear_device_into(x, w, y, 1, in_f, out_f)
17959 }
17960
17961 pub fn linear_decode_exact(
17968 &self,
17969 x: &CudaSlice<f32>,
17970 w: &CudaSlice<f32>,
17971 m_tokens: usize,
17972 in_f: usize,
17973 out_f: usize,
17974 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17975 if m_tokens == 1 {
17976 return self.linear(x, w, 1, in_f, out_f);
17977 }
17978 let xv = self.view(x, m_tokens * in_f);
17979 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
17980 for t in 0..m_tokens {
17981 let row = xv.slice(t * in_f..(t + 1) * in_f);
17982 let mut xr = self.alloc_uninit::<f32>(in_f)?;
17983 self.copy_view_into(&mut xr, 0, &row, in_f)?;
17984 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
17985 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
17986 }
17987 Ok(y)
17988 }
17989
17990 pub fn linear(
17991 &self,
17992 x: &CudaSlice<f32>,
17993 w: &CudaSlice<f32>,
17994 m_tokens: usize,
17995 in_f: usize,
17996 out_f: usize,
17997 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17998 self.linear_device(x, w, m_tokens, in_f, out_f)
17999 }
18000
18001 fn linear_device<I>(
18002 &self,
18003 x: &I,
18004 w: &I,
18005 m_tokens: usize,
18006 in_f: usize,
18007 out_f: usize,
18008 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
18009 where
18010 I: cudarc::driver::DevicePtr<f32>,
18011 {
18012 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)?;
18014 Ok(c)
18015 }
18016
18017 fn linear_device_into<I, O>(
18018 &self,
18019 x: &I,
18020 w: &I,
18021 c: &mut O,
18022 m_tokens: usize,
18023 in_f: usize,
18024 out_f: usize,
18025 ) -> Result<(), Box<dyn std::error::Error>>
18026 where
18027 I: cudarc::driver::DevicePtr<f32>,
18028 O: cudarc::driver::DevicePtrMut<f32>,
18029 {
18030 use cudarc::cublaslt::{Matmul, MatmulConfig};
18031 let cfg = MatmulConfig {
18032 transa: true,
18033 transb: false,
18034 transc: false,
18035 m: out_f as u64,
18036 n: m_tokens as u64,
18037 k: in_f as u64,
18038 alpha: 1.0,
18039 lda: in_f as i64,
18040 ldb: in_f as i64,
18041 beta: 0.0,
18042 ldc: out_f as i64,
18043 stride_a: None,
18044 stride_b: None,
18045 stride_c: None,
18046 stride_bias: None,
18047 batch_size: None,
18048 };
18049 let blas = self.gpu.blas();
18050 unsafe {
18051 blas.matmul(cfg, w, x, c, None, None)?;
18052 }
18053 Ok(())
18054 }
18055
18056 pub fn sdpa_naive(
18058 &self,
18059 q: &CudaSlice<f32>,
18060 k: &CudaSlice<f32>,
18061 v: &CudaSlice<f32>,
18062 o: &mut CudaSlice<f32>,
18063 head_dim: usize,
18064 n_head: usize,
18065 n_head_kv: usize,
18066 t: usize,
18067 t_kv: usize,
18068 scale: f32,
18069 causal: bool,
18070 ) -> Result<(), Box<dyn std::error::Error>> {
18071 let f = self.func("sdpa_naive_f32");
18072 let cfg = LaunchConfig {
18073 grid_dim: (n_head as u32, t as u32, 1),
18074 block_dim: (128, 1, 1),
18075 shared_mem_bytes: (t_kv * 4) as u32,
18076 };
18077 let (hd, nh, nhkv, ti, tkvi, cz) = (
18078 head_dim as i32,
18079 n_head as i32,
18080 n_head_kv as i32,
18081 t as i32,
18082 t_kv as i32,
18083 causal as i32,
18084 );
18085 let __s_b = self.gpu.stream();
18086 let mut b = __s_b.launch_builder(&f);
18087 b.arg(q)
18088 .arg(k)
18089 .arg(v)
18090 .arg(o)
18091 .arg(&hd)
18092 .arg(&nh)
18093 .arg(&nhkv)
18094 .arg(&ti)
18095 .arg(&tkvi)
18096 .arg(&scale)
18097 .arg(&cz);
18098 unsafe {
18099 b.launch(cfg)?;
18100 }
18101 Ok(())
18102 }
18103
18104 #[allow(clippy::too_many_arguments)]
18109 pub fn sdpa_naive_island(
18110 &self,
18111 q: &CudaSlice<f32>,
18112 k: &CudaSlice<f32>,
18113 v: &CudaSlice<f32>,
18114 o: &mut CudaSlice<f32>,
18115 span_id: &CudaSlice<i32>,
18116 head_dim: usize,
18117 n_head: usize,
18118 n_head_kv: usize,
18119 t: usize,
18120 t_kv: usize,
18121 scale: f32,
18122 window: usize,
18123 ) -> Result<(), Box<dyn std::error::Error>> {
18124 let f = self.func("sdpa_naive_island_f32");
18125 let cfg = LaunchConfig {
18126 grid_dim: (n_head as u32, t as u32, 1),
18127 block_dim: (128, 1, 1),
18128 shared_mem_bytes: (t_kv * 4) as u32,
18129 };
18130 let (hd, nh, nhkv, ti, tkvi, wi) = (
18131 head_dim as i32,
18132 n_head as i32,
18133 n_head_kv as i32,
18134 t as i32,
18135 t_kv as i32,
18136 window as i32,
18137 );
18138 let __s_b = self.gpu.stream();
18139 let mut b = __s_b.launch_builder(&f);
18140 b.arg(q)
18141 .arg(k)
18142 .arg(v)
18143 .arg(o)
18144 .arg(span_id)
18145 .arg(&hd)
18146 .arg(&nh)
18147 .arg(&nhkv)
18148 .arg(&ti)
18149 .arg(&tkvi)
18150 .arg(&scale)
18151 .arg(&wi);
18152 unsafe {
18153 b.launch(cfg)?;
18154 }
18155 Ok(())
18156 }
18157
18158 #[allow(clippy::too_many_arguments)]
18160 pub fn sdpa_naive_w(
18161 &self,
18162 q: &CudaSlice<f32>,
18163 k: &CudaSlice<f32>,
18164 v: &CudaSlice<f32>,
18165 o: &mut CudaSlice<f32>,
18166 head_dim: usize,
18167 n_head: usize,
18168 n_head_kv: usize,
18169 t: usize,
18170 t_kv: usize,
18171 scale: f32,
18172 causal: bool,
18173 window: usize,
18174 ) -> Result<(), Box<dyn std::error::Error>> {
18175 let f = self.func("sdpa_naive_w_f32");
18176 let cfg = LaunchConfig {
18177 grid_dim: (n_head as u32, t as u32, 1),
18178 block_dim: (128, 1, 1),
18179 shared_mem_bytes: (t_kv * 4) as u32,
18180 };
18181 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18182 head_dim as i32,
18183 n_head as i32,
18184 n_head_kv as i32,
18185 t as i32,
18186 t_kv as i32,
18187 causal as i32,
18188 window as i32,
18189 );
18190 let __s_b = self.gpu.stream();
18191 let mut b = __s_b.launch_builder(&f);
18192 b.arg(q)
18193 .arg(k)
18194 .arg(v)
18195 .arg(o)
18196 .arg(&hd)
18197 .arg(&nh)
18198 .arg(&nhkv)
18199 .arg(&ti)
18200 .arg(&tkvi)
18201 .arg(&scale)
18202 .arg(&cz)
18203 .arg(&wi);
18204 unsafe {
18205 b.launch(cfg)?;
18206 }
18207 Ok(())
18208 }
18209
18210 #[allow(clippy::too_many_arguments)]
18220 pub fn sdpa_naive_w_lo(
18221 &self,
18222 q: &CudaSlice<f32>,
18223 k: &CudaSlice<f32>,
18224 v: &CudaSlice<f32>,
18225 o: &mut CudaSlice<f32>,
18226 head_dim: usize,
18227 n_head: usize,
18228 n_head_kv: usize,
18229 t: usize,
18230 t_kv: usize,
18231 scale: f32,
18232 causal: bool,
18233 window: usize,
18234 ) -> Result<(), Box<dyn std::error::Error>> {
18235 let kv_lo = if window > 0 {
18236 (t_kv - t + 1).saturating_sub(window)
18237 } else {
18238 0
18239 };
18240 let smem = (t_kv - kv_lo) * 4;
18241 if smem > 48 * 1024 {
18242 return Err(format!(
18243 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
18244 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
18245 a window this wide needs the multi-pass long-ctx kernel"
18246 )
18247 .into());
18248 }
18249 let f = self.func("sdpa_naive_w_lo_f32");
18250 let cfg = LaunchConfig {
18251 grid_dim: (n_head as u32, t as u32, 1),
18252 block_dim: (128, 1, 1),
18253 shared_mem_bytes: smem as u32,
18254 };
18255 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
18256 head_dim as i32,
18257 n_head as i32,
18258 n_head_kv as i32,
18259 t as i32,
18260 t_kv as i32,
18261 causal as i32,
18262 window as i32,
18263 kv_lo as i32,
18264 );
18265 let __s_b = self.gpu.stream();
18266 let mut b = __s_b.launch_builder(&f);
18267 b.arg(q)
18268 .arg(k)
18269 .arg(v)
18270 .arg(o)
18271 .arg(&hd)
18272 .arg(&nh)
18273 .arg(&nhkv)
18274 .arg(&ti)
18275 .arg(&tkvi)
18276 .arg(&scale)
18277 .arg(&cz)
18278 .arg(&wi)
18279 .arg(&lo);
18280 unsafe {
18281 b.launch(cfg)?;
18282 }
18283 Ok(())
18284 }
18285
18286 pub fn sdpa_naive_view(
18288 &self,
18289 q: &CudaSlice<f32>,
18290 k: &cudarc::driver::CudaView<f32>,
18291 v: &cudarc::driver::CudaView<f32>,
18292 o: &mut CudaSlice<f32>,
18293 head_dim: usize,
18294 n_head: usize,
18295 n_head_kv: usize,
18296 t: usize,
18297 t_kv: usize,
18298 scale: f32,
18299 causal: bool,
18300 ) -> Result<(), Box<dyn std::error::Error>> {
18301 let f = self.func("sdpa_naive_f32");
18302 let cfg = LaunchConfig {
18303 grid_dim: (n_head as u32, t as u32, 1),
18304 block_dim: (128, 1, 1),
18305 shared_mem_bytes: (t_kv * 4) as u32,
18306 };
18307 let (hd, nh, nhkv, ti, tkvi, cz) = (
18308 head_dim as i32,
18309 n_head as i32,
18310 n_head_kv as i32,
18311 t as i32,
18312 t_kv as i32,
18313 causal as i32,
18314 );
18315 let __s_b = self.gpu.stream();
18316 let mut b = __s_b.launch_builder(&f);
18317 b.arg(q)
18318 .arg(k)
18319 .arg(v)
18320 .arg(o)
18321 .arg(&hd)
18322 .arg(&nh)
18323 .arg(&nhkv)
18324 .arg(&ti)
18325 .arg(&tkvi)
18326 .arg(&scale)
18327 .arg(&cz);
18328 unsafe {
18329 b.launch(cfg)?;
18330 }
18331 Ok(())
18332 }
18333
18334 #[allow(clippy::too_many_arguments)]
18342 pub fn fa_dequant_kv_view_f32(
18343 &self,
18344 k: &cudarc::driver::CudaView<u8>,
18345 v: &cudarc::driver::CudaView<u8>,
18346 kf: &mut CudaSlice<f32>,
18347 vf: &mut CudaSlice<f32>,
18348 kv_dim_k: usize,
18349 kv_dim_v: usize,
18350 t_kv: usize,
18351 k_tok_bytes: usize,
18352 v_tok_bytes: usize,
18353 g: bool,
18354 ) -> Result<(), Box<dyn std::error::Error>> {
18355 let f = if g {
18356 self.func_g("fa_dequant_kv_ws_f32")
18357 } else {
18358 self.func("fa_dequant_kv_ws_f32")
18359 };
18360 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
18361 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18362 let cfg = LaunchConfig {
18363 grid_dim: (nblk.max(1), 1, 1),
18364 block_dim: (256, 1, 1),
18365 shared_mem_bytes: 0,
18366 };
18367 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
18368 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
18369 let __s_b = self.gpu.stream();
18370 let mut b = __s_b.launch_builder(&f);
18371 b.arg(k)
18372 .arg(v)
18373 .arg(&mut *kf)
18374 .arg(&mut *vf)
18375 .arg(&kdk)
18376 .arg(&kdv)
18377 .arg(&tkvi)
18378 .arg(&ktb)
18379 .arg(&vtb);
18380 unsafe {
18381 b.launch(cfg)?;
18382 }
18383 Ok(())
18384 }
18385
18386 #[allow(clippy::too_many_arguments)]
18387 pub fn sdpa_naive_quantized_view(
18388 &self,
18389 q: &CudaSlice<f32>,
18390 k: &cudarc::driver::CudaView<u8>,
18391 v: &cudarc::driver::CudaView<u8>,
18392 o: &mut CudaSlice<f32>,
18393 head_dim: usize,
18394 n_head: usize,
18395 n_head_kv: usize,
18396 t: usize,
18397 t_kv: usize,
18398 scale: f32,
18399 causal: bool,
18400 k_tok_bytes: usize,
18401 v_tok_bytes: usize,
18402 ) -> Result<(), Box<dyn std::error::Error>> {
18403 let kv_dim = n_head_kv * head_dim;
18404 let mut kf = self.uninit(t_kv * kv_dim)?;
18405 let mut vf = self.uninit(t_kv * kv_dim)?;
18406 let f = self.func("fa_dequant_kv_ws_f32");
18407 let total = (2 * t_kv * kv_dim) as u64;
18408 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18409 let cfg = LaunchConfig {
18410 grid_dim: (nblk.max(1), 1, 1),
18411 block_dim: (256, 1, 1),
18412 shared_mem_bytes: 0,
18413 };
18414 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18415 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18416 let __s_b = self.gpu.stream();
18417 let mut b = __s_b.launch_builder(&f);
18418 b.arg(k)
18419 .arg(v)
18420 .arg(&mut kf)
18421 .arg(&mut vf)
18422 .arg(&kv_dim_i)
18423 .arg(&kv_dim_i)
18424 .arg(&t_kv_i)
18425 .arg(&k_tok_bytes_i)
18426 .arg(&v_tok_bytes_i);
18427 unsafe { b.launch(cfg)? };
18428 self.sdpa_naive(
18429 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18430 )
18431 }
18432
18433 #[allow(clippy::too_many_arguments)]
18445 pub fn sdpa_naive_w_quantized_view(
18446 &self,
18447 q: &CudaSlice<f32>,
18448 k: &cudarc::driver::CudaView<u8>,
18449 v: &cudarc::driver::CudaView<u8>,
18450 o: &mut CudaSlice<f32>,
18451 head_dim: usize,
18452 n_head: usize,
18453 n_head_kv: usize,
18454 t: usize,
18455 t_kv: usize,
18456 scale: f32,
18457 causal: bool,
18458 window: usize,
18459 k_tok_bytes: usize,
18460 v_tok_bytes: usize,
18461 ) -> Result<(), Box<dyn std::error::Error>> {
18462 let kv_dim = n_head_kv * head_dim;
18463 let mut kf = self.uninit(t_kv * kv_dim)?;
18464 let mut vf = self.uninit(t_kv * kv_dim)?;
18465 let f = self.func("fa_dequant_kv_ws_f32");
18466 let total = (2 * t_kv * kv_dim) as u64;
18467 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
18468 let cfg = LaunchConfig {
18469 grid_dim: (nblk.max(1), 1, 1),
18470 block_dim: (256, 1, 1),
18471 shared_mem_bytes: 0,
18472 };
18473 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
18474 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
18475 let __s_b = self.gpu.stream();
18476 let mut b = __s_b.launch_builder(&f);
18477 b.arg(k)
18478 .arg(v)
18479 .arg(&mut kf)
18480 .arg(&mut vf)
18481 .arg(&kv_dim_i)
18482 .arg(&kv_dim_i)
18483 .arg(&t_kv_i)
18484 .arg(&k_tok_bytes_i)
18485 .arg(&v_tok_bytes_i);
18486 unsafe { b.launch(cfg)? };
18487 self.sdpa_naive_w(
18488 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18489 )
18490 }
18491
18492 pub fn fa_prefill(
18496 &self,
18497 q: &CudaSlice<f32>,
18498 k: &CudaSlice<f32>,
18499 v: &CudaSlice<f32>,
18500 o: &mut CudaSlice<f32>,
18501 head_dim: usize,
18502 n_head: usize,
18503 n_head_kv: usize,
18504 t: usize,
18505 t_kv: usize,
18506 scale: f32,
18507 causal: bool,
18508 ) -> Result<(), Box<dyn std::error::Error>> {
18509 if portable_mma_gated() {
18510 return self.sdpa_naive(
18511 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
18512 );
18513 }
18514 let fa3_on = head_dim == 256
18522 && causal
18523 && t == t_kv
18524 && match std::env::var("MEMRA_FA3").as_deref() {
18525 Ok("0") => false,
18526 Ok("1") => true,
18527 _ => cfg!(memra_hopper_mma),
18528 };
18529 if fa3_on {
18530 let n = t * n_head * head_dim;
18531 let nkv = t * n_head_kv * head_dim;
18532 let mut q16 = self.alloc_u8_uninit(n * 2)?;
18533 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
18534 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
18535 self.f32_to_bf16_into(q, &mut q16, n)?;
18536 self.f32_to_bf16_into(k, &mut k16, nkv)?;
18537 self.f32_to_bf16_into(v, &mut v16, nkv)?;
18538 let rc = {
18539 use cudarc::driver::{DevicePtr, DevicePtrMut};
18540 let stream = self.gpu.stream();
18541 let (qp, _g1) = q16.device_ptr(&stream);
18542 let (kp, _g2) = k16.device_ptr(&stream);
18543 let (vp, _g3) = v16.device_ptr(&stream);
18544 let (op, _g4) = o.device_ptr_mut(&stream);
18545 unsafe {
18546 memra_fa3_prefill(
18547 qp as *const core::ffi::c_void,
18548 kp as *const core::ffi::c_void,
18549 vp as *const core::ffi::c_void,
18550 op as *mut f32,
18551 t as i32,
18552 n_head as i32,
18553 n_head_kv as i32,
18554 head_dim as i32,
18555 scale,
18556 stream.cu_stream() as *mut core::ffi::c_void,
18557 )
18558 }
18559 };
18560 if rc != 0 {
18561 return Err(format!("memra_fa3_prefill rc={rc}").into());
18562 }
18563 return Ok(());
18564 }
18565 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18570 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
18571 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
18572 const BLOCK_Q: usize = 64;
18573 const BKX: usize = 32;
18574 let f = self.func("fa_prefill_bf16_p1");
18575 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
18576 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
18577 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18578 f.set_attribute(
18579 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18580 shmem as i32,
18581 )?;
18582 let cfg = LaunchConfig {
18583 grid_dim: (
18584 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18585 n_head as u32,
18586 1,
18587 ),
18588 block_dim: (32, 4, 1),
18589 shared_mem_bytes: shmem,
18590 };
18591 let (hd, nh, nhkv, ti, tkvi, cz) = (
18592 head_dim as i32,
18593 n_head as i32,
18594 n_head_kv as i32,
18595 t as i32,
18596 t_kv as i32,
18597 causal as i32,
18598 );
18599 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18600 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18601 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
18602 let __s_b = self.gpu.stream();
18603 let mut b = __s_b.launch_builder(&f);
18604 b.arg(&qb)
18605 .arg(&kb)
18606 .arg(&vb)
18607 .arg(o)
18608 .arg(&hd)
18609 .arg(&nh)
18610 .arg(&nhkv)
18611 .arg(&ti)
18612 .arg(&tkvi)
18613 .arg(&scale)
18614 .arg(&cz);
18615 unsafe {
18616 b.launch(cfg)?;
18617 }
18618 return Ok(());
18619 }
18620 const BK: usize = 32;
18626 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
18629 let (block_q, warps, w2_sfx): (usize, u32, &str) =
18630 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
18631 let hd_sfx = fa_hd_suffix(head_dim)?;
18635 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18636 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
18641 let (kb16, vb16) = if bf16kv {
18642 let n = t_kv * n_head_kv * head_dim;
18643 let mut kb = self.alloc_u8_uninit(n * 2)?;
18644 let mut vb = self.alloc_u8_uninit(n * 2)?;
18645 let fcv = self.func("f32_to_bf16_bulk");
18646 let ni = n as i64;
18647 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
18648 let __s_b = self.gpu.stream();
18649 let mut b = __s_b.launch_builder(&fcv);
18650 b.arg(k).arg(&mut kb).arg(&ni);
18651 unsafe {
18652 b.launch(cfgc)?;
18653 }
18654 let __s_b = self.gpu.stream();
18655 let mut b = __s_b.launch_builder(&fcv);
18656 b.arg(v).arg(&mut vb).arg(&ni);
18657 unsafe {
18658 b.launch(cfgc)?;
18659 }
18660 (Some(kb), Some(vb))
18661 } else {
18662 (None, None)
18663 };
18664 let f = self.func(&if bf16kv {
18665 format!("fa_prefill_bf16kv_pp{hd_sfx}")
18666 } else {
18667 format!(
18668 "fa_prefill_f32{}{}{hd_sfx}",
18669 if floor { "" } else { "_pp" },
18670 if floor { "" } else { w2_sfx }
18671 )
18672 });
18673 let kv_stages = if bf16kv { 2 } else { 1 };
18676 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
18677 + 4 * (block_q * BK + 2 * block_q)) as u32;
18678 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18679 f.set_attribute(
18680 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18681 shmem as i32,
18682 )?;
18683 let cfg = LaunchConfig {
18684 grid_dim: (
18685 (t as u32 + block_q as u32 - 1) / block_q as u32,
18686 n_head as u32,
18687 1,
18688 ),
18689 block_dim: (32, warps, 1),
18690 shared_mem_bytes: shmem,
18691 };
18692 let (hd, nh, nhkv, ti, tkvi, cz) = (
18693 head_dim as i32,
18694 n_head as i32,
18695 n_head_kv as i32,
18696 t as i32,
18697 t_kv as i32,
18698 causal as i32,
18699 );
18700 let __s_b = self.gpu.stream();
18701 let mut b = __s_b.launch_builder(&f);
18702 b.arg(q);
18703 match (&kb16, &vb16) {
18704 (Some(kb), Some(vb)) => {
18705 b.arg(kb).arg(vb);
18706 }
18707 _ => {
18708 b.arg(k).arg(v);
18709 }
18710 }
18711 b.arg(o)
18712 .arg(&hd)
18713 .arg(&nh)
18714 .arg(&nhkv)
18715 .arg(&ti)
18716 .arg(&tkvi)
18717 .arg(&scale)
18718 .arg(&cz);
18719 unsafe {
18720 b.launch(cfg)?;
18721 }
18722 Ok(())
18723 }
18724
18725 #[allow(clippy::too_many_arguments)]
18729 pub fn fa_prefill_w(
18730 &self,
18731 q: &CudaSlice<f32>,
18732 k: &CudaSlice<f32>,
18733 v: &CudaSlice<f32>,
18734 o: &mut CudaSlice<f32>,
18735 head_dim: usize,
18736 n_head: usize,
18737 n_head_kv: usize,
18738 t: usize,
18739 t_kv: usize,
18740 scale: f32,
18741 causal: bool,
18742 window: usize,
18743 ) -> Result<(), Box<dyn std::error::Error>> {
18744 if portable_mma_gated() {
18747 return self.sdpa_naive_w(
18748 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
18749 );
18750 }
18751 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18755 let faw_f32 =
18756 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
18757 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
18758 self.fa_prefill_w_arm(
18759 q,
18760 k,
18761 v,
18762 o,
18763 head_dim,
18764 n_head,
18765 n_head_kv,
18766 t,
18767 t_kv,
18768 scale,
18769 causal,
18770 window,
18771 floor || faw_f32,
18772 floor,
18773 )
18774 }
18775
18776 #[allow(clippy::too_many_arguments)]
18779 pub fn fa_prefill_w_pre(
18780 &self,
18781 qb: &CudaSlice<u8>,
18782 kb: &CudaSlice<u8>,
18783 vb: &CudaSlice<u8>,
18784 o: &mut CudaSlice<f32>,
18785 head_dim: usize,
18786 n_head: usize,
18787 n_head_kv: usize,
18788 t: usize,
18789 t_kv: usize,
18790 scale: f32,
18791 causal: bool,
18792 window: usize,
18793 v_f16: bool,
18794 ) -> Result<(), Box<dyn std::error::Error>> {
18795 const BLOCK_Q: usize = 64;
18796 const BK: usize = 32;
18797 debug_assert_eq!(head_dim, 256);
18798 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
18799 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
18800 if hp {
18801 const BLOCK_QH: usize = 32;
18802 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
18805 let vh: &CudaSlice<u8> = if v_f16 {
18806 vb
18807 } else {
18808 let n = t_kv * n_head_kv * head_dim;
18809 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
18810 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
18811 }
18812 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
18813 vguard.as_ref().unwrap()
18814 };
18815 let f = self.func("fa_prefill_w_bf16_p1h2");
18816 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
18817 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18818 f.set_attribute(
18819 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18820 shmem as i32,
18821 )?;
18822 let cfg = LaunchConfig {
18823 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
18824 block_dim: (32, 4, 1),
18825 shared_mem_bytes: shmem,
18826 };
18827 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18828 head_dim as i32,
18829 n_head as i32,
18830 n_head_kv as i32,
18831 t as i32,
18832 t_kv as i32,
18833 causal as i32,
18834 window as i32,
18835 );
18836 let __s_b = self.gpu.stream();
18837 let mut b = __s_b.launch_builder(&f);
18838 b.arg(qb)
18839 .arg(kb)
18840 .arg(vh)
18841 .arg(o)
18842 .arg(&hd)
18843 .arg(&nh)
18844 .arg(&nhkv)
18845 .arg(&ti)
18846 .arg(&tkvi)
18847 .arg(&scale)
18848 .arg(&cz)
18849 .arg(&wi);
18850 unsafe {
18851 b.launch(cfg)?;
18852 }
18853 return Ok(());
18854 }
18855 let f = self.func("fa_prefill_w_bf16_p1");
18856 let shmem =
18857 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
18858 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18859 f.set_attribute(
18860 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18861 shmem as i32,
18862 )?;
18863 let cfg = LaunchConfig {
18864 grid_dim: (
18865 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18866 n_head as u32,
18867 1,
18868 ),
18869 block_dim: (32, 4, 1),
18870 shared_mem_bytes: shmem,
18871 };
18872 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18873 head_dim as i32,
18874 n_head as i32,
18875 n_head_kv as i32,
18876 t as i32,
18877 t_kv as i32,
18878 causal as i32,
18879 window as i32,
18880 );
18881 let __s_b = self.gpu.stream();
18882 let mut b = __s_b.launch_builder(&f);
18883 b.arg(qb)
18884 .arg(kb)
18885 .arg(vb)
18886 .arg(o)
18887 .arg(&hd)
18888 .arg(&nh)
18889 .arg(&nhkv)
18890 .arg(&ti)
18891 .arg(&tkvi)
18892 .arg(&scale)
18893 .arg(&cz)
18894 .arg(&wi);
18895 unsafe {
18896 b.launch(cfg)?;
18897 }
18898 Ok(())
18899 }
18900
18901 #[allow(clippy::too_many_arguments)]
18903 pub fn fa_prefill_w_arm(
18904 &self,
18905 q: &CudaSlice<f32>,
18906 k: &CudaSlice<f32>,
18907 v: &CudaSlice<f32>,
18908 o: &mut CudaSlice<f32>,
18909 head_dim: usize,
18910 n_head: usize,
18911 n_head_kv: usize,
18912 t: usize,
18913 t_kv: usize,
18914 scale: f32,
18915 causal: bool,
18916 window: usize,
18917 f32_stage: bool,
18918 floor: bool,
18919 ) -> Result<(), Box<dyn std::error::Error>> {
18920 const BLOCK_Q: usize = 64;
18921 const BK: usize = 32;
18922 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
18923 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18927 let p1 = !floor
18928 && !f32_stage
18929 && *P1_ON.get_or_init(|| {
18930 std::env::var("MEMRA_FAW_P1")
18931 .map(|v| v != "0")
18932 .unwrap_or(true)
18933 });
18934 let hp =
18935 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
18936 if hp {
18937 const BLOCK_QH: usize = 32;
18938 let f = self.func("fa_prefill_w_bf16_p1h2");
18939 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
18940 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18941 f.set_attribute(
18942 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18943 shmem as i32,
18944 )?;
18945 let cfg = LaunchConfig {
18946 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
18947 block_dim: (32, 4, 1),
18948 shared_mem_bytes: shmem,
18949 };
18950 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
18951 head_dim as i32,
18952 n_head as i32,
18953 n_head_kv as i32,
18954 t as i32,
18955 t_kv as i32,
18956 causal as i32,
18957 window as i32,
18958 );
18959 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
18960 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
18961 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
18962 let __s_b = self.gpu.stream();
18963 let mut b = __s_b.launch_builder(&f);
18964 b.arg(&qb)
18965 .arg(&kb)
18966 .arg(&vh)
18967 .arg(o)
18968 .arg(&hd)
18969 .arg(&nh)
18970 .arg(&nhkv)
18971 .arg(&ti)
18972 .arg(&tkvi)
18973 .arg(&scale)
18974 .arg(&cz)
18975 .arg(&wi);
18976 unsafe {
18977 b.launch(cfg)?;
18978 }
18979 return Ok(());
18980 }
18981 if p1 {
18982 let f = self.func("fa_prefill_w_bf16_p1");
18983 let shmem =
18984 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
18985 use cudarc::driver::sys::CUfunction_attribute_enum as A;
18986 f.set_attribute(
18987 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
18988 shmem as i32,
18989 )?;
18990 let cfg = LaunchConfig {
18991 grid_dim: (
18992 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
18993 n_head as u32,
18994 1,
18995 ),
18996 block_dim: (32, 4, 1),
18997 shared_mem_bytes: shmem,
18998 };
18999 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19000 head_dim as i32,
19001 n_head as i32,
19002 n_head_kv as i32,
19003 t as i32,
19004 t_kv as i32,
19005 causal as i32,
19006 window as i32,
19007 );
19008 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19009 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19010 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19011 let __s_b = self.gpu.stream();
19012 let mut b = __s_b.launch_builder(&f);
19013 b.arg(&qb)
19014 .arg(&kb)
19015 .arg(&vb)
19016 .arg(o)
19017 .arg(&hd)
19018 .arg(&nh)
19019 .arg(&nhkv)
19020 .arg(&ti)
19021 .arg(&tkvi)
19022 .arg(&scale)
19023 .arg(&cz)
19024 .arg(&wi);
19025 unsafe {
19026 b.launch(cfg)?;
19027 }
19028 return Ok(());
19029 }
19030 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19033 let g4 = !floor
19034 && !f32_stage
19035 && n_head_kv == 1
19036 && n_head % 4 == 0
19037 && *G4_ON.get_or_init(|| {
19038 std::env::var("MEMRA_FAW_G4")
19039 .map(|v| v != "0")
19040 .unwrap_or(true)
19041 });
19042 if g4 {
19043 const SP_M: usize = 16;
19044 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19047 let o2 = *O2_ON.get_or_init(|| {
19048 std::env::var("MEMRA_FAW_O2")
19049 .map(|v| v != "0")
19050 .unwrap_or(true)
19051 });
19052 let f = self.func(if o2 {
19053 "fa_prefill_w_bf16_g4o2"
19054 } else {
19055 "fa_prefill_w_bf16_g4"
19056 });
19057 let shmem = if o2 {
19058 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
19059 } else {
19060 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
19061 as u32
19062 };
19063 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19064 f.set_attribute(
19065 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19066 shmem as i32,
19067 )?;
19068 let cfg = LaunchConfig {
19069 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
19070 block_dim: (32, 4, 1),
19071 shared_mem_bytes: shmem,
19072 };
19073 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19074 head_dim as i32,
19075 n_head as i32,
19076 n_head_kv as i32,
19077 t as i32,
19078 t_kv as i32,
19079 causal as i32,
19080 window as i32,
19081 );
19082 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19083 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19084 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19085 let __s_b = self.gpu.stream();
19086 let mut b = __s_b.launch_builder(&f);
19087 b.arg(&qb)
19088 .arg(&kb)
19089 .arg(&vb)
19090 .arg(o)
19091 .arg(&hd)
19092 .arg(&nh)
19093 .arg(&nhkv)
19094 .arg(&ti)
19095 .arg(&tkvi)
19096 .arg(&scale)
19097 .arg(&cz)
19098 .arg(&wi);
19099 unsafe {
19100 b.launch(cfg)?;
19101 }
19102 return Ok(());
19103 }
19104 let f = self.func(if floor {
19105 "fa_prefill_w_f32"
19106 } else if f32_stage {
19107 "fa_prefill_w_f32_pp"
19108 } else {
19109 "fa_prefill_w_bf16_pp"
19110 });
19111 let shmem =
19112 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19113 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19114 f.set_attribute(
19115 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19116 shmem as i32,
19117 )?;
19118 let cfg = LaunchConfig {
19119 grid_dim: (
19120 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19121 n_head as u32,
19122 1,
19123 ),
19124 block_dim: (32, 4, 1),
19125 shared_mem_bytes: shmem,
19126 };
19127 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
19128 head_dim as i32,
19129 n_head as i32,
19130 n_head_kv as i32,
19131 t as i32,
19132 t_kv as i32,
19133 causal as i32,
19134 window as i32,
19135 );
19136 if f32_stage {
19137 let __s_b = self.gpu.stream();
19138 let mut b = __s_b.launch_builder(&f);
19139 b.arg(q)
19140 .arg(k)
19141 .arg(v)
19142 .arg(o)
19143 .arg(&hd)
19144 .arg(&nh)
19145 .arg(&nhkv)
19146 .arg(&ti)
19147 .arg(&tkvi)
19148 .arg(&scale)
19149 .arg(&cz)
19150 .arg(&wi);
19151 unsafe {
19152 b.launch(cfg)?;
19153 }
19154 } else {
19155 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19156 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19157 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19158 let __s_b = self.gpu.stream();
19159 let mut b = __s_b.launch_builder(&f);
19160 b.arg(&qb)
19161 .arg(&kb)
19162 .arg(&vb)
19163 .arg(o)
19164 .arg(&hd)
19165 .arg(&nh)
19166 .arg(&nhkv)
19167 .arg(&ti)
19168 .arg(&tkvi)
19169 .arg(&scale)
19170 .arg(&cz)
19171 .arg(&wi);
19172 unsafe {
19173 b.launch(cfg)?;
19174 }
19175 }
19176 Ok(())
19177 }
19178
19179 #[allow(clippy::too_many_arguments)]
19183 pub fn fa_prefill_hd512(
19184 &self,
19185 q: &CudaSlice<f32>,
19186 k: &CudaSlice<f32>,
19187 v: &CudaSlice<f32>,
19188 o: &mut CudaSlice<f32>,
19189 head_dim: usize,
19190 n_head: usize,
19191 n_head_kv: usize,
19192 t: usize,
19193 t_kv: usize,
19194 scale: f32,
19195 causal: bool,
19196 ) -> Result<(), Box<dyn std::error::Error>> {
19197 if portable_mma_gated() {
19199 return self.sdpa_naive(
19200 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
19201 );
19202 }
19203 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19209 let f32_stage =
19210 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
19211 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19215 let sp = !f32_stage
19216 && *SP_ON.get_or_init(|| {
19217 std::env::var("MEMRA_FA512_SP")
19218 .map(|v| v != "0")
19219 .unwrap_or(true)
19220 });
19221 self.fa_prefill_hd512_arm(
19222 q,
19223 k,
19224 v,
19225 o,
19226 head_dim,
19227 n_head,
19228 n_head_kv,
19229 t,
19230 t_kv,
19231 scale,
19232 causal,
19233 f32_stage,
19234 sp,
19235 sp && fa_f16pv_on(),
19236 )
19237 }
19238
19239 #[allow(clippy::too_many_arguments)]
19241 pub fn fa_prefill_hd512_pre(
19242 &self,
19243 qb: &CudaSlice<u8>,
19244 kb: &CudaSlice<u8>,
19245 vb: &CudaSlice<u8>,
19246 o: &mut CudaSlice<f32>,
19247 head_dim: usize,
19248 n_head: usize,
19249 n_head_kv: usize,
19250 t: usize,
19251 t_kv: usize,
19252 scale: f32,
19253 causal: bool,
19254 v_f16: bool,
19255 ) -> Result<(), Box<dyn std::error::Error>> {
19256 debug_assert_eq!(head_dim, 512);
19257 const SP_M: usize = 16;
19258 const BKS: usize = 32;
19259 let f16pv = fa_f16pv_on();
19263 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19264 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19265 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
19266 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
19267 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
19268 let n = t_kv * n_head_kv * head_dim;
19270 let need = n * 2;
19271 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
19272 *vguard = Some(self.alloc_uninit::<u8>(need)?);
19273 }
19274 let dst = vguard.as_mut().unwrap();
19275 self.bf16_to_f16_into(vb, n, dst)?;
19276 vguard.as_ref().unwrap()
19277 } else {
19278 vb
19279 };
19280 let f = self.func(if hp {
19281 "fa_prefill_bf16_hd512_sp16h2"
19282 } else {
19283 match (f16pv, nw) {
19284 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19285 (true, _) => "fa_prefill_bf16_hd512_sp16",
19286 _ => "fa_prefill_bf16_hd512_sp",
19287 }
19288 });
19289 let (nwarp, npart) = if hp {
19290 (4usize, 4usize)
19291 } else if nw > 2 {
19292 (nw, nw)
19293 } else {
19294 (2, 1)
19295 };
19296 let shmem = if hp {
19298 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
19299 as u32
19300 } else {
19301 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19302 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19303 };
19304 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19305 f.set_attribute(
19306 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19307 shmem as i32,
19308 )?;
19309 let grid_y = if hp {
19310 (n_head / 2) as u32
19311 } else {
19312 n_head as u32
19313 };
19314 let cfg = LaunchConfig {
19315 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19316 block_dim: (32, nwarp as u32, 1),
19317 shared_mem_bytes: shmem,
19318 };
19319 let (hd, nh, nhkv, ti, tkvi, cz) = (
19320 head_dim as i32,
19321 n_head as i32,
19322 n_head_kv as i32,
19323 t as i32,
19324 t_kv as i32,
19325 causal as i32,
19326 );
19327 let __s_b = self.gpu.stream();
19328 let mut b = __s_b.launch_builder(&f);
19329 b.arg(qb)
19330 .arg(kb)
19331 .arg(vref)
19332 .arg(o)
19333 .arg(&hd)
19334 .arg(&nh)
19335 .arg(&nhkv)
19336 .arg(&ti)
19337 .arg(&tkvi)
19338 .arg(&scale)
19339 .arg(&cz);
19340 unsafe {
19341 b.launch(cfg)?;
19342 }
19343 Ok(())
19344 }
19345
19346 #[allow(clippy::too_many_arguments)]
19349 pub fn fa_prefill_hd512_arm(
19350 &self,
19351 q: &CudaSlice<f32>,
19352 k: &CudaSlice<f32>,
19353 v: &CudaSlice<f32>,
19354 o: &mut CudaSlice<f32>,
19355 head_dim: usize,
19356 n_head: usize,
19357 n_head_kv: usize,
19358 t: usize,
19359 t_kv: usize,
19360 scale: f32,
19361 causal: bool,
19362 f32_stage: bool,
19363 sp: bool,
19364 f16pv: bool,
19365 ) -> Result<(), Box<dyn std::error::Error>> {
19366 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
19367 if sp && !f32_stage {
19368 const SP_M: usize = 16;
19372 const BKS: usize = 32;
19373 let nw = if f16pv { fa512_wide_warps() } else { 2 };
19374 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
19375 let f = self.func(if hp {
19376 "fa_prefill_bf16_hd512_sp16h2"
19377 } else {
19378 match (f16pv, nw) {
19379 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
19380 (true, _) => "fa_prefill_bf16_hd512_sp16",
19381 _ => "fa_prefill_bf16_hd512_sp",
19382 }
19383 });
19384 let (nwarp, npart) = if hp {
19385 (4usize, 4usize)
19386 } else if nw > 2 {
19387 (nw, nw)
19388 } else {
19389 (2, 1)
19390 };
19391 let shmem = if hp {
19392 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
19393 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
19394 } else {
19395 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
19396 + 4 * (npart * SP_M * BKS + SP_M)) as u32
19397 };
19398 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19399 f.set_attribute(
19400 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19401 shmem as i32,
19402 )?;
19403 let grid_y = if hp {
19404 (n_head / 2) as u32
19405 } else {
19406 n_head as u32
19407 };
19408 let cfg = LaunchConfig {
19409 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
19410 block_dim: (32, nwarp as u32, 1),
19411 shared_mem_bytes: shmem,
19412 };
19413 let (hd, nh, nhkv, ti, tkvi, cz) = (
19414 head_dim as i32,
19415 n_head as i32,
19416 n_head_kv as i32,
19417 t as i32,
19418 t_kv as i32,
19419 causal as i32,
19420 );
19421 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19422 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19423 let vb = if f16pv {
19424 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
19425 } else {
19426 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
19427 };
19428 let __s_b = self.gpu.stream();
19429 let mut b = __s_b.launch_builder(&f);
19430 b.arg(&qb)
19431 .arg(&kb)
19432 .arg(&vb)
19433 .arg(o)
19434 .arg(&hd)
19435 .arg(&nh)
19436 .arg(&nhkv)
19437 .arg(&ti)
19438 .arg(&tkvi)
19439 .arg(&scale)
19440 .arg(&cz);
19441 unsafe {
19442 b.launch(cfg)?;
19443 }
19444 return Ok(());
19445 }
19446 const BLOCK_Q: usize = 32;
19447 const BK: usize = 32;
19448 const HALF: usize = 256;
19449 let f = self.func(if f32_stage {
19450 "fa_prefill_f32_hd512"
19451 } else {
19452 "fa_prefill_bf16_hd512"
19453 });
19454 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
19456 + 4 * BLOCK_Q) as u32;
19457 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19458 f.set_attribute(
19459 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19460 shmem as i32,
19461 )?;
19462 let cfg = LaunchConfig {
19463 grid_dim: (
19464 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19465 n_head as u32,
19466 2,
19467 ),
19468 block_dim: (32, 2, 1),
19469 shared_mem_bytes: shmem,
19470 };
19471 let (hd, nh, nhkv, ti, tkvi, cz) = (
19472 head_dim as i32,
19473 n_head as i32,
19474 n_head_kv as i32,
19475 t as i32,
19476 t_kv as i32,
19477 causal as i32,
19478 );
19479 if f32_stage {
19480 let __s_b = self.gpu.stream();
19481 let mut b = __s_b.launch_builder(&f);
19482 b.arg(q)
19483 .arg(k)
19484 .arg(v)
19485 .arg(o)
19486 .arg(&hd)
19487 .arg(&nh)
19488 .arg(&nhkv)
19489 .arg(&ti)
19490 .arg(&tkvi)
19491 .arg(&scale)
19492 .arg(&cz);
19493 unsafe {
19494 b.launch(cfg)?;
19495 }
19496 } else {
19497 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
19498 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
19499 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
19500 let __s_b = self.gpu.stream();
19501 let mut b = __s_b.launch_builder(&f);
19502 b.arg(&qb)
19503 .arg(&kb)
19504 .arg(&vb)
19505 .arg(o)
19506 .arg(&hd)
19507 .arg(&nh)
19508 .arg(&nhkv)
19509 .arg(&ti)
19510 .arg(&tkvi)
19511 .arg(&scale)
19512 .arg(&cz);
19513 unsafe {
19514 b.launch(cfg)?;
19515 }
19516 }
19517 Ok(())
19518 }
19519
19520 #[allow(clippy::too_many_arguments)]
19524 pub fn rope_neox2_bf16e(
19525 &self,
19526 q: &mut CudaSlice<f32>,
19527 k: &mut CudaSlice<f32>,
19528 qb: &mut CudaSlice<u8>,
19529 kb: &mut CudaSlice<u8>,
19530 pos: &CudaSlice<i32>,
19531 head_dim: usize,
19532 n_dims: usize,
19533 nh_q: usize,
19534 nh_k: usize,
19535 n_tokens: usize,
19536 base: f32,
19537 freq_scale: f32,
19538 ff: Option<&CudaSlice<f32>>,
19539 ) -> Result<(), Box<dyn std::error::Error>> {
19540 let f = self.func("rope_neox2_bf16e_f32");
19541 let rows = ((nh_q + nh_k) * n_tokens) as u32;
19542 let cfg = LaunchConfig {
19543 grid_dim: (rows, 1, 1),
19544 block_dim: ((head_dim / 2) as u32, 1, 1),
19545 shared_mem_bytes: 0,
19546 };
19547 let theta_scale = base.powf(-2.0 / n_dims as f32);
19548 let (hd, nd, nhq, nhk, nt) = (
19549 head_dim as i32,
19550 n_dims as i32,
19551 nh_q as i32,
19552 nh_k as i32,
19553 n_tokens as i32,
19554 );
19555 let __s_b = self.gpu.stream();
19556 let mut b = __s_b.launch_builder(&f);
19557 match ff {
19558 Some(t) => {
19559 b.arg(&mut *q)
19560 .arg(&mut *k)
19561 .arg(&mut *qb)
19562 .arg(&mut *kb)
19563 .arg(pos)
19564 .arg(&hd)
19565 .arg(&nd)
19566 .arg(&nhq)
19567 .arg(&nhk)
19568 .arg(&nt)
19569 .arg(&theta_scale)
19570 .arg(&freq_scale)
19571 .arg(t);
19572 unsafe {
19573 b.launch(cfg)?;
19574 }
19575 }
19576 None => {
19577 let null: u64 = 0;
19578 b.arg(&mut *q)
19579 .arg(&mut *k)
19580 .arg(&mut *qb)
19581 .arg(&mut *kb)
19582 .arg(pos)
19583 .arg(&hd)
19584 .arg(&nd)
19585 .arg(&nhq)
19586 .arg(&nhk)
19587 .arg(&nt)
19588 .arg(&theta_scale)
19589 .arg(&freq_scale)
19590 .arg(&null);
19591 unsafe {
19592 b.launch(cfg)?;
19593 }
19594 }
19595 }
19596 Ok(())
19597 }
19598
19599 pub fn f32_to_bf16(
19602 &self,
19603 x: &CudaSlice<f32>,
19604 n: usize,
19605 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19606 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
19607 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19608 let f = self.func("f32_to_bf16_flat");
19609 let n_i = n as i64;
19610 let cfg = LaunchConfig {
19611 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19612 block_dim: (256, 1, 1),
19613 shared_mem_bytes: 0,
19614 };
19615 let __s_b = self.gpu.stream();
19616 let mut b = __s_b.launch_builder(&f);
19617 b.arg(x).arg(&mut y).arg(&n_i);
19618 unsafe {
19619 b.launch(cfg)?;
19620 }
19621 Ok(y)
19622 }
19623
19624 pub fn f32_to_f16(
19625 &self,
19626 x: &CudaSlice<f32>,
19627 n: usize,
19628 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19629 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
19630 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19631 let f = self.func("f32_to_f16_flat");
19632 let n_i = n as i64;
19633 let cfg = LaunchConfig {
19634 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
19635 block_dim: (256, 1, 1),
19636 shared_mem_bytes: 0,
19637 };
19638 let __s_b = self.gpu.stream();
19639 let mut b = __s_b.launch_builder(&f);
19640 b.arg(x).arg(&mut y).arg(&n_i);
19641 unsafe {
19642 b.launch(cfg)?;
19643 }
19644 Ok(y)
19645 }
19646
19647 pub fn bf16_to_f16(
19649 &self,
19650 xb: &CudaSlice<u8>,
19651 n: usize,
19652 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
19653 let mut y = self.alloc_uninit::<u8>(n * 2)?;
19654 self.bf16_to_f16_into(xb, n, &mut y)?;
19655 Ok(y)
19656 }
19657
19658 pub fn bf16_to_f16_into(
19660 &self,
19661 xb: &CudaSlice<u8>,
19662 n: usize,
19663 y: &mut CudaSlice<u8>,
19664 ) -> Result<(), Box<dyn std::error::Error>> {
19665 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
19666 assert!(y.len() >= n * 2);
19667 let f = self.func("bf16_to_f16_flat");
19668 let n2 = (n / 2) as i64;
19669 let cfg = LaunchConfig {
19670 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
19671 block_dim: (256, 1, 1),
19672 shared_mem_bytes: 0,
19673 };
19674 let __s_b = self.gpu.stream();
19675 let mut b = __s_b.launch_builder(&f);
19676 b.arg(xb).arg(y).arg(&n2);
19677 unsafe {
19678 b.launch(cfg)?;
19679 }
19680 Ok(())
19681 }
19682
19683 #[allow(clippy::too_many_arguments)]
19688 pub fn fa_prefill_vl8(
19689 &self,
19690 seqs: &[FaSeqVl],
19691 head_dim: usize,
19692 n_head: usize,
19693 n_head_kv: usize,
19694 scale: f32,
19695 ) -> Result<(), Box<dyn std::error::Error>> {
19696 const BK: usize = 32;
19697 let b = seqs.len();
19698 assert!(b >= 1 && b <= 8);
19699 let mut packed = [FaSeqVl::default(); 8];
19700 packed[..b].copy_from_slice(seqs);
19701 let v = FaVl8(packed);
19702 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19703 let ept = (n_head_kv * head_dim) as i32;
19704 {
19705 let f = self.func("fa_mirror_vl");
19706 let max_n = (max_t as i64) * ept as i64;
19707 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
19708 for which in 0..2i32 {
19709 let cfg = LaunchConfig {
19710 grid_dim: (blocks, 1, b as u32),
19711 block_dim: (256, 1, 1),
19712 shared_mem_bytes: 0,
19713 };
19714 let __s_lb = self.gpu.stream();
19715 let mut lb = __s_lb.launch_builder(&f);
19716 lb.arg(&v).arg(&ept).arg(&which);
19717 unsafe {
19718 lb.launch(cfg)?;
19719 }
19720 }
19721 }
19722 let hd_sfx = fa_hd_suffix(head_dim)?;
19723 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
19724 let block_q = 64usize;
19725 let kv_stages = 2usize;
19726 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
19727 + 4 * (block_q * BK + 2 * block_q)) as u32;
19728 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19729 f.set_attribute(
19730 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19731 shmem as i32,
19732 )?;
19733 let cfg = LaunchConfig {
19734 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
19735 block_dim: (32, 4, 1),
19736 shared_mem_bytes: shmem,
19737 };
19738 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19739 let __s_lb = self.gpu.stream();
19740 let mut lb = __s_lb.launch_builder(&f);
19741 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
19742 unsafe {
19743 lb.launch(cfg)?;
19744 }
19745 Ok(())
19746 }
19747
19748 #[allow(clippy::too_many_arguments)]
19752 pub fn attn_pre_vl8(
19753 &self,
19754 seqs: &[AttnPreVl],
19755 wq: &CudaSlice<f32>,
19756 wk: &CudaSlice<f32>,
19757 head_dim: usize,
19758 rope_dims: usize,
19759 n_head: usize,
19760 n_head_kv: usize,
19761 eps: f32,
19762 freq_base: f32,
19763 freq_scale: f32,
19764 kv_dim_k: usize,
19765 kv_dim_v: usize,
19766 k_tok_bytes: usize,
19767 v_tok_bytes: usize,
19768 ) -> Result<(), Box<dyn std::error::Error>> {
19769 let b = seqs.len();
19770 assert!(b >= 1 && b <= 8);
19771 let mut packed = [AttnPreVl::default(); 8];
19772 packed[..b].copy_from_slice(seqs);
19773 let v = AttnPreVl8(packed);
19774 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
19775 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
19776 {
19777 let f = self.func("q_gate_split_vl");
19778 let n = max_t * (n_head * head_dim) as u32;
19779 let cfg = LaunchConfig {
19780 grid_dim: (n.div_ceil(256), 1, b as u32),
19781 block_dim: (256, 1, 1),
19782 shared_mem_bytes: 0,
19783 };
19784 let __s_lb = self.gpu.stream();
19785 let mut lb = __s_lb.launch_builder(&f);
19786 lb.arg(&v).arg(&hd).arg(&nh);
19787 unsafe {
19788 lb.launch(cfg)?;
19789 }
19790 }
19791 {
19792 let f = self.func("attn_rms_vl");
19793 let cfg = LaunchConfig {
19794 grid_dim: (max_t * n_head as u32, 2, b as u32),
19795 block_dim: (rms_block(), 1, 1),
19796 shared_mem_bytes: 0,
19797 };
19798 let __s_lb = self.gpu.stream();
19799 let mut lb = __s_lb.launch_builder(&f);
19800 lb.arg(&v)
19801 .arg(wq)
19802 .arg(wk)
19803 .arg(&hd)
19804 .arg(&nh)
19805 .arg(&nhkv)
19806 .arg(&eps);
19807 unsafe {
19808 lb.launch(cfg)?;
19809 }
19810 }
19811 {
19812 let f = self.func("attn_rope_vl");
19813 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
19814 let nd = rope_dims as i32;
19815 let cfg = LaunchConfig {
19816 grid_dim: (max_t * n_head as u32, 2, b as u32),
19817 block_dim: ((head_dim / 2) as u32, 1, 1),
19818 shared_mem_bytes: 0,
19819 };
19820 let __s_lb = self.gpu.stream();
19821 let mut lb = __s_lb.launch_builder(&f);
19822 lb.arg(&v)
19823 .arg(&hd)
19824 .arg(&nd)
19825 .arg(&nh)
19826 .arg(&nhkv)
19827 .arg(&theta_scale)
19828 .arg(&freq_scale);
19829 unsafe {
19830 lb.launch(cfg)?;
19831 }
19832 }
19833 {
19834 let f = self.func("append_kv_vl");
19835 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
19836 let cfg = LaunchConfig {
19837 grid_dim: (nblk, max_t, b as u32),
19838 block_dim: (32, 1, 1),
19839 shared_mem_bytes: 0,
19840 };
19841 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
19842 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19843 let __s_lb = self.gpu.stream();
19844 let mut lb = __s_lb.launch_builder(&f);
19845 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
19846 unsafe {
19847 lb.launch(cfg)?;
19848 }
19849 }
19850 Ok(())
19851 }
19852
19853 pub fn fa_prefill_view(
19858 &self,
19859 q: &CudaSlice<f32>,
19860 k: &cudarc::driver::CudaView<u8>,
19861 v: &cudarc::driver::CudaView<u8>,
19862 o: &mut CudaSlice<f32>,
19863 head_dim: usize,
19864 n_head: usize,
19865 n_head_kv: usize,
19866 t: usize,
19867 t_kv: usize,
19868 scale: f32,
19869 causal: bool,
19870 k_tok_bytes: usize,
19871 v_tok_bytes: usize,
19872 g: bool,
19873 ) -> Result<(), Box<dyn std::error::Error>> {
19874 if portable_mma_gated() {
19875 return self.sdpa_naive_quantized_view(
19876 q,
19877 k,
19878 v,
19879 o,
19880 head_dim,
19881 n_head,
19882 n_head_kv,
19883 t,
19884 t_kv,
19885 scale,
19886 causal,
19887 k_tok_bytes,
19888 v_tok_bytes,
19889 );
19890 }
19891 const BLOCK_Q: usize = 64;
19892 const BK: usize = 32;
19893 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
19896 let f = if g {
19897 self.func_g(&name)
19898 } else {
19899 self.func(&name)
19900 };
19901 let shmem =
19902 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
19903 use cudarc::driver::sys::CUfunction_attribute_enum as A;
19904 f.set_attribute(
19905 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
19906 shmem as i32,
19907 )?;
19908 let cfg = LaunchConfig {
19909 grid_dim: (
19910 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
19911 n_head as u32,
19912 1,
19913 ),
19914 block_dim: (32, 4, 1),
19915 shared_mem_bytes: shmem,
19916 };
19917 let (hd, nh, nhkv, ti, tkvi, cz) = (
19918 head_dim as i32,
19919 n_head as i32,
19920 n_head_kv as i32,
19921 t as i32,
19922 t_kv as i32,
19923 causal as i32,
19924 );
19925 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
19926 let __s_b = self.gpu.stream();
19927 let mut b = __s_b.launch_builder(&f);
19928 b.arg(q)
19929 .arg(k)
19930 .arg(v)
19931 .arg(o)
19932 .arg(&hd)
19933 .arg(&nh)
19934 .arg(&nhkv)
19935 .arg(&ti)
19936 .arg(&tkvi)
19937 .arg(&scale)
19938 .arg(&cz)
19939 .arg(&ktb)
19940 .arg(&vtb);
19941 unsafe {
19942 b.launch(cfg)?;
19943 }
19944 Ok(())
19945 }
19946
19947 #[allow(clippy::too_many_arguments)]
19957 pub fn fa_prefill_view_ws(
19958 &self,
19959 q: &CudaSlice<f32>,
19960 k: &cudarc::driver::CudaView<u8>,
19961 v: &cudarc::driver::CudaView<u8>,
19962 o: &mut CudaSlice<f32>,
19963 head_dim: usize,
19964 n_head: usize,
19965 n_head_kv: usize,
19966 t: usize,
19967 t_kv: usize,
19968 scale: f32,
19969 causal: bool,
19970 k_tok_bytes: usize,
19971 v_tok_bytes: usize,
19972 g: bool,
19973 ) -> Result<(), Box<dyn std::error::Error>> {
19974 if portable_mma_gated() {
19975 return self.sdpa_naive_quantized_view(
19976 q,
19977 k,
19978 v,
19979 o,
19980 head_dim,
19981 n_head,
19982 n_head_kv,
19983 t,
19984 t_kv,
19985 scale,
19986 causal,
19987 k_tok_bytes,
19988 v_tok_bytes,
19989 );
19990 }
19991 const BLOCK_Q: usize = 64;
19992 const BK: usize = 32;
19993 let kv_dim_k = n_head_kv * head_dim;
19994 let kv_dim_v = n_head_kv * head_dim;
19995 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
19997 let mut guard = self.prime_deqw_ws.lock().unwrap();
19999 let need_grow = match guard.as_ref() {
20000 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20001 None => true,
20002 };
20003 if need_grow {
20004 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20005 let (ck, cv) = guard
20006 .as_ref()
20007 .map(|(a, b)| (a.len(), b.len()))
20008 .unwrap_or((0, 0));
20009 *guard = Some((
20010 self.alloc_u8(grow(ck, k_ws_bytes))?,
20011 self.alloc_u8(grow(cv, v_ws_bytes))?,
20012 ));
20013 }
20014 let (kw, vw) = guard.as_mut().unwrap();
20015 {
20017 let f = if g {
20019 self.func_g("fa_dequant_kv_ws_bf16")
20020 } else {
20021 self.func("fa_dequant_kv_ws_bf16")
20022 };
20023 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20024 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20025 let cfg = LaunchConfig {
20026 grid_dim: (nblk.max(1), 1, 1),
20027 block_dim: (256, 1, 1),
20028 shared_mem_bytes: 0,
20029 };
20030 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20031 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20032 let __s_b = self.gpu.stream();
20033 let mut b = __s_b.launch_builder(&f);
20034 b.arg(k)
20035 .arg(v)
20036 .arg(&mut *kw)
20037 .arg(&mut *vw)
20038 .arg(&kdk)
20039 .arg(&kdv)
20040 .arg(&tkvi)
20041 .arg(&ktb)
20042 .arg(&vtb);
20043 unsafe {
20044 b.launch(cfg)?;
20045 }
20046 }
20047 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20055 .map(|v| v != "0")
20056 .unwrap_or(true);
20057 {
20058 let hd_sfx = fa_hd_suffix(head_dim)?;
20059 let f = self.func(&format!(
20060 "fa_prefill_qw{}{hd_sfx}",
20061 if db { "_db" } else { "" }
20062 ));
20063 let shmem = if db {
20064 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20066 } else {
20067 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20068 };
20069 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20070 f.set_attribute(
20071 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20072 shmem as i32,
20073 )?;
20074 let cfg = LaunchConfig {
20075 grid_dim: (
20076 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20077 n_head as u32,
20078 1,
20079 ),
20080 block_dim: (32, 4, 1),
20081 shared_mem_bytes: shmem,
20082 };
20083 let (hd, nh, nhkv, ti, tkvi, cz) = (
20084 head_dim as i32,
20085 n_head as i32,
20086 n_head_kv as i32,
20087 t as i32,
20088 t_kv as i32,
20089 causal as i32,
20090 );
20091 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
20092 let __s_b = self.gpu.stream();
20093 let mut b = __s_b.launch_builder(&f);
20094 b.arg(q)
20095 .arg(&*kw)
20096 .arg(&*vw)
20097 .arg(o)
20098 .arg(&hd)
20099 .arg(&nh)
20100 .arg(&nhkv)
20101 .arg(&ti)
20102 .arg(&tkvi)
20103 .arg(&scale)
20104 .arg(&cz)
20105 .arg(&kdk)
20106 .arg(&kdv);
20107 unsafe {
20108 b.launch(cfg)?;
20109 }
20110 }
20111 Ok(())
20112 }
20113
20114 #[allow(clippy::too_many_arguments)]
20130 pub fn fa_prefill_view_ws_w_hd128(
20131 &self,
20132 q: &CudaSlice<f32>,
20133 k: &cudarc::driver::CudaView<u8>,
20134 v: &cudarc::driver::CudaView<u8>,
20135 o: &mut CudaSlice<f32>,
20136 head_dim: usize,
20137 n_head: usize,
20138 n_head_kv: usize,
20139 t: usize,
20140 t_kv: usize,
20141 scale: f32,
20142 causal: bool,
20143 window: usize,
20144 k_tok_bytes: usize,
20145 v_tok_bytes: usize,
20146 ) -> Result<(), Box<dyn std::error::Error>> {
20147 assert_eq!(
20148 head_dim, 128,
20149 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
20150 );
20151 if portable_mma_gated() {
20152 return self.sdpa_naive_w_quantized_view(
20153 q,
20154 k,
20155 v,
20156 o,
20157 head_dim,
20158 n_head,
20159 n_head_kv,
20160 t,
20161 t_kv,
20162 scale,
20163 causal,
20164 window,
20165 k_tok_bytes,
20166 v_tok_bytes,
20167 );
20168 }
20169 const BLOCK_Q: usize = 64;
20170 const BK: usize = 32;
20171 let kv_dim_k = n_head_kv * head_dim;
20172 let kv_dim_v = n_head_kv * head_dim;
20173 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
20175 let mut guard = self.prime_deqw_ws.lock().unwrap();
20176 let need_grow = match guard.as_ref() {
20177 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
20178 None => true,
20179 };
20180 if need_grow {
20181 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
20182 let (ck, cv) = guard
20183 .as_ref()
20184 .map(|(a, b)| (a.len(), b.len()))
20185 .unwrap_or((0, 0));
20186 *guard = Some((
20187 self.alloc_u8(grow(ck, k_ws_bytes))?,
20188 self.alloc_u8(grow(cv, v_ws_bytes))?,
20189 ));
20190 }
20191 let (kw, vw) = guard.as_mut().unwrap();
20192 {
20195 let f = self.func("fa_dequant_kv_ws_bf16");
20196 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20197 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20198 let cfg = LaunchConfig {
20199 grid_dim: (nblk.max(1), 1, 1),
20200 block_dim: (256, 1, 1),
20201 shared_mem_bytes: 0,
20202 };
20203 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20204 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20205 let __s_b = self.gpu.stream();
20206 let mut b = __s_b.launch_builder(&f);
20207 b.arg(k)
20208 .arg(v)
20209 .arg(&mut *kw)
20210 .arg(&mut *vw)
20211 .arg(&kdk)
20212 .arg(&kdv)
20213 .arg(&tkvi)
20214 .arg(&ktb)
20215 .arg(&vtb);
20216 unsafe {
20217 b.launch(cfg)?;
20218 }
20219 }
20220 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
20222 .map(|v| v != "0")
20223 .unwrap_or(true);
20224 {
20225 let f = self.func(if db {
20226 "fa_prefill_qw_db_w_hd128"
20227 } else {
20228 "fa_prefill_qw_w_hd128"
20229 });
20230 let shmem = if db {
20231 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
20232 } else {
20233 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
20234 };
20235 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20236 f.set_attribute(
20237 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20238 shmem as i32,
20239 )?;
20240 let cfg = LaunchConfig {
20241 grid_dim: (
20242 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20243 n_head as u32,
20244 1,
20245 ),
20246 block_dim: (32, 4, 1),
20247 shared_mem_bytes: shmem,
20248 };
20249 let (hd, nh, nhkv, ti, tkvi, cz) = (
20250 head_dim as i32,
20251 n_head as i32,
20252 n_head_kv as i32,
20253 t as i32,
20254 t_kv as i32,
20255 causal as i32,
20256 );
20257 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
20258 let __s_b = self.gpu.stream();
20259 let mut b = __s_b.launch_builder(&f);
20260 b.arg(q)
20261 .arg(&*kw)
20262 .arg(&*vw)
20263 .arg(o)
20264 .arg(&hd)
20265 .arg(&nh)
20266 .arg(&nhkv)
20267 .arg(&ti)
20268 .arg(&tkvi)
20269 .arg(&scale)
20270 .arg(&cz)
20271 .arg(&kdk)
20272 .arg(&kdv)
20273 .arg(&wnd);
20274 unsafe {
20275 b.launch(cfg)?;
20276 }
20277 }
20278 Ok(())
20279 }
20280
20281 pub fn fa_decode(
20285 &self,
20286 q: &CudaSlice<f32>,
20287 k: &cudarc::driver::CudaView<u8>,
20288 v: &cudarc::driver::CudaView<u8>,
20289 o: &mut CudaSlice<f32>,
20290 head_dim: usize,
20291 n_head: usize,
20292 n_head_kv: usize,
20293 t_kv: usize,
20294 scale: f32,
20295 k_tok_bytes: usize,
20296 v_tok_bytes: usize,
20297 ) -> Result<(), Box<dyn std::error::Error>> {
20298 self.fa_decode_kvmod(
20299 q,
20300 k,
20301 v,
20302 o,
20303 head_dim,
20304 n_head,
20305 n_head_kv,
20306 t_kv,
20307 scale,
20308 k_tok_bytes,
20309 v_tok_bytes,
20310 false,
20311 )
20312 }
20313
20314 #[allow(clippy::too_many_arguments)]
20318 #[allow(clippy::too_many_arguments)]
20322 #[allow(clippy::too_many_arguments)]
20323 fn fa_decode_scalar_unified(
20324 &self,
20325 q: &cudarc::driver::CudaView<f32>,
20326 k: &cudarc::driver::CudaView<u8>,
20327 v: &cudarc::driver::CudaView<u8>,
20328 o: &mut cudarc::driver::CudaViewMut<f32>,
20329 head_dim: usize,
20330 n_head: usize,
20331 n_head_kv: usize,
20332 t_kv_host: usize,
20333 t_kv_dev: Option<&CudaSlice<i32>>,
20334 scale: f32,
20335 n_splits: usize,
20336 split_keys: usize,
20337 k_tok_bytes: usize,
20338 v_tok_bytes: usize,
20339 g: bool,
20340 part_o: &mut CudaSlice<f32>,
20341 part_m: &mut CudaSlice<f32>,
20342 part_l: &mut CudaSlice<f32>,
20343 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
20344 ) -> Result<(), Box<dyn std::error::Error>> {
20345 let f = if g {
20346 self.func_g("fa_decode_f32")
20347 } else {
20348 self.fa_func("fa_decode_f32", head_dim)
20349 };
20350 let cfg = LaunchConfig {
20351 grid_dim: (n_head as u32, n_splits as u32, 1),
20352 block_dim: (head_dim as u32, 1, 1),
20353 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
20354 };
20355 let (hd, nh, nhkv, nsp) = (
20356 head_dim as i32,
20357 n_head as i32,
20358 n_head_kv as i32,
20359 n_splits as i32,
20360 );
20361 let (ktb, vtb, tkvi, ski) = (
20362 k_tok_bytes as i64,
20363 v_tok_bytes as i64,
20364 t_kv_host as i32,
20365 split_keys as i32,
20366 );
20367 let __s_b = self.gpu.stream();
20368 let mut b = __s_b.launch_builder(&f);
20369 match t_kv_dev {
20370 Some(d) => {
20371 b.arg(q)
20372 .arg(k)
20373 .arg(v)
20374 .arg(&mut *part_o)
20375 .arg(&mut *part_m)
20376 .arg(&mut *part_l)
20377 .arg(&hd)
20378 .arg(&nh)
20379 .arg(&nhkv)
20380 .arg(&tkvi)
20381 .arg(d)
20382 .arg(&scale)
20383 .arg(&nsp)
20384 .arg(&ski)
20385 .arg(&ktb)
20386 .arg(&vtb);
20387 unsafe {
20388 b.launch(cfg)?;
20389 }
20390 }
20391 None => {
20392 let null: u64 = 0;
20393 b.arg(q)
20394 .arg(k)
20395 .arg(v)
20396 .arg(&mut *part_o)
20397 .arg(&mut *part_m)
20398 .arg(&mut *part_l)
20399 .arg(&hd)
20400 .arg(&nh)
20401 .arg(&nhkv)
20402 .arg(&tkvi)
20403 .arg(&null)
20404 .arg(&scale)
20405 .arg(&nsp)
20406 .arg(&ski)
20407 .arg(&ktb)
20408 .arg(&vtb);
20409 unsafe {
20410 b.launch(cfg)?;
20411 }
20412 }
20413 }
20414 let cfg2 = LaunchConfig {
20415 grid_dim: (n_head as u32, 1, 1),
20416 block_dim: (head_dim as u32, 1, 1),
20417 shared_mem_bytes: 0,
20418 };
20419 if let Some((oq, od)) = q8_out {
20420 let fc = if g {
20422 self.func_g("fa_decode_combine_q8_1")
20423 } else {
20424 self.fa_func("fa_decode_combine_q8_1", head_dim)
20425 };
20426 let __s_b2 = self.gpu.stream();
20427 let mut b2 = __s_b2.launch_builder(&fc);
20428 b2.arg(&*part_o)
20429 .arg(&*part_m)
20430 .arg(&*part_l)
20431 .arg(oq)
20432 .arg(od)
20433 .arg(&hd)
20434 .arg(&nh)
20435 .arg(&nsp);
20436 unsafe {
20437 b2.launch(cfg2)?;
20438 }
20439 return Ok(());
20440 }
20441 let fc = if g {
20442 self.func_g("fa_decode_combine_f32")
20443 } else {
20444 self.fa_func("fa_decode_combine_f32", head_dim)
20445 };
20446 let __s_b2 = self.gpu.stream();
20447 let mut b2 = __s_b2.launch_builder(&fc);
20448 b2.arg(&*part_o)
20449 .arg(&*part_m)
20450 .arg(&*part_l)
20451 .arg(o)
20452 .arg(&hd)
20453 .arg(&nh)
20454 .arg(&nsp);
20455 unsafe {
20456 b2.launch(cfg2)?;
20457 }
20458 Ok(())
20459 }
20460
20461 pub fn fa_decode_kvmod(
20462 &self,
20463 q: &CudaSlice<f32>,
20464 k: &cudarc::driver::CudaView<u8>,
20465 v: &cudarc::driver::CudaView<u8>,
20466 o: &mut CudaSlice<f32>,
20467 head_dim: usize,
20468 n_head: usize,
20469 n_head_kv: usize,
20470 t_kv: usize,
20471 scale: f32,
20472 k_tok_bytes: usize,
20473 v_tok_bytes: usize,
20474 g: bool,
20475 ) -> Result<(), Box<dyn std::error::Error>> {
20476 let q_view = q.as_view();
20477 let mut o_view = o.as_view_mut();
20478 self.fa_decode_kvmod_view(
20479 &q_view,
20480 k,
20481 v,
20482 &mut o_view,
20483 head_dim,
20484 n_head,
20485 n_head_kv,
20486 t_kv,
20487 scale,
20488 k_tok_bytes,
20489 v_tok_bytes,
20490 g,
20491 )
20492 }
20493
20494 #[allow(clippy::too_many_arguments)]
20499 pub fn fa_decode_kvmod_view(
20500 &self,
20501 q: &cudarc::driver::CudaView<f32>,
20502 k: &cudarc::driver::CudaView<u8>,
20503 v: &cudarc::driver::CudaView<u8>,
20504 o: &mut cudarc::driver::CudaViewMut<f32>,
20505 head_dim: usize,
20506 n_head: usize,
20507 n_head_kv: usize,
20508 t_kv: usize,
20509 scale: f32,
20510 k_tok_bytes: usize,
20511 v_tok_bytes: usize,
20512 g: bool,
20513 ) -> Result<(), Box<dyn std::error::Error>> {
20514 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
20535 if g && head_dim == 256 && !fa_v4_at(t_kv) {
20539 fa_vec = false;
20540 }
20541 let sp = fa_split_keys(t_kv, n_head_kv);
20542 let n_splits = if fa_vec {
20543 ((t_kv + sp - 1) / sp).max(1)
20544 } else {
20545 ((t_kv + 255) / 256).max(1)
20546 };
20547 let o_len = n_head * n_splits * head_dim;
20548 let ml_len = n_head * n_splits;
20549 let mut part_guard = self.fa_part_pool.lock().unwrap();
20550 if part_guard
20551 .as_ref()
20552 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20553 .unwrap_or(true)
20554 {
20555 let old = part_guard.take();
20566 let (co, cm) = old
20567 .as_ref()
20568 .map(|pp| (pp.0.len(), pp.1.len()))
20569 .unwrap_or((0, 0));
20570 if let Some(old) = old {
20571 self.fa_part_retired.lock().unwrap().push(old);
20572 }
20573 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20574 eprintln!(
20575 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20576 co, o_len, cm, ml_len
20577 );
20578 }
20579 *part_guard = Some((
20580 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20581 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20582 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20583 ));
20584 }
20585 let pg = part_guard.as_mut().unwrap();
20586 self.gpu
20587 .stream()
20588 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20589 self.gpu
20590 .stream()
20591 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20592 self.gpu
20593 .stream()
20594 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20595 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20596 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
20597 let (hd, nh, nhkv, tkvi, nsp) = (
20598 head_dim as i32,
20599 n_head as i32,
20600 n_head_kv as i32,
20601 t_kv as i32,
20602 n_splits as i32,
20603 );
20604 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20605 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
20609 let fa512_min = fa512_min_tkv();
20614 let deep = fa_vec
20617 && head_dim == 256
20618 && fa_v4_at(t_kv)
20619 && !g
20620 && fa_deep_at(t_kv)
20621 && !matches!(fa_v4_mode(), "noB3" | "stage");
20622 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
20623 let gqa = (n_head / n_head_kv).max(1) as u32;
20626 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
20627 (
20628 fv,
20629 LaunchConfig {
20630 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20631 block_dim: (32, gqa, 1),
20632 shared_mem_bytes: 0,
20633 },
20634 )
20635 } else if fa_vec && head_dim <= 256 {
20636 let gqa = (n_head / n_head_kv).max(1) as u32;
20637 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
20648 let smem_tkv = *SMEM_TKV.get_or_init(|| {
20649 std::env::var("MEMRA_FA_SMEM_TKV")
20650 .ok()
20651 .and_then(|v| v.parse().ok())
20652 .unwrap_or_else(|| {
20653 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
20654 })
20655 });
20656 if fa_v4_at(t_kv) && head_dim == 256 {
20657 let v4name = match fa_v4_mode() {
20661 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
20664 _ => "fa_decode_vec_q_v4",
20665 };
20666 let fv = if g {
20667 self.func_g(v4name)
20668 } else {
20669 self.func(v4name)
20670 };
20671 let shmem = (if deep { 12160 } else { 11520 }
20674 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
20675 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20676 fv.set_attribute(
20677 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20678 shmem as i32,
20679 )?;
20680 (
20681 fv,
20682 LaunchConfig {
20683 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20684 block_dim: (32, gqa, 1),
20685 shared_mem_bytes: shmem,
20686 },
20687 )
20688 } else if fa_v3_active(head_dim) {
20689 let fv = if g {
20692 self.func_g("fa_decode_vec_q_v3")
20693 } else {
20694 self.func("fa_decode_vec_q_v3")
20695 };
20696 let shmem = (32 * head_dim * 2) as u32; (
20698 fv,
20699 LaunchConfig {
20700 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20701 block_dim: (32, gqa, 1),
20702 shared_mem_bytes: shmem,
20703 },
20704 )
20705 } else if fa_v2_on() {
20706 let fv = if g {
20710 self.func_g("fa_decode_vec_q_v2")
20711 } else {
20712 self.func("fa_decode_vec_q_v2")
20713 };
20714 let shmem = (2 * 32 * head_dim * 2) as u32; (
20716 fv,
20717 LaunchConfig {
20718 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20719 block_dim: (32, gqa, 1),
20720 shared_mem_bytes: shmem,
20721 },
20722 )
20723 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
20724 {
20725 let fv = if g {
20729 self.func_g("fa_decode_vec_q_smem")
20730 } else {
20731 self.func("fa_decode_vec_q_smem")
20732 };
20733 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
20735 fv.set_attribute(
20736 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20737 shmem as i32,
20738 )?;
20739 (
20740 fv,
20741 LaunchConfig {
20742 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20743 block_dim: (32, gqa, 1),
20744 shared_mem_bytes: shmem,
20745 },
20746 )
20747 } else {
20748 let fv = if g {
20751 self.func_g("fa_decode_vec_q")
20752 } else {
20753 self.func("fa_decode_vec_q")
20754 };
20755 (
20756 fv,
20757 LaunchConfig {
20758 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
20759 block_dim: (32, gqa, 1),
20760 shared_mem_bytes: 0,
20761 },
20762 )
20763 }
20764 } else {
20765 return self.fa_decode_scalar_unified(
20768 q,
20769 k,
20770 v,
20771 o,
20772 head_dim,
20773 n_head,
20774 n_head_kv,
20775 t_kv,
20776 None,
20777 scale,
20778 n_splits,
20779 if fa_vec { sp } else { 256 },
20780 k_tok_bytes,
20781 v_tok_bytes,
20782 g,
20783 part_o,
20784 part_m,
20785 part_l,
20786 None,
20787 );
20788 };
20789 let __s_b = self.gpu.stream();
20790 let mut b = __s_b.launch_builder(&f);
20791 b.arg(q)
20792 .arg(k)
20793 .arg(v)
20794 .arg(&mut *part_o)
20795 .arg(&mut *part_m)
20796 .arg(&mut *part_l)
20797 .arg(&hd)
20798 .arg(&nh)
20799 .arg(&nhkv)
20800 .arg(&tkvi)
20801 .arg(&scale)
20802 .arg(&nsp)
20803 .arg(&ktb)
20804 .arg(&vtb);
20805 unsafe {
20806 b.launch(cfg)?;
20807 }
20808 let (fc, cfg2) = (
20811 if g {
20812 self.func_g("fa_decode_combine_f32")
20813 } else {
20814 self.fa_func("fa_decode_combine_f32", head_dim)
20815 },
20816 LaunchConfig {
20817 grid_dim: (n_head as u32, 1, 1),
20818 block_dim: (head_dim as u32, 1, 1),
20819 shared_mem_bytes: 0,
20820 },
20821 );
20822 let __s_b2 = self.gpu.stream();
20823 let mut b2 = __s_b2.launch_builder(&fc);
20824 b2.arg(&*part_o)
20825 .arg(&*part_m)
20826 .arg(&*part_l)
20827 .arg(o)
20828 .arg(&hd)
20829 .arg(&nh)
20830 .arg(&nsp);
20831 unsafe {
20832 b2.launch(cfg2)?;
20833 }
20834 Ok(())
20835 }
20836
20837 #[allow(clippy::too_many_arguments)]
20848 pub fn fa_decode_batch_seqs_v4(
20849 &self,
20850 q: &CudaSlice<f32>,
20851 kv_ptrs: &cudarc::driver::CudaView<u64>,
20852 pos_seq: &CudaSlice<i32>,
20853 o: &mut CudaSlice<f32>,
20854 head_dim: usize,
20855 n_head: usize,
20856 n_head_kv: usize,
20857 b_n: usize,
20858 t_kv_max: usize,
20859 scale: f32,
20860 split_keys: usize,
20861 k_tok_bytes: usize,
20862 v_tok_bytes: usize,
20863 ) -> Result<(), Box<dyn std::error::Error>> {
20864 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
20865 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
20866 let o_len = b_n * n_head * n_splits_max * head_dim;
20867 let ml_len = b_n * n_head * n_splits_max;
20868 let mut part_guard = self.fa_part_pool.lock().unwrap();
20869 if part_guard
20870 .as_ref()
20871 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
20872 .unwrap_or(true)
20873 {
20874 let old = part_guard.take();
20885 let (co, cm) = old
20886 .as_ref()
20887 .map(|pp| (pp.0.len(), pp.1.len()))
20888 .unwrap_or((0, 0));
20889 if let Some(old) = old {
20890 self.fa_part_retired.lock().unwrap().push(old);
20891 }
20892 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
20893 eprintln!(
20894 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
20895 co, o_len, cm, ml_len
20896 );
20897 }
20898 *part_guard = Some((
20899 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
20900 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20901 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
20902 ));
20903 }
20904 let pg = part_guard.as_mut().unwrap();
20905 self.gpu
20906 .stream()
20907 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
20908 self.gpu
20909 .stream()
20910 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
20911 self.gpu
20912 .stream()
20913 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
20914 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
20915 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
20916 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
20917 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20918 let gqa = (n_head / n_head_kv).max(1) as u32;
20919 let f = self.func("fa_decode_vec_q_seqs_v4");
20920 let shmem = (11520 + 32 * head_dim * 2) as u32;
20922 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20923 f.set_attribute(
20924 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20925 shmem as i32,
20926 )?;
20927 let cfg = LaunchConfig {
20928 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
20929 block_dim: (32, gqa, 1),
20930 shared_mem_bytes: shmem,
20931 };
20932 {
20933 let __s_b = self.gpu.stream();
20934 let mut b = __s_b.launch_builder(&f);
20935 b.arg(q)
20936 .arg(kv_ptrs)
20937 .arg(pos_seq)
20938 .arg(&mut *part_o)
20939 .arg(&mut *part_m)
20940 .arg(&mut *part_l)
20941 .arg(&hd)
20942 .arg(&nh)
20943 .arg(&nhkv)
20944 .arg(&scale)
20945 .arg(&nspm)
20946 .arg(&spk)
20947 .arg(&ktb)
20948 .arg(&vtb);
20949 unsafe {
20950 b.launch(cfg)?;
20951 }
20952 }
20953 let fc = self.func("fa_decode_combine_seqs");
20954 let cfg2 = LaunchConfig {
20955 grid_dim: (n_head as u32, b_n as u32, 1),
20956 block_dim: (head_dim as u32, 1, 1),
20957 shared_mem_bytes: 0,
20958 };
20959 let __s_b2 = self.gpu.stream();
20960 let mut b2 = __s_b2.launch_builder(&fc);
20961 b2.arg(&*part_o)
20962 .arg(&*part_m)
20963 .arg(&*part_l)
20964 .arg(o)
20965 .arg(&hd)
20966 .arg(&nh)
20967 .arg(pos_seq)
20968 .arg(&nspm)
20969 .arg(&spk);
20970 unsafe {
20971 b2.launch(cfg2)?;
20972 }
20973 Ok(())
20974 }
20975
20976 #[allow(clippy::too_many_arguments)]
20983 pub fn append_kv_quantized_seqs(
20984 &self,
20985 k_rows: &CudaSlice<f32>,
20986 v_rows: &CudaSlice<f32>,
20987 kv_ptrs: &cudarc::driver::CudaView<u64>,
20988 pos_seq: &CudaSlice<i32>,
20989 b_n: usize,
20990 kv_dim_k: usize,
20991 kv_dim_v: usize,
20992 k_tok_bytes: usize,
20993 v_tok_bytes: usize,
20994 ) -> Result<(), Box<dyn std::error::Error>> {
20995 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
20996 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
20997 let cfg = LaunchConfig {
20998 grid_dim: (nblk, b_n as u32, 1),
20999 block_dim: (32, 1, 1),
21000 shared_mem_bytes: 0,
21001 };
21002 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21003 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21004 let __s_b = self.gpu.stream();
21005 let mut b = __s_b.launch_builder(&f);
21006 b.arg(k_rows)
21007 .arg(v_rows)
21008 .arg(kv_ptrs)
21009 .arg(pos_seq)
21010 .arg(&kdk)
21011 .arg(&kdv)
21012 .arg(&ktb)
21013 .arg(&vtb);
21014 unsafe {
21015 b.launch(cfg)?;
21016 }
21017 Ok(())
21018 }
21019
21020 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
21026 std::env::var("MEMRA_NO_FA_VEC").is_err()
21027 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
21028 && base_len + 1 >= fa_vec_min_tkv()
21029 && head_dim <= 256
21030 && head_dim % 32 == 0
21031 }
21032
21033 #[allow(clippy::too_many_arguments)]
21042 pub fn fa_decode_rows(
21043 &self,
21044 q: &CudaSlice<f32>,
21045 k: &cudarc::driver::CudaView<u8>,
21046 v: &cudarc::driver::CudaView<u8>,
21047 o: &mut CudaSlice<f32>,
21048 head_dim: usize,
21049 n_head: usize,
21050 n_head_kv: usize,
21051 base_len: usize,
21052 t: usize,
21053 scale: f32,
21054 k_tok_bytes: usize,
21055 v_tok_bytes: usize,
21056 base_dev: Option<(&CudaSlice<i32>, i32)>,
21060 kv_shared: bool,
21063 g: bool,
21067 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21070 ) -> Result<(), Box<dyn std::error::Error>> {
21071 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
21072 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
21079 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21080 let v = *SP512.get_or_init(|| {
21083 std::env::var("MEMRA_FA_SP512")
21084 .ok()
21085 .and_then(|x| x.parse().ok())
21086 .unwrap_or(0)
21087 });
21088 sp = if v >= 8 {
21089 v
21090 } else {
21091 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21092 };
21093 }
21094 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21095 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21096 let gqa = (n_head / n_head_kv).max(1) as u32;
21097 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
21108 groups.push((0, t, sp));
21109 } else {
21110 let mut r0 = 0usize;
21111 while r0 < t {
21112 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
21113 let mut r1 = r0 + 1;
21114 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
21115 r1 += 1;
21116 }
21117 groups.push((r0, r1 - r0, sp_g));
21118 r0 = r1;
21119 }
21120 }
21121 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21125 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
21126 std::env::var("MEMRA_FA_SMEM_TKV")
21127 .ok()
21128 .and_then(|v| v.parse().ok())
21129 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21130 });
21131 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
21132 let v3 = fa_v3_active(head_dim);
21133 let smem_rows =
21134 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
21135 let _ = kv_shared;
21140 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
21143 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21157 let tb512 = head_dim == 512
21159 && sp <= 32
21160 && n_head / n_head_kv.max(1) <= 16
21161 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
21162 let fname = if tb512 {
21163 "fa_decode_vec_q_rows_v4_512_tb"
21164 } else if i2 {
21165 "fa_decode_vec_q_rows_dpl16_i2"
21166 } else if head_dim == 512 {
21167 "fa_decode_vec_q_rows_dpl16"
21168 }
21169 else if v4 {
21171 "fa_decode_vec_q_rows_v4"
21172 } else if v3 {
21173 "fa_decode_vec_q_rows_v3"
21174 } else if fa_v2_on() {
21175 "fa_decode_vec_q_rows_v2"
21176 } else if smem_rows {
21177 "fa_decode_vec_q_rows_smem"
21178 } else {
21179 "fa_decode_vec_q_rows"
21180 };
21181 let f = if head_dim == 512 {
21182 self.fa_func(fname, head_dim)
21183 } else if g {
21184 self.func_g(if smem_rows {
21192 "fa_decode_vec_q_rows"
21193 } else {
21194 fname
21195 })
21196 } else {
21197 self.func(fname)
21198 };
21199 let shmem = if tb512 {
21200 let gk = Self::gkv_on();
21202 let sh =
21203 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
21204 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21205 f.set_attribute(
21206 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21207 sh as i32,
21208 )?;
21209 sh
21210 } else if v4 || v3 || smem_rows || fa_v2_on() {
21211 let sh = (if v4 {
21213 11520 + 32 * head_dim * if g { 1 } else { 2 }
21214 } else if v3 {
21215 32 * head_dim * 2
21216 } else {
21217 2 * 32 * head_dim * 2
21218 }) as u32;
21219 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21220 f.set_attribute(
21221 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21222 sh as i32,
21223 )?;
21224 sh
21225 } else {
21226 0
21227 };
21228 for &(r0, t_g, sp_g) in &groups {
21232 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
21233 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
21234 let base_i = (base_len + r0) as i32;
21235 let o_len = t_g * n_head * n_splits_g * head_dim;
21236 let ml_len = t_g * n_head * n_splits_g;
21237 let mut part_guard = self.fa_part_pool.lock().unwrap();
21238 if part_guard
21239 .as_ref()
21240 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21241 .unwrap_or(true)
21242 {
21243 let old = part_guard.take();
21254 let (co, cm) = old
21255 .as_ref()
21256 .map(|pp| (pp.0.len(), pp.1.len()))
21257 .unwrap_or((0, 0));
21258 if let Some(old) = old {
21259 self.fa_part_retired.lock().unwrap().push(old);
21260 }
21261 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21262 eprintln!(
21263 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21264 co, o_len, cm, ml_len
21265 );
21266 }
21267 *part_guard = Some((
21268 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21269 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21270 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21271 ));
21272 }
21273 let pg = part_guard.as_mut().unwrap();
21274 self.gpu
21275 .stream()
21276 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21277 self.gpu
21278 .stream()
21279 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21280 self.gpu
21281 .stream()
21282 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21283 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21284 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
21285 let qv = self.view(q, t * n_head * head_dim);
21286 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21287 let cfg = LaunchConfig {
21288 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
21289 block_dim: (32, gqa, 1),
21290 shared_mem_bytes: shmem,
21291 };
21292 {
21293 let __s_b = self.gpu.stream();
21294 let mut b = __s_b.launch_builder(&f);
21295 if tb512 {
21296 let (bd, plus) =
21298 base_dev.expect("hd512 rows twin requires a device base counter");
21299 let plus_g = plus + r0 as i32;
21300 let nr = t_g as i32;
21301 if Self::pdl_on() && Self::pdl_wb_on() {
21302 use cudarc::driver::{DevicePtr, DevicePtrMut};
21304 let s = &self.gpu.stream();
21305 let (pq, _b0) = q_g.device_ptr(s);
21306 let (pk, _b1) = k.device_ptr(s);
21307 let (pv, _b2) = v.device_ptr(s);
21308 let (po, _b3) = part_o.device_ptr_mut(s);
21309 let (pm, _b4) = part_m.device_ptr_mut(s);
21310 let (pl, _b5) = part_l.device_ptr_mut(s);
21311 let (pb, _b6) = bd.device_ptr(s);
21312 let mut ps = [
21313 &pq as *const _ as *mut std::ffi::c_void,
21314 &pk as *const _ as *mut _,
21315 &pv as *const _ as *mut _,
21316 &po as *const _ as *mut _,
21317 &pm as *const _ as *mut _,
21318 &pl as *const _ as *mut _,
21319 &hd as *const _ as *mut _,
21320 &nh as *const _ as *mut _,
21321 &nhkv as *const _ as *mut _,
21322 &pb as *const _ as *mut _,
21323 &plus_g as *const _ as *mut _,
21324 &scale as *const _ as *mut _,
21325 &nspm as *const _ as *mut _,
21326 &spk as *const _ as *mut _,
21327 &ktb as *const _ as *mut _,
21328 &vtb as *const _ as *mut _,
21329 &nr as *const _ as *mut _,
21330 ];
21331 unsafe {
21332 self.launch_pdl_flash(
21333 Self::gkv_on(),
21334 "fa_decode_vec_q_rows_v4_512_tb",
21335 (n_head_kv as u32, n_splits_g as u32, 1),
21336 (32, gqa, 1),
21337 shmem,
21338 &mut ps,
21339 )?;
21340 }
21341 } else {
21342 let cfg_tb = LaunchConfig {
21343 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
21344 block_dim: (32, gqa, 1),
21345 shared_mem_bytes: shmem,
21346 };
21347 b.arg(&q_g)
21348 .arg(k)
21349 .arg(v)
21350 .arg(&mut *part_o)
21351 .arg(&mut *part_m)
21352 .arg(&mut *part_l)
21353 .arg(&hd)
21354 .arg(&nh)
21355 .arg(&nhkv)
21356 .arg(bd)
21357 .arg(&plus_g)
21358 .arg(&scale)
21359 .arg(&nspm)
21360 .arg(&spk)
21361 .arg(&ktb)
21362 .arg(&vtb)
21363 .arg(&nr);
21364 unsafe {
21365 b.launch(cfg_tb)?;
21366 }
21367 }
21368 } else if head_dim == 512 {
21369 let (bd, plus) =
21370 base_dev.expect("hd512 rows twin requires a device base counter");
21371 let plus_g = plus + r0 as i32;
21372 b.arg(&q_g)
21373 .arg(k)
21374 .arg(v)
21375 .arg(&mut *part_o)
21376 .arg(&mut *part_m)
21377 .arg(&mut *part_l)
21378 .arg(&hd)
21379 .arg(&nh)
21380 .arg(&nhkv)
21381 .arg(bd)
21382 .arg(&plus_g)
21383 .arg(&scale)
21384 .arg(&nspm)
21385 .arg(&spk)
21386 .arg(&ktb)
21387 .arg(&vtb);
21388 unsafe {
21389 b.launch(cfg)?;
21390 }
21391 } else {
21392 b.arg(&q_g)
21393 .arg(k)
21394 .arg(v)
21395 .arg(&mut *part_o)
21396 .arg(&mut *part_m)
21397 .arg(&mut *part_l)
21398 .arg(&hd)
21399 .arg(&nh)
21400 .arg(&nhkv)
21401 .arg(&base_i)
21402 .arg(&scale)
21403 .arg(&nspm)
21404 .arg(&spk)
21405 .arg(&ktb)
21406 .arg(&vtb);
21407 unsafe {
21408 b.launch(cfg)?;
21409 }
21410 }
21411 }
21412 let cfg2 = LaunchConfig {
21413 grid_dim: (n_head as u32, t_g as u32, 1),
21414 block_dim: (head_dim as u32, 1, 1),
21415 shared_mem_bytes: 0,
21416 };
21417 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
21418 if head_dim == 512 {
21419 let (bd, plus) = base_dev.unwrap();
21422 let plus_g = plus + r0 as i32;
21423 if let Some((oq, od)) = q8_out.as_mut() {
21424 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
21426 if Self::pdl_on() && Self::pdl_wb_on() {
21427 use cudarc::driver::{DevicePtr, DevicePtrMut};
21429 let s = &self.gpu.stream();
21430 let (po, _g0) = part_o.device_ptr(s);
21431 let (pm, _g1) = part_m.device_ptr(s);
21432 let (pl, _g2) = part_l.device_ptr(s);
21433 let (pq, _g3) = oq.device_ptr_mut(s);
21434 let (pd, _g4) = od.device_ptr_mut(s);
21435 let (pb, _g5) = bd.device_ptr(s);
21436 let mut ps = [
21437 &po as *const _ as *mut std::ffi::c_void,
21438 &pm as *const _ as *mut _,
21439 &pl as *const _ as *mut _,
21440 &pq as *const _ as *mut _,
21441 &pd as *const _ as *mut _,
21442 &hd as *const _ as *mut _,
21443 &nh as *const _ as *mut _,
21444 &pb as *const _ as *mut _,
21445 &plus_g as *const _ as *mut _,
21446 &nspm as *const _ as *mut _,
21447 &spk as *const _ as *mut _,
21448 ];
21449 unsafe {
21450 self.launch_pdl_flash(
21451 Self::gkv_on(),
21452 "fa_decode_combine_rows_dc_q8_1",
21453 cfg2.grid_dim,
21454 cfg2.block_dim,
21455 0,
21456 &mut ps,
21457 )?;
21458 }
21459 continue;
21460 }
21461 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
21462 let __s_b2 = self.gpu.stream();
21463 let mut b2 = __s_b2.launch_builder(&fc);
21464 b2.arg(&*part_o)
21465 .arg(&*part_m)
21466 .arg(&*part_l)
21467 .arg(&mut **oq)
21468 .arg(&mut **od)
21469 .arg(&hd)
21470 .arg(&nh)
21471 .arg(bd)
21472 .arg(&plus_g)
21473 .arg(&nspm)
21474 .arg(&spk);
21475 unsafe {
21476 b2.launch(cfg2)?;
21477 }
21478 continue;
21479 }
21480 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
21481 let __s_b2 = self.gpu.stream();
21482 let mut b2 = __s_b2.launch_builder(&fc);
21483 b2.arg(&*part_o)
21484 .arg(&*part_m)
21485 .arg(&*part_l)
21486 .arg(&mut o_g)
21487 .arg(&hd)
21488 .arg(&nh)
21489 .arg(bd)
21490 .arg(&plus_g)
21491 .arg(&nspm)
21492 .arg(&spk);
21493 unsafe {
21494 b2.launch(cfg2)?;
21495 }
21496 } else {
21497 assert!(
21500 q8_out.is_none(),
21501 "rows q8 emit requires the hd512 dc combine"
21502 );
21503 let fc = self.func("fa_decode_combine_rows");
21504 let __s_b2 = self.gpu.stream();
21505 let mut b2 = __s_b2.launch_builder(&fc);
21506 b2.arg(&*part_o)
21507 .arg(&*part_m)
21508 .arg(&*part_l)
21509 .arg(&mut o_g)
21510 .arg(&hd)
21511 .arg(&nh)
21512 .arg(&base_i)
21513 .arg(&nspm)
21514 .arg(&spk);
21515 unsafe {
21516 b2.launch(cfg2)?;
21517 }
21518 }
21519 }
21520 Ok(())
21521 }
21522
21523 #[allow(clippy::too_many_arguments)]
21527 pub fn fa_decode_rows_w(
21528 &self,
21529 q: &CudaSlice<f32>,
21530 k: &cudarc::driver::CudaView<u8>,
21531 v: &cudarc::driver::CudaView<u8>,
21532 o: &mut CudaSlice<f32>,
21533 head_dim: usize,
21534 n_head: usize,
21535 n_head_kv: usize,
21536 base_dev: &CudaSlice<i32>,
21537 base_plus: i32,
21538 t: usize,
21539 scale: f32,
21540 window: usize,
21541 k_tok_bytes: usize,
21542 v_tok_bytes: usize,
21543 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
21544 ) -> Result<(), Box<dyn std::error::Error>> {
21545 debug_assert!(head_dim == 256);
21550 let sp = {
21558 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21559 let v = *SPW.get_or_init(|| {
21560 std::env::var("MEMRA_FA_SPW")
21561 .ok()
21562 .and_then(|x| x.parse().ok())
21563 .unwrap_or(0)
21564 });
21565 if v >= 8 {
21566 v
21567 } else {
21568 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
21569 }
21570 };
21571 let n_splits_max = (window + sp - 1) / sp;
21572 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21573 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
21574 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21575 let gqa = (n_head / n_head_kv).max(1) as u32;
21576 let o_len = t * n_head * n_splits_max * head_dim;
21577 let ml_len = t * n_head * n_splits_max;
21578 let mut part_guard = self.fa_part_pool.lock().unwrap();
21579 if part_guard
21580 .as_ref()
21581 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21582 .unwrap_or(true)
21583 {
21584 let old = part_guard.take();
21595 let (co, cm) = old
21596 .as_ref()
21597 .map(|pp| (pp.0.len(), pp.1.len()))
21598 .unwrap_or((0, 0));
21599 if let Some(old) = old {
21600 self.fa_part_retired.lock().unwrap().push(old);
21601 }
21602 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21603 eprintln!(
21604 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21605 co, o_len, cm, ml_len
21606 );
21607 }
21608 *part_guard = Some((
21609 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21610 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21611 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21612 ));
21613 }
21614 let pg = part_guard.as_mut().unwrap();
21615 self.gpu
21616 .stream()
21617 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21618 self.gpu
21619 .stream()
21620 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21621 self.gpu
21622 .stream()
21623 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21624 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21625 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
21631 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
21632 std::env::var("MEMRA_FA_SMEM_TKV")
21633 .ok()
21634 .and_then(|v| v.parse().ok())
21635 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
21636 });
21637 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21643 let wg = Self::wkv_on();
21648 let sp2 =
21651 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
21652 if sp2 {
21653 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21654 if Self::pdl_on() && Self::pdl_wb_on() {
21655 use cudarc::driver::{DevicePtr, DevicePtrMut};
21657 let s = &self.gpu.stream();
21658 let (pq, _b0) = q.device_ptr(s);
21659 let (pk, _b1) = k.device_ptr(s);
21660 let (pv, _b2) = v.device_ptr(s);
21661 let (po, _b3) = part_o.device_ptr_mut(s);
21662 let (pm, _b4) = part_m.device_ptr_mut(s);
21663 let (pl, _b5) = part_l.device_ptr_mut(s);
21664 let (pb, _b6) = base_dev.device_ptr(s);
21665 let mut ps = [
21666 &pq as *const _ as *mut std::ffi::c_void,
21667 &pk as *const _ as *mut _,
21668 &pv as *const _ as *mut _,
21669 &po as *const _ as *mut _,
21670 &pm as *const _ as *mut _,
21671 &pl as *const _ as *mut _,
21672 &hd as *const _ as *mut _,
21673 &nh as *const _ as *mut _,
21674 &nhkv as *const _ as *mut _,
21675 &pb as *const _ as *mut _,
21676 &base_plus as *const _ as *mut _,
21677 &scale as *const _ as *mut _,
21678 &nspm as *const _ as *mut _,
21679 &spk as *const _ as *mut _,
21680 &ktb as *const _ as *mut _,
21681 &vtb as *const _ as *mut _,
21682 &wini as *const _ as *mut _,
21683 ];
21684 unsafe {
21685 self.launch_pdl_flash(
21686 wg,
21687 "fa_decode_vec_q_rows_v4_w_sp",
21688 (n_head_kv as u32, n_splits_max as u32, t as u32),
21689 (32, gqa + 1, 1),
21690 sh,
21691 &mut ps,
21692 )?;
21693 }
21694 } else {
21695 let f = if wg {
21696 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
21697 } else {
21698 self.func("fa_decode_vec_q_rows_v4_w_sp")
21699 };
21700 f.set_attribute(
21701 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21702 sh as i32,
21703 )?;
21704 let cfg = LaunchConfig {
21705 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21706 block_dim: (32, gqa + 1, 1),
21707 shared_mem_bytes: sh,
21708 };
21709 let __s_b = self.gpu.stream();
21710 let mut b = __s_b.launch_builder(&f);
21711 b.arg(q)
21712 .arg(k)
21713 .arg(v)
21714 .arg(&mut *part_o)
21715 .arg(&mut *part_m)
21716 .arg(&mut *part_l)
21717 .arg(&hd)
21718 .arg(&nh)
21719 .arg(&nhkv)
21720 .arg(base_dev)
21721 .arg(&base_plus)
21722 .arg(&scale)
21723 .arg(&nspm)
21724 .arg(&spk)
21725 .arg(&ktb)
21726 .arg(&vtb)
21727 .arg(&wini);
21728 unsafe {
21729 b.launch(cfg)?;
21730 }
21731 }
21732 } else {
21733 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
21734 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
21736 use cudarc::driver::{DevicePtr, DevicePtrMut};
21737 let s = &self.gpu.stream();
21738 let (pq, _b0) = q.device_ptr(s);
21739 let (pk, _b1) = k.device_ptr(s);
21740 let (pv, _b2) = v.device_ptr(s);
21741 let (po, _b3) = part_o.device_ptr_mut(s);
21742 let (pm, _b4) = part_m.device_ptr_mut(s);
21743 let (pl, _b5) = part_l.device_ptr_mut(s);
21744 let (pb, _b6) = base_dev.device_ptr(s);
21745 let mut ps = [
21746 &pq as *const _ as *mut std::ffi::c_void,
21747 &pk as *const _ as *mut _,
21748 &pv as *const _ as *mut _,
21749 &po as *const _ as *mut _,
21750 &pm as *const _ as *mut _,
21751 &pl as *const _ as *mut _,
21752 &hd as *const _ as *mut _,
21753 &nh as *const _ as *mut _,
21754 &nhkv as *const _ as *mut _,
21755 &pb as *const _ as *mut _,
21756 &base_plus as *const _ as *mut _,
21757 &scale as *const _ as *mut _,
21758 &nspm as *const _ as *mut _,
21759 &spk as *const _ as *mut _,
21760 &ktb as *const _ as *mut _,
21761 &vtb as *const _ as *mut _,
21762 &wini as *const _ as *mut _,
21763 ];
21764 unsafe {
21765 self.launch_pdl_flash(
21766 wg,
21767 "fa_decode_vec_q_rows_v4_w",
21768 (n_head_kv as u32, n_splits_max as u32, t as u32),
21769 (32, gqa, 1),
21770 sh,
21771 &mut ps,
21772 )?;
21773 }
21774 } else {
21775 let pick = |name: &str| {
21776 if wg {
21777 self.func_g(name)
21778 } else {
21779 self.func(name)
21780 }
21781 };
21782 let (f, sh) = if fa_v4_at(window) {
21783 let f = pick("fa_decode_vec_q_rows_v4_w");
21784 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
21785 } else if smem_tkv > 0 && window >= smem_tkv {
21786 (
21789 pick("fa_decode_vec_q_rows_smem_w"),
21790 (2 * 32 * head_dim * 2) as u32,
21791 )
21792 } else {
21793 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
21794 };
21795 f.set_attribute(
21796 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21797 sh as i32,
21798 )?;
21799 let cfg = LaunchConfig {
21800 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
21801 block_dim: (32, gqa, 1),
21802 shared_mem_bytes: sh,
21803 };
21804 let __s_b = self.gpu.stream();
21805 let mut b = __s_b.launch_builder(&f);
21806 b.arg(q)
21807 .arg(k)
21808 .arg(v)
21809 .arg(&mut *part_o)
21810 .arg(&mut *part_m)
21811 .arg(&mut *part_l)
21812 .arg(&hd)
21813 .arg(&nh)
21814 .arg(&nhkv)
21815 .arg(base_dev)
21816 .arg(&base_plus)
21817 .arg(&scale)
21818 .arg(&nspm)
21819 .arg(&spk)
21820 .arg(&ktb)
21821 .arg(&vtb)
21822 .arg(&wini);
21823 unsafe {
21824 b.launch(cfg)?;
21825 }
21826 }
21827 }
21828 let cfg2 = LaunchConfig {
21829 grid_dim: (n_head as u32, t as u32, 1),
21830 block_dim: (head_dim as u32, 1, 1),
21831 shared_mem_bytes: 0,
21832 };
21833 if let Some((oq, od)) = q8_out {
21834 if Self::pdl_on() && Self::pdl_wb_on() {
21837 use cudarc::driver::{DevicePtr, DevicePtrMut};
21839 let s = &self.gpu.stream();
21840 let (po, _g0) = part_o.device_ptr(s);
21841 let (pm, _g1) = part_m.device_ptr(s);
21842 let (pl, _g2) = part_l.device_ptr(s);
21843 let (pq, _g3) = oq.device_ptr_mut(s);
21844 let (pd, _g4) = od.device_ptr_mut(s);
21845 let mut ps = [
21846 &po as *const _ as *mut std::ffi::c_void,
21847 &pm as *const _ as *mut _,
21848 &pl as *const _ as *mut _,
21849 &pq as *const _ as *mut _,
21850 &pd as *const _ as *mut _,
21851 &hd as *const _ as *mut _,
21852 &nh as *const _ as *mut _,
21853 &nspm as *const _ as *mut _,
21854 &spk as *const _ as *mut _,
21855 &wini as *const _ as *mut _,
21856 ];
21857 unsafe {
21858 self.launch_pdl_flash(
21859 wg,
21860 "fa_decode_combine_rows_w_q8_1",
21861 cfg2.grid_dim,
21862 cfg2.block_dim,
21863 0,
21864 &mut ps,
21865 )?;
21866 }
21867 return Ok(());
21868 }
21869 let fc = if wg {
21870 self.func_g("fa_decode_combine_rows_w_q8_1")
21871 } else {
21872 self.func("fa_decode_combine_rows_w_q8_1")
21873 };
21874 let __s_b2 = self.gpu.stream();
21875 let mut b2 = __s_b2.launch_builder(&fc);
21876 b2.arg(&*part_o)
21877 .arg(&*part_m)
21878 .arg(&*part_l)
21879 .arg(oq)
21880 .arg(od)
21881 .arg(&hd)
21882 .arg(&nh)
21883 .arg(&nspm)
21884 .arg(&spk)
21885 .arg(&wini);
21886 unsafe {
21887 b2.launch(cfg2)?;
21888 }
21889 return Ok(());
21890 }
21891 let fc = if wg {
21892 self.func_g("fa_decode_combine_rows_w")
21893 } else {
21894 self.func("fa_decode_combine_rows_w")
21895 };
21896 let __s_b2 = self.gpu.stream();
21897 let mut b2 = __s_b2.launch_builder(&fc);
21898 b2.arg(&*part_o)
21899 .arg(&*part_m)
21900 .arg(&*part_l)
21901 .arg(o)
21902 .arg(&hd)
21903 .arg(&nh)
21904 .arg(&nspm)
21905 .arg(&spk)
21906 .arg(&wini);
21907 unsafe {
21908 b2.launch(cfg2)?;
21909 }
21910 Ok(())
21911 }
21912
21913 #[allow(clippy::too_many_arguments)]
21919 pub fn fa_decode_rows_dc(
21920 &self,
21921 q: &CudaSlice<f32>,
21922 k: &cudarc::driver::CudaView<u8>,
21923 v: &cudarc::driver::CudaView<u8>,
21924 o: &mut CudaSlice<f32>,
21925 head_dim: usize,
21926 n_head: usize,
21927 n_head_kv: usize,
21928 base_dev: &CudaSlice<i32>,
21929 t_kv_upper: usize,
21930 t: usize,
21931 scale: f32,
21932 k_tok_bytes: usize,
21933 v_tok_bytes: usize,
21934 base_plus: i32,
21935 g: bool,
21936 ) -> Result<(), Box<dyn std::error::Error>> {
21937 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
21938 assert!(
21939 v4 || fa_v3_active(head_dim),
21940 "stream fa rows requires the v3 or v4 lane"
21941 );
21942 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
21943 if v4 {
21944 let sp = fa_split_keys(t_kv_upper, n_head_kv);
21945 let n_splits_max = (t_kv_upper + sp - 1) / sp;
21946 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21947 let (nspm, spk) = (n_splits_max as i32, sp as i32);
21948 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21949 let gqa = (n_head / n_head_kv).max(1) as u32;
21950 let o_len = t * n_head * n_splits_max * head_dim;
21951 let ml_len = t * n_head * n_splits_max;
21952 let mut part_guard = self.fa_part_pool.lock().unwrap();
21953 if part_guard
21954 .as_ref()
21955 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
21956 .unwrap_or(true)
21957 {
21958 let old = part_guard.take();
21969 let (co, cm) = old
21970 .as_ref()
21971 .map(|pp| (pp.0.len(), pp.1.len()))
21972 .unwrap_or((0, 0));
21973 if let Some(old) = old {
21974 self.fa_part_retired.lock().unwrap().push(old);
21975 }
21976 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
21977 eprintln!(
21978 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
21979 co, o_len, cm, ml_len
21980 );
21981 }
21982 *part_guard = Some((
21983 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
21984 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21985 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
21986 ));
21987 }
21988 let pg = part_guard.as_mut().unwrap();
21989 self.gpu
21990 .stream()
21991 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
21992 self.gpu
21993 .stream()
21994 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
21995 self.gpu
21996 .stream()
21997 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
21998 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
21999 let f = if g {
22000 self.func_g("fa_decode_vec_q_rows_v4_dc")
22001 } else {
22002 self.func("fa_decode_vec_q_rows_v4_dc")
22003 };
22004 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22005 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22006 f.set_attribute(
22007 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22008 sh as i32,
22009 )?;
22010 let cfg = LaunchConfig {
22011 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22012 block_dim: (32, gqa, 1),
22013 shared_mem_bytes: sh,
22014 };
22015 let __s_b = self.gpu.stream();
22016 let mut b = __s_b.launch_builder(&f);
22017 b.arg(q)
22018 .arg(k)
22019 .arg(v)
22020 .arg(&mut *part_o)
22021 .arg(&mut *part_m)
22022 .arg(&mut *part_l)
22023 .arg(&hd)
22024 .arg(&nh)
22025 .arg(&nhkv)
22026 .arg(base_dev)
22027 .arg(&base_plus)
22028 .arg(&scale)
22029 .arg(&nspm)
22030 .arg(&spk)
22031 .arg(&ktb)
22032 .arg(&vtb);
22033 unsafe {
22034 b.launch(cfg)?;
22035 }
22036 let fc = self.func("fa_decode_combine_rows_dc");
22037 let cfg2 = LaunchConfig {
22038 grid_dim: (n_head as u32, t as u32, 1),
22039 block_dim: (head_dim as u32, 1, 1),
22040 shared_mem_bytes: 0,
22041 };
22042 let __s_b2 = self.gpu.stream();
22043 let mut b2 = __s_b2.launch_builder(&fc);
22044 b2.arg(&*part_o)
22045 .arg(&*part_m)
22046 .arg(&*part_l)
22047 .arg(o)
22048 .arg(&hd)
22049 .arg(&nh)
22050 .arg(base_dev)
22051 .arg(&base_plus)
22052 .arg(&nspm)
22053 .arg(&spk);
22054 unsafe {
22055 b2.launch(cfg2)?;
22056 }
22057 return Ok(());
22058 }
22059 let sp = fa_split_keys(t_kv_upper, n_head_kv);
22060 let n_splits_max = (t_kv_upper + sp - 1) / sp;
22061 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22062 let (nspm, spk) = (n_splits_max as i32, sp as i32);
22063 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22064 let gqa = (n_head / n_head_kv).max(1) as u32;
22065 let o_len = t * n_head * n_splits_max * head_dim;
22066 let ml_len = t * n_head * n_splits_max;
22067 let mut part_guard = self.fa_part_pool.lock().unwrap();
22068 if part_guard
22069 .as_ref()
22070 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22071 .unwrap_or(true)
22072 {
22073 let old = part_guard.take();
22084 let (co, cm) = old
22085 .as_ref()
22086 .map(|pp| (pp.0.len(), pp.1.len()))
22087 .unwrap_or((0, 0));
22088 if let Some(old) = old {
22089 self.fa_part_retired.lock().unwrap().push(old);
22090 }
22091 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22092 eprintln!(
22093 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22094 co, o_len, cm, ml_len
22095 );
22096 }
22097 *part_guard = Some((
22098 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22099 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22100 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22101 ));
22102 }
22103 let pg = part_guard.as_mut().unwrap();
22104 self.gpu
22105 .stream()
22106 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22107 self.gpu
22108 .stream()
22109 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22110 self.gpu
22111 .stream()
22112 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22113 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22114 let f = self.func("fa_decode_vec_q_rows_v3_dc");
22115 let sh = (32 * head_dim * 2) as u32;
22116 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22117 f.set_attribute(
22118 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22119 sh as i32,
22120 )?;
22121 let cfg = LaunchConfig {
22122 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
22123 block_dim: (32, gqa, 1),
22124 shared_mem_bytes: sh,
22125 };
22126 let __s_b = self.gpu.stream();
22127 let mut b = __s_b.launch_builder(&f);
22128 b.arg(q)
22129 .arg(k)
22130 .arg(v)
22131 .arg(&mut *part_o)
22132 .arg(&mut *part_m)
22133 .arg(&mut *part_l)
22134 .arg(&hd)
22135 .arg(&nh)
22136 .arg(&nhkv)
22137 .arg(base_dev)
22138 .arg(&scale)
22139 .arg(&nspm)
22140 .arg(&spk)
22141 .arg(&ktb)
22142 .arg(&vtb);
22143 unsafe {
22144 b.launch(cfg)?;
22145 }
22146 let fc = self.func("fa_decode_combine_rows_dc");
22147 let cfg2 = LaunchConfig {
22148 grid_dim: (n_head as u32, t as u32, 1),
22149 block_dim: (head_dim as u32, 1, 1),
22150 shared_mem_bytes: 0,
22151 };
22152 let plus0 = 0i32;
22153 let __s_b2 = self.gpu.stream();
22154 let mut b2 = __s_b2.launch_builder(&fc);
22155 b2.arg(&*part_o)
22156 .arg(&*part_m)
22157 .arg(&*part_l)
22158 .arg(o)
22159 .arg(&hd)
22160 .arg(&nh)
22161 .arg(base_dev)
22162 .arg(&plus0)
22163 .arg(&nspm)
22164 .arg(&spk);
22165 unsafe {
22166 b2.launch(cfg2)?;
22167 }
22168 Ok(())
22169 }
22170
22171 pub fn fa_decode_dc(
22182 &self,
22183 q: &CudaSlice<f32>,
22184 k: &cudarc::driver::CudaView<u8>,
22185 v: &cudarc::driver::CudaView<u8>,
22186 o: &mut CudaSlice<f32>,
22187 head_dim: usize,
22188 n_head: usize,
22189 n_head_kv: usize,
22190 t_kv_dev: &CudaSlice<i32>,
22191 bucket_max: usize,
22192 scale: f32,
22193 k_tok_bytes: usize,
22194 v_tok_bytes: usize,
22195 g: bool,
22196 ) -> Result<(), Box<dyn std::error::Error>> {
22197 self.fa_decode_dc_q8(
22198 q,
22199 k,
22200 v,
22201 o,
22202 head_dim,
22203 n_head,
22204 n_head_kv,
22205 t_kv_dev,
22206 bucket_max,
22207 scale,
22208 k_tok_bytes,
22209 v_tok_bytes,
22210 g,
22211 None,
22212 )
22213 }
22214
22215 #[allow(clippy::too_many_arguments)]
22218 pub fn fa_decode_dc_q8(
22219 &self,
22220 q: &CudaSlice<f32>,
22221 k: &cudarc::driver::CudaView<u8>,
22222 v: &cudarc::driver::CudaView<u8>,
22223 o: &mut CudaSlice<f32>,
22224 head_dim: usize,
22225 n_head: usize,
22226 n_head_kv: usize,
22227 t_kv_dev: &CudaSlice<i32>,
22228 bucket_max: usize,
22229 scale: f32,
22230 k_tok_bytes: usize,
22231 v_tok_bytes: usize,
22232 g: bool,
22233 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22234 ) -> Result<(), Box<dyn std::error::Error>> {
22235 let mut fa_vec =
22243 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22244 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
22245 fa_vec = false;
22246 } let sp = fa_split_keys(bucket_max, n_head_kv);
22248 let n_splits = if fa_vec {
22249 ((bucket_max + sp - 1) / sp).max(1)
22250 } else {
22251 ((bucket_max + 255) / 256).max(1)
22252 };
22253 let o_len = n_head * n_splits * head_dim;
22254 let ml_len = n_head * n_splits;
22255 let mut part_guard = self.fa_part_pool.lock().unwrap();
22256 if part_guard
22257 .as_ref()
22258 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22259 .unwrap_or(true)
22260 {
22261 let old = part_guard.take();
22272 let (co, cm) = old
22273 .as_ref()
22274 .map(|pp| (pp.0.len(), pp.1.len()))
22275 .unwrap_or((0, 0));
22276 if let Some(old) = old {
22277 self.fa_part_retired.lock().unwrap().push(old);
22278 }
22279 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22280 eprintln!(
22281 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22282 co, o_len, cm, ml_len
22283 );
22284 }
22285 *part_guard = Some((
22286 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22287 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22288 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22289 ));
22290 }
22291 let pg = part_guard.as_mut().unwrap();
22292 self.gpu
22293 .stream()
22294 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22295 self.gpu
22296 .stream()
22297 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22298 self.gpu
22299 .stream()
22300 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22301 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22302 let (hd, nh, nhkv, nsp) = (
22303 head_dim as i32,
22304 n_head as i32,
22305 n_head_kv as i32,
22306 n_splits as i32,
22307 );
22308 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22309 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22310 let deep = fa_vec
22313 && head_dim == 256
22314 && fa_v4_at(bucket_max)
22315 && !g
22316 && fa_deep_at(bucket_max)
22317 && !matches!(fa_v4_mode(), "noB3" | "stage");
22318 let (f, cfg) = if fa_vec
22319 && head_dim == 512
22320 && bucket_max >= {
22321 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22322 *FA512_MIN_DC.get_or_init(|| {
22323 std::env::var("MEMRA_FA512_MIN")
22324 .ok()
22325 .and_then(|v| v.parse().ok())
22326 .unwrap_or(512)
22327 })
22328 } {
22329 let gqa = (n_head / n_head_kv).max(1) as u32;
22331 (
22332 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
22333 LaunchConfig {
22334 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22335 block_dim: (32, gqa, 1),
22336 shared_mem_bytes: 0,
22337 },
22338 )
22339 } else if fa_vec && head_dim == 512 {
22340 let q_view = q.as_view();
22343 let mut o_view = o.as_view_mut();
22344 return self.fa_decode_scalar_unified(
22345 &q_view,
22346 k,
22347 v,
22348 &mut o_view,
22349 head_dim,
22350 n_head,
22351 n_head_kv,
22352 0,
22353 Some(t_kv_dev),
22354 scale,
22355 n_splits,
22356 sp,
22357 k_tok_bytes,
22358 v_tok_bytes,
22359 g,
22360 &mut *part_o,
22361 &mut *part_m,
22362 &mut *part_l,
22363 q8_out,
22364 );
22365 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
22366 let gqa = (n_head / n_head_kv).max(1) as u32;
22369 let fv = if g {
22370 self.func_g("fa_decode_vec_q_v4_dc")
22371 } else if deep {
22372 self.func("fa_decode_vec_q_v4_deep_dc")
22373 } else {
22374 self.func("fa_decode_vec_q_v4_dc")
22375 };
22376 let shmem =
22377 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22378 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22379 fv.set_attribute(
22380 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22381 shmem as i32,
22382 )?;
22383 (
22384 fv,
22385 LaunchConfig {
22386 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22387 block_dim: (32, gqa, 1),
22388 shared_mem_bytes: shmem,
22389 },
22390 )
22391 } else if fa_vec && fa_v3_active(head_dim) {
22392 let gqa = (n_head / n_head_kv).max(1) as u32;
22395 let fv = if g {
22396 self.func_g("fa_decode_vec_q_v3_dc")
22397 } else {
22398 self.func("fa_decode_vec_q_v3_dc")
22399 };
22400 let shmem = (32 * head_dim * 2) as u32; (
22402 fv,
22403 LaunchConfig {
22404 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22405 block_dim: (32, gqa, 1),
22406 shared_mem_bytes: shmem,
22407 },
22408 )
22409 } else if fa_vec && fa_v2_on() {
22410 let gqa = (n_head / n_head_kv).max(1) as u32;
22414 let fv = if g {
22415 self.func_g("fa_decode_vec_q_v2_dc")
22416 } else {
22417 self.func("fa_decode_vec_q_v2_dc")
22418 };
22419 let shmem = (2 * 32 * head_dim * 2) as u32; (
22421 fv,
22422 LaunchConfig {
22423 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22424 block_dim: (32, gqa, 1),
22425 shared_mem_bytes: shmem,
22426 },
22427 )
22428 } else if fa_vec {
22429 let gqa = (n_head / n_head_kv).max(1) as u32;
22430 let fv = if g {
22432 self.func_g("fa_decode_vec_q_dc")
22433 } else {
22434 self.func("fa_decode_vec_q_dc")
22435 };
22436 (
22437 fv,
22438 LaunchConfig {
22439 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22440 block_dim: (32, gqa, 1),
22441 shared_mem_bytes: 0,
22442 },
22443 )
22444 } else {
22445 let q_view = q.as_view();
22446 let mut o_view = o.as_view_mut();
22447 return self.fa_decode_scalar_unified(
22448 &q_view,
22449 k,
22450 v,
22451 &mut o_view,
22452 head_dim,
22453 n_head,
22454 n_head_kv,
22455 0,
22456 Some(t_kv_dev),
22457 scale,
22458 n_splits,
22459 if fa_vec { sp } else { 256 },
22460 k_tok_bytes,
22461 v_tok_bytes,
22462 g,
22463 &mut *part_o,
22464 &mut *part_m,
22465 &mut *part_l,
22466 q8_out,
22467 );
22468 };
22469 let ski = sp as i32; let __s_b = self.gpu.stream();
22471 let mut b = __s_b.launch_builder(&f);
22472 b.arg(q)
22473 .arg(k)
22474 .arg(v)
22475 .arg(&mut *part_o)
22476 .arg(&mut *part_m)
22477 .arg(&mut *part_l)
22478 .arg(&hd)
22479 .arg(&nh)
22480 .arg(&nhkv)
22481 .arg(t_kv_dev)
22482 .arg(&scale)
22483 .arg(&nsp)
22484 .arg(&ski)
22485 .arg(&ktb)
22486 .arg(&vtb);
22487 unsafe {
22488 b.launch(cfg)?;
22489 }
22490 let cfg2 = LaunchConfig {
22491 grid_dim: (n_head as u32, 1, 1),
22492 block_dim: (head_dim as u32, 1, 1),
22493 shared_mem_bytes: 0,
22494 };
22495 if let Some((oq, od)) = q8_out {
22496 let fc = if g {
22497 self.func_g("fa_decode_combine_q8_1")
22498 } else {
22499 self.fa_func("fa_decode_combine_q8_1", head_dim)
22500 };
22501 let __s_b2 = self.gpu.stream();
22502 let mut b2 = __s_b2.launch_builder(&fc);
22503 b2.arg(&*part_o)
22504 .arg(&*part_m)
22505 .arg(&*part_l)
22506 .arg(oq)
22507 .arg(od)
22508 .arg(&hd)
22509 .arg(&nh)
22510 .arg(&nsp);
22511 unsafe {
22512 b2.launch(cfg2)?;
22513 }
22514 return Ok(());
22515 }
22516 let fc = if g {
22517 self.func_g("fa_decode_combine_f32")
22518 } else {
22519 self.fa_func("fa_decode_combine_f32", head_dim)
22520 };
22521 let __s_b2 = self.gpu.stream();
22522 let mut b2 = __s_b2.launch_builder(&fc);
22523 b2.arg(&*part_o)
22524 .arg(&*part_m)
22525 .arg(&*part_l)
22526 .arg(o)
22527 .arg(&hd)
22528 .arg(&nh)
22529 .arg(&nsp);
22530 unsafe {
22531 b2.launch(cfg2)?;
22532 }
22533 Ok(())
22534 }
22535
22536 #[allow(clippy::too_many_arguments)]
22540 pub fn append_kv_quantized_dcw(
22541 &self,
22542 k_row: &CudaSlice<f32>,
22543 v_row: &CudaSlice<f32>,
22544 kc: &mut CudaSlice<u8>,
22545 vc: &mut CudaSlice<u8>,
22546 len_dev: &CudaSlice<i32>,
22547 base_dev: Option<&CudaSlice<i32>>,
22548 kv_dim_k: usize,
22549 kv_dim_v: usize,
22550 k_tok_bytes: usize,
22551 v_tok_bytes: usize,
22552 ) -> Result<(), Box<dyn std::error::Error>> {
22553 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
22554 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
22555 let cfg = LaunchConfig {
22556 grid_dim: (nblk, 1, 1),
22557 block_dim: (32, 1, 1),
22558 shared_mem_bytes: 0,
22559 };
22560 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22561 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22562 let null: u64 = 0;
22563 let __s_b = self.gpu.stream();
22564 let mut b = __s_b.launch_builder(&f);
22565 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
22566 match base_dev {
22567 Some(base) => {
22568 b.arg(base);
22569 }
22570 None => {
22571 b.arg(&null);
22572 }
22573 }
22574 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
22575 unsafe {
22576 b.launch(cfg)?;
22577 }
22578 Ok(())
22579 }
22580
22581 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
22583 let f = self.func("inc_i32");
22584 let cfg = LaunchConfig {
22585 grid_dim: (1, 1, 1),
22586 block_dim: (1, 1, 1),
22587 shared_mem_bytes: 0,
22588 };
22589 let __s_b = self.gpu.stream();
22590 let mut b = __s_b.launch_builder(&f);
22591 b.arg(counter);
22592 unsafe {
22593 b.launch(cfg)?;
22594 }
22595 Ok(())
22596 }
22597
22598 #[allow(clippy::too_many_arguments)]
22607 fn fa_part_pool_grow(
22611 &self,
22612 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
22613 o_len: usize,
22614 ml_len: usize,
22615 ) -> Result<(), Box<dyn std::error::Error>> {
22616 if part_guard
22617 .as_ref()
22618 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22619 .unwrap_or(true)
22620 {
22621 let old = part_guard.take();
22622 let (co, cm) = old
22623 .as_ref()
22624 .map(|pp| (pp.0.len(), pp.1.len()))
22625 .unwrap_or((0, 0));
22626 if let Some(old) = old {
22627 self.fa_part_retired.lock().unwrap().push(old);
22628 }
22629 *part_guard = Some((
22630 self.alloc_uninit::<f32>(o_len.max(2 * co))?,
22631 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22632 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
22633 ));
22634 }
22635 Ok(())
22636 }
22637
22638 pub fn fa_dcw_pool_ensure(
22641 &self,
22642 head_dim: usize,
22643 n_head: usize,
22644 n_head_kv: usize,
22645 bucket_max: usize,
22646 ) -> Result<(), Box<dyn std::error::Error>> {
22647 let sp = fa_split_keys(bucket_max, n_head_kv);
22648 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22649 let o_len = n_head * n_splits * head_dim;
22650 let ml_len = n_head * n_splits;
22651 let mut part_guard = self.fa_part_pool.lock().unwrap();
22652 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
22653 }
22654
22655 pub fn fa_decode_dcw(
22656 &self,
22657 q: &CudaSlice<f32>,
22658 k_ring: &cudarc::driver::CudaView<u8>,
22659 v_ring: &cudarc::driver::CudaView<u8>,
22660 o: &mut CudaSlice<f32>,
22661 head_dim: usize,
22662 n_head: usize,
22663 n_head_kv: usize,
22664 len_dev: &CudaSlice<i32>,
22665 base_dev: Option<&CudaSlice<i32>>,
22666 window: usize,
22667 bucket_max: usize,
22668 scale: f32,
22669 k_tok_bytes: usize,
22670 v_tok_bytes: usize,
22671 fused_gate: Option<&CudaSlice<f32>>,
22675 ) -> Result<(), Box<dyn std::error::Error>> {
22676 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
22677 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
22678 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"
22679 .into());
22680 }
22681 let sp = fa_split_keys(bucket_max, n_head_kv);
22682 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
22683 let o_len = n_head * n_splits * head_dim;
22684 let ml_len = n_head * n_splits;
22685 let mut part_guard = self.fa_part_pool.lock().unwrap();
22686 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
22687 let pg = part_guard.as_mut().unwrap();
22688 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
22693 let memset_on = *MEMSET_ON
22698 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
22699 || crate::tp::token_graph_building();
22700 if memset_on {
22701 self.gpu
22702 .stream()
22703 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22704 self.gpu
22705 .stream()
22706 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22707 self.gpu
22708 .stream()
22709 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22710 }
22711 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22712 let (hd, nh, nhkv, nsp) = (
22713 head_dim as i32,
22714 n_head as i32,
22715 n_head_kv as i32,
22716 n_splits as i32,
22717 );
22718 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22719 let (ski, win) = (sp as i32, window as i32);
22720 let gqa = (n_head / n_head_kv).max(1) as u32;
22721 let smem = (32 * head_dim * 2) as u32; let f = self.func("fa_decode_vec_q_v3_dcw");
22723 let cfg = LaunchConfig {
22724 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22725 block_dim: (32, gqa, 1),
22726 shared_mem_bytes: smem,
22727 };
22728 let null: u64 = 0;
22729 let __s_b = self.gpu.stream();
22730 let mut b = __s_b.launch_builder(&f);
22731 b.arg(q)
22732 .arg(k_ring)
22733 .arg(v_ring)
22734 .arg(&mut *part_o)
22735 .arg(&mut *part_m)
22736 .arg(&mut *part_l)
22737 .arg(&hd)
22738 .arg(&nh)
22739 .arg(&nhkv)
22740 .arg(len_dev);
22741 match base_dev {
22742 Some(base) => {
22743 b.arg(base);
22744 }
22745 None => {
22746 b.arg(&null);
22747 }
22748 }
22749 b.arg(&win)
22750 .arg(&scale)
22751 .arg(&nsp)
22752 .arg(&ski)
22753 .arg(&ktb)
22754 .arg(&vtb);
22755 unsafe {
22756 b.launch(cfg)?;
22757 }
22758 let fc = if fused_gate.is_some() {
22759 self.func("fa_decode_combine_gate_f32")
22760 } else {
22761 self.fa_func("fa_decode_combine_f32", head_dim)
22762 };
22763 let cfg2 = LaunchConfig {
22764 grid_dim: (n_head as u32, 1, 1),
22765 block_dim: (head_dim as u32, 1, 1),
22766 shared_mem_bytes: 0,
22767 };
22768 let __s_b2 = self.gpu.stream();
22769 let mut b2 = __s_b2.launch_builder(&fc);
22770 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
22771 if let Some(gate_row) = fused_gate {
22772 b2.arg(gate_row);
22773 }
22774 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
22775 unsafe {
22776 b2.launch(cfg2)?;
22777 }
22778 Ok(())
22779 }
22780
22781 pub fn fa_geom_eager(
22787 &self,
22788 t_kv: usize,
22789 head_dim: usize,
22790 n_head_kv: usize,
22791 g: bool,
22792 ) -> (bool, usize) {
22793 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
22797 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
22803 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
22804 if g && head_dim == 256 && !fa_v4_at(t_kv) {
22810 fa_vec = false;
22811 }
22812 let sp = fa_split_keys(t_kv, n_head_kv);
22813 let n_splits = if fa_vec {
22814 ((t_kv + sp - 1) / sp).max(1)
22815 } else {
22816 ((t_kv + 255) / 256).max(1)
22817 };
22818 (fa_vec, n_splits)
22819 }
22820
22821 pub fn fa_bucket_key(
22827 &self,
22828 t_kv: usize,
22829 head_dim: usize,
22830 n_head_kv: usize,
22831 g: bool,
22832 ) -> (bool, usize) {
22833 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
22834 }
22835
22836 pub fn capture_graph_retained<F>(
22848 &self,
22849 step: F,
22850 ) -> Result<
22851 (
22852 cudarc::driver::CudaGraph,
22853 Vec<Box<dyn std::any::Any + Send>>,
22854 ),
22855 Box<dyn std::error::Error>,
22856 >
22857 where
22858 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22859 {
22860 use cudarc::driver::sys::CUgraphInstantiate_flags;
22861 self.capture_graph_retained_flags(
22862 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
22863 step,
22864 )
22865 }
22866
22867 pub fn capture_graph_retained_flags<F>(
22872 &self,
22873 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
22874 mut step: F,
22875 ) -> Result<
22876 (
22877 cudarc::driver::CudaGraph,
22878 Vec<Box<dyn std::any::Any + Send>>,
22879 ),
22880 Box<dyn std::error::Error>,
22881 >
22882 where
22883 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22884 {
22885 use cudarc::driver::sys::CUstreamCaptureMode;
22886 self.capture_keep.lock().unwrap().clear();
22894 let was_tracking = self.gpu.ctx.is_event_tracking();
22895 if was_tracking {
22896 unsafe {
22897 self.gpu.ctx.disable_event_tracking();
22898 }
22899 }
22900 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
22901 self.capture_keep_on
22902 .store(true, std::sync::atomic::Ordering::Relaxed);
22903 let w = (|| {
22904 step(self)?;
22905 step(self)
22906 })();
22907 self.capture_keep_on
22908 .store(false, std::sync::atomic::Ordering::Relaxed);
22909 w?;
22910 self.gpu.stream().synchronize()?;
22911 self.gpu
22912 .stream()
22913 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
22914 let r = step(self);
22915 let g = self.gpu.stream().end_capture(flags);
22916 r?;
22917 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
22918 graph.upload()?;
22919 Ok(graph)
22920 };
22921 let result = run();
22922 self.capture_keep_on
22923 .store(false, std::sync::atomic::Ordering::Relaxed);
22924 if was_tracking {
22925 unsafe {
22926 self.gpu.ctx.enable_event_tracking();
22927 }
22928 }
22929 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
22930 Ok((result?, keeper))
22931 }
22932
22933 pub fn capture_graph_retained_nowarm<F>(
22939 &self,
22940 mut step: F,
22941 ) -> Result<
22942 (
22943 cudarc::driver::CudaGraph,
22944 Vec<Box<dyn std::any::Any + Send>>,
22945 ),
22946 Box<dyn std::error::Error>,
22947 >
22948 where
22949 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22950 {
22951 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
22952 let was_tracking = self.gpu.ctx.is_event_tracking();
22953 if was_tracking {
22954 unsafe {
22955 self.gpu.ctx.disable_event_tracking();
22956 }
22957 }
22958 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
22959 self.gpu.stream().synchronize()?;
22960 self.gpu
22961 .stream()
22962 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
22963 let r = step(self);
22964 let g = self.gpu.stream().end_capture(
22965 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
22966 );
22967 r?;
22968 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
22969 graph.upload()?;
22970 Ok(graph)
22971 };
22972 let result = run();
22973 if was_tracking {
22974 unsafe {
22975 self.gpu.ctx.enable_event_tracking();
22976 }
22977 }
22978 Ok((result?, Vec::new()))
22979 }
22980
22981 pub fn capture_graph<F>(
22982 &self,
22983 mut step: F,
22984 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
22985 where
22986 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
22987 {
22988 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
22989 let was_tracking = self.gpu.ctx.is_event_tracking();
22997 if was_tracking {
22998 unsafe {
22999 self.gpu.ctx.disable_event_tracking();
23000 }
23001 }
23002 let iflag = {
23009 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
23010 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
23011 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
23014 Ok("priority") => {
23015 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
23016 }
23017 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
23018 })
23019 };
23020 let ct = {
23027 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23028 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
23029 };
23030 let warmups = {
23053 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23054 *W.get_or_init(|| {
23055 std::env::var("MEMRA_GRAPH_WARMUPS")
23056 .ok()
23057 .and_then(|v| v.parse().ok())
23058 .filter(|n| *n >= 1)
23059 .unwrap_or(1)
23060 })
23061 };
23062 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
23063 let t_w = std::time::Instant::now();
23064 for _ in 0..warmups {
23066 step(self)?;
23067 }
23068 self.gpu.stream().synchronize()?;
23069 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
23070 let t_c = std::time::Instant::now();
23072 self.gpu
23073 .stream()
23074 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
23075 let r = step(self);
23078 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
23079 let t_i = std::time::Instant::now();
23080 let g = self.gpu.stream().end_capture(iflag);
23081 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
23082 r?;
23083 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
23084 let t_u = std::time::Instant::now();
23085 graph.upload()?;
23086 if ct {
23087 println!(
23088 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
23089 instantiate {ms_inst:.2} ms upload {:.2} ms",
23090 t_u.elapsed().as_secs_f64() * 1e3
23091 );
23092 }
23093 Ok(graph)
23094 };
23095 let result = run();
23096 if was_tracking {
23097 unsafe {
23098 self.gpu.ctx.enable_event_tracking();
23099 }
23100 }
23101 result
23102 }
23103
23104 pub fn gdn_scan_s128_view(
23106 &self,
23107 q: &CudaSlice<f32>,
23108 k: &CudaSlice<f32>,
23109 v: &CudaSlice<f32>,
23110 g: &CudaSlice<f32>,
23111 beta: &CudaSlice<f32>,
23112 state_in: &cudarc::driver::CudaView<f32>,
23113 state_out: &mut cudarc::driver::CudaViewMut<f32>,
23114 o: &mut CudaSlice<f32>,
23115 n_head: usize,
23116 t: usize,
23117 scale: f32,
23118 ) -> Result<(), Box<dyn std::error::Error>> {
23119 let f = self.func("gdn_scan_s128");
23120 const S_V: u32 = 128;
23121 const WARP: u32 = 32;
23122 const COLS: u32 = 4;
23123 let cfg = LaunchConfig {
23124 grid_dim: (n_head as u32, 1, S_V / COLS),
23125 block_dim: (WARP, COLS, 1),
23126 shared_mem_bytes: 0,
23127 };
23128 let (h, ti) = (n_head as i32, t as i32);
23129 let __s_b = self.gpu.stream();
23130 let mut b = __s_b.launch_builder(&f);
23131 b.arg(q)
23132 .arg(k)
23133 .arg(v)
23134 .arg(g)
23135 .arg(beta)
23136 .arg(state_in)
23137 .arg(state_out)
23138 .arg(o)
23139 .arg(&h)
23140 .arg(&ti)
23141 .arg(&scale);
23142 unsafe {
23143 b.launch(cfg)?;
23144 }
23145 Ok(())
23146 }
23147
23148 pub fn ssm_conv1d_view(
23150 &self,
23151 x: &cudarc::driver::CudaView<f32>,
23152 w: &CudaSlice<f32>,
23153 y: &mut CudaSlice<f32>,
23154 conv_dim: usize,
23155 t: usize,
23156 d_conv: usize,
23157 silu: bool,
23158 ) -> Result<(), Box<dyn std::error::Error>> {
23159 let f = self.func("ssm_conv1d_silu_f32");
23160 let cfg = LaunchConfig {
23162 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23163 block_dim: (256, 1, 1),
23164 shared_mem_bytes: 0,
23165 };
23166 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23167 let __s_b = self.gpu.stream();
23168 let mut b = __s_b.launch_builder(&f);
23169 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23170 unsafe {
23171 b.launch(cfg)?;
23172 }
23173 Ok(())
23174 }
23175
23176 pub fn ssm_conv1d_tm(
23183 &self,
23184 qkv_tm: &CudaSlice<f32>,
23185 w: &CudaSlice<f32>,
23186 y: &mut CudaSlice<f32>,
23187 conv_dim: usize,
23188 t: usize,
23189 d_conv: usize,
23190 ) -> Result<(), Box<dyn std::error::Error>> {
23191 let f = self.func("ssm_conv1d_tm_f32");
23192 let cfg = LaunchConfig {
23193 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23194 block_dim: (256, 1, 1),
23195 shared_mem_bytes: 0,
23196 };
23197 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23198 let __s_b = self.gpu.stream();
23199 let mut b = __s_b.launch_builder(&f);
23200 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
23201 unsafe {
23202 b.launch(cfg)?;
23203 }
23204 Ok(())
23205 }
23206
23207 pub fn ssm_conv1d_tm_state(
23215 &self,
23216 qkv_tm: &CudaSlice<f32>,
23217 conv_state: &mut CudaSlice<f32>,
23218 w: &CudaSlice<f32>,
23219 y: &mut CudaSlice<f32>,
23220 conv_dim: usize,
23221 t: usize,
23222 d_conv: usize,
23223 ) -> Result<(), Box<dyn std::error::Error>> {
23224 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
23225 }
23226
23227 #[allow(clippy::too_many_arguments)]
23230 pub fn ssm_conv1d_tm_state_pad(
23231 &self,
23232 qkv_tm: &CudaSlice<f32>,
23233 conv_state: &mut CudaSlice<f32>,
23234 w: &CudaSlice<f32>,
23235 y: &mut CudaSlice<f32>,
23236 conv_dim: usize,
23237 t: usize,
23238 d_conv: usize,
23239 pad_len: Option<&CudaSlice<i32>>,
23240 ) -> Result<(), Box<dyn std::error::Error>> {
23241 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23242 let ring_old = if t < d_conv - 1 {
23246 Some(self.clone_dtod(conv_state)?)
23247 } else {
23248 None
23249 };
23250 {
23251 let f = self.func("ssm_conv1d_tm_state_f32");
23252 let cfg = LaunchConfig {
23253 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23254 block_dim: (256, 1, 1),
23255 shared_mem_bytes: 0,
23256 };
23257 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23258 let __s_b = self.gpu.stream();
23259 let mut b = __s_b.launch_builder(&f);
23260 b.arg(qkv_tm)
23261 .arg(&*conv_state)
23262 .arg(w)
23263 .arg(y)
23264 .arg(&cd)
23265 .arg(&ti)
23266 .arg(&dc);
23267 unsafe {
23268 b.launch(cfg)?;
23269 }
23270 }
23271 match (ring_old, pad_len) {
23272 (None, Some(len_d)) => {
23273 let f = self.func("ssm_conv_ring_update_dev_f32");
23274 let n = conv_dim * (d_conv - 1);
23275 let cfg = LaunchConfig::for_num_elems(n as u32);
23276 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23277 let __s_b = self.gpu.stream();
23278 let mut b = __s_b.launch_builder(&f);
23279 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23280 unsafe {
23281 b.launch(cfg)?;
23282 }
23283 }
23284 (None, None) => {
23285 let f = self.func("ssm_conv_ring_update_f32");
23286 let n = conv_dim * (d_conv - 1);
23287 let cfg = LaunchConfig::for_num_elems(n as u32);
23288 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23289 let __s_b = self.gpu.stream();
23290 let mut b = __s_b.launch_builder(&f);
23291 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23292 unsafe {
23293 b.launch(cfg)?;
23294 }
23295 }
23296 (Some(old), _) => {
23297 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
23298 }
23299 }
23300 Ok(())
23301 }
23302
23303 pub fn ssm_conv1d_tm_state_pad_v(
23305 &self,
23306 qkv_tm: &cudarc::driver::CudaView<f32>,
23307 conv_state: &mut CudaSlice<f32>,
23308 w: &CudaSlice<f32>,
23309 y: &mut CudaSlice<f32>,
23310 conv_dim: usize,
23311 t: usize,
23312 d_conv: usize,
23313 pad_len: Option<&CudaSlice<i32>>,
23314 ) -> Result<(), Box<dyn std::error::Error>> {
23315 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
23316 let ring_old = if t < d_conv - 1 {
23320 Some(self.clone_dtod(conv_state)?)
23321 } else {
23322 None
23323 };
23324 {
23325 let f = self.func("ssm_conv1d_tm_state_f32");
23326 let cfg = LaunchConfig {
23327 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23328 block_dim: (256, 1, 1),
23329 shared_mem_bytes: 0,
23330 };
23331 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23332 let __s_b = self.gpu.stream();
23333 let mut b = __s_b.launch_builder(&f);
23334 b.arg(qkv_tm)
23335 .arg(&*conv_state)
23336 .arg(w)
23337 .arg(y)
23338 .arg(&cd)
23339 .arg(&ti)
23340 .arg(&dc);
23341 unsafe {
23342 b.launch(cfg)?;
23343 }
23344 }
23345 match (ring_old, pad_len) {
23346 (None, Some(len_d)) => {
23347 let f = self.func("ssm_conv_ring_update_dev_f32");
23348 let n = conv_dim * (d_conv - 1);
23349 let cfg = LaunchConfig::for_num_elems(n as u32);
23350 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23351 let __s_b = self.gpu.stream();
23352 let mut b = __s_b.launch_builder(&f);
23353 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
23354 unsafe {
23355 b.launch(cfg)?;
23356 }
23357 }
23358 (None, None) => {
23359 let f = self.func("ssm_conv_ring_update_f32");
23360 let n = conv_dim * (d_conv - 1);
23361 let cfg = LaunchConfig::for_num_elems(n as u32);
23362 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23363 let __s_b = self.gpu.stream();
23364 let mut b = __s_b.launch_builder(&f);
23365 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
23366 unsafe {
23367 b.launch(cfg)?;
23368 }
23369 }
23370 (Some(_), _) => unreachable!(
23371 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
23372 ),
23373 }
23374 Ok(())
23375 }
23376
23377 pub fn ssm_conv_ring_rebuild(
23382 &self,
23383 qkv_tm: &CudaSlice<f32>,
23384 ring_old: &CudaSlice<f32>,
23385 conv_state: &mut CudaSlice<f32>,
23386 conv_dim: usize,
23387 tc: usize,
23388 d_conv: usize,
23389 ) -> Result<(), Box<dyn std::error::Error>> {
23390 let f = self.func("ssm_conv_ring_rebuild_f32");
23391 let n = conv_dim * (d_conv - 1);
23392 let cfg = LaunchConfig::for_num_elems(n as u32);
23393 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
23394 let __s_b = self.gpu.stream();
23395 let mut b = __s_b.launch_builder(&f);
23396 b.arg(qkv_tm)
23397 .arg(ring_old)
23398 .arg(conv_state)
23399 .arg(&cd)
23400 .arg(&ti)
23401 .arg(&dc);
23402 unsafe {
23403 b.launch(cfg)?;
23404 }
23405 Ok(())
23406 }
23407
23408 #[allow(clippy::too_many_arguments)]
23413 pub fn gdn_prep_decode(
23414 &self,
23415 conv_out: &CudaSlice<f32>,
23416 beta_raw: &CudaSlice<f32>,
23417 alpha: &CudaSlice<f32>,
23418 dt_bias: &CudaSlice<f32>,
23419 a: &CudaSlice<f32>,
23420 q_l2: &mut CudaSlice<f32>,
23421 k_l2: &mut CudaSlice<f32>,
23422 v_g: &mut CudaSlice<f32>,
23423 beta: &mut CudaSlice<f32>,
23424 g_log: &mut CudaSlice<f32>,
23425 d_state: usize,
23426 num_v: usize,
23427 num_k: usize,
23428 key_dim: usize,
23429 eps: f32,
23430 ) -> Result<(), Box<dyn std::error::Error>> {
23431 let f = self.func("gdn_prep_decode_f32");
23432 let cfg = LaunchConfig {
23433 grid_dim: (num_v as u32, 1, 1),
23434 block_dim: (32, 4, 1),
23435 shared_mem_bytes: 0,
23436 };
23437 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23438 let __s_b = self.gpu.stream();
23439 let mut b = __s_b.launch_builder(&f);
23440 b.arg(conv_out)
23441 .arg(beta_raw)
23442 .arg(alpha)
23443 .arg(dt_bias)
23444 .arg(a)
23445 .arg(q_l2)
23446 .arg(k_l2)
23447 .arg(v_g)
23448 .arg(beta)
23449 .arg(g_log)
23450 .arg(&ds)
23451 .arg(&nv)
23452 .arg(&nk)
23453 .arg(&kd)
23454 .arg(&eps);
23455 unsafe {
23456 b.launch(cfg)?;
23457 }
23458 Ok(())
23459 }
23460
23461 #[allow(clippy::too_many_arguments)]
23465 pub fn ssm_conv1d_gdn(
23466 &self,
23467 qkv_tm: &CudaSlice<f32>,
23468 w: &CudaSlice<f32>,
23469 q_g: &mut CudaSlice<f32>,
23470 k_g: &mut CudaSlice<f32>,
23471 v_g: &mut CudaSlice<f32>,
23472 conv_dim: usize,
23473 t: usize,
23474 d_conv: usize,
23475 d_state: usize,
23476 num_v: usize,
23477 num_k: usize,
23478 key_dim: usize,
23479 ) -> Result<(), Box<dyn std::error::Error>> {
23480 let f = self.func("ssm_conv1d_gdn_f32");
23481 let cfg = LaunchConfig {
23482 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
23483 block_dim: (256, 1, 1),
23484 shared_mem_bytes: 0,
23485 };
23486 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
23487 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
23488 let __s_b = self.gpu.stream();
23489 let mut b = __s_b.launch_builder(&f);
23490 b.arg(qkv_tm)
23491 .arg(w)
23492 .arg(q_g)
23493 .arg(k_g)
23494 .arg(v_g)
23495 .arg(&cd)
23496 .arg(&ti)
23497 .arg(&dc)
23498 .arg(&ds)
23499 .arg(&nv)
23500 .arg(&nk)
23501 .arg(&kd);
23502 unsafe {
23503 b.launch(cfg)?;
23504 }
23505 Ok(())
23506 }
23507
23508 pub fn ssm_conv1d(
23509 &self,
23510 x: &CudaSlice<f32>,
23511 w: &CudaSlice<f32>,
23512 y: &mut CudaSlice<f32>,
23513 conv_dim: usize,
23514 t: usize,
23515 d_conv: usize,
23516 silu: bool,
23517 ) -> Result<(), Box<dyn std::error::Error>> {
23518 let f = self.func("ssm_conv1d_silu_f32");
23519 let cfg = LaunchConfig {
23520 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
23521 block_dim: (256, 1, 1),
23522 shared_mem_bytes: 0,
23523 };
23524 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
23525 let __s_b = self.gpu.stream();
23526 let mut b = __s_b.launch_builder(&f);
23527 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
23528 unsafe {
23529 b.launch(cfg)?;
23530 }
23531 Ok(())
23532 }
23533
23534 pub fn gdn_scan_s128(
23537 &self,
23538 q: &CudaSlice<f32>,
23539 k: &CudaSlice<f32>,
23540 v: &CudaSlice<f32>,
23541 g: &CudaSlice<f32>,
23542 beta: &CudaSlice<f32>,
23543 state_in: &CudaSlice<f32>,
23544 state_out: &mut CudaSlice<f32>,
23545 o: &mut CudaSlice<f32>,
23546 n_head: usize,
23547 t: usize,
23548 scale: f32,
23549 ) -> Result<(), Box<dyn std::error::Error>> {
23550 let f = self.func("gdn_scan_s128");
23551 const S_V: u32 = 128;
23552 const WARP: u32 = 32;
23553 const COLS_PER_BLOCK: u32 = 4;
23554 let cfg = LaunchConfig {
23555 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
23556 block_dim: (WARP, COLS_PER_BLOCK, 1),
23557 shared_mem_bytes: 0,
23558 };
23559 let (h, ti) = (n_head as i32, t as i32);
23560 let __s_b = self.gpu.stream();
23561 let mut b = __s_b.launch_builder(&f);
23562 b.arg(q)
23563 .arg(k)
23564 .arg(v)
23565 .arg(g)
23566 .arg(beta)
23567 .arg(state_in)
23568 .arg(state_out)
23569 .arg(o)
23570 .arg(&h)
23571 .arg(&ti)
23572 .arg(&scale);
23573 unsafe {
23574 b.launch(cfg)?;
23575 }
23576 Ok(())
23577 }
23578
23579 #[allow(clippy::too_many_arguments)]
23584 pub fn ssm_conv1d_fused_decode_b(
23585 &self,
23586 qkv_cols: &CudaSlice<f32>,
23587 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
23588 w: &CudaSlice<f32>,
23589 conv_outs: &mut CudaSlice<f32>,
23590 conv_dim: usize,
23591 d_conv: usize,
23592 b_n: usize,
23593 ) -> Result<(), Box<dyn std::error::Error>> {
23594 let f = self.func("ssm_conv1d_fused_decode_b_f32");
23595 let cfg = LaunchConfig {
23596 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
23597 block_dim: (256, 1, 1),
23598 shared_mem_bytes: 0,
23599 };
23600 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23601 let __s_b = self.gpu.stream();
23602 let mut b = __s_b.launch_builder(&f);
23603 b.arg(qkv_cols)
23604 .arg(conv_state_ptrs)
23605 .arg(w)
23606 .arg(conv_outs)
23607 .arg(&cd)
23608 .arg(&dc);
23609 unsafe {
23610 b.launch(cfg)?;
23611 }
23612 Ok(())
23613 }
23614
23615 #[allow(clippy::too_many_arguments)]
23616 pub fn gdn_prep_decode_b(
23617 &self,
23618 conv_outs: &CudaSlice<f32>,
23619 beta_raws: &CudaSlice<f32>,
23620 alphas: &CudaSlice<f32>,
23621 dt_bias: &CudaSlice<f32>,
23622 a: &CudaSlice<f32>,
23623 q_l2: &mut CudaSlice<f32>,
23624 k_l2: &mut CudaSlice<f32>,
23625 v_g: &mut CudaSlice<f32>,
23626 beta: &mut CudaSlice<f32>,
23627 g_log: &mut CudaSlice<f32>,
23628 d_state: usize,
23629 num_v: usize,
23630 num_k: usize,
23631 key_dim: usize,
23632 eps: f32,
23633 conv_dim: usize,
23634 b_n: usize,
23635 ) -> Result<(), Box<dyn std::error::Error>> {
23636 let f = self.func("gdn_prep_decode_b_f32");
23637 let cfg = LaunchConfig {
23638 grid_dim: (num_v as u32, 1, b_n as u32),
23639 block_dim: (32, 4, 1),
23640 shared_mem_bytes: 0,
23641 };
23642 let (ds, nv, nk, kd, cd) = (
23643 d_state as i32,
23644 num_v as i32,
23645 num_k as i32,
23646 key_dim as i32,
23647 conv_dim as i32,
23648 );
23649 let __s_b = self.gpu.stream();
23650 let mut b = __s_b.launch_builder(&f);
23651 b.arg(conv_outs)
23652 .arg(beta_raws)
23653 .arg(alphas)
23654 .arg(dt_bias)
23655 .arg(a)
23656 .arg(q_l2)
23657 .arg(k_l2)
23658 .arg(v_g)
23659 .arg(beta)
23660 .arg(g_log)
23661 .arg(&ds)
23662 .arg(&nv)
23663 .arg(&nk)
23664 .arg(&kd)
23665 .arg(&eps)
23666 .arg(&cd);
23667 unsafe {
23668 b.launch(cfg)?;
23669 }
23670 Ok(())
23671 }
23672
23673 #[allow(clippy::too_many_arguments)]
23674 pub fn gdn_scan_s128_batched(
23675 &self,
23676 q: &CudaSlice<f32>,
23677 k: &CudaSlice<f32>,
23678 v: &CudaSlice<f32>,
23679 g: &CudaSlice<f32>,
23680 beta: &CudaSlice<f32>,
23681 state_in_ptrs: &cudarc::driver::CudaView<u64>,
23682 state_out_ptrs: &cudarc::driver::CudaView<u64>,
23683 o: &mut CudaSlice<f32>,
23684 n_head: usize,
23685 b_n: usize,
23686 scale: f32,
23687 ) -> Result<(), Box<dyn std::error::Error>> {
23688 let f = self.func("gdn_scan_s128_b");
23689 const S_V: u32 = 128;
23690 const WARP: u32 = 32;
23691 const COLS_PER_BLOCK: u32 = 4;
23692 let cfg = LaunchConfig {
23693 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
23694 block_dim: (WARP, COLS_PER_BLOCK, 1),
23695 shared_mem_bytes: 0,
23696 };
23697 let h = n_head as i32;
23698 let __s_b = self.gpu.stream();
23699 let mut b = __s_b.launch_builder(&f);
23700 b.arg(q)
23701 .arg(k)
23702 .arg(v)
23703 .arg(g)
23704 .arg(beta)
23705 .arg(state_in_ptrs)
23706 .arg(state_out_ptrs)
23707 .arg(o)
23708 .arg(&h)
23709 .arg(&scale);
23710 unsafe {
23711 b.launch(cfg)?;
23712 }
23713 Ok(())
23714 }
23715
23716 #[allow(clippy::too_many_arguments)]
23722 pub fn ssm_conv1d_fused_decode_b_view(
23723 &self,
23724 qkv_cols: &cudarc::driver::CudaView<f32>,
23725 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
23726 w: &CudaSlice<f32>,
23727 conv_outs: &mut CudaSlice<f32>,
23728 conv_dim: usize,
23729 d_conv: usize,
23730 b_n: usize,
23731 ) -> Result<(), Box<dyn std::error::Error>> {
23732 let f = self.func("ssm_conv1d_fused_decode_b_f32");
23733 let cfg = LaunchConfig {
23734 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
23735 block_dim: (256, 1, 1),
23736 shared_mem_bytes: 0,
23737 };
23738 let (cd, dc) = (conv_dim as i32, d_conv as i32);
23739 let __s_b = self.gpu.stream();
23740 let mut b = __s_b.launch_builder(&f);
23741 b.arg(qkv_cols)
23742 .arg(conv_state_ptrs)
23743 .arg(w)
23744 .arg(conv_outs)
23745 .arg(&cd)
23746 .arg(&dc);
23747 unsafe {
23748 b.launch(cfg)?;
23749 }
23750 Ok(())
23751 }
23752
23753 #[allow(clippy::too_many_arguments)]
23754 pub fn gdn_prep_decode_b_view(
23755 &self,
23756 conv_outs: &CudaSlice<f32>,
23757 beta_raws: &cudarc::driver::CudaView<f32>,
23758 alphas: &cudarc::driver::CudaView<f32>,
23759 dt_bias: &CudaSlice<f32>,
23760 a: &CudaSlice<f32>,
23761 q_l2: &mut CudaSlice<f32>,
23762 k_l2: &mut CudaSlice<f32>,
23763 v_g: &mut CudaSlice<f32>,
23764 beta: &mut CudaSlice<f32>,
23765 g_log: &mut CudaSlice<f32>,
23766 d_state: usize,
23767 num_v: usize,
23768 num_k: usize,
23769 key_dim: usize,
23770 eps: f32,
23771 conv_dim: usize,
23772 b_n: usize,
23773 ) -> Result<(), Box<dyn std::error::Error>> {
23774 let f = self.func("gdn_prep_decode_b_f32");
23775 let cfg = LaunchConfig {
23776 grid_dim: (num_v as u32, 1, b_n as u32),
23777 block_dim: (32, 4, 1),
23778 shared_mem_bytes: 0,
23779 };
23780 let (ds, nv, nk, kd, cd) = (
23781 d_state as i32,
23782 num_v as i32,
23783 num_k as i32,
23784 key_dim as i32,
23785 conv_dim as i32,
23786 );
23787 let __s_b = self.gpu.stream();
23788 let mut b = __s_b.launch_builder(&f);
23789 b.arg(conv_outs)
23790 .arg(beta_raws)
23791 .arg(alphas)
23792 .arg(dt_bias)
23793 .arg(a)
23794 .arg(q_l2)
23795 .arg(k_l2)
23796 .arg(v_g)
23797 .arg(beta)
23798 .arg(g_log)
23799 .arg(&ds)
23800 .arg(&nv)
23801 .arg(&nk)
23802 .arg(&kd)
23803 .arg(&eps)
23804 .arg(&cd);
23805 unsafe {
23806 b.launch(cfg)?;
23807 }
23808 Ok(())
23809 }
23810
23811 #[allow(clippy::too_many_arguments)]
23812 pub fn gdn_scan_s128_batched_view(
23813 &self,
23814 q: &CudaSlice<f32>,
23815 k: &CudaSlice<f32>,
23816 v: &CudaSlice<f32>,
23817 g: &CudaSlice<f32>,
23818 beta: &CudaSlice<f32>,
23819 state_in_ptrs: &cudarc::driver::CudaView<u64>,
23820 state_out_ptrs: &cudarc::driver::CudaView<u64>,
23821 o: &mut cudarc::driver::CudaViewMut<f32>,
23822 n_head: usize,
23823 b_n: usize,
23824 scale: f32,
23825 ) -> Result<(), Box<dyn std::error::Error>> {
23826 let f = self.func("gdn_scan_s128_b");
23827 const S_V: u32 = 128;
23828 const WARP: u32 = 32;
23829 const COLS_PER_BLOCK: u32 = 4;
23830 let cfg = LaunchConfig {
23831 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
23832 block_dim: (WARP, COLS_PER_BLOCK, 1),
23833 shared_mem_bytes: 0,
23834 };
23835 let h = n_head as i32;
23836 let __s_b = self.gpu.stream();
23837 let mut b = __s_b.launch_builder(&f);
23838 b.arg(q)
23839 .arg(k)
23840 .arg(v)
23841 .arg(g)
23842 .arg(beta)
23843 .arg(state_in_ptrs)
23844 .arg(state_out_ptrs)
23845 .arg(o)
23846 .arg(&h)
23847 .arg(&scale);
23848 unsafe {
23849 b.launch(cfg)?;
23850 }
23851 Ok(())
23852 }
23853
23854 pub fn gdn_chunked_enabled() -> bool {
23863 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23864 *E.get_or_init(|| {
23865 std::env::var("MEMRA_GDN_CHUNKED")
23866 .map(|v| v != "0")
23867 .unwrap_or(true)
23868 })
23869 }
23870
23871 pub fn gdn_chunk_size() -> usize {
23876 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23877 *C.get_or_init(|| {
23878 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
23879 .ok()
23880 .and_then(|v| v.parse().ok())
23881 .unwrap_or(32);
23882 c.clamp(32, 128) / 32 * 32
23883 })
23884 }
23885
23886 #[allow(clippy::too_many_arguments)]
23891 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
23894 #[allow(clippy::too_many_arguments)]
23895 pub fn gdn_chunk_k123(
23896 &self,
23897 q: &CudaSlice<f32>,
23898 k: &CudaSlice<f32>,
23899 v: &CudaSlice<f32>,
23900 g: &CudaSlice<f32>,
23901 beta: &CudaSlice<f32>,
23902 wb16: Option<&mut CudaSlice<u8>>,
23903 n_head: usize,
23904 t: usize,
23905 c: usize,
23906 hk: usize,
23907 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
23908 ) -> Result<
23909 (
23910 CudaSlice<f32>,
23911 CudaSlice<f32>,
23912 CudaSlice<f32>,
23913 CudaSlice<f32>,
23914 ),
23915 Box<dyn std::error::Error>,
23916 > {
23917 const D: usize = 128;
23918 let h = n_head;
23919 let nc = (t + c - 1) / c;
23920 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
23921 let mut gcum = self.uninit(t * h)?;
23922 let mut a = self.uninit(nc * h * c * c)?;
23923 let mut p = self.uninit(nc * h * c * c)?;
23924 let mut u = self.uninit(nc * h * c * D)?;
23925 let mut w = self.uninit(nc * h * c * D)?;
23926 {
23927 let f = self.func("gdn_chunk_cumgate_f32");
23929 let cfg = LaunchConfig {
23930 grid_dim: (nc as u32, h as u32, 1),
23931 block_dim: (32, 1, 1),
23932 shared_mem_bytes: 0,
23933 };
23934 let __s_b = self.gpu.stream();
23935 let mut b = __s_b.launch_builder(&f);
23936 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
23937 unsafe {
23938 b.launch(cfg)?;
23939 }
23940 }
23941 if let Some((qb, kb, pb)) = k2w {
23942 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
23945 let f = self.func("gdn_k2_wgmma");
23946 let cfg = LaunchConfig {
23947 grid_dim: (nc as u32, h as u32, 1),
23948 block_dim: (128, 1, 1),
23949 shared_mem_bytes: 0,
23950 };
23951 let hki = hk as i32;
23952 let __s_b = self.gpu.stream();
23953 let mut b = __s_b.launch_builder(&f);
23954 b.arg(qb)
23955 .arg(kb)
23956 .arg(&gcum)
23957 .arg(beta)
23958 .arg(&mut a)
23959 .arg(&mut *pb)
23960 .arg(&hi)
23961 .arg(&ti)
23962 .arg(&ci)
23963 .arg(&hki);
23964 unsafe {
23965 b.launch(cfg)?;
23966 }
23967 } else if c <= 64 && !portable_mma_gated() {
23968 let f = self.func("gdn_chunk_attn_f32");
23970 f.set_attribute(
23971 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23972 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
23973 )?;
23974 let jt = ((c + 31) / 32) as u32;
23975 let cfg = LaunchConfig {
23976 grid_dim: (nc as u32, h as u32, jt),
23977 block_dim: (256, 1, 1),
23978 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
23979 };
23980 let hki = hk as i32;
23981 let __s_b = self.gpu.stream();
23982 let mut b = __s_b.launch_builder(&f);
23983 b.arg(q)
23984 .arg(k)
23985 .arg(&gcum)
23986 .arg(beta)
23987 .arg(&mut a)
23988 .arg(&mut p)
23989 .arg(&hi)
23990 .arg(&ti)
23991 .arg(&ci)
23992 .arg(&hki);
23993 unsafe {
23994 b.launch(cfg)?;
23995 }
23996 } else {
23997 assert!(
23999 hk == h,
24000 "generic K2 is broadcast-only (de-broadcast rides C==32)"
24001 );
24002 let f = self.func("gdn_chunk_attn_g_f32");
24003 let cfg = LaunchConfig {
24004 grid_dim: (nc as u32, h as u32, 1),
24005 block_dim: (32, 8, 1),
24006 shared_mem_bytes: 0,
24007 };
24008 let __s_b = self.gpu.stream();
24009 let mut b = __s_b.launch_builder(&f);
24010 b.arg(q)
24011 .arg(k)
24012 .arg(&gcum)
24013 .arg(beta)
24014 .arg(&mut a)
24015 .arg(&mut p)
24016 .arg(&hi)
24017 .arg(&ti)
24018 .arg(&ci);
24019 unsafe {
24020 b.launch(cfg)?;
24021 }
24022 }
24023 {
24024 let cfg = LaunchConfig {
24026 grid_dim: (nc as u32, h as u32, 1),
24027 block_dim: (256, 1, 1),
24028 shared_mem_bytes: 0,
24029 };
24030 match c {
24031 32 | 64 => {
24032 let f = self.func(if c == 32 {
24033 "gdn_chunk_solve32_f32"
24034 } else {
24035 "gdn_chunk_solve64_f32"
24036 });
24037 let wb: u64 = match wb16 {
24039 Some(d) => self.addr_u8(d),
24040 None => 0,
24041 };
24042 let hki = hk as i32;
24043 let __s_b = self.gpu.stream();
24044 let mut b = __s_b.launch_builder(&f);
24045 b.arg(v)
24046 .arg(k)
24047 .arg(&a)
24048 .arg(&gcum)
24049 .arg(&mut u)
24050 .arg(&mut w)
24051 .arg(&wb)
24052 .arg(&hi)
24053 .arg(&ti)
24054 .arg(&hki);
24055 unsafe {
24056 b.launch(cfg)?;
24057 }
24058 }
24059 _ => {
24060 assert!(hk == h, "generic K3 is broadcast-only");
24061 let f = self.func("gdn_chunk_solve_f32");
24062 let __s_b = self.gpu.stream();
24063 let mut b = __s_b.launch_builder(&f);
24064 b.arg(v)
24065 .arg(k)
24066 .arg(&a)
24067 .arg(&gcum)
24068 .arg(&mut u)
24069 .arg(&mut w)
24070 .arg(&hi)
24071 .arg(&ti)
24072 .arg(&ci);
24073 unsafe {
24074 b.launch(cfg)?;
24075 }
24076 }
24077 }
24078 }
24079 Ok((gcum, p, u, w))
24080 }
24081
24082 pub fn gdn_db_on() -> bool {
24086 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
24087 }
24088
24089 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
24098 !portable_mma_gated()
24099 && c == 32
24100 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24101 Ok("1") => true,
24102 Ok("0") => false,
24103 _ => gdn_mma_default_on(),
24104 }
24105 }
24106
24107 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
24114 cfg!(memra_hopper_mma)
24115 && self.gdn_mma_enabled(c)
24116 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
24117 }
24118
24119 #[allow(clippy::too_many_arguments)]
24124 pub fn ssm_conv1d_gdn_state_pad(
24125 &self,
24126 qkv_tm: &cudarc::driver::CudaView<f32>,
24127 conv_state: &mut CudaSlice<f32>,
24128 w: &CudaSlice<f32>,
24129 q_g: &mut CudaSlice<f32>,
24130 k_g: &mut CudaSlice<f32>,
24131 v_g: &mut CudaSlice<f32>,
24132 conv_dim: usize,
24133 t: usize,
24134 d_conv: usize,
24135 d_state: usize,
24136 num_v: usize,
24137 num_k: usize,
24138 key_dim: usize,
24139 hk: usize,
24140 pad_len: Option<&CudaSlice<i32>>,
24141 ) -> Result<(), Box<dyn std::error::Error>> {
24142 assert!(
24143 t >= d_conv - 1,
24144 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
24145 );
24146 {
24147 let f = self.func("ssm_conv1d_gdn_state_f32");
24148 let cfg = LaunchConfig {
24149 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
24150 block_dim: (256, 1, 1),
24151 shared_mem_bytes: 0,
24152 };
24153 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24154 let (ds, nv, nk, kd, hki) = (
24155 d_state as i32,
24156 num_v as i32,
24157 num_k as i32,
24158 key_dim as i32,
24159 hk as i32,
24160 );
24161 let __s_b = self.gpu.stream();
24162 let mut b = __s_b.launch_builder(&f);
24163 b.arg(qkv_tm)
24164 .arg(&*conv_state)
24165 .arg(w)
24166 .arg(q_g)
24167 .arg(k_g)
24168 .arg(v_g)
24169 .arg(&cd)
24170 .arg(&ti)
24171 .arg(&dc)
24172 .arg(&ds)
24173 .arg(&nv)
24174 .arg(&nk)
24175 .arg(&kd)
24176 .arg(&hki);
24177 unsafe {
24178 b.launch(cfg)?;
24179 }
24180 }
24181 match pad_len {
24182 Some(len_d) => {
24183 let f = self.func("ssm_conv_ring_update_dev_f32");
24184 let n = conv_dim * (d_conv - 1);
24185 let cfg = LaunchConfig::for_num_elems(n as u32);
24186 let (cd, dc) = (conv_dim as i32, d_conv as i32);
24187 let __s_b = self.gpu.stream();
24188 let mut b = __s_b.launch_builder(&f);
24189 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
24190 unsafe {
24191 b.launch(cfg)?;
24192 }
24193 }
24194 None => {
24195 let f = self.func("ssm_conv_ring_update_f32");
24196 let n = conv_dim * (d_conv - 1);
24197 let cfg = LaunchConfig::for_num_elems(n as u32);
24198 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
24199 let __s_b = self.gpu.stream();
24200 let mut b = __s_b.launch_builder(&f);
24201 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
24202 unsafe {
24203 b.launch(cfg)?;
24204 }
24205 }
24206 }
24207 Ok(())
24208 }
24209
24210 pub fn gdn_chunk_alloc(
24214 &self,
24215 n_head: usize,
24216 t: usize,
24217 c: usize,
24218 hk: usize,
24219 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
24220 const D: usize = 128;
24221 assert!(
24222 c == 32,
24223 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
24224 );
24225 let h = n_head;
24226 let nc = (t + c - 1) / c;
24227 Ok(GdnChunkBufs {
24228 gcum: self.uninit(t * h)?,
24229 a: self.uninit(nc * h * c * c)?,
24230 p: self.uninit(nc * h * c * c)?,
24231 u: self.uninit(nc * h * c * D)?,
24232 w: self.uninit(nc * h * c * D)?,
24233 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24234 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24235 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
24236 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
24237 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
24238 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
24239 o: self.uninit(D * h * t)?,
24240 t,
24241 nc,
24242 })
24243 }
24244
24245 pub fn f32_to_bf16_v(
24247 &self,
24248 x: &cudarc::driver::CudaView<f32>,
24249 dst: &mut CudaSlice<u8>,
24250 n: usize,
24251 ) -> Result<(), Box<dyn std::error::Error>> {
24252 let f = self.func("f32_to_bf16_bulk");
24253 let ni = n as i64;
24254 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24255 let __s_b = self.gpu.stream();
24256 let mut b = __s_b.launch_builder(&f);
24257 b.arg(x).arg(dst).arg(&ni);
24258 unsafe {
24259 b.launch(cfg)?;
24260 }
24261 Ok(())
24262 }
24263
24264 pub fn f32_to_bf16_into(
24266 &self,
24267 x: &CudaSlice<f32>,
24268 dst: &mut CudaSlice<u8>,
24269 n: usize,
24270 ) -> Result<(), Box<dyn std::error::Error>> {
24271 let f = self.func("f32_to_bf16_bulk");
24272 let ni = n as i64;
24273 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
24274 let __s_b = self.gpu.stream();
24275 let mut b = __s_b.launch_builder(&f);
24276 b.arg(x).arg(dst).arg(&ni);
24277 unsafe {
24278 b.launch(cfg)?;
24279 }
24280 Ok(())
24281 }
24282
24283 pub fn gdn_chunk_k123_vl8(
24286 &self,
24287 seqs: &[GdnSeqVl],
24288 n_head: usize,
24289 hk: usize,
24290 wq: Option<&GdnWVl8>,
24291 ) -> Result<(), Box<dyn std::error::Error>> {
24292 let b = seqs.len();
24293 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
24294 let mut packed = [GdnSeqVl::default(); 8];
24295 packed[..b].copy_from_slice(seqs);
24296 let v = GdnVl8(packed);
24297 let (hi, ci) = (n_head as i32, 32i32);
24298 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24299 {
24300 let f = self.func("gdn_chunk_cumgate_vl");
24301 let cfg = LaunchConfig {
24302 grid_dim: (max_nc, n_head as u32, b as u32),
24303 block_dim: (32, 1, 1),
24304 shared_mem_bytes: 0,
24305 };
24306 let __s_lb = self.gpu.stream();
24307 let mut lb = __s_lb.launch_builder(&f);
24308 lb.arg(&v).arg(&hi).arg(&ci);
24309 unsafe {
24310 lb.launch(cfg)?;
24311 }
24312 }
24313 let hki = hk as i32;
24314 if let Some(w) = wq {
24315 let f = self.func("gdn_k2_wgmma_vl");
24317 let cfg = LaunchConfig {
24318 grid_dim: (max_nc, n_head as u32, b as u32),
24319 block_dim: (128, 1, 1),
24320 shared_mem_bytes: 0,
24321 };
24322 let __s_lb = self.gpu.stream();
24323 let mut lb = __s_lb.launch_builder(&f);
24324 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
24325 unsafe {
24326 lb.launch(cfg)?;
24327 }
24328 } else {
24329 let f = self.func("gdn_chunk_attn_vl");
24330 f.set_attribute(
24331 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24332 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
24333 )?;
24334 let cfg = LaunchConfig {
24335 grid_dim: (max_nc, n_head as u32, b as u32),
24336 block_dim: (256, 1, 1),
24337 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
24338 };
24339 let __s_lb = self.gpu.stream();
24340 let mut lb = __s_lb.launch_builder(&f);
24341 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24342 unsafe {
24343 lb.launch(cfg)?;
24344 }
24345 }
24346 {
24347 let f = self.func("gdn_chunk_solve32_vl");
24348 let cfg = LaunchConfig {
24349 grid_dim: (max_nc, n_head as u32, b as u32),
24350 block_dim: (256, 1, 1),
24351 shared_mem_bytes: 0,
24352 };
24353 let __s_lb = self.gpu.stream();
24354 let mut lb = __s_lb.launch_builder(&f);
24355 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24356 unsafe {
24357 lb.launch(cfg)?;
24358 }
24359 }
24360 Ok(())
24361 }
24362
24363 #[allow(clippy::too_many_arguments)]
24367 pub fn gdn_prep_vl8(
24368 &self,
24369 seqs: &[GdnPrepVl],
24370 conv_w: &CudaSlice<f32>,
24371 dt_bias: &CudaSlice<f32>,
24372 a: &CudaSlice<f32>,
24373 conv_dim: usize,
24374 d_conv: usize,
24375 d_state: usize,
24376 num_v: usize,
24377 num_k: usize,
24378 key_dim: usize,
24379 hk: usize,
24380 eps: f32,
24381 ) -> Result<(), Box<dyn std::error::Error>> {
24382 let b = seqs.len();
24383 assert!(b >= 1 && b <= 8);
24384 let mut packed = [GdnPrepVl::default(); 8];
24385 packed[..b].copy_from_slice(seqs);
24386 let v = GdnPrepVl8(packed);
24387 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24388 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
24389 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
24390 assert!(
24391 conv_fuse || hk == num_v,
24392 "de-broadcast requires the fused conv"
24393 );
24394 if conv_fuse {
24395 let f = self.func("ssm_conv1d_gdn_state_vl");
24396 let cfg = LaunchConfig {
24397 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24398 block_dim: (256, 1, 1),
24399 shared_mem_bytes: 0,
24400 };
24401 let (dsi, nvi, nki, kdi, hki) = (
24402 d_state as i32,
24403 num_v as i32,
24404 num_k as i32,
24405 key_dim as i32,
24406 hk as i32,
24407 );
24408 let __s_lb = self.gpu.stream();
24409 let mut lb = __s_lb.launch_builder(&f);
24410 lb.arg(&v)
24411 .arg(conv_w)
24412 .arg(&cdi)
24413 .arg(&dci)
24414 .arg(&dsi)
24415 .arg(&nvi)
24416 .arg(&nki)
24417 .arg(&kdi)
24418 .arg(&hki);
24419 unsafe {
24420 lb.launch(cfg)?;
24421 }
24422 } else {
24423 let f = self.func("ssm_conv1d_tm_state_vl");
24424 let cfg = LaunchConfig {
24425 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
24426 block_dim: (256, 1, 1),
24427 shared_mem_bytes: 0,
24428 };
24429 let __s_lb = self.gpu.stream();
24430 let mut lb = __s_lb.launch_builder(&f);
24431 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
24432 unsafe {
24433 lb.launch(cfg)?;
24434 }
24435 }
24436 {
24437 let f = self.func("ssm_conv_ring_update_vl");
24438 let n = (conv_dim * (d_conv - 1)) as u32;
24439 let cfg = LaunchConfig {
24440 grid_dim: (n.div_ceil(256), 1, b as u32),
24441 block_dim: (256, 1, 1),
24442 shared_mem_bytes: 0,
24443 };
24444 let __s_lb = self.gpu.stream();
24445 let mut lb = __s_lb.launch_builder(&f);
24446 lb.arg(&v).arg(&cdi).arg(&dci);
24447 unsafe {
24448 lb.launch(cfg)?;
24449 }
24450 }
24451 if !conv_fuse {
24452 let f = self.func("qkv_to_gdn_repack_vl");
24453 let n = max_t * (num_v * d_state) as u32;
24454 let cfg = LaunchConfig {
24455 grid_dim: (n.div_ceil(256), 1, b as u32),
24456 block_dim: (256, 1, 1),
24457 shared_mem_bytes: 0,
24458 };
24459 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
24460 let __s_lb = self.gpu.stream();
24461 let mut lb = __s_lb.launch_builder(&f);
24462 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
24463 unsafe {
24464 lb.launch(cfg)?;
24465 }
24466 }
24467 if Self::l2_v2_on(d_state) {
24468 let f = self.func("gdn_l2_v2_vl");
24469 let cfg = LaunchConfig {
24470 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
24471 block_dim: (256, 1, 1),
24472 shared_mem_bytes: 0,
24473 };
24474 let (dsi, nvi) = (d_state as i32, hk as i32);
24475 let __s_lb = self.gpu.stream();
24476 let mut lb = __s_lb.launch_builder(&f);
24477 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24478 unsafe {
24479 lb.launch(cfg)?;
24480 }
24481 } else {
24482 let f = self.func("gdn_l2_vl");
24483 let cfg = LaunchConfig {
24484 grid_dim: (max_t * hk as u32, 2, b as u32),
24485 block_dim: (256, 1, 1),
24486 shared_mem_bytes: 0,
24487 };
24488 let (dsi, nvi) = (d_state as i32, hk as i32);
24489 let __s_lb = self.gpu.stream();
24490 let mut lb = __s_lb.launch_builder(&f);
24491 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
24492 unsafe {
24493 lb.launch(cfg)?;
24494 }
24495 }
24496 {
24497 let f = self.func("gdn_gate_prep_vl");
24498 let n = max_t * num_v as u32;
24499 let cfg = LaunchConfig {
24500 grid_dim: (n.div_ceil(256), 1, b as u32),
24501 block_dim: (256, 1, 1),
24502 shared_mem_bytes: 0,
24503 };
24504 let nvi = num_v as i32;
24505 let __s_lb = self.gpu.stream();
24506 let mut lb = __s_lb.launch_builder(&f);
24507 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
24508 unsafe {
24509 lb.launch(cfg)?;
24510 }
24511 }
24512 Ok(())
24513 }
24514
24515 pub fn gdn_mirror_vl8(
24517 &self,
24518 seqs: &[GdnSeqVl],
24519 n_head: usize,
24520 which: i32,
24521 hk: usize,
24522 ) -> Result<(), Box<dyn std::error::Error>> {
24523 let b = seqs.len();
24524 assert!(b >= 1 && b <= 8);
24525 let mut packed = [GdnSeqVl::default(); 8];
24526 packed[..b].copy_from_slice(seqs);
24527 let v = GdnVl8(packed);
24528 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
24529 let max_n = seqs
24530 .iter()
24531 .map(|s| {
24532 if which == 0 {
24533 s.t as i64 * ept as i64
24534 } else {
24535 s.nc as i64 * ept as i64 * 32
24536 }
24537 })
24538 .max()
24539 .unwrap();
24540 let f = self.func("gdn_mirror_vl");
24541 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
24542 let cfg = LaunchConfig {
24543 grid_dim: (blocks, 1, b as u32),
24544 block_dim: (256, 1, 1),
24545 shared_mem_bytes: 0,
24546 };
24547 let __s_lb = self.gpu.stream();
24548 let mut lb = __s_lb.launch_builder(&f);
24549 lb.arg(&v).arg(&ept).arg(&which);
24550 unsafe {
24551 lb.launch(cfg)?;
24552 }
24553 Ok(())
24554 }
24555
24556 pub fn gdn_tail_vl8(
24558 &self,
24559 seqs: &[GdnPrepVl],
24560 norm_w: &CudaSlice<f32>,
24561 d_state: usize,
24562 num_v: usize,
24563 eps: f32,
24564 ) -> Result<(), Box<dyn std::error::Error>> {
24565 let b = seqs.len();
24566 assert!(b >= 1 && b <= 8);
24567 let mut packed = [GdnPrepVl::default(); 8];
24568 packed[..b].copy_from_slice(seqs);
24569 let v = GdnPrepVl8(packed);
24570 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
24571 let f = self.func("gated_rmsnorm_f16out_vl");
24572 let cfg = LaunchConfig {
24574 grid_dim: (max_t * num_v as u32, 1, b as u32),
24575 block_dim: (128, 1, 1),
24576 shared_mem_bytes: 0,
24577 };
24578 let (dsi, nvi) = (d_state as i32, num_v as i32);
24579 let __s_lb = self.gpu.stream();
24580 let mut lb = __s_lb.launch_builder(&f);
24581 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
24582 unsafe {
24583 lb.launch(cfg)?;
24584 }
24585 Ok(())
24586 }
24587
24588 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
24591 use cudarc::driver::DevicePtr;
24592 let s = self.gpu.stream();
24593 let (p, _g) = x.device_ptr(&s);
24594 p as u64
24595 }
24596 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
24597 use cudarc::driver::DevicePtrMut;
24598 let s = self.gpu.stream();
24599 let (p, _g) = x.device_ptr_mut(&s);
24600 p as u64
24601 }
24602 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
24603 use cudarc::driver::DevicePtr;
24604 let s = self.gpu.stream();
24605 let (p, _g) = x.device_ptr(&s);
24606 p as u64
24607 }
24608 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
24609 use cudarc::driver::DevicePtr;
24610 let s = self.gpu.stream();
24611 let (p, _g) = x.device_ptr(&s);
24612 p as u64
24613 }
24614
24615 pub fn gdn_chunk_vl8(
24619 &self,
24620 seqs: &[GdnSeqVl],
24621 n_head: usize,
24622 scale: f32,
24623 hk: usize,
24624 wq: Option<&GdnWVl8>,
24625 ) -> Result<(), Box<dyn std::error::Error>> {
24626 const NSPLIT: u32 = 4;
24627 let b = seqs.len();
24628 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
24629 let mut packed = [GdnSeqVl::default(); 8];
24630 packed[..b].copy_from_slice(seqs);
24631 let v = GdnVl8(packed);
24632 let (hi, ci) = (n_head as i32, 32i32);
24633 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
24634 let hki = hk as i32;
24635 if let Some(w) = wq {
24636 let f = self.func("gdn_k45_wgmma_vl");
24638 let cfg = LaunchConfig {
24639 grid_dim: (n_head as u32, NSPLIT, b as u32),
24640 block_dim: (256, 1, 1),
24641 shared_mem_bytes: 0,
24642 };
24643 let __s_lb = self.gpu.stream();
24644 let mut lb = __s_lb.launch_builder(&f);
24645 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
24646 unsafe {
24647 lb.launch(cfg)?;
24648 }
24649 let _ = max_nc;
24650 return Ok(());
24651 }
24652 {
24653 let f = self.func("gdn_chunk_state_mma_vl");
24654 let cfg = LaunchConfig {
24655 grid_dim: (n_head as u32, NSPLIT, b as u32),
24656 block_dim: (256, 1, 1),
24657 shared_mem_bytes: 0,
24658 };
24659 let __s_lb = self.gpu.stream();
24660 let mut lb = __s_lb.launch_builder(&f);
24661 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
24662 unsafe {
24663 lb.launch(cfg)?;
24664 }
24665 }
24666 {
24667 let f = self.func("gdn_chunk_output_mma_vl");
24668 let cfg = LaunchConfig {
24669 grid_dim: (max_nc, n_head as u32, b as u32),
24670 block_dim: (256, 1, 1),
24671 shared_mem_bytes: 0,
24672 };
24673 let __s_lb = self.gpu.stream();
24674 let mut lb = __s_lb.launch_builder(&f);
24675 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
24676 unsafe {
24677 lb.launch(cfg)?;
24678 }
24679 }
24680 Ok(())
24681 }
24682 pub fn gdn_scan_chunked(
24683 &self,
24684 q: &CudaSlice<f32>,
24685 k: &CudaSlice<f32>,
24686 v: &CudaSlice<f32>,
24687 g: &CudaSlice<f32>,
24688 beta: &CudaSlice<f32>,
24689 kb16_pre: Option<&CudaSlice<u8>>,
24690 qb16_pre: Option<&CudaSlice<u8>>,
24691 state_in: &CudaSlice<f32>,
24692 state_out: &mut CudaSlice<f32>,
24693 o: &mut CudaSlice<f32>,
24694 n_head: usize,
24695 t: usize,
24696 scale: f32,
24697 c: usize,
24698 hk: usize,
24699 ) -> Result<(), Box<dyn std::error::Error>> {
24700 const D: usize = 128;
24701 const NSPLIT: u32 = 4;
24702 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
24703 let h = n_head;
24704 let nc = (t + c - 1) / c;
24705 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
24706 let gdn_mma_pre = !portable_mma_gated()
24711 && c == 32
24712 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24713 Ok("1") => true,
24714 Ok("0") => false,
24715 _ => gdn_mma_default_on(),
24716 };
24717 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
24718 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
24719 } else {
24720 None
24721 };
24722 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
24727 && gdn_mma_pre
24728 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
24729 let nk = t * hk * D;
24730 let mut kb16_local: Option<CudaSlice<u8>> = None;
24731 if gdn_mma_pre && kb16_pre.is_none() {
24732 let mut kb = self.alloc_u8_uninit(nk * 2)?;
24733 let f = self.func("f32_to_bf16_bulk");
24734 let n2 = nk as i64;
24735 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
24736 let __s_b = self.gpu.stream();
24737 let mut b = __s_b.launch_builder(&f);
24738 b.arg(k).arg(&mut kb).arg(&n2);
24739 unsafe {
24740 b.launch(cfg2)?;
24741 }
24742 kb16_local = Some(kb);
24743 }
24744 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
24745 if let Some(kb) = kb16_pre {
24746 assert!(kb.len() >= nk * 2, "kb16_pre too small");
24747 }
24748 let mut qb16: Option<CudaSlice<u8>> = None;
24749 let mut pb16: Option<CudaSlice<u8>> = None;
24750 if gdn_wgmma_pre {
24751 if qb16_pre.is_none() {
24754 let mut qb = self.alloc_u8_uninit(nk * 2)?;
24755 let f = self.func("f32_to_bf16_bulk");
24756 let n2 = nk as i64;
24757 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
24758 let __s_b = self.gpu.stream();
24759 let mut b = __s_b.launch_builder(&f);
24760 b.arg(q).arg(&mut qb).arg(&n2);
24761 unsafe {
24762 b.launch(cfg2)?;
24763 }
24764 qb16 = Some(qb);
24765 } else if let Some(qb) = qb16_pre {
24766 assert!(qb.len() >= nk * 2, "qb16_pre too small");
24767 }
24768 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
24769 }
24770 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
24771 let k2w = if gdn_wgmma_pre {
24772 Some((
24773 *qb16_ref0.as_ref().unwrap(),
24774 *kb16_ref0.as_ref().unwrap(),
24775 pb16.as_mut().unwrap(),
24776 ))
24777 } else {
24778 None
24779 };
24780 let (gcum, p, u, w) =
24781 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
24782 let _ = &w;
24783 let mut y = self.uninit(nc * h * c * D)?;
24784 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
24798 && c == 32
24799 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
24800 Ok("1") => true,
24801 Ok("0") => false,
24802 _ => gdn_mma_default_on(),
24803 };
24804 if gdn_mma {
24805 let wb16 = wb16_pre
24806 .take()
24807 .expect("mma path pre-allocates wb16 (K3 store fold)");
24808 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
24809 if gdn_wgmma_pre {
24821 let qb16 = qb16_ref0.unwrap();
24823 let pb16 = pb16.as_ref().unwrap();
24824 {
24825 let f = self.func("gdn_k45_wgmma");
24826 let cfg = LaunchConfig {
24827 grid_dim: (h as u32, 4, 1),
24828 block_dim: (256, 1, 1),
24829 shared_mem_bytes: 0,
24830 };
24831 let hki = hk as i32;
24832 let __s_b = self.gpu.stream();
24833 let mut b = __s_b.launch_builder(&f);
24834 b.arg(kb16_ref)
24835 .arg(&gcum)
24836 .arg(beta)
24837 .arg(&u)
24838 .arg(&wb16)
24839 .arg(qb16)
24840 .arg(pb16)
24841 .arg(o)
24842 .arg(&scale)
24843 .arg(state_in)
24844 .arg(&mut *state_out)
24845 .arg(&hi)
24846 .arg(&ti)
24847 .arg(&ci)
24848 .arg(&hki);
24849 unsafe {
24850 b.launch(cfg)?;
24851 }
24852 }
24853 return Ok(());
24854 }
24855 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
24859 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
24860 {
24861 let f = self.func("gdn_chunk_state_mma");
24862 let cfg = LaunchConfig {
24863 grid_dim: (h as u32, NSPLIT, 1),
24864 block_dim: (256, 1, 1),
24865 shared_mem_bytes: 0,
24866 };
24867 let hki = hk as i32;
24868 let __s_b = self.gpu.stream();
24869 let mut b = __s_b.launch_builder(&f);
24870 b.arg(kb16_ref)
24871 .arg(&gcum)
24872 .arg(beta)
24873 .arg(&u)
24874 .arg(&wb16)
24875 .arg(&mut y16)
24876 .arg(&mut ssnap16)
24877 .arg(state_in)
24878 .arg(&mut *state_out)
24879 .arg(&hi)
24880 .arg(&ti)
24881 .arg(&ci)
24882 .arg(&hki);
24883 unsafe {
24884 b.launch(cfg)?;
24885 }
24886 }
24887 {
24888 let f = self.func("gdn_chunk_output_mma");
24890 let jt = ((c + 31) / 32) as u32;
24891 let cfg = LaunchConfig {
24892 grid_dim: (nc as u32, h as u32, jt),
24893 block_dim: (256, 1, 1),
24894 shared_mem_bytes: 0,
24895 };
24896 let hki = hk as i32;
24897 let __s_b = self.gpu.stream();
24898 let mut b = __s_b.launch_builder(&f);
24899 b.arg(q)
24900 .arg(&gcum)
24901 .arg(&p)
24902 .arg(&y16)
24903 .arg(&ssnap16)
24904 .arg(o)
24905 .arg(&hi)
24906 .arg(&ti)
24907 .arg(&ci)
24908 .arg(&scale)
24909 .arg(&hki);
24910 unsafe {
24911 b.launch(cfg)?;
24912 }
24913 }
24914 return Ok(());
24915 }
24916 {
24917 let f = self.func("gdn_chunk_state_f32");
24919 let cfg = LaunchConfig {
24920 grid_dim: (h as u32, NSPLIT, 1),
24921 block_dim: (256, 1, 1),
24922 shared_mem_bytes: 0,
24923 };
24924 let __s_b = self.gpu.stream();
24925 let mut b = __s_b.launch_builder(&f);
24926 b.arg(k)
24927 .arg(&gcum)
24928 .arg(beta)
24929 .arg(&u)
24930 .arg(&w)
24931 .arg(&mut y)
24932 .arg(&mut ssnap)
24933 .arg(state_in)
24934 .arg(&mut *state_out)
24935 .arg(&hi)
24936 .arg(&ti)
24937 .arg(&ci);
24938 unsafe {
24939 b.launch(cfg)?;
24940 }
24941 }
24942 {
24943 let f = self.func("gdn_chunk_output_f32");
24945 let jt = ((c + 31) / 32) as u32;
24946 let cfg = LaunchConfig {
24947 grid_dim: (nc as u32, h as u32, jt),
24948 block_dim: (256, 1, 1),
24949 shared_mem_bytes: 0,
24950 };
24951 let __s_b = self.gpu.stream();
24952 let mut b = __s_b.launch_builder(&f);
24953 b.arg(q)
24954 .arg(&gcum)
24955 .arg(&p)
24956 .arg(&y)
24957 .arg(&ssnap)
24958 .arg(o)
24959 .arg(&hi)
24960 .arg(&ti)
24961 .arg(&ci)
24962 .arg(&scale);
24963 unsafe {
24964 b.launch(cfg)?;
24965 }
24966 }
24967 Ok(())
24968 }
24969
24970 #[allow(clippy::too_many_arguments)]
24979 #[allow(clippy::too_many_arguments)]
24980 pub fn gdn_scan_prefill(
24981 &self,
24982 q: &CudaSlice<f32>,
24983 k: &CudaSlice<f32>,
24984 v: &CudaSlice<f32>,
24985 g: &CudaSlice<f32>,
24986 beta: &CudaSlice<f32>,
24987 kb16_pre: Option<&CudaSlice<u8>>,
24988 qb16_pre: Option<&CudaSlice<u8>>,
24989 state_in: &CudaSlice<f32>,
24990 state_out: &mut CudaSlice<f32>,
24991 o: &mut CudaSlice<f32>,
24992 n_head: usize,
24993 t: usize,
24994 scale: f32,
24995 hk: usize,
24996 ) -> Result<(), Box<dyn std::error::Error>> {
24997 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
24998 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
24999 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
25000 }
25001 if Self::gdn_chunked_enabled() && t >= 16 {
25002 self.gdn_scan_chunked(
25003 q,
25004 k,
25005 v,
25006 g,
25007 beta,
25008 kb16_pre,
25009 qb16_pre,
25010 state_in,
25011 state_out,
25012 o,
25013 n_head,
25014 t,
25015 scale,
25016 Self::gdn_chunk_size(),
25017 hk,
25018 )
25019 } else {
25020 assert!(
25021 hk == n_head,
25022 "s128 scan is broadcast-only (prep guarantees by predicate)"
25023 );
25024 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
25025 }
25026 }
25027
25028 #[allow(clippy::too_many_arguments)]
25030 fn gdn_scan_diff(
25031 &self,
25032 q: &CudaSlice<f32>,
25033 k: &CudaSlice<f32>,
25034 v: &CudaSlice<f32>,
25035 g: &CudaSlice<f32>,
25036 beta: &CudaSlice<f32>,
25037 state_in: &CudaSlice<f32>,
25038 state_out: &mut CudaSlice<f32>,
25039 o: &mut CudaSlice<f32>,
25040 n_head: usize,
25041 t: usize,
25042 scale: f32,
25043 ) -> Result<(), Box<dyn std::error::Error>> {
25044 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
25045 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
25046 let mut o_c = self.uninit(o.len())?;
25047 let mut st_c = self.uninit(state_out.len())?;
25048 self.gdn_scan_chunked(
25049 q,
25050 k,
25051 v,
25052 g,
25053 beta,
25054 None,
25055 None,
25056 state_in,
25057 &mut st_c,
25058 &mut o_c,
25059 n_head,
25060 t,
25061 scale,
25062 Self::gdn_chunk_size(),
25063 n_head,
25064 )?;
25065 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
25066 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
25067 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
25068 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
25069 let mut max_abs = 0f32;
25070 let mut max_rel = 0f32;
25071 let mut sum_rel = 0f64;
25072 for (x, y) in a.iter().zip(b) {
25073 let ad = (x - y).abs();
25074 let rel = ad / x.abs().max(y.abs()).max(1e-3);
25075 if ad > max_abs {
25076 max_abs = ad;
25077 }
25078 if rel > max_rel {
25079 max_rel = rel;
25080 }
25081 sum_rel += rel as f64;
25082 }
25083 (max_abs, max_rel, sum_rel / a.len() as f64)
25084 };
25085 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
25086 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
25087 println!(
25088 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
25089 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
25090 Self::gdn_chunk_size()
25091 );
25092 Ok(())
25093 }
25094
25095 pub fn gdn_glog(
25097 &self,
25098 alpha: &CudaSlice<f32>,
25099 dt_bias: &CudaSlice<f32>,
25100 a: &CudaSlice<f32>,
25101 g_log: &mut CudaSlice<f32>,
25102 n_head: usize,
25103 t: usize,
25104 ) -> Result<(), Box<dyn std::error::Error>> {
25105 let f = self.func("gdn_glog_f32");
25106 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25107 let (h, ti) = (n_head as i32, t as i32);
25108 let __s_b = self.gpu.stream();
25109 let mut b = __s_b.launch_builder(&f);
25110 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25111 unsafe {
25112 b.launch(cfg)?;
25113 }
25114 Ok(())
25115 }
25116
25117 pub fn sigmoid_v(
25120 &self,
25121 x: &cudarc::driver::CudaView<f32>,
25122 y: &mut CudaSlice<f32>,
25123 n: usize,
25124 ) -> Result<(), Box<dyn std::error::Error>> {
25125 let f = self.func("sigmoid_f32");
25126 let cfg = LaunchConfig::for_num_elems(n as u32);
25127 let ni = n as i32;
25128 let __s_b = self.gpu.stream();
25129 let mut b = __s_b.launch_builder(&f);
25130 b.arg(x).arg(y).arg(&ni);
25131 unsafe {
25132 b.launch(cfg)?;
25133 }
25134 Ok(())
25135 }
25136
25137 pub fn gdn_glog_v(
25138 &self,
25139 alpha: &cudarc::driver::CudaView<f32>,
25140 dt_bias: &CudaSlice<f32>,
25141 a: &CudaSlice<f32>,
25142 g_log: &mut CudaSlice<f32>,
25143 n_head: usize,
25144 t: usize,
25145 ) -> Result<(), Box<dyn std::error::Error>> {
25146 let f = self.func("gdn_glog_f32");
25147 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
25148 let (h, ti) = (n_head as i32, t as i32);
25149 let __s_b = self.gpu.stream();
25150 let mut b = __s_b.launch_builder(&f);
25151 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
25152 unsafe {
25153 b.launch(cfg)?;
25154 }
25155 Ok(())
25156 }
25157
25158 pub fn sigmoid(
25159 &self,
25160 x: &CudaSlice<f32>,
25161 y: &mut CudaSlice<f32>,
25162 n: usize,
25163 ) -> Result<(), Box<dyn std::error::Error>> {
25164 let f = self.func("sigmoid_f32");
25165 let cfg = LaunchConfig::for_num_elems(n as u32);
25166 let ni = n as i32;
25167 let __s_b = self.gpu.stream();
25168 let mut b = __s_b.launch_builder(&f);
25169 b.arg(x).arg(y).arg(&ni);
25170 unsafe {
25171 b.launch(cfg)?;
25172 }
25173 Ok(())
25174 }
25175
25176 pub fn sig_mul_f16out(
25179 &self,
25180 a: &CudaSlice<f32>,
25181 g: &CudaSlice<f32>,
25182 dst: &mut CudaSlice<f32>,
25183 dst16: &mut CudaSlice<u8>,
25184 n: usize,
25185 ) -> Result<(), Box<dyn std::error::Error>> {
25186 let f = self.func("sig_mul_f16out_f32");
25187 let cfg = LaunchConfig::for_num_elems(n as u32);
25188 let ni = n as i32;
25189 let __s_b = self.gpu.stream();
25190 let mut b = __s_b.launch_builder(&f);
25191 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
25192 unsafe {
25193 b.launch(cfg)?;
25194 }
25195 Ok(())
25196 }
25197
25198 #[allow(clippy::too_many_arguments)]
25207 pub fn attn_head_gate(
25208 &self,
25209 a: &CudaSlice<f32>,
25210 g: &CudaSlice<f32>,
25211 dst: &mut CudaSlice<f32>,
25212 dst16: Option<&mut CudaSlice<u8>>,
25213 head_dim: usize,
25214 n_head: usize,
25215 t: usize,
25216 ) -> Result<(), Box<dyn std::error::Error>> {
25217 let f = self.func("attn_head_gate_f32");
25218 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25219 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25220 let d16: u64 = match dst16 {
25222 Some(d) => self.addr_u8(d),
25223 None => 0,
25224 };
25225 let __s_b = self.gpu.stream();
25226 let mut b = __s_b.launch_builder(&f);
25227 b.arg(a)
25228 .arg(g)
25229 .arg(dst)
25230 .arg(&d16)
25231 .arg(&hd)
25232 .arg(&nh)
25233 .arg(&ti);
25234 unsafe {
25235 b.launch(cfg)?;
25236 }
25237 Ok(())
25238 }
25239
25240 #[allow(clippy::too_many_arguments)]
25249 pub fn swiglu_clamped_mul_scaled(
25250 &self,
25251 gate: &CudaSlice<f32>,
25252 up: &CudaSlice<f32>,
25253 gs: f32,
25254 us: f32,
25255 limit: f32,
25256 dst: &mut CudaSlice<f32>,
25257 n: usize,
25258 ) -> Result<(), Box<dyn std::error::Error>> {
25259 debug_assert!(
25260 limit > 1e-6,
25261 "swiglu_clamped needs a live limit; use silu_mul_scaled"
25262 );
25263 let f = self.func("swiglu_clamped_mul_scaled_f32");
25264 let cfg = LaunchConfig::for_num_elems(n as u32);
25265 let ni = n as i32;
25266 let __s_b = self.gpu.stream();
25267 let mut b = __s_b.launch_builder(&f);
25268 b.arg(gate)
25269 .arg(up)
25270 .arg(&gs)
25271 .arg(&us)
25272 .arg(&limit)
25273 .arg(dst)
25274 .arg(&ni);
25275 unsafe {
25276 b.launch(cfg)?;
25277 }
25278 Ok(())
25279 }
25280
25281 pub fn gated_rmsnorm(
25283 &self,
25284 o: &CudaSlice<f32>,
25285 w: &CudaSlice<f32>,
25286 z: &CudaSlice<f32>,
25287 dst: &mut CudaSlice<f32>,
25288 ncols: usize,
25289 nrows: usize,
25290 eps: f32,
25291 ) -> Result<(), Box<dyn std::error::Error>> {
25292 let f = self.func("gated_rmsnorm_f32");
25293 let cfg = LaunchConfig {
25294 grid_dim: (nrows as u32, 1, 1),
25295 block_dim: (128, 1, 1),
25296 shared_mem_bytes: 0,
25297 };
25298 let (nc, e) = (ncols as i32, eps);
25299 let __s_b = self.gpu.stream();
25300 let mut b = __s_b.launch_builder(&f);
25301 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25302 unsafe {
25303 b.launch(cfg)?;
25304 }
25305 Ok(())
25306 }
25307
25308 pub fn gated_rmsnorm_f16out(
25311 &self,
25312 o: &CudaSlice<f32>,
25313 w: &CudaSlice<f32>,
25314 z: &CudaSlice<f32>,
25315 dst: &mut CudaSlice<f32>,
25316 dst16: &mut CudaSlice<u8>,
25317 ncols: usize,
25318 nrows: usize,
25319 eps: f32,
25320 ) -> Result<(), Box<dyn std::error::Error>> {
25321 let f = self.func("gated_rmsnorm_f16out_f32");
25322 let cfg = LaunchConfig {
25324 grid_dim: (nrows as u32, 1, 1),
25325 block_dim: (128, 1, 1),
25326 shared_mem_bytes: 0,
25327 };
25328 let (nc, e) = (ncols as i32, eps);
25329 let __s_b = self.gpu.stream();
25330 let mut b = __s_b.launch_builder(&f);
25331 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25332 unsafe {
25333 b.launch(cfg)?;
25334 }
25335 Ok(())
25336 }
25337
25338 #[allow(clippy::too_many_arguments)]
25342 pub fn add_rms_norm_zq8(
25343 &self,
25344 a: &CudaSlice<f32>,
25345 b_in: &CudaSlice<f32>,
25346 w: &CudaSlice<f32>,
25347 res: &mut CudaSlice<f32>,
25348 z: &mut CudaSlice<f32>,
25349 ncols: usize,
25350 nrows: usize,
25351 eps: f32,
25352 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25353 assert!(ncols % 32 == 0);
25354 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
25355 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25356 let f = self.func("add_rms_norm_zq8");
25357 let cfg = LaunchConfig {
25358 grid_dim: (nrows as u32, 1, 1),
25359 block_dim: (1024, 1, 1),
25360 shared_mem_bytes: 0,
25361 };
25362 let (nc, ep) = (ncols as i32, eps);
25363 let __s_b = self.gpu.stream();
25364 let mut b = __s_b.launch_builder(&f);
25365 b.arg(a)
25366 .arg(b_in)
25367 .arg(w)
25368 .arg(res)
25369 .arg(z)
25370 .arg(&mut q)
25371 .arg(&mut d)
25372 .arg(&nc)
25373 .arg(&ep);
25374 unsafe {
25375 b.launch(cfg)?;
25376 }
25377 Ok((q, d))
25378 }
25379
25380 pub fn gated_rmsnorm_zv(
25385 &self,
25386 o: &CudaSlice<f32>,
25387 w: &CudaSlice<f32>,
25388 z: &cudarc::driver::CudaView<f32>,
25389 dst: &mut CudaSlice<f32>,
25390 ncols: usize,
25391 nrows: usize,
25392 eps: f32,
25393 ) -> Result<(), Box<dyn std::error::Error>> {
25394 let f = self.func("gated_rmsnorm_f32");
25395 let cfg = LaunchConfig {
25396 grid_dim: (nrows as u32, 1, 1),
25397 block_dim: (128, 1, 1),
25398 shared_mem_bytes: 0,
25399 };
25400 let (nc, e) = (ncols as i32, eps);
25401 let __s_b = self.gpu.stream();
25402 let mut b = __s_b.launch_builder(&f);
25403 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
25404 unsafe {
25405 b.launch(cfg)?;
25406 }
25407 Ok(())
25408 }
25409
25410 pub fn gated_rmsnorm_f16out_zv(
25411 &self,
25412 o: &CudaSlice<f32>,
25413 w: &CudaSlice<f32>,
25414 z: &cudarc::driver::CudaView<f32>,
25415 dst: &mut CudaSlice<f32>,
25416 dst16: &mut CudaSlice<u8>,
25417 ncols: usize,
25418 nrows: usize,
25419 eps: f32,
25420 ) -> Result<(), Box<dyn std::error::Error>> {
25421 let f = self.func("gated_rmsnorm_f16out_f32");
25422 let cfg = LaunchConfig {
25424 grid_dim: (nrows as u32, 1, 1),
25425 block_dim: (128, 1, 1),
25426 shared_mem_bytes: 0,
25427 };
25428 let (nc, e) = (ncols as i32, eps);
25429 let __s_b = self.gpu.stream();
25430 let mut b = __s_b.launch_builder(&f);
25431 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
25432 unsafe {
25433 b.launch(cfg)?;
25434 }
25435 Ok(())
25436 }
25437
25438 pub fn gated_rmsnorm_q8_1(
25439 &self,
25440 o: &CudaSlice<f32>,
25441 w: &CudaSlice<f32>,
25442 z: &CudaSlice<f32>,
25443 ncols: usize,
25444 nrows: usize,
25445 eps: f32,
25446 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
25447 assert!(ncols % 32 == 0);
25448 let f = self.func("gated_rmsnorm_q8_1");
25449 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
25450 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
25451 let cfg = LaunchConfig {
25452 grid_dim: (nrows as u32, 1, 1),
25453 block_dim: (128, 1, 1),
25454 shared_mem_bytes: 0,
25455 };
25456 let (nc, ep) = (ncols as i32, eps);
25457 let __s_b = self.gpu.stream();
25458 let mut b = __s_b.launch_builder(&f);
25459 b.arg(o)
25460 .arg(w)
25461 .arg(z)
25462 .arg(&mut out_q)
25463 .arg(&mut out_d)
25464 .arg(&nc)
25465 .arg(&ep);
25466 unsafe {
25467 b.launch(cfg)?;
25468 }
25469 Ok((out_q, out_d))
25470 }
25471
25472 pub fn transpose(
25474 &self,
25475 inp: &CudaSlice<f32>,
25476 rows: usize,
25477 cols: usize,
25478 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25479 let f = self.func("transpose_f32");
25480 let mut out = self.zeros(rows * cols)?;
25481 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
25482 let (r, c) = (rows as i32, cols as i32);
25483 let __s_b = self.gpu.stream();
25484 let mut b = __s_b.launch_builder(&f);
25485 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
25486 unsafe {
25487 b.launch(cfg)?;
25488 }
25489 Ok(out)
25490 }
25491
25492 pub fn repeat_heads(
25494 &self,
25495 inp: &CudaSlice<f32>,
25496 out: &mut CudaSlice<f32>,
25497 head_dim: usize,
25498 n_in: usize,
25499 n_out: usize,
25500 t: usize,
25501 ) -> Result<(), Box<dyn std::error::Error>> {
25502 let f = self.func("repeat_heads_f32");
25503 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
25504 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
25505 let __s_b = self.gpu.stream();
25506 let mut b = __s_b.launch_builder(&f);
25507 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
25508 unsafe {
25509 b.launch(cfg)?;
25510 }
25511 Ok(())
25512 }
25513
25514 pub fn q_gate_split(
25521 &self,
25522 qf: &CudaSlice<f32>,
25523 q_out: &mut CudaSlice<f32>,
25524 gate_out: &mut CudaSlice<f32>,
25525 head_dim: usize,
25526 n_head: usize,
25527 t: usize,
25528 ) -> Result<(), Box<dyn std::error::Error>> {
25529 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
25530 let out_need = head_dim * n_head * t;
25531 if q_out.len() < out_need || gate_out.len() < out_need {
25532 return Err(format!(
25533 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
25534 q_out.len(),
25535 gate_out.len()
25536 )
25537 .into());
25538 }
25539 let f = self.func("q_gate_split_f32");
25540 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
25541 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
25542 let __s_b = self.gpu.stream();
25543 let mut b = __s_b.launch_builder(&f);
25544 b.arg(qf)
25545 .arg(q_out)
25546 .arg(gate_out)
25547 .arg(&hd)
25548 .arg(&nh)
25549 .arg(&ti);
25550 unsafe {
25551 b.launch(cfg)?;
25552 }
25553 Ok(())
25554 }
25555
25556 pub fn qkv_to_gdn_repack(
25560 &self,
25561 conv_out: &CudaSlice<f32>,
25562 q_g: &mut CudaSlice<f32>,
25563 k_g: &mut CudaSlice<f32>,
25564 v_g: &mut CudaSlice<f32>,
25565 d_state: usize,
25566 num_v: usize,
25567 num_k: usize,
25568 key_dim: usize,
25569 t: usize,
25570 ) -> Result<(), Box<dyn std::error::Error>> {
25571 let f = self.func("qkv_to_gdn_repack_f32");
25572 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
25573 let (ds, nv, nk, kd, ti) = (
25574 d_state as i32,
25575 num_v as i32,
25576 num_k as i32,
25577 key_dim as i32,
25578 t as i32,
25579 );
25580 let __s_b = self.gpu.stream();
25581 let mut b = __s_b.launch_builder(&f);
25582 b.arg(conv_out)
25583 .arg(q_g)
25584 .arg(k_g)
25585 .arg(v_g)
25586 .arg(&ds)
25587 .arg(&nv)
25588 .arg(&nk)
25589 .arg(&kd)
25590 .arg(&ti);
25591 unsafe {
25592 b.launch(cfg)?;
25593 }
25594 Ok(())
25595 }
25596
25597 pub fn conv_left_pad(
25600 &self,
25601 src: &CudaSlice<f32>,
25602 dst: &mut CudaSlice<f32>,
25603 conv_dim: usize,
25604 t: usize,
25605 pad: usize,
25606 ) -> Result<(), Box<dyn std::error::Error>> {
25607 let f = self.func("conv_left_pad_f32");
25608 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
25609 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
25610 let __s_b = self.gpu.stream();
25611 let mut b = __s_b.launch_builder(&f);
25612 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
25613 unsafe {
25614 b.launch(cfg)?;
25615 }
25616 Ok(())
25617 }
25618
25619 pub fn conv_assemble_and_roll(
25623 &self,
25624 qkv_col: &CudaSlice<f32>,
25625 conv_state: &mut CudaSlice<f32>,
25626 conv_in: &mut CudaSlice<f32>,
25627 conv_dim: usize,
25628 pad: usize,
25629 ) -> Result<(), Box<dyn std::error::Error>> {
25630 let f = self.func("conv_assemble_and_roll_f32");
25631 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
25632 let (cd, p) = (conv_dim as i32, pad as i32);
25633 let __s_b = self.gpu.stream();
25634 let mut b = __s_b.launch_builder(&f);
25635 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
25636 unsafe {
25637 b.launch(cfg)?;
25638 }
25639 Ok(())
25640 }
25641
25642 pub fn ssm_conv1d_fused_decode(
25648 &self,
25649 qkv_col: &CudaSlice<f32>,
25650 conv_state: &mut CudaSlice<f32>,
25651 w: &CudaSlice<f32>,
25652 conv_out: &mut CudaSlice<f32>,
25653 conv_dim: usize,
25654 d_conv: usize,
25655 ) -> Result<(), Box<dyn std::error::Error>> {
25656 let f = self.func("ssm_conv1d_fused_decode_f32");
25657 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
25658 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25659 let __s_b = self.gpu.stream();
25660 let mut b = __s_b.launch_builder(&f);
25661 b.arg(qkv_col)
25662 .arg(conv_state)
25663 .arg(w)
25664 .arg(conv_out)
25665 .arg(&cd)
25666 .arg(&dc);
25667 unsafe {
25668 b.launch(cfg)?;
25669 }
25670 Ok(())
25671 }
25672
25673 pub fn slice_range(
25676 &self,
25677 src: &CudaSlice<f32>,
25678 start: usize,
25679 len: usize,
25680 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25681 let host = self.gpu.stream().clone_dtoh(src)?;
25682 self.gpu.stream().synchronize()?;
25683 Ok(self.htod(&host[start..start + len])?)
25684 }
25685}
25686
25687#[cfg(test)]
25688mod target_dispatch_tests {
25689 use super::legacy_quant_gemm_allowed;
25690
25691 #[test]
25692 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
25693 assert!(legacy_quant_gemm_allowed(false, false, false));
25695 assert!(!legacy_quant_gemm_allowed(false, false, true));
25696 assert!(!legacy_quant_gemm_allowed(true, false, false));
25698 assert!(!legacy_quant_gemm_allowed(true, false, true));
25699 assert!(legacy_quant_gemm_allowed(true, true, false));
25701 assert!(!legacy_quant_gemm_allowed(true, true, true));
25702 }
25703
25704 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
25705 #[test]
25706 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
25707 assert!(!legacy_quant_gemm_allowed(
25708 cfg!(memra_portable_cuda),
25709 cfg!(memra_hopper_mma),
25710 false
25711 ));
25712 }
25713
25714 #[cfg(memra_hopper_mma)]
25715 #[test]
25716 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
25717 assert!(legacy_quant_gemm_allowed(
25718 cfg!(memra_portable_cuda),
25719 cfg!(memra_hopper_mma),
25720 false
25721 ));
25722 assert!(super::portable_mma_gated() == false);
25723 }
25724}
25725
25726impl memra_kv::KvDev for Engine {
25729 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25730 Engine::zeros(self, n)
25731 }
25732 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25733 Engine::uninit(self, n)
25734 }
25735 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
25736 Engine::alloc_u8(self, n)
25737 }
25738 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
25739 Engine::htod_i32(self, v)
25740 }
25741 fn clone_dtod(
25742 &self,
25743 src: &CudaSlice<f32>,
25744 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
25745 Engine::clone_dtod(self, src)
25746 }
25747 fn copy_into(
25748 &self,
25749 dst: &mut CudaSlice<f32>,
25750 off: usize,
25751 src: &CudaSlice<f32>,
25752 len: usize,
25753 ) -> Result<(), Box<dyn std::error::Error>> {
25754 Engine::copy_into(self, dst, off, src, len)
25755 }
25756 fn set_i32_one(
25757 &self,
25758 d: &mut CudaSlice<i32>,
25759 v: i32,
25760 ) -> Result<(), Box<dyn std::error::Error>> {
25761 Engine::set_i32_one(self, d, v)
25762 }
25763}
25764
25765#[cfg(test)]
25766mod fused_gate_bounds_tests {
25767 use super::*;
25768
25769 #[test]
25782 #[ignore = "requires a CUDA GPU"]
25783 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
25784 let e = Engine::new(0).unwrap();
25785 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
25786 let fused = 2 * head_dim * n_head * t;
25787 let out_n = head_dim * n_head * t;
25788
25789 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
25791 let mut q = e.uninit(out_n).unwrap();
25792 let mut gate = e.uninit(out_n).unwrap();
25793 let err = e
25794 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
25795 .expect_err("half-width wq must be refused, not read past")
25796 .to_string();
25797 assert!(err.contains("NO fused gate"), "{err}");
25798 assert!(err.contains(&format!("{fused}")), "{err}");
25799
25800 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
25803 let wide = e.htod(&host).unwrap();
25804 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
25805 .expect("full-width wq splits");
25806 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
25807 for tok in 0..t {
25808 for hh in 0..n_head {
25809 for d in 0..head_dim {
25810 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
25811 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
25812 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
25813 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
25814 }
25815 }
25816 }
25817
25818 let mut small = e.uninit(out_n - 1).unwrap();
25820 assert!(
25821 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
25822 .is_err()
25823 );
25824 }
25825}
25826
25827#[cfg(test)]
25831mod fused_rope_width_tests {
25832 use super::Engine;
25833
25834 #[test]
25837 fn full_width_is_accepted() {
25838 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
25839 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
25840 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
25841 }
25842
25843 #[test]
25856 fn gemma4_official_artifact_widths_pass() {
25857 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
25858 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
25859 }
25860
25861 #[test]
25864 fn partial_rotary_is_refused_with_the_geometry_named() {
25865 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
25867 .expect_err("partial rotary must refuse");
25868 let msg = err.to_string();
25869 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
25870 assert!(msg.contains("n_rot 64"), "{msg}");
25871 assert!(msg.contains("head_dim 256"), "{msg}");
25872 assert!(
25873 msg.contains("64..256"),
25874 "names the band it would corrupt: {msg}"
25875 );
25876 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
25878 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
25880 }
25881}