1use cudarc::driver::sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES;
4use cudarc::driver::{
5 CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DeviceSlice, LaunchConfig,
6 PushKernelArg,
7};
8use cudarc::nvrtc::Ptx;
9use std::sync::{Arc, Mutex};
10
11const GDN_K2_DYNAMIC_SHARED_BYTES: u32 = 67_072;
12
13const SDPA_NAIVE_SMEM_MAX: usize = 48 * 1024;
18
19const SDPA_NAIVE_GMEM_WS_MAX: usize = 1 << 30;
23
24#[cfg(debug_assertions)]
25pub(crate) fn debug_assert_tensor_stream_device<T>(
26 tensor: &CudaSlice<T>,
27 stream: &CudaStream,
28 site: &str,
29) {
30 let tensor_dev = tensor.ordinal();
31 let stream_dev = stream.context().ordinal();
32 assert_eq!(
33 tensor_dev, stream_dev,
34 "PP cross-device tensor read at {site}: tensor on dev{tensor_dev}, stream on dev{stream_dev}"
35 );
36}
37
38fn ensure_tensor_stream_device<T>(
39 tensor: &impl DeviceSlice<T>,
40 stream: &CudaStream,
41 site: &str,
42) -> Result<(), Box<dyn std::error::Error>> {
43 let tensor_dev = tensor.stream().context().ordinal();
44 let stream_dev = stream.context().ordinal();
45 if tensor_dev != stream_dev {
46 return Err(format!(
47 "PP cross-device tensor access at {site}: tensor on dev{tensor_dev}, \
48 stream on dev{stream_dev}"
49 )
50 .into());
51 }
52 Ok(())
53}
54
55pub use memra_gguf;
56pub use memra_runtime;
57
58pub mod forward;
59pub mod hybrid;
60pub mod hybrid_forward;
61pub mod model;
62pub mod sigrouter_contract;
63pub mod vision;
64pub mod vision_gemma;
65pub mod vision_pre;
66pub mod cache {
69 pub use memra_kv::*;
70}
71pub mod decode;
72pub mod decode_batch;
73pub mod dflash;
74pub mod eagle;
75pub mod gemma_spec;
76pub mod graph_update;
77pub mod mla;
81pub mod moesd;
82pub mod parallel;
83pub mod plan_backend;
84pub mod pp;
85pub mod round_stream;
86pub mod spec;
87pub mod tp;
88pub use memra_sampling as sampler;
89
90pub fn moe_f16g_mode() -> u8 {
134 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
135 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
136 Ok("0") => 0,
137 Ok("2") => 2,
138 Ok("3") => 3,
139 Ok(_) => 1,
140 Err(_) => 2,
143 })
144}
145pub fn moe_f16g_sk_params() -> (i32, i32) {
159 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
160 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
161 Ok("0") => (-1, 0),
162 Ok("32") => (0, i32::MAX),
163 Ok("128") => (0, 1),
164 _ => {
165 let cross = std::env::var("MEMRA_F16G_SK_CROSS")
166 .ok()
167 .and_then(|v| v.parse().ok())
168 .unwrap_or(64);
169 (0, cross)
170 }
171 })
172}
173pub fn moe_f16g_direct_on(qtype: i32) -> bool {
184 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
185 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
186 Ok("0") => 0,
187 Ok("kq") => 1,
188 _ => 2,
189 });
190 match m {
191 0 => false,
192 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
193 _ => true,
194 }
195}
196pub fn moe_f16g_tail_on() -> bool {
205 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
206 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
207}
208
209pub fn moe_f16g_gemma_on() -> bool {
216 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
217 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
218}
219
220pub fn moe_fuse_actq_on() -> bool {
224 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
225 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
226}
227
228pub fn router_prefill_exact_on() -> bool {
238 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
239 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
240}
241
242pub fn router_kernel_on() -> bool {
243 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244 *ON.get_or_init(|| {
245 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
246 if !on {
247 eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)");
248 }
249 on
250 })
251}
252
253pub const ROUTER_BATCH_MIN_T: usize = 8;
268pub fn router_batch_on() -> bool {
269 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
270 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
271}
272mod cpu_experts;
273#[cfg(memra_cutlass)]
274pub mod cutlass_ffi;
275pub mod dsv4_ffi;
276pub mod dsv4_gpu;
277pub mod f16_ffi;
278pub mod fp8_ffi;
279pub mod mmq_ffi;
280pub mod moe_cache;
281pub mod prime_graph;
282pub mod spill;
283mod spill_pread;
284
285const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
292const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
293const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
294const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
295const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
296const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
297const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
299
300fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
306 assert!(
307 !(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
308 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane"
309 );
310 match std::env::var("MEMRA_GEMM_FATBIN") {
311 Ok(path) => std::borrow::Cow::Owned(
312 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}")),
313 ),
314 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
315 }
316}
317
318pub(crate) const fn portable_mma_gated() -> bool {
325 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
326}
327
328#[track_caller]
345pub(crate) fn refuse_portable_force(var: &str, needs: &str) {
346 assert!(
347 !portable_mma_gated(),
348 "{var} forces a kernel path this build does not contain: it needs {needs}, and this is a \
349 portable-CUDA build (sm_89). Unset {var} — the default path serves this arch."
350 );
351}
352
353pub(crate) const fn gdn_mma_default_on() -> bool {
362 cfg!(memra_hopper_mma) || konst_eq(env!("MEMRA_BUILT_CUDA_ARCH"), "120a")
363}
364
365const fn konst_eq(a: &str, b: &str) -> bool {
367 let (a, b) = (a.as_bytes(), b.as_bytes());
368 if a.len() != b.len() {
369 return false;
370 }
371 let mut i = 0;
372 while i < a.len() {
373 if a[i] != b[i] {
374 return false;
375 }
376 i += 1;
377 }
378 true
379}
380
381const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
386 (!portable_cuda || hopper_mma) && !no_gemm
387}
388
389const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
397const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
398const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
399const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
400const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
401
402pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
405
406fn flash_fatbin_bytes() -> &'static [u8] {
408 match kv_cache_formats() {
409 ("q8_0", "q5_1") => FLASH_FATBIN,
410 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
411 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
412 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
413 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
414 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
415 other => unreachable!("kv_cache_formats returned {other:?}"),
416 }
417}
418
419fn k1_launch_override() -> Option<(u32, u32, u32)> {
426 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
427 *K1.get_or_init(|| {
428 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
429 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
430 match p.as_slice() {
431 [bm, bn, w] => Some((*bm, *bn, *w)),
432 _ => None,
433 }
434 })
435}
436
437pub(crate) fn wgmma_gemm_enabled() -> bool {
444 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
445 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
446}
447
448pub const FA_VEC_MIN_TKV: usize = 96;
463pub fn fa_vec_min_tkv() -> usize {
467 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
468 *V.get_or_init(|| {
469 std::env::var("MEMRA_FA_VEC_MIN")
470 .ok()
471 .and_then(|v| v.parse().ok())
472 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
473 })
474}
475
476pub fn fa_f16pv_on() -> bool {
487 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
488 *ON.get_or_init(|| {
489 std::env::var("MEMRA_FA_F16PV")
490 .map(|v| v != "0")
491 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err())
492 })
493}
494
495pub fn fa512_hp_on() -> bool {
499 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
500 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
501}
502
503pub fn faw_hp_on() -> bool {
507 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
508 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
509}
510
511pub fn fa512_wide_warps() -> usize {
515 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
516 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
517 Ok("1") => 4,
518 _ => 2,
519 })
520}
521
522pub fn fa512_min_tkv() -> usize {
525 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
526 *FA512_MIN.get_or_init(|| {
527 std::env::var("MEMRA_FA512_MIN")
528 .ok()
529 .and_then(|v| v.parse().ok())
530 .unwrap_or(512)
531 })
532}
533pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
537 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
538pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(32);
542pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
548 std::sync::atomic::AtomicBool::new(false);
549pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
556 std::sync::atomic::AtomicBool::new(true);
557pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
558 std::sync::atomic::AtomicUsize::new(16);
559pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
564pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
566pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
571pub use memra_kv::KV_FP8_FORCE;
574pub(crate) fn mmv_block() -> u32 {
579 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
580 *V.get_or_init(|| {
581 std::env::var("MEMRA_MMV_BLOCK")
582 .ok()
583 .and_then(|v| v.parse().ok())
584 .filter(|&b: &u32| (64..=256).contains(&b) && b % 32 == 0)
585 .unwrap_or(128)
586 })
587}
588
589static STEP37_SERVING_DEFAULTS: std::sync::atomic::AtomicBool =
615 std::sync::atomic::AtomicBool::new(false);
616
617pub fn arm_step37_serving_defaults() {
618 STEP37_SERVING_DEFAULTS.store(true, std::sync::atomic::Ordering::Relaxed);
619 crate::cache::set_swa_ring_default(true);
620 eprintln!(
621 "[step37-defaults] serving doors armed ON for the SlidingGatedMoe program \
622 (per-flag =0 kills, =1 forces; owner flip 2026-08-27)"
623 );
624}
625
626pub(crate) fn step37_defaults_armed() -> bool {
627 STEP37_SERVING_DEFAULTS.load(std::sync::atomic::Ordering::Relaxed)
628}
629
630pub(crate) fn step37_door(cell: &'static std::sync::OnceLock<Option<bool>>, name: &str) -> bool {
634 match *cell.get_or_init(|| match std::env::var(name).ok().as_deref() {
635 Some("1") => Some(true),
636 Some("0") => Some(false),
637 _ => None,
638 }) {
639 Some(forced) => forced,
640 None => step37_defaults_armed(),
641 }
642}
643
644pub(crate) fn w8_hybrid_on() -> bool {
645 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
646 step37_door(&ENV, "MEMRA_W8_HYBRID")
647}
648
649pub(crate) fn step_tp_w8_on() -> bool {
650 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
651 step37_door(&ENV, "MEMRA_STEP_TP_W8")
652}
653
654pub(crate) fn w8_view_on() -> bool {
659 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
660 *ON.get_or_init(|| std::env::var("MEMRA_W8_VIEW").as_deref() == Ok("1"))
661}
662
663pub(crate) fn step_gemm_prime_on() -> bool {
677 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
678 step37_door(&ENV, "MEMRA_STEP_GEMM_PRIME")
679}
680
681pub(crate) fn q8t_wonce_on() -> bool {
682 static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
683 step37_door(&ENV, "MEMRA_Q8T_WONCE")
684}
685
686pub(crate) fn sig_expf_dev_on() -> bool {
691 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
692 *ON.get_or_init(|| std::env::var("MEMRA_SIG_EXPF_DEV").as_deref() == Ok("1"))
693}
694
695pub(crate) fn topk_fast_on() -> bool {
696 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
697 *ON.get_or_init(|| std::env::var("MEMRA_TOPK_FAST").as_deref() == Ok("1"))
698}
699
700fn sigmoid_topk_kernel(sig_expf: bool, fast: bool, n_used: usize) -> &'static str {
705 match (sig_expf, fast && n_used <= 8) {
706 (true, true) => "moe_router_sigmoid_topk_f32_dexp_fast",
707 (true, false) => "moe_router_sigmoid_topk_f32_dexp",
708 (false, true) => "moe_router_sigmoid_topk_f32_fast",
709 (false, false) => "moe_router_sigmoid_topk_f32",
710 }
711}
712
713#[cfg(test)]
714mod sigmoid_topk_dispatch_tests {
715 #[test]
716 fn fast_kernel_refuses_wide_topk_and_composes_with_dexp() {
717 use super::sigmoid_topk_kernel;
718
719 assert_eq!(
720 sigmoid_topk_kernel(false, true, 8),
721 "moe_router_sigmoid_topk_f32_fast"
722 );
723 assert_eq!(
724 sigmoid_topk_kernel(true, true, 8),
725 "moe_router_sigmoid_topk_f32_dexp_fast"
726 );
727 assert_eq!(
728 sigmoid_topk_kernel(false, true, 9),
729 "moe_router_sigmoid_topk_f32"
730 );
731 assert_eq!(
732 sigmoid_topk_kernel(true, true, 9),
733 "moe_router_sigmoid_topk_f32_dexp"
734 );
735 }
736}
737
738pub(crate) fn rms_block() -> u32 {
739 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
740 *V.get_or_init(|| {
741 std::env::var("MEMRA_RMS_BLOCK")
742 .ok()
743 .and_then(|v| v.parse().ok())
744 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
745 })
746}
747
748pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
749 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
750 if let Some(forced) = *S.get_or_init(|| {
751 std::env::var("MEMRA_FA_SPLIT")
752 .ok()
753 .and_then(|v| v.parse().ok())
754 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
755 }) {
756 return forced;
757 }
758 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
776 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1")
777 {
778 return if t_kv <= 8192 {
779 16
780 } else if t_kv <= 16384 {
781 64
782 } else {
783 128
784 };
785 }
786 let big_rig = fa_sm_count() >= 128;
787 if big_rig {
788 let _ = n_head_kv;
789 if t_kv <= 2048 {
790 static SHORT: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
799 if let Some(sp) = *SHORT.get_or_init(|| {
800 std::env::var("MEMRA_FA_SP_SHORT")
801 .ok()
802 .and_then(|v| v.parse().ok())
803 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
804 }) {
805 return sp;
806 }
807 16
808 } else if t_kv <= 16384 {
809 64
810 } else {
811 128
812 }
813 } else if n_head_kv <= 4 {
814 if t_kv <= 512 {
835 8
836 } else if t_kv <= 16384 {
837 64
838 } else {
839 128
840 }
841 } else {
842 if t_kv <= 8192 {
843 32
844 } else if t_kv <= 16384 {
845 64
846 } else {
847 128
848 }
849 }
850}
851
852pub(crate) fn fa_sm_count() -> i32 {
855 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
856 *N.get_or_init(|| {
857 cudarc::driver::result::init().ok();
858 cudarc::driver::result::device::get(0)
859 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
860 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
861 .unwrap_or(82)
862 })
863}
864
865fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
869 match head_dim {
870 256 => Ok(""),
871 128 => Ok("_hd128"),
872 d => Err(format!(
873 "fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
874 callers must gate to sdpa_naive"
875 )
876 .into()),
877 }
878}
879
880pub const QT_Q8_0: i32 = 0;
882pub const QT_Q4_K: i32 = 1;
883pub const QT_Q6_K: i32 = 2;
884pub const QT_Q5_K: i32 = 3;
885pub const QT_Q3_K: i32 = 4;
886pub const QT_IQ4_XS: i32 = 5;
887pub const QT_IQ3_S: i32 = 6;
888pub const QT_NVFP4: i32 = 7;
889pub const QT_NVFP4_V2: i32 = 107;
892pub const QT_F8_E4M3: i32 = 10;
898pub const QT_NVFP4_RP: i32 = 9;
901pub const QT_F32: i32 = 8;
903pub const QT_BF16: i32 = 11;
904pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
909pub const QT_F8_E4M3_BLK: i32 = 14;
925
926pub struct Engine {
928 pub gpu: memra_runtime::Gpu,
929 module: Arc<CudaModule>,
930 hybrid: Arc<CudaModule>,
931 qmatvec: Arc<CudaModule>,
932 flash: Arc<CudaModule>,
933 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
937 gemm: Arc<CudaModule>,
938 router: Arc<CudaModule>,
939 sample: Arc<CudaModule>,
941 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
945 w8_mirrors: Mutex<std::collections::HashMap<(u64, u32, u32), CudaSlice<u8>>>,
954 w8_act: Mutex<std::collections::HashMap<usize, (CudaSlice<i8>, CudaSlice<f32>)>>,
957 moe_cache_layout: Mutex<Option<Vec<usize>>>,
961 capture_keep_on: std::sync::atomic::AtomicBool,
967 verify_exact: std::sync::atomic::AtomicBool,
972 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
973 pub copy_stream: Arc<CudaStream>,
975 #[cfg(memra_cutlass)]
982 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
983 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
987 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
990 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
994 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
998 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
1000 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
1001 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
1006 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
1011 router_stage: Mutex<Option<PinnedStage>>,
1015}
1016
1017fn fa_v2_on() -> bool {
1027 std::env::var("MEMRA_FA_V2")
1033 .map(|v| v != "0")
1034 .unwrap_or(true)
1035}
1036
1037pub(crate) fn fa_part_zero_on() -> bool {
1047 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1048 *ON.get_or_init(|| std::env::var("MEMRA_FA_PART_ZERO").as_deref() == Ok("1"))
1049}
1050
1051pub(crate) fn fa_v3_on() -> bool {
1052 std::env::var("MEMRA_FA_V3")
1056 .map(|v| v != "0")
1057 .unwrap_or(true)
1058}
1059
1060fn fa_v4_mode() -> &'static str {
1065 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
1066 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
1067}
1068fn fa_v4_on() -> bool {
1069 fa_v4_mode() != "0"
1070} pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
1079 std::sync::atomic::AtomicUsize::new(1024);
1080pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
1081 std::sync::atomic::AtomicUsize::new(usize::MAX);
1082pub fn fa_v4_at_pub(t_kv: usize) -> bool {
1083 fa_v4_at(t_kv)
1084}
1085fn fa_v4_at(t_kv: usize) -> bool {
1086 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
1087 let mx = *M.get_or_init(|| {
1088 std::env::var("MEMRA_FA_V4_MAX")
1089 .ok()
1090 .and_then(|v| v.parse().ok())
1091 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
1092 });
1093 fa_v4_on() && t_kv < mx
1094}
1095pub const FA_DEEP_MIN_DEFAULT: usize = 0;
1109fn fa_deep_at(t_kv: usize) -> bool {
1110 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") {
1111 return false;
1112 }
1113 let min = std::env::var("MEMRA_FA_DEEP_MIN")
1114 .ok()
1115 .and_then(|v| v.parse().ok())
1116 .unwrap_or(FA_DEEP_MIN_DEFAULT);
1117 t_kv >= min
1118}
1119pub fn fa_deep_at_pub(t_kv: usize) -> bool {
1121 fa_deep_at(t_kv)
1122}
1123
1124fn fa_v3_active(head_dim: usize) -> bool {
1125 fa_v3_on()
1128 && head_dim % 128 == 0
1129 && kv_cache_formats() == ("q8_0", "q5_1")
1130 && !Engine::kv_fp8_on()
1131}
1132
1133pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
1141 std::env::var("MEMRA_NO_FA_VEC").is_err()
1142 && t_kv >= fa_vec_min_tkv()
1143 && head_dim == 256
1144 && fa_v4_at(t_kv)
1145 && !matches!(fa_v4_mode(), "noB3" | "stage")
1146 && !Engine::kv_fp8_on()
1147}
1148pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize {
1150 fa_split_keys(t_kv, n_head_kv)
1151}
1152
1153struct PinnedStage {
1158 ptr: *mut u8,
1159 cap: usize,
1160}
1161unsafe impl Send for PinnedStage {}
1162impl PinnedStage {
1163 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
1164 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
1165 Ok(PinnedStage { ptr, cap })
1166 }
1167}
1168impl Drop for PinnedStage {
1169 fn drop(&mut self) {
1170 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
1171 }
1172}
1173
1174pub const ARGMAX_NB: usize = 256;
1177
1178pub(crate) use memra_fa3_vl as fa3_vl_raw;
1180
1181unsafe extern "C" {
1182 fn memra_fa3_prefill(
1184 q16: *const core::ffi::c_void,
1185 k16: *const core::ffi::c_void,
1186 v16: *const core::ffi::c_void,
1187 o: *mut f32,
1188 t: i32,
1189 h: i32,
1190 hkv: i32,
1191 d: i32,
1192 scale: f32,
1193 stream: *mut core::ffi::c_void,
1194 ) -> i32;
1195 pub(crate) fn memra_fa3_vl(
1197 q16s: *const *const core::ffi::c_void,
1198 k16s: *const *const core::ffi::c_void,
1199 v16s: *const *const core::ffi::c_void,
1200 os: *const *mut f32,
1201 ts: *const i32,
1202 b: i32,
1203 h: i32,
1204 hkv: i32,
1205 d: i32,
1206 scale: f32,
1207 stream: *mut core::ffi::c_void,
1208 ) -> i32;
1209}
1210
1211#[repr(C)]
1216#[derive(Clone, Copy)]
1217pub struct WPtr8(pub [u64; 8]);
1218unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
1219
1220#[repr(C)]
1225#[derive(Clone, Copy, Default)]
1226pub struct GdnSeqVl {
1227 pub kb16: u64,
1228 pub gcum: u64,
1229 pub beta: u64,
1230 pub u: u64,
1231 pub wb16: u64,
1232 pub y: u64,
1233 pub ssnap: u64,
1234 pub state_in: u64,
1235 pub state_out: u64,
1236 pub q: u64,
1237 pub p: u64,
1238 pub o: u64,
1239 pub k: u64,
1240 pub v: u64,
1241 pub g: u64,
1242 pub a: u64,
1243 pub w: u64,
1244 pub t: i32,
1245 pub nc: i32,
1246}
1247unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
1248#[repr(C)]
1249#[derive(Clone, Copy)]
1250pub struct GdnVl8(pub [GdnSeqVl; 8]);
1251unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
1252
1253#[repr(C)]
1256#[derive(Clone, Copy, Default)]
1257pub struct GdnWVl {
1258 pub qb16: u64,
1259 pub pb16: u64,
1260}
1261unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
1262#[repr(C)]
1263#[derive(Clone, Copy)]
1264pub struct GdnWVl8(pub [GdnWVl; 8]);
1265unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
1266
1267#[repr(C)]
1269#[derive(Clone, Copy, Default)]
1270pub struct GdnPrepVl {
1271 pub qkv: u64,
1272 pub conv_state: u64,
1273 pub conv_out: u64,
1274 pub q_g: u64,
1275 pub k_g: u64,
1276 pub v_g: u64,
1277 pub q_l2: u64,
1278 pub k_l2: u64,
1279 pub beta_raw: u64,
1280 pub alpha: u64,
1281 pub beta: u64,
1282 pub g_log: u64,
1283 pub o: u64,
1284 pub z: u64,
1285 pub gn: u64,
1286 pub gn16: u64,
1287 pub kb16: u64,
1288 pub qb16: u64,
1289 pub t: i32,
1290 pub pad: i32,
1291}
1292unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
1293#[repr(C)]
1294#[derive(Clone, Copy)]
1295pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
1296unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
1297
1298#[repr(C)]
1300#[derive(Clone, Copy, Default)]
1301pub struct FaSeqVl {
1302 pub q: u64,
1303 pub k16: u64,
1304 pub v16: u64,
1305 pub o: u64,
1306 pub kf: u64,
1307 pub vf: u64,
1308 pub t: i32,
1309 pub pad: i32,
1310}
1311unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
1312#[repr(C)]
1313#[derive(Clone, Copy)]
1314pub struct FaVl8(pub [FaSeqVl; 8]);
1315unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
1316
1317#[repr(C)]
1319#[derive(Clone, Copy, Default)]
1320pub struct AttnPreVl {
1321 pub qf: u64,
1322 pub kf: u64,
1323 pub vf: u64,
1324 pub q: u64,
1325 pub gate: u64,
1326 pub qn: u64,
1327 pub kn: u64,
1328 pub kc: u64,
1329 pub vc: u64,
1330 pub t: i32,
1331 pub pad: i32,
1332}
1333unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
1334#[repr(C)]
1335#[derive(Clone, Copy)]
1336pub struct AttnPreVl8(pub [AttnPreVl; 8]);
1337unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
1338
1339pub struct GdnChunkBufs {
1342 pub gcum: CudaSlice<f32>,
1343 pub a: CudaSlice<f32>,
1344 pub p: CudaSlice<f32>,
1345 pub u: CudaSlice<f32>,
1346 pub w: CudaSlice<f32>,
1347 pub kb16: CudaSlice<u8>,
1348 pub wb16: CudaSlice<u8>,
1349 pub y16: CudaSlice<u8>,
1350 pub ssnap16: CudaSlice<u8>,
1351 pub qb16: CudaSlice<u8>,
1352 pub pb16: CudaSlice<u8>,
1353 pub o: CudaSlice<f32>,
1354 pub t: usize,
1355 pub nc: usize,
1356}
1357
1358#[repr(C)]
1360#[derive(Clone, Copy)]
1361pub struct F32x8(pub [f32; 8]);
1362unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
1363
1364pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1368
1369#[must_use = "dropping immediately ends the exact scope"]
1374pub struct ExactScope<'a> {
1375 flag: &'a std::sync::atomic::AtomicBool,
1376 prev: bool,
1377}
1378
1379impl<'a> ExactScope<'a> {
1380 pub(crate) fn set(flag: &'a std::sync::atomic::AtomicBool, on: bool) -> Self {
1381 let prev = flag.load(std::sync::atomic::Ordering::Relaxed);
1382 flag.store(on, std::sync::atomic::Ordering::Relaxed);
1383 ExactScope { flag, prev }
1384 }
1385}
1386
1387impl Drop for ExactScope<'_> {
1388 fn drop(&mut self) {
1389 self.flag
1390 .store(self.prev, std::sync::atomic::Ordering::Relaxed);
1391 }
1392}
1393
1394#[cfg(test)]
1395mod exact_scope_tests {
1396 use std::sync::atomic::{AtomicBool, Ordering};
1397
1398 #[test]
1399 fn error_path_restores_verify_exact() {
1400 let flag = AtomicBool::new(false);
1405 let failing = |flag: &AtomicBool| -> Result<(), &'static str> {
1406 let _scope = super::ExactScope::set(flag, true);
1407 assert!(flag.load(Ordering::Relaxed), "scope arms the flag");
1408 Err("draft forward failed")? };
1410 assert!(failing(&flag).is_err());
1411 assert!(
1412 !flag.load(Ordering::Relaxed),
1413 "error propagation must restore the pre-scope value"
1414 );
1415 let flag = AtomicBool::new(true);
1417 {
1418 let _scope = super::ExactScope::set(&flag, true);
1419 }
1420 assert!(flag.load(Ordering::Relaxed));
1421 let flag = AtomicBool::new(false);
1423 let scope = super::ExactScope::set(&flag, true);
1424 drop(scope);
1425 assert!(!flag.load(Ordering::Relaxed));
1426 }
1427}
1428
1429impl Engine {
1430 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
1431 let gpu = memra_runtime::Gpu::new(ordinal)?;
1432 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
1436 use cudarc::driver::sys::CUdevice_attribute_enum as A;
1437 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
1438 .and_then(|d| unsafe {
1439 Ok((
1440 cudarc::driver::result::device::get_attribute(
1441 d,
1442 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
1443 )?,
1444 cudarc::driver::result::device::get_attribute(
1445 d,
1446 A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
1447 )?,
1448 ))
1449 })
1450 .unwrap_or((0, 0));
1451 let built = env!("MEMRA_BUILT_CUDA_ARCH");
1452 let ok = matches!(
1453 (built, maj, min),
1454 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9)
1455 );
1456 if !ok {
1457 return Err(format!(
1458 "memra was built for sm_{built} but device {ordinal} reports compute \
1459 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
1460 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass."
1461 )
1462 .into());
1463 }
1464 }
1465 unsafe {
1470 use cudarc::driver::sys;
1471 let dev: sys::CUdevice = ordinal as sys::CUdevice;
1472 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1473 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
1474 let mut thresh: u64 = u64::MAX;
1475 let _ = sys::cuMemPoolSetAttribute(
1476 pool,
1477 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
1478 &mut thresh as *mut u64 as *mut core::ffi::c_void,
1479 );
1480 }
1481 }
1482 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
1483 let hybrid = gpu
1484 .ctx
1485 .load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
1486 let qmatvec = gpu
1487 .ctx
1488 .load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
1489 let flash = gpu
1490 .ctx
1491 .load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
1492 let gemm = gpu
1493 .ctx
1494 .load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
1495 let router = gpu
1496 .ctx
1497 .load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
1498 let sample = gpu
1499 .ctx
1500 .load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
1501 let copy_stream = gpu.ctx.new_stream()?;
1502 if std::env::var("MEMRA_EVT")
1518 .map(|v| v == "1")
1519 .unwrap_or(false)
1520 {
1521 } else {
1523 unsafe {
1524 gpu.ctx.disable_event_tracking();
1525 }
1526 }
1527 Ok(Self {
1528 gpu,
1529 module,
1530 hybrid,
1531 qmatvec,
1532 flash,
1533 flash_g: std::sync::OnceLock::new(),
1534 gemm,
1535 router,
1536 sample,
1537 moe_cache: Mutex::new(None),
1538 w8_mirrors: Mutex::new(std::collections::HashMap::new()),
1539 w8_act: Mutex::new(std::collections::HashMap::new()),
1540 moe_cache_layout: Mutex::new(None),
1541 copy_stream,
1542 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
1543 verify_exact: std::sync::atomic::AtomicBool::new(false),
1544 capture_keep: Mutex::new(Vec::new()),
1545 argmax_partials: Mutex::new(None),
1546 prime_deqw_ws: Mutex::new(None),
1547 router_stage: Mutex::new(None),
1548 fp8_scratch: Mutex::new(None),
1549 fa_vf16_scratch: Mutex::new(None),
1550 fa_part_pool: Mutex::new(None),
1551 fa_part_retired: Mutex::new(Vec::new()),
1552 fn_cache: Mutex::new(Default::default()),
1553 f16_scratch: Mutex::new(None),
1554 #[cfg(memra_cutlass)]
1555 cutlass_scratch: Mutex::new(None),
1556 })
1557 }
1558
1559 pub fn ctx(&self) -> &Arc<CudaContext> {
1560 &self.gpu.ctx
1561 }
1562
1563 pub fn pool_cached_bytes(&self) -> usize {
1581 let (reserved, used) = self.pool_reserved_used();
1582 reserved.saturating_sub(used)
1583 }
1584
1585 pub fn device_graph_mem_reserved(&self) -> usize {
1595 use cudarc::driver::sys as cus;
1596 let Ok(dev) = cudarc::driver::result::device::get(self.gpu.ctx.ordinal() as i32) else {
1597 return 0;
1598 };
1599 let mut bytes: u64 = 0;
1600 let rc = unsafe {
1601 cus::cuDeviceGetGraphMemAttribute(
1602 dev,
1603 cus::CUgraphMem_attribute::CU_GRAPH_MEM_ATTR_RESERVED_MEM_CURRENT,
1604 &mut bytes as *mut u64 as *mut std::ffi::c_void,
1605 )
1606 };
1607 if rc == cus::cudaError_enum::CUDA_SUCCESS {
1608 bytes as usize
1609 } else {
1610 0
1611 }
1612 }
1613
1614 pub fn pool_trim_to_zero(&self) -> usize {
1628 use cudarc::driver::sys;
1629 let (before, _) = self.pool_reserved_used();
1630 unsafe {
1631 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1632 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1633 != sys::CUresult::CUDA_SUCCESS
1634 {
1635 return 0;
1636 }
1637 let _ = sys::cuMemPoolTrimTo(pool, 0);
1638 }
1639 let (after, _) = self.pool_reserved_used();
1640 before.saturating_sub(after)
1641 }
1642
1643 pub fn pool_reserved_used(&self) -> (usize, usize) {
1644 use cudarc::driver::sys;
1645 unsafe {
1646 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
1647 if sys::cuDeviceGetDefaultMemPool(&mut pool, self.gpu.ctx.ordinal() as sys::CUdevice)
1648 != sys::CUresult::CUDA_SUCCESS
1649 {
1650 return (0, 0);
1651 }
1652 let (mut reserved, mut used) = (0u64, 0u64);
1653 if sys::cuMemPoolGetAttribute(
1654 pool,
1655 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
1656 &mut reserved as *mut u64 as *mut core::ffi::c_void,
1657 ) != sys::CUresult::CUDA_SUCCESS
1658 {
1659 return (0, 0);
1660 }
1661 if sys::cuMemPoolGetAttribute(
1662 pool,
1663 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
1664 &mut used as *mut u64 as *mut core::ffi::c_void,
1665 ) != sys::CUresult::CUDA_SUCCESS
1666 {
1667 return (0, 0);
1668 }
1669 (reserved as usize, used as usize)
1670 }
1671 }
1672
1673 pub fn stream(&self) -> Arc<CudaStream> {
1676 self.gpu.stream()
1677 }
1678 pub fn gkv_on() -> bool {
1681 memra_kv::gkv_on()
1682 }
1683
1684 pub fn wkv_on() -> bool {
1696 memra_kv::wkv_on()
1697 }
1698
1699 pub fn kv_fp8_on() -> bool {
1705 memra_kv::kv_fp8_on()
1706 }
1707
1708 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1711 if head_dim == 512 && Self::gkv_on() {
1712 self.func_g(name)
1713 } else {
1714 self.func(name)
1715 }
1716 }
1717
1718 fn func_g(&self, name: &str) -> CudaFunction {
1722 let m = self.flash_g.get_or_init(|| {
1723 self.gpu
1724 .ctx
1725 .load_module(cudarc::nvrtc::Ptx::from_binary(
1726 FLASH_FATBIN_KF8VF8.to_vec(),
1727 ))
1728 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1729 });
1730 let key = format!("g:{name}");
1731 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) {
1732 return f.clone();
1733 }
1734 let f = match m.load_function(name) {
1735 Ok(f) => f,
1736 Err(_) => self.func(name),
1737 };
1738 self.fn_cache.lock().unwrap().insert(key, f.clone());
1739 f
1740 }
1741
1742 fn func(&self, name: &str) -> CudaFunction {
1743 if let Some(f) = self.fn_cache.lock().unwrap().get(name) {
1746 return f.clone();
1747 }
1748 let f = self
1749 .module
1750 .load_function(name)
1751 .or_else(|_| self.hybrid.load_function(name))
1752 .or_else(|_| self.qmatvec.load_function(name))
1753 .or_else(|_| self.flash.load_function(name))
1754 .or_else(|_| self.gemm.load_function(name))
1755 .or_else(|_| self.router.load_function(name))
1756 .or_else(|_| self.sample.load_function(name))
1757 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1758 self.fn_cache
1759 .lock()
1760 .unwrap()
1761 .insert(name.to_string(), f.clone());
1762 f
1763 }
1764
1765 pub fn scatter_trim_logits(
1768 &self,
1769 src: &CudaSlice<f32>,
1770 d2t: &CudaSlice<u32>,
1771 dst: &mut CudaSlice<f32>,
1772 d_vocab: usize,
1773 n_vocab: usize,
1774 ) -> Result<(), Box<dyn std::error::Error>> {
1775 let f1 = self.func("scatter_trim_logits_f32");
1776 let f2 = self.func("scatter_trim_logits_pass2_f32");
1777 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1778 let cfg1 = LaunchConfig {
1779 grid_dim: (256, 1, 1),
1780 block_dim: (256, 1, 1),
1781 shared_mem_bytes: 0,
1782 };
1783 let __s_b1 = self.gpu.stream();
1784 let mut b1 = __s_b1.launch_builder(&f1);
1785 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1786 unsafe {
1787 b1.launch(cfg1)?;
1788 }
1789 let cfg2 = LaunchConfig {
1790 grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1),
1791 block_dim: (256, 1, 1),
1792 shared_mem_bytes: 0,
1793 };
1794 let __s_b2 = self.gpu.stream();
1795 let mut b2 = __s_b2.launch_builder(&f2);
1796 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1797 unsafe {
1798 b2.launch(cfg2)?;
1799 }
1800 Ok(())
1801 }
1802
1803 #[allow(clippy::too_many_arguments)]
1809 pub fn filter_stats(
1810 &self,
1811 x: &CudaSlice<f32>,
1812 row_stride: usize,
1813 rows: &CudaSlice<i32>,
1814 out_th: &mut CudaSlice<f32>,
1815 out_z: &mut CudaSlice<f32>,
1816 out_max: &mut CudaSlice<f32>,
1817 n: usize,
1818 nrow: usize,
1819 temp: f32,
1820 top_k: i32,
1821 top_p: f32,
1822 min_p: f32,
1823 ) -> Result<(), Box<dyn std::error::Error>> {
1824 static COOP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1851 let coop_on =
1852 *COOP_ON.get_or_init(|| std::env::var("MEMRA_FILTER_COOP").as_deref() != Ok("0"));
1853 if coop_on && self.sm_count() >= 16 {
1854 let cap = self.sm_count() as usize / 16;
1855 let mut done = 0usize;
1856 while done < nrow {
1857 let chunk = cap.min(nrow - done);
1858 self.filter_stats_coop_chunk(
1859 x, row_stride, rows, done, out_th, out_z, out_max, n, chunk, temp, top_k,
1860 top_p, min_p,
1861 )?;
1862 done += chunk;
1863 }
1864 return Ok(());
1865 }
1866 self.filter_stats_plain_program(
1867 x, row_stride, rows, out_th, out_z, out_max, n, nrow, temp, top_k, top_p, min_p,
1868 )
1869 }
1870
1871 #[allow(clippy::too_many_arguments)]
1876 pub fn filter_stats_coop_chunk(
1877 &self,
1878 x: &CudaSlice<f32>,
1879 row_stride: usize,
1880 rows: &CudaSlice<i32>,
1881 row0: usize,
1882 out_th: &mut CudaSlice<f32>,
1883 out_z: &mut CudaSlice<f32>,
1884 out_max: &mut CudaSlice<f32>,
1885 n: usize,
1886 chunk: usize,
1887 temp: f32,
1888 top_k: i32,
1889 top_p: f32,
1890 min_p: f32,
1891 ) -> Result<(), Box<dyn std::error::Error>> {
1892 let (ni, nr, rs) = (n as i32, chunk as i32, row_stride as i64);
1893 let f = self.func("filter_stats_coop_f32");
1894 let mut ws = self.alloc_uninit::<f32>(chunk * (2 * 16 + 2))?;
1895 let cfg = LaunchConfig {
1896 grid_dim: (16, chunk as u32, 1),
1897 block_dim: (512, 1, 1),
1898 shared_mem_bytes: 0,
1899 };
1900 let rows_v = rows.slice(row0..row0 + chunk);
1901 let mut th_v = out_th.slice_mut(row0..row0 + chunk);
1902 let mut z_v = out_z.slice_mut(row0..row0 + chunk);
1903 let mut mx_v = out_max.slice_mut(row0..row0 + chunk);
1904 let __s_b = self.gpu.stream();
1905 let mut b = __s_b.launch_builder(&f);
1906 b.arg(x)
1907 .arg(&rs)
1908 .arg(&rows_v)
1909 .arg(&mut th_v)
1910 .arg(&mut z_v)
1911 .arg(&mut mx_v)
1912 .arg(&mut ws)
1913 .arg(&ni)
1914 .arg(&nr)
1915 .arg(&temp)
1916 .arg(&top_k)
1917 .arg(&top_p)
1918 .arg(&min_p);
1919 unsafe {
1920 b.launch_cooperative(cfg)?;
1921 }
1922 Ok(())
1923 }
1924
1925 #[allow(clippy::too_many_arguments)]
1929 pub fn filter_stats_plain_program(
1930 &self,
1931 x: &CudaSlice<f32>,
1932 row_stride: usize,
1933 rows: &CudaSlice<i32>,
1934 out_th: &mut CudaSlice<f32>,
1935 out_z: &mut CudaSlice<f32>,
1936 out_max: &mut CudaSlice<f32>,
1937 n: usize,
1938 nrow: usize,
1939 temp: f32,
1940 top_k: i32,
1941 top_p: f32,
1942 min_p: f32,
1943 ) -> Result<(), Box<dyn std::error::Error>> {
1944 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1945 let f = self.func("filter_stats_f32");
1946 let cfg = LaunchConfig {
1947 grid_dim: (nrow as u32, 1, 1),
1948 block_dim: (1024, 1, 1),
1949 shared_mem_bytes: 0,
1950 };
1951 let __s_b = self.gpu.stream();
1952 let mut b = __s_b.launch_builder(&f);
1953 b.arg(x)
1954 .arg(&rs)
1955 .arg(rows)
1956 .arg(&mut *out_th)
1957 .arg(&mut *out_z)
1958 .arg(&mut *out_max)
1959 .arg(&ni)
1960 .arg(&nr)
1961 .arg(&temp)
1962 .arg(&top_k)
1963 .arg(&top_p)
1964 .arg(&min_p);
1965 unsafe {
1966 b.launch(cfg)?;
1967 }
1968 Ok(())
1969 }
1970
1971 #[allow(clippy::too_many_arguments)]
1973 pub fn softmax_gather_filtered(
1974 &self,
1975 x: &CudaSlice<f32>,
1976 row_stride: usize,
1977 ids: &CudaSlice<u32>,
1978 rows: &CudaSlice<i32>,
1979 th: &CudaSlice<f32>,
1980 z: &CudaSlice<f32>,
1981 out: &mut CudaSlice<f32>,
1982 n: usize,
1983 npair: usize,
1984 temp: f32,
1985 ) -> Result<(), Box<dyn std::error::Error>> {
1986 let f = self.func("softmax_gather_filtered_f32");
1987 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1988 let cfg = LaunchConfig {
1989 grid_dim: (npair as u32, 1, 1),
1990 block_dim: (256, 1, 1),
1991 shared_mem_bytes: 0,
1992 };
1993 let __s_b = self.gpu.stream();
1994 let mut b = __s_b.launch_builder(&f);
1995 b.arg(x)
1996 .arg(&rs)
1997 .arg(ids)
1998 .arg(rows)
1999 .arg(th)
2000 .arg(z)
2001 .arg(&mut *out)
2002 .arg(&ni)
2003 .arg(&np)
2004 .arg(&temp);
2005 unsafe {
2006 b.launch(cfg)?;
2007 }
2008 Ok(())
2009 }
2010
2011 #[allow(clippy::too_many_arguments)]
2013 pub fn residual_sample_filtered(
2014 &self,
2015 p: &CudaSlice<f32>,
2016 q: Option<&CudaSlice<f32>>,
2017 n: usize,
2018 temp: f32,
2019 seed: u64,
2020 stream_pos: u32,
2021 p_stats: (f32, f32, f32),
2022 q_stats: (f32, f32, f32),
2023 out_tok: &mut CudaSlice<u32>,
2024 ) -> Result<(), Box<dyn std::error::Error>> {
2025 let f = self.func("residual_sample_filtered_f32");
2026 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2027 let has_q: i32 = q.is_some() as i32;
2028 let qbuf = q.unwrap_or(p);
2029 let (pm, pth, pz) = p_stats;
2030 let (qm, qth, qz) = q_stats;
2031 let cfg = LaunchConfig {
2032 grid_dim: (1, 1, 1),
2033 block_dim: (1024, 1, 1),
2034 shared_mem_bytes: 0,
2035 };
2036 let __s_b = self.gpu.stream();
2037 let mut b = __s_b.launch_builder(&f);
2038 b.arg(p)
2039 .arg(qbuf)
2040 .arg(&has_q)
2041 .arg(&ni)
2042 .arg(&temp)
2043 .arg(&slo)
2044 .arg(&shi)
2045 .arg(&stream_pos)
2046 .arg(&pm)
2047 .arg(&pth)
2048 .arg(&pz)
2049 .arg(&qm)
2050 .arg(&qth)
2051 .arg(&qz)
2052 .arg(&mut *out_tok);
2053 unsafe {
2054 b.launch(cfg)?;
2055 }
2056 Ok(())
2057 }
2058
2059 #[allow(clippy::too_many_arguments)]
2065 pub fn residual_sample_sparse_q(
2066 &self,
2067 p: &CudaSlice<f32>,
2068 cand_ids: &CudaSlice<u32>,
2069 q_probs: &CudaSlice<f32>,
2070 n_cand: usize,
2071 n: usize,
2072 temp: f32,
2073 seed: u64,
2074 stream_pos: u32,
2075 p_stats: (f32, f32, f32),
2076 out_tok: &mut CudaSlice<u32>,
2077 ) -> Result<(), Box<dyn std::error::Error>> {
2078 assert!(
2079 n_cand >= 1 && n_cand <= 32,
2080 "residual_sample_sparse_q supports 1..=32 candidates, got {n_cand}"
2081 );
2082 let f = self.func("residual_sample_sparse_q_f32");
2083 let (ni, nc) = (n as i32, n_cand as i32);
2084 let (slo, shi) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2085 let (pm, pth, pz) = p_stats;
2086 let cfg = LaunchConfig {
2087 grid_dim: (1, 1, 1),
2088 block_dim: (1024, 1, 1),
2089 shared_mem_bytes: 0,
2090 };
2091 let __s_b = self.gpu.stream();
2092 let mut b = __s_b.launch_builder(&f);
2093 b.arg(p)
2094 .arg(cand_ids)
2095 .arg(q_probs)
2096 .arg(&nc)
2097 .arg(&ni)
2098 .arg(&temp)
2099 .arg(&slo)
2100 .arg(&shi)
2101 .arg(&stream_pos)
2102 .arg(&pm)
2103 .arg(&pth)
2104 .arg(&pz)
2105 .arg(&mut *out_tok);
2106 unsafe {
2107 b.launch(cfg)?;
2108 }
2109 Ok(())
2110 }
2111
2112 #[allow(clippy::too_many_arguments)]
2114 pub fn gumbel_perturb_filtered(
2115 &self,
2116 x: &CudaSlice<f32>,
2117 y: &mut CudaSlice<f32>,
2118 n: usize,
2119 seed: u64,
2120 stream_pos: u32,
2121 temp: f32,
2122 row_max: f32,
2123 th: f32,
2124 ) -> Result<(), Box<dyn std::error::Error>> {
2125 let f = self.func("gumbel_perturb_filtered_f32");
2126 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2127 let cfg = LaunchConfig {
2128 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2129 block_dim: (256, 1, 1),
2130 shared_mem_bytes: 0,
2131 };
2132 let __s_b = self.gpu.stream();
2133 let mut b = __s_b.launch_builder(&f);
2134 b.arg(x)
2135 .arg(&mut *y)
2136 .arg(&ni)
2137 .arg(&slo)
2138 .arg(&shi)
2139 .arg(&stream_pos)
2140 .arg(&temp)
2141 .arg(&row_max)
2142 .arg(&th);
2143 unsafe {
2144 b.launch(cfg)?;
2145 }
2146 Ok(())
2147 }
2148
2149 #[allow(clippy::too_many_arguments)]
2153 pub fn penalize_logits(
2154 &self,
2155 x: &mut CudaSlice<f32>,
2156 hist: &CudaSlice<u32>,
2157 n_hist: usize,
2158 rep: f32,
2159 freq: f32,
2160 present: f32,
2161 n: usize,
2162 ) -> Result<(), Box<dyn std::error::Error>> {
2163 if n_hist == 0 {
2164 return Ok(());
2165 }
2166 let f = self.func("penalize_logits_f32");
2167 let (nh, ni) = (n_hist as i32, n as i32);
2168 let cfg = LaunchConfig {
2169 grid_dim: (n_hist.div_ceil(128) as u32, 1, 1),
2170 block_dim: (128, 1, 1),
2171 shared_mem_bytes: 0,
2172 };
2173 let __s_b = self.gpu.stream();
2174 let mut b = __s_b.launch_builder(&f);
2175 b.arg(&mut *x)
2176 .arg(hist)
2177 .arg(&nh)
2178 .arg(&rep)
2179 .arg(&freq)
2180 .arg(&present)
2181 .arg(&ni);
2182 unsafe {
2183 b.launch(cfg)?;
2184 }
2185 Ok(())
2186 }
2187
2188 #[allow(clippy::too_many_arguments)]
2190 pub fn penalize_logits_rows(
2191 &self,
2192 x: &mut CudaSlice<f32>,
2193 hist: &CudaSlice<u32>,
2194 n_hist: usize,
2195 rep: f32,
2196 freq: f32,
2197 present: f32,
2198 n: usize,
2199 nrow: usize,
2200 ) -> Result<(), Box<dyn std::error::Error>> {
2201 if n_hist == 0 || nrow == 0 {
2202 return Ok(());
2203 }
2204 let f = self.func("penalize_logits_rows_f32");
2205 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
2206 let cfg = LaunchConfig {
2207 grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1),
2208 block_dim: (128, 1, 1),
2209 shared_mem_bytes: 0,
2210 };
2211 let __s_b = self.gpu.stream();
2212 let mut b = __s_b.launch_builder(&f);
2213 b.arg(&mut *x)
2214 .arg(hist)
2215 .arg(&nh)
2216 .arg(&rep)
2217 .arg(&freq)
2218 .arg(&present)
2219 .arg(&ni)
2220 .arg(&nr);
2221 unsafe {
2222 b.launch(cfg)?;
2223 }
2224 Ok(())
2225 }
2226
2227 #[allow(clippy::too_many_arguments)]
2233 pub fn penalize_logits_sparse_rows(
2234 &self,
2235 x: &mut CudaSlice<f32>,
2236 ids: &[u32],
2237 counts: &[u32],
2238 offsets: &[i32],
2239 rows: &[i32],
2240 reps: &[f32],
2241 freqs: &[f32],
2242 presents: &[f32],
2243 n: usize,
2244 ) -> Result<(), Box<dyn std::error::Error>> {
2245 let nrow = rows.len();
2246 if nrow == 0 {
2247 return Ok(());
2248 }
2249 let _ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2250 let _nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2251 let entry_count =
2252 i32::try_from(ids.len()).map_err(|_| "sparse penalty entry count must fit CUDA i32")?;
2253 if ids.len() != counts.len()
2254 || offsets.len() != nrow + 1
2255 || reps.len() != nrow
2256 || freqs.len() != nrow
2257 || presents.len() != nrow
2258 || offsets.first().copied() != Some(0)
2259 || offsets.last().copied() != Some(entry_count)
2260 {
2261 return Err("sparse penalty row metadata shape mismatch".into());
2262 }
2263 if counts.contains(&0) {
2264 return Err("sparse penalty counts must be positive".into());
2265 }
2266 let mut max_len = 0usize;
2267 for pair in offsets.windows(2) {
2268 if pair[0] < 0 || pair[1] < pair[0] {
2269 return Err("sparse penalty offsets must be monotonic".into());
2270 }
2271 max_len = max_len.max((pair[1] - pair[0]) as usize);
2272 }
2273 if max_len == 0 {
2274 return Ok(());
2275 }
2276
2277 let mut seen = std::collections::HashSet::with_capacity(ids.len());
2278 for (r, &row) in rows.iter().enumerate() {
2279 if row < 0 || (row as usize + 1).saturating_mul(n) > x.len() {
2280 return Err("sparse penalty row index exceeds logits shape".into());
2281 }
2282 let begin = offsets[r] as usize;
2283 let end = offsets[r + 1] as usize;
2284 for &id in &ids[begin..end] {
2285 if id as usize >= n {
2286 return Err("sparse penalty token id exceeds logits row".into());
2287 }
2288 if !seen.insert((row, id)) {
2289 return Err("sparse penalty entries must be unique per logits row".into());
2290 }
2291 }
2292 }
2293
2294 unsafe {
2296 self.penalize_logits_sparse_rows_unchecked(
2297 x, ids, counts, offsets, rows, reps, freqs, presents, n,
2298 )
2299 }
2300 }
2301
2302 #[allow(clippy::too_many_arguments)]
2311 pub(crate) unsafe fn penalize_logits_sparse_rows_unchecked(
2312 &self,
2313 x: &mut CudaSlice<f32>,
2314 ids: &[u32],
2315 counts: &[u32],
2316 offsets: &[i32],
2317 rows: &[i32],
2318 reps: &[f32],
2319 freqs: &[f32],
2320 presents: &[f32],
2321 n: usize,
2322 ) -> Result<(), Box<dyn std::error::Error>> {
2323 let nrow = rows.len();
2324 if nrow == 0 {
2325 return Ok(());
2326 }
2327 let max_len = offsets
2328 .windows(2)
2329 .map(|pair| (pair[1] - pair[0]) as usize)
2330 .max()
2331 .unwrap_or(0);
2332 if max_len == 0 {
2333 return Ok(());
2334 }
2335 let ids_d = self.htod_u32_v(ids)?;
2336 let counts_d = self.htod_u32_v(counts)?;
2337 let offsets_d = self.htod_i32(offsets)?;
2338 let rows_d = self.htod_i32(rows)?;
2339 let reps_d = self.htod(reps)?;
2340 let freqs_d = self.htod(freqs)?;
2341 let presents_d = self.htod(presents)?;
2342 let f = self.func("penalize_logits_sparse_rows_f32");
2343 let ni = i32::try_from(n).map_err(|_| "sparse penalty logits width must fit CUDA i32")?;
2344 let nr = i32::try_from(nrow).map_err(|_| "sparse penalty row count must fit CUDA i32")?;
2345 let cfg = LaunchConfig {
2346 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2347 block_dim: (128, 1, 1),
2348 shared_mem_bytes: 0,
2349 };
2350 let __s_b = self.gpu.stream();
2351 let mut b = __s_b.launch_builder(&f);
2352 b.arg(&mut *x)
2353 .arg(&ids_d)
2354 .arg(&counts_d)
2355 .arg(&offsets_d)
2356 .arg(&rows_d)
2357 .arg(&reps_d)
2358 .arg(&freqs_d)
2359 .arg(&presents_d)
2360 .arg(&ni)
2361 .arg(&nr);
2362 unsafe {
2363 b.launch(cfg)?;
2364 }
2365 Ok(())
2366 }
2367
2368 #[allow(clippy::too_many_arguments)]
2376 pub fn penalize_logits_rows_inc(
2377 &self,
2378 x: &mut CudaSlice<f32>,
2379 hist: &CudaSlice<u32>,
2380 n_hist0: usize,
2381 rep: f32,
2382 freq: f32,
2383 present: f32,
2384 n: usize,
2385 nrow: usize,
2386 win: usize,
2387 ) -> Result<(), Box<dyn std::error::Error>> {
2388 if nrow == 0 || win == 0 || (n_hist0 == 0 && nrow == 1) {
2389 return Ok(());
2390 }
2391 debug_assert!(
2392 hist.len() >= n_hist0 + nrow - 1,
2393 "rows-inc hist must carry n_hist0 + nrow - 1 ids"
2394 );
2395 let f = self.func("penalize_logits_rows_inc_f32");
2396 let max_len = win.min(n_hist0 + nrow - 1).max(1);
2397 let (nh, ni, nr, wi) = (n_hist0 as i32, n as i32, nrow as i32, win as i32);
2398 let cfg = LaunchConfig {
2399 grid_dim: (max_len.div_ceil(128) as u32, nrow as u32, 1),
2400 block_dim: (128, 1, 1),
2401 shared_mem_bytes: 0,
2402 };
2403 let __s_b = self.gpu.stream();
2404 let mut b = __s_b.launch_builder(&f);
2405 b.arg(&mut *x)
2406 .arg(hist)
2407 .arg(&nh)
2408 .arg(&rep)
2409 .arg(&freq)
2410 .arg(&present)
2411 .arg(&ni)
2412 .arg(&nr)
2413 .arg(&wi);
2414 unsafe {
2415 b.launch(cfg)?;
2416 }
2417 Ok(())
2418 }
2419
2420 pub fn wpf_level() -> u32 {
2428 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
2429 *ON.get_or_init(|| {
2430 std::env::var("MEMRA_WPF")
2431 .ok()
2432 .and_then(|v| v.parse().ok())
2433 .unwrap_or(1)
2434 })
2435 }
2436
2437 pub fn set_verify_exact(&self, on: bool) {
2453 self.verify_exact
2454 .store(on, std::sync::atomic::Ordering::Relaxed);
2455 }
2456 pub(crate) fn verify_exact_on(&self) -> bool {
2457 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
2458 }
2459
2460 pub fn exact_scope(&self, on: bool) -> ExactScope<'_> {
2466 ExactScope::set(&self.verify_exact, on)
2467 }
2468
2469 pub fn qkv_append_on() -> bool {
2472 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2473 *ON.get_or_init(|| {
2474 std::env::var("MEMRA_QKV_APPEND")
2475 .map(|v| v != "0")
2476 .unwrap_or(true)
2477 })
2478 }
2479
2480 pub fn pdl_wb_on() -> bool {
2483 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2484 *ON.get_or_init(|| {
2485 std::env::var("MEMRA_PDL_WB")
2486 .map(|v| v != "0")
2487 .unwrap_or(true)
2488 })
2489 }
2490
2491 pub fn norm_ilp_on() -> bool {
2499 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2500 *ON.get_or_init(|| {
2501 std::env::var("MEMRA_NORM_ILP")
2502 .map(|v| v != "0")
2503 .unwrap_or(true)
2504 })
2505 }
2506
2507 pub fn tk_ffn_dual_on() -> bool {
2515 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2516 *ON.get_or_init(|| {
2517 std::env::var("MEMRA_TK_FFN_DUAL")
2518 .map(|v| v != "0")
2519 .unwrap_or(true)
2520 })
2521 }
2522
2523 pub fn pdl_mmvq_on() -> bool {
2527 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2528 *ON.get_or_init(|| {
2529 std::env::var("MEMRA_PDL_MMVQ")
2530 .map(|v| v != "0")
2531 .unwrap_or(true)
2532 })
2533 }
2534
2535 pub fn pdl_on() -> bool {
2536 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2537 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
2538 }
2539
2540 pub fn pdl_nvfp4q8_on() -> bool {
2546 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2547 *ON.get_or_init(|| {
2548 std::env::var("MEMRA_PDL_NVFP4")
2549 .map(|v| v != "0")
2550 .unwrap_or(true)
2551 })
2552 }
2553
2554 fn q40_mr1_on() -> bool {
2560 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
2561 match *Q40MR.get_or_init(|| {
2562 std::env::var("MEMRA_Q40_MR")
2563 .ok()
2564 .and_then(|v| v.parse().ok())
2565 }) {
2566 Some(v) => v == 1,
2567 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
2568 }
2569 }
2570
2571 fn pdl_func_flash(
2576 &self,
2577 g: bool,
2578 name: &'static str,
2579 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2580 use cudarc::driver::sys as cu;
2581 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
2588 std::sync::Mutex::new(None);
2589 static FNS: std::sync::Mutex<
2590 Option<std::collections::HashMap<(usize, bool, &'static str), usize>>,
2591 > = std::sync::Mutex::new(None);
2592 let ctx_key = self.ctx().cu_ctx() as usize;
2593 if let Some(&f) = FNS
2594 .lock()
2595 .unwrap()
2596 .get_or_insert_with(Default::default)
2597 .get(&(ctx_key, g, name))
2598 {
2599 return Ok(f as cu::CUfunction);
2600 }
2601 let module = {
2602 let mut mods = MODS.lock().unwrap();
2603 let map = mods.get_or_insert_with(Default::default);
2604 match map.get(&(ctx_key, g)) {
2605 Some(&m) => m,
2606 None => {
2607 let m = self.pdl_load_module_in_ctx(if g {
2608 FLASH_FATBIN_KF8VF8
2609 } else {
2610 FLASH_FATBIN
2611 })?;
2612 map.insert((ctx_key, g), m);
2613 m
2614 }
2615 }
2616 };
2617 let cname = std::ffi::CString::new(name)?;
2618 let mut f: cu::CUfunction = std::ptr::null_mut();
2619 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2620 if r != cu::CUresult::CUDA_SUCCESS {
2621 return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into());
2622 }
2623 FNS.lock()
2624 .unwrap()
2625 .get_or_insert_with(Default::default)
2626 .insert((ctx_key, g, name), f as usize);
2627 Ok(f)
2628 }
2629
2630 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
2635 use cudarc::driver::sys as cu;
2636 let mut prev: cu::CUcontext = std::ptr::null_mut();
2637 unsafe {
2638 cu::cuCtxGetCurrent(&mut prev).result()?;
2639 }
2640 self.ctx().bind_to_thread()?;
2641 let mut m: cu::CUmodule = std::ptr::null_mut();
2642 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
2643 let restore = if prev.is_null() {
2644 cu::CUresult::CUDA_SUCCESS
2645 } else {
2646 unsafe { cu::cuCtxSetCurrent(prev) }
2647 };
2648 if r != cu::CUresult::CUDA_SUCCESS {
2649 return Err(format!("pdl module load: {r:?}").into());
2650 }
2651 if restore != cu::CUresult::CUDA_SUCCESS {
2652 return Err(format!("pdl module load: ctx restore {restore:?}").into());
2653 }
2654 Ok(m as usize)
2655 }
2656
2657 pub fn raw_kernel_function(
2660 &self,
2661 name: &'static str,
2662 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2663 self.pdl_func(name)
2664 }
2665
2666 fn pdl_func(
2667 &self,
2668 name: &'static str,
2669 ) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
2670 use cudarc::driver::sys as cu;
2671 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2674 std::sync::Mutex::new(None);
2675 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
2678 std::sync::Mutex::new(None);
2679 static FNS: std::sync::Mutex<
2680 Option<std::collections::HashMap<(usize, &'static str), usize>>,
2681 > = std::sync::Mutex::new(None);
2682 let ctx_key = self.ctx().cu_ctx() as usize;
2683 if let Some(&f) = FNS
2684 .lock()
2685 .unwrap()
2686 .get_or_insert_with(Default::default)
2687 .get(&(ctx_key, name))
2688 {
2689 return Ok(f as cu::CUfunction);
2690 }
2691 let module = {
2692 let mut mods = MODULES.lock().unwrap();
2693 let map = mods.get_or_insert_with(Default::default);
2694 match map.get(&ctx_key) {
2695 Some(&m) => m,
2696 None => {
2697 let m = self.pdl_load_module_in_ctx(FATBIN)?;
2698 map.insert(ctx_key, m);
2699 m
2700 }
2701 }
2702 };
2703 let cname = std::ffi::CString::new(name)?;
2704 let mut f: cu::CUfunction = std::ptr::null_mut();
2705 let mut r =
2706 unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
2707 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
2708 let qmodule = {
2709 let mut mods = QMODULES.lock().unwrap();
2710 let map = mods.get_or_insert_with(Default::default);
2711 match map.get(&ctx_key) {
2712 Some(&m) => m,
2713 None => {
2714 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
2715 map.insert(ctx_key, m);
2716 m
2717 }
2718 }
2719 };
2720 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
2721 }
2722 if r != cu::CUresult::CUDA_SUCCESS {
2723 return Err(format!("pdl_func {name}: {r:?}").into());
2724 }
2725 FNS.lock()
2726 .unwrap()
2727 .get_or_insert_with(Default::default)
2728 .insert((ctx_key, name), f as usize);
2729 Ok(f)
2730 }
2731
2732 unsafe fn launch_pdl_flash(
2744 &self,
2745 g: bool,
2746 name: &'static str,
2747 grid: (u32, u32, u32),
2748 block: (u32, u32, u32),
2749 smem: u32,
2750 params: &mut [*mut std::ffi::c_void],
2751 ) -> Result<(), Box<dyn std::error::Error>> {
2752 use cudarc::driver::sys as cu;
2753 let f = self.pdl_func_flash(g, name)?;
2754 if smem > 0 {
2755 let r =
2757 unsafe {
2758 cu::cuFuncSetAttribute(f,
2759 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
2760 smem as i32)
2761 };
2762 if r != cu::CUresult::CUDA_SUCCESS {
2763 return Err(format!("pdl smem attr {name}: {r:?}").into());
2764 }
2765 }
2766 let mut attr = cu::CUlaunchAttribute {
2767 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2768 pad: [0; 4],
2769 value: cu::CUlaunchAttributeValue {
2770 programmaticStreamSerializationAllowed: 1,
2771 },
2772 };
2773 let cfg = cu::CUlaunchConfig {
2774 gridDimX: grid.0,
2775 gridDimY: grid.1,
2776 gridDimZ: grid.2,
2777 blockDimX: block.0,
2778 blockDimY: block.1,
2779 blockDimZ: block.2,
2780 sharedMemBytes: smem,
2781 hStream: self.gpu.stream().cu_stream(),
2782 attrs: &mut attr,
2783 numAttrs: 1,
2784 };
2785 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2786 if r != cu::CUresult::CUDA_SUCCESS {
2787 return Err(format!("launch_pdl_flash {name}: {r:?}").into());
2788 }
2789 Ok(())
2790 }
2791
2792 unsafe fn launch_pdl(
2793 &self,
2794 name: &'static str,
2795 grid: (u32, u32, u32),
2796 block: (u32, u32, u32),
2797 params: &mut [*mut std::ffi::c_void],
2798 ) -> Result<(), Box<dyn std::error::Error>> {
2799 use cudarc::driver::sys as cu;
2800 let f = self.pdl_func(name)?;
2801 let mut attr = cu::CUlaunchAttribute {
2802 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
2803 pad: [0; 4],
2804 value: cu::CUlaunchAttributeValue {
2805 programmaticStreamSerializationAllowed: 1,
2806 },
2807 };
2808 let cfg = cu::CUlaunchConfig {
2809 gridDimX: grid.0,
2810 gridDimY: grid.1,
2811 gridDimZ: grid.2,
2812 blockDimX: block.0,
2813 blockDimY: block.1,
2814 blockDimZ: block.2,
2815 sharedMemBytes: 0,
2816 hStream: self.gpu.stream().cu_stream(),
2817 attrs: &mut attr,
2818 numAttrs: 1,
2819 };
2820 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
2821 if r != cu::CUresult::CUDA_SUCCESS {
2822 return Err(format!("launch_pdl {name}: {r:?}").into());
2823 }
2824 Ok(())
2825 }
2826
2827 pub fn prefetch_weight_l2(
2830 &self,
2831 w: &crate::model::GpuTensor,
2832 ) -> Result<(), Box<dyn std::error::Error>> {
2833 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
2834 let p = rp4.as_ref().unwrap_or(bytes);
2835 self.prefetch_l2(p, p.len())?;
2836 }
2837 Ok(())
2838 }
2839
2840 pub fn gather_row_bf16(
2843 &self,
2844 table: &CudaSlice<u8>,
2845 tok: &CudaSlice<u32>,
2846 idx: usize,
2847 dst: &mut CudaSlice<f32>,
2848 ncols: usize,
2849 ) -> Result<(), Box<dyn std::error::Error>> {
2850 let f = self.func("gather_row_bf16_f32");
2851 let cfg = LaunchConfig {
2852 grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
2853 block_dim: (256, 1, 1),
2854 shared_mem_bytes: 0,
2855 };
2856 let (nc, ix) = (ncols as i32, idx as i32);
2857 let __s_b = self.gpu.stream();
2858 let mut b = __s_b.launch_builder(&f);
2859 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
2860 unsafe {
2861 b.launch(cfg)?;
2862 }
2863 Ok(())
2864 }
2865
2866 #[allow(clippy::too_many_arguments)]
2872 pub fn dflash2_dynconv(
2873 &self,
2874 x: &CudaSlice<f32>,
2875 dyn_: &CudaSlice<f32>,
2876 base: &CudaSlice<f32>,
2877 out: &mut CudaSlice<f32>,
2878 rows: usize,
2879 hidden: usize,
2880 group_size: usize,
2881 ksize: usize,
2882 half: usize,
2883 ) -> Result<(), Box<dyn std::error::Error>> {
2884 assert_eq!(hidden % group_size, 0, "hidden % group_size != 0");
2885 let f = self.func("dflash2_dynconv_f32");
2886 let n = rows * hidden;
2887 let cfg = LaunchConfig {
2888 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2889 block_dim: (256, 1, 1),
2890 shared_mem_bytes: 0,
2891 };
2892 let (ri, hi, gi, ki, hf) = (
2893 rows as i32,
2894 hidden as i32,
2895 group_size as i32,
2896 ksize as i32,
2897 half as i32,
2898 );
2899 let __s_b = self.gpu.stream();
2900 let mut b = __s_b.launch_builder(&f);
2901 b.arg(x)
2902 .arg(dyn_)
2903 .arg(base)
2904 .arg(out)
2905 .arg(&ri)
2906 .arg(&hi)
2907 .arg(&gi)
2908 .arg(&ki)
2909 .arg(&hf);
2910 unsafe {
2911 b.launch(cfg)?;
2912 }
2913 Ok(())
2914 }
2915
2916 pub fn topk_rows(
2920 &self,
2921 logits: &CudaSlice<f32>,
2922 n_rows: usize,
2923 n_cols: usize,
2924 k: usize,
2925 ) -> Result<(CudaSlice<f32>, CudaSlice<u32>), Box<dyn std::error::Error>> {
2926 assert!(k <= 32 && k >= 1, "topk_rows supports 1..=32, got {k}");
2927 assert!(k <= n_cols, "topk_rows: k {k} > n_cols {n_cols}");
2928 let f = self.func("topk_rows_f32");
2929 let nth = 256usize;
2930 let mut vals = self.uninit(n_rows * k)?;
2931 let mut idxs = self.gpu.stream().alloc_zeros::<u32>(n_rows * k)?;
2932 let cfg = LaunchConfig {
2933 grid_dim: (n_rows as u32, 1, 1),
2934 block_dim: (nth as u32, 1, 1),
2935 shared_mem_bytes: (nth * k * 8) as u32,
2936 };
2937 let (nr, nc, ki) = (n_rows as i32, n_cols as i32, k as i32);
2938 let __s_b = self.gpu.stream();
2939 let mut b = __s_b.launch_builder(&f);
2940 b.arg(logits)
2941 .arg(&nr)
2942 .arg(&nc)
2943 .arg(&ki)
2944 .arg(&mut vals)
2945 .arg(&mut idxs);
2946 unsafe {
2947 b.launch(cfg)?;
2948 }
2949 Ok((vals, idxs))
2950 }
2951
2952 pub fn add_row_inplace(
2954 &self,
2955 logits: &mut CudaSlice<f32>,
2956 bias: &CudaSlice<f32>,
2957 n: usize,
2958 row_off: usize,
2959 ) -> Result<(), Box<dyn std::error::Error>> {
2960 let f = self.func("add_row_inplace_f32");
2961 let cfg = LaunchConfig {
2962 grid_dim: (n.div_ceil(256) as u32, 1, 1),
2963 block_dim: (256, 1, 1),
2964 shared_mem_bytes: 0,
2965 };
2966 let (ni, off) = (n as i32, row_off as i64);
2967 let __s_b = self.gpu.stream();
2968 let mut b = __s_b.launch_builder(&f);
2969 b.arg(logits).arg(bias).arg(&ni).arg(&off);
2970 unsafe {
2971 b.launch(cfg)?;
2972 }
2973 Ok(())
2974 }
2975
2976 pub fn prefetch_l2(
2978 &self,
2979 p: &CudaSlice<u8>,
2980 n: usize,
2981 ) -> Result<(), Box<dyn std::error::Error>> {
2982 let f = self.func("prefetch_l2_bytes");
2983 let lines = n.div_ceil(128);
2984 let ni = n as i64;
2985 let cfg = LaunchConfig {
2986 grid_dim: (lines.div_ceil(256) as u32, 1, 1),
2987 block_dim: (256, 1, 1),
2988 shared_mem_bytes: 0,
2989 };
2990 let __s_b = self.gpu.stream();
2991 let mut b = __s_b.launch_builder(&f);
2992 b.arg(p).arg(&ni);
2993 unsafe {
2994 b.launch(cfg)?;
2995 }
2996 Ok(())
2997 }
2998
2999 pub fn router_gemv(
3002 &self,
3003 w: &CudaSlice<f32>,
3004 x: &CudaSlice<f32>,
3005 n_embd: usize,
3006 n_experts: usize,
3007 t: usize,
3008 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3009 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
3015 Ok("0") => false,
3016 Ok(_) => true,
3017 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3018 };
3019 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
3029 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
3030 }
3031
3032 pub fn router_gemv_form(
3035 &self,
3036 w: &CudaSlice<f32>,
3037 x: &CudaSlice<f32>,
3038 n_embd: usize,
3039 n_experts: usize,
3040 t: usize,
3041 w8: bool,
3042 batch: bool,
3043 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3044 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
3045 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
3046 let f = if batch {
3047 self.func("router_gemv_f32_w8_batch")
3048 } else if w8 {
3049 self.func("router_gemv_f32_w8")
3050 } else {
3051 self.func("router_gemv_f32")
3052 };
3053 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3054 let cfg = if batch {
3055 LaunchConfig {
3056 grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
3057 block_dim: (32, 8, 1),
3058 shared_mem_bytes: 0,
3059 }
3060 } else {
3061 LaunchConfig {
3062 grid_dim: (n_experts as u32, t as u32, 1),
3063 block_dim: (32, if w8 { 8 } else { 1 }, 1),
3064 shared_mem_bytes: 0,
3065 }
3066 };
3067 let __s_b = self.gpu.stream();
3068 let mut b = __s_b.launch_builder(&f);
3069 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
3070 unsafe {
3071 b.launch(cfg)?;
3072 }
3073 Ok(y)
3074 }
3075
3076 pub fn router_gemv_into(
3079 &self,
3080 w: &CudaSlice<f32>,
3081 x: &CudaSlice<f32>,
3082 y: &mut CudaSlice<f32>,
3083 n_embd: usize,
3084 n_experts: usize,
3085 t: usize,
3086 ) -> Result<(), Box<dyn std::error::Error>> {
3087 if y.len() < t * n_experts {
3088 return Err("router_gemv_into output too small".into());
3089 }
3090 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
3091 Ok("0") => false,
3092 Ok(_) => true,
3093 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
3094 };
3095 let f = if w8 {
3096 self.func("router_gemv_f32_w8")
3097 } else {
3098 self.func("router_gemv_f32")
3099 };
3100 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
3101 let cfg = LaunchConfig {
3102 grid_dim: (n_experts as u32, t as u32, 1),
3103 block_dim: (32, if w8 { 8 } else { 1 }, 1),
3104 shared_mem_bytes: 0,
3105 };
3106 let __s_b = self.gpu.stream();
3107 let mut b = __s_b.launch_builder(&f);
3108 b.arg(w).arg(x).arg(&mut *y).arg(&ne).arg(&nx).arg(&ti);
3109 unsafe {
3110 b.launch(cfg)?;
3111 }
3112 Ok(())
3113 }
3114
3115 pub fn rows_permute(
3117 &self,
3118 src: &CudaSlice<f32>,
3119 idx: &CudaSlice<i32>,
3120 nrows: usize,
3121 ncols: usize,
3122 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3123 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
3124 let f = self.func("rows_permute_f32");
3125 let (nc, nr) = (ncols as i32, nrows as i32);
3126 let cfg = LaunchConfig {
3127 grid_dim: (nrows as u32, 1, 1),
3128 block_dim: (256, 1, 1),
3129 shared_mem_bytes: 0,
3130 };
3131 let __s_b = self.gpu.stream();
3132 let mut b = __s_b.launch_builder(&f);
3133 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
3134 unsafe {
3135 b.launch(cfg)?;
3136 }
3137 Ok(dst)
3138 }
3139
3140 pub fn sigmoid_dot_rows(
3145 &self,
3146 x: &CudaSlice<f32>,
3147 w: &CudaSlice<f32>,
3148 n_embd: usize,
3149 t: usize,
3150 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3151 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3154 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
3155 let gs = self.linear(x, w, t, n_embd, 1)?;
3156 let mut g = self.uninit(t)?;
3157 self.sigmoid(&gs, &mut g, t)?;
3158 return Ok(g);
3159 }
3160 let mut g = self.alloc_uninit::<f32>(t)?;
3166 let f = self.func("sigmoid_dot_rows_f32");
3167 let (ne, ti) = (n_embd as i32, t as i32);
3168 let cfg = LaunchConfig {
3169 grid_dim: (t as u32, 1, 1),
3170 block_dim: (32, 8, 1),
3171 shared_mem_bytes: 0,
3172 };
3173 let __s_b = self.gpu.stream();
3174 let mut b = __s_b.launch_builder(&f);
3175 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
3176 unsafe {
3177 b.launch(cfg)?;
3178 }
3179 Ok(g)
3180 }
3181
3182 pub fn sigmoid_dot_rows_into(
3184 &self,
3185 x: &CudaSlice<f32>,
3186 w: &CudaSlice<f32>,
3187 g: &mut CudaSlice<f32>,
3188 n_embd: usize,
3189 t: usize,
3190 ) -> Result<(), Box<dyn std::error::Error>> {
3191 if g.len() < t {
3192 return Err("sigmoid_dot_rows_into output too small".into());
3193 }
3194 let f = self.func("sigmoid_dot_rows_f32");
3195 let (ne, ti) = (n_embd as i32, t as i32);
3196 let cfg = LaunchConfig {
3197 grid_dim: (t as u32, 1, 1),
3198 block_dim: (32, 8, 1),
3199 shared_mem_bytes: 0,
3200 };
3201 let __s_b = self.gpu.stream();
3202 let mut b = __s_b.launch_builder(&f);
3203 b.arg(x).arg(w).arg(&mut *g).arg(&ne).arg(&ti);
3204 unsafe {
3205 b.launch(cfg)?;
3206 }
3207 Ok(())
3208 }
3209
3210 pub fn spec_rollback_stream(
3212 &self,
3213 len_ptrs: &CudaSlice<u64>,
3214 pos_start: &CudaSlice<i32>,
3215 acc: &CudaSlice<u32>,
3216 base: usize,
3217 n_rows: usize,
3218 ) -> Result<(), Box<dyn std::error::Error>> {
3219 let f = self.func("spec_rollback_stream");
3220 let (b, nr) = (base as i32, n_rows as i32);
3221 let cfg = LaunchConfig {
3222 grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
3223 block_dim: (64, 1, 1),
3224 shared_mem_bytes: 0,
3225 };
3226 let __s_bl = self.gpu.stream();
3227 let mut bl = __s_bl.launch_builder(&f);
3228 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
3229 unsafe {
3230 bl.launch(cfg)?;
3231 }
3232 Ok(())
3233 }
3234
3235 pub fn plain_tok_ring(
3237 &self,
3238 vam: &CudaSlice<u32>,
3239 pos_start: &CudaSlice<i32>,
3240 base: usize,
3241 ring: &mut CudaSlice<u32>,
3242 ) -> Result<(), Box<dyn std::error::Error>> {
3243 let f = self.func("plain_tok_ring");
3244 let (b, cap) = (base as i32, ring.len() as i32);
3245 let cfg = LaunchConfig {
3246 grid_dim: (1, 1, 1),
3247 block_dim: (32, 1, 1),
3248 shared_mem_bytes: 0,
3249 };
3250 let __s_bl = self.gpu.stream();
3251 let mut bl = __s_bl.launch_builder(&f);
3252 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
3253 unsafe {
3254 bl.launch(cfg)?;
3255 }
3256 Ok(())
3257 }
3258
3259 pub fn spec_ring_commit(
3261 &self,
3262 vtok: &CudaSlice<u32>,
3263 acc: &CudaSlice<u32>,
3264 brk: &CudaSlice<u32>,
3265 ring: &mut CudaSlice<u32>,
3266 pend: &mut CudaSlice<u32>,
3267 ) -> Result<(), Box<dyn std::error::Error>> {
3268 let f = self.func("spec_ring_commit");
3269 let cfg = LaunchConfig {
3270 grid_dim: (1, 1, 1),
3271 block_dim: (32, 1, 1),
3272 shared_mem_bytes: 0,
3273 };
3274 let __s_b = self.gpu.stream();
3275 let mut b = __s_b.launch_builder(&f);
3276 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
3277 unsafe {
3278 b.launch(cfg)?;
3279 }
3280 Ok(())
3281 }
3282 pub fn i32_copy_add(
3283 &self,
3284 src: &CudaSlice<i32>,
3285 dst: &mut CudaSlice<i32>,
3286 delta: i32,
3287 ) -> Result<(), Box<dyn std::error::Error>> {
3288 let f = self.func("i32_copy_add");
3289 let cfg = LaunchConfig {
3290 grid_dim: (1, 1, 1),
3291 block_dim: (32, 1, 1),
3292 shared_mem_bytes: 0,
3293 };
3294 let __s_b = self.gpu.stream();
3295 let mut b = __s_b.launch_builder(&f);
3296 b.arg(src).arg(dst).arg(&delta);
3297 unsafe {
3298 b.launch(cfg)?;
3299 }
3300 Ok(())
3301 }
3302 pub fn u32_copy(
3303 &self,
3304 src: &CudaSlice<u32>,
3305 dst: &mut CudaSlice<u32>,
3306 ) -> Result<(), Box<dyn std::error::Error>> {
3307 let f = self.func("u32_copy");
3308 let cfg = LaunchConfig {
3309 grid_dim: (1, 1, 1),
3310 block_dim: (32, 1, 1),
3311 shared_mem_bytes: 0,
3312 };
3313 let __s_b = self.gpu.stream();
3314 let mut b = __s_b.launch_builder(&f);
3315 b.arg(src).arg(dst);
3316 unsafe {
3317 b.launch(cfg)?;
3318 }
3319 Ok(())
3320 }
3321
3322 pub fn spec_adapt_k(
3326 &self,
3327 acc: &CudaSlice<u32>,
3328 brk: &mut CudaSlice<u32>,
3329 floor: usize,
3330 cap: usize,
3331 ) -> Result<(), Box<dyn std::error::Error>> {
3332 let f = self.func("spec_adapt_k");
3333 let (fl, cp) = (floor as i32, cap as i32);
3334 let cfg = LaunchConfig {
3335 grid_dim: (1, 1, 1),
3336 block_dim: (32, 1, 1),
3337 shared_mem_bytes: 0,
3338 };
3339 let __s_b = self.gpu.stream();
3340 let mut b = __s_b.launch_builder(&f);
3341 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
3342 unsafe {
3343 b.launch(cfg)?;
3344 }
3345 Ok(())
3346 }
3347
3348 pub fn spec_accept_greedy_dc(
3350 &self,
3351 preds: &CudaSlice<u32>,
3352 vtok: &CudaSlice<u32>,
3353 last_pred: &CudaSlice<u32>,
3354 brk: &CudaSlice<u32>,
3355 out: &mut CudaSlice<u32>,
3356 ) -> Result<(), Box<dyn std::error::Error>> {
3357 let f = self.func("spec_accept_greedy_dc");
3358 let cfg = LaunchConfig {
3359 grid_dim: (1, 1, 1),
3360 block_dim: (32, 1, 1),
3361 shared_mem_bytes: 0,
3362 };
3363 let __s_b = self.gpu.stream();
3364 let mut b = __s_b.launch_builder(&f);
3365 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
3366 unsafe {
3367 b.launch(cfg)?;
3368 }
3369 Ok(())
3370 }
3371
3372 pub fn pos_iota(
3374 &self,
3375 pos0: &CudaSlice<i32>,
3376 out: &mut CudaSlice<i32>,
3377 t: usize,
3378 ) -> Result<(), Box<dyn std::error::Error>> {
3379 let f = self.func("pos_iota_i32");
3380 let ti = t as i32;
3381 let cfg = LaunchConfig {
3382 grid_dim: (1, 1, 1),
3383 block_dim: (t.max(1) as u32, 1, 1),
3384 shared_mem_bytes: 0,
3385 };
3386 let __s_b = self.gpu.stream();
3387 let mut b = __s_b.launch_builder(&f);
3388 b.arg(pos0).arg(out).arg(&ti);
3389 unsafe {
3390 b.launch(cfg)?;
3391 }
3392 Ok(())
3393 }
3394 #[allow(clippy::too_many_arguments)]
3395 pub fn append_kv_quantized_rows_dc(
3396 &self,
3397 k_rows: &CudaSlice<f32>,
3398 v_rows: &CudaSlice<f32>,
3399 kc: &mut CudaSlice<u8>,
3400 vc: &mut CudaSlice<u8>,
3401 t0_dev: &CudaSlice<i32>,
3402 t: usize,
3403 kv_dim_k: usize,
3404 kv_dim_v: usize,
3405 k_tok_bytes: usize,
3406 v_tok_bytes: usize,
3407 g: bool,
3408 ) -> Result<(), Box<dyn std::error::Error>> {
3409 let f = if g {
3410 self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc")
3411 } else {
3412 self.func("append_quantize_kv_q8_0_q5_1_rows_dc")
3413 };
3414 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
3415 let cfg = LaunchConfig {
3416 grid_dim: (nblk, t as u32, 1),
3417 block_dim: (32, 1, 1),
3418 shared_mem_bytes: 0,
3419 };
3420 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3421 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3422 let __s_b = self.gpu.stream();
3423 let mut b = __s_b.launch_builder(&f);
3424 b.arg(k_rows)
3425 .arg(v_rows)
3426 .arg(kc)
3427 .arg(vc)
3428 .arg(t0_dev)
3429 .arg(&kdk)
3430 .arg(&kdv)
3431 .arg(&ktb)
3432 .arg(&vtb);
3433 unsafe {
3434 b.launch(cfg)?;
3435 }
3436 Ok(())
3437 }
3438
3439 #[allow(clippy::too_many_arguments)]
3442 pub fn append_kv_quantized_row_dc_inc(
3443 &self,
3444 k_row: &CudaSlice<f32>,
3445 v_row: &CudaSlice<f32>,
3446 kc: &mut CudaSlice<u8>,
3447 vc: &mut CudaSlice<u8>,
3448 t0_dev: &mut CudaSlice<i32>,
3449 kv_dim_k: usize,
3450 kv_dim_v: usize,
3451 k_tok_bytes: usize,
3452 v_tok_bytes: usize,
3453 g: bool,
3454 ) -> Result<(), Box<dyn std::error::Error>> {
3455 let f = if g {
3456 self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc")
3457 } else {
3458 self.func("append_quantize_kv_q8_0_q5_1_dc_inc")
3459 };
3460 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
3461 let cfg = LaunchConfig {
3462 grid_dim: (1, 1, 1),
3463 block_dim: (nthreads, 1, 1),
3464 shared_mem_bytes: 0,
3465 };
3466 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
3467 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
3468 let __s_b = self.gpu.stream();
3469 let mut b = __s_b.launch_builder(&f);
3470 b.arg(k_row)
3471 .arg(v_row)
3472 .arg(kc)
3473 .arg(vc)
3474 .arg(t0_dev)
3475 .arg(&kdk)
3476 .arg(&kdv)
3477 .arg(&ktb)
3478 .arg(&vtb);
3479 unsafe {
3480 b.launch(cfg)?;
3481 }
3482 Ok(())
3483 }
3484
3485 pub fn pack_tok_p(
3487 &self,
3488 tok: &CudaSlice<u32>,
3489 p: &CudaSlice<f32>,
3490 out: &mut CudaSlice<u32>,
3491 slot: usize,
3492 ) -> Result<(), Box<dyn std::error::Error>> {
3493 let f = self.func("pack_tok_p");
3494 let sl = slot as i32;
3495 let cfg = LaunchConfig {
3496 grid_dim: (1, 1, 1),
3497 block_dim: (32, 1, 1),
3498 shared_mem_bytes: 0,
3499 };
3500 let __s_b = self.gpu.stream();
3501 let mut b = __s_b.launch_builder(&f);
3502 b.arg(tok).arg(p).arg(out).arg(&sl);
3503 unsafe {
3504 b.launch(cfg)?;
3505 }
3506 Ok(())
3507 }
3508 pub fn tok_map_u32(
3509 &self,
3510 tok: &mut CudaSlice<u32>,
3511 map: &CudaSlice<u32>,
3512 ) -> Result<(), Box<dyn std::error::Error>> {
3513 let f = self.func("tok_map_u32");
3514 let cfg = LaunchConfig {
3515 grid_dim: (1, 1, 1),
3516 block_dim: (32, 1, 1),
3517 shared_mem_bytes: 0,
3518 };
3519 let __s_b = self.gpu.stream();
3520 let mut b = __s_b.launch_builder(&f);
3521 b.arg(tok).arg(map);
3522 unsafe {
3523 b.launch(cfg)?;
3524 }
3525 Ok(())
3526 }
3527
3528 #[allow(clippy::too_many_arguments)]
3530 pub fn spec_assemble_verify(
3531 &self,
3532 tokp: &CudaSlice<u32>,
3533 pend: &CudaSlice<u32>,
3534 d2t: Option<&CudaSlice<u32>>,
3535 vtok: &mut CudaSlice<u32>,
3536 brk: &mut CudaSlice<u32>,
3537 p_min: f32,
3538 k: usize,
3539 pmin0: bool,
3540 ) -> Result<(), Box<dyn std::error::Error>> {
3541 let f = self.func("spec_assemble_verify");
3542 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
3543 let cfg = LaunchConfig {
3544 grid_dim: (1, 1, 1),
3545 block_dim: (32, 1, 1),
3546 shared_mem_bytes: 0,
3547 };
3548 let __s_b = self.gpu.stream();
3549 let mut b = __s_b.launch_builder(&f);
3550 match d2t {
3551 Some(m) => {
3552 b.arg(tokp)
3553 .arg(pend)
3554 .arg(m)
3555 .arg(vtok)
3556 .arg(brk)
3557 .arg(&p_min)
3558 .arg(&ki)
3559 .arg(&pm);
3560 unsafe {
3561 b.launch(cfg)?;
3562 }
3563 }
3564 None => {
3565 let null: u64 = 0;
3566 b.arg(tokp)
3567 .arg(pend)
3568 .arg(&null)
3569 .arg(vtok)
3570 .arg(brk)
3571 .arg(&p_min)
3572 .arg(&ki)
3573 .arg(&pm);
3574 unsafe {
3575 b.launch(cfg)?;
3576 }
3577 }
3578 }
3579 Ok(())
3580 }
3581
3582 #[allow(clippy::too_many_arguments)]
3584 pub fn ssm_conv_ring_rebuild_dc(
3585 &self,
3586 qkv_tm: &CudaSlice<f32>,
3587 ring_old: &CudaSlice<f32>,
3588 conv_state: &mut CudaSlice<f32>,
3589 conv_dim: usize,
3590 acc: &CudaSlice<u32>,
3591 base: usize,
3592 t_v: usize,
3593 d_conv: usize,
3594 ) -> Result<(), Box<dyn std::error::Error>> {
3595 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
3596 let n = conv_dim * (d_conv - 1);
3597 let cfg = LaunchConfig::for_num_elems(n as u32);
3598 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
3599 let __s_b = self.gpu.stream();
3600 let mut b = __s_b.launch_builder(&f);
3601 b.arg(qkv_tm)
3602 .arg(ring_old)
3603 .arg(conv_state)
3604 .arg(&cd)
3605 .arg(acc)
3606 .arg(&b0)
3607 .arg(&tv)
3608 .arg(&dc);
3609 unsafe {
3610 b.launch(cfg)?;
3611 }
3612 Ok(())
3613 }
3614 #[allow(clippy::too_many_arguments)]
3615 pub fn gdn_scan_s128_dc(
3616 &self,
3617 q: &CudaSlice<f32>,
3618 k: &CudaSlice<f32>,
3619 v: &CudaSlice<f32>,
3620 g: &CudaSlice<f32>,
3621 beta: &CudaSlice<f32>,
3622 state_in: &CudaSlice<f32>,
3623 state_out: &mut CudaSlice<f32>,
3624 o: &mut CudaSlice<f32>,
3625 n_head: usize,
3626 acc: &CudaSlice<u32>,
3627 base: usize,
3628 t_v: usize,
3629 scale: f32,
3630 ) -> Result<(), Box<dyn std::error::Error>> {
3631 let f = self.func("gdn_scan_s128_dc");
3632 const S_V: u32 = 128;
3633 const WARP: u32 = 32;
3634 const COLS_PER_BLOCK: u32 = 4;
3635 let cfg = LaunchConfig {
3636 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
3637 block_dim: (WARP, COLS_PER_BLOCK, 1),
3638 shared_mem_bytes: 0,
3639 };
3640 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
3641 let __s_b = self.gpu.stream();
3642 let mut b = __s_b.launch_builder(&f);
3643 b.arg(q)
3644 .arg(k)
3645 .arg(v)
3646 .arg(g)
3647 .arg(beta)
3648 .arg(state_in)
3649 .arg(state_out)
3650 .arg(o)
3651 .arg(&h)
3652 .arg(acc)
3653 .arg(&b0)
3654 .arg(&tv)
3655 .arg(&scale);
3656 unsafe {
3657 b.launch(cfg)?;
3658 }
3659 Ok(())
3660 }
3661
3662 pub fn spec_rollback_kv(
3664 &self,
3665 len_ptrs: &CudaSlice<u64>,
3666 saved: &CudaSlice<i32>,
3667 acc: &CudaSlice<u32>,
3668 base: usize,
3669 n_layer: usize,
3670 ) -> Result<(), Box<dyn std::error::Error>> {
3671 let f = self.func("spec_rollback_kv");
3672 let (b, nl) = (base as i32, n_layer as i32);
3673 let cfg = LaunchConfig {
3674 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3675 block_dim: (64, 1, 1),
3676 shared_mem_bytes: 0,
3677 };
3678 let __s_bl = self.gpu.stream();
3679 let mut bl = __s_bl.launch_builder(&f);
3680 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
3681 unsafe {
3682 bl.launch(cfg)?;
3683 }
3684 Ok(())
3685 }
3686
3687 pub fn spec_fork_valid(
3689 &self,
3690 acc: &CudaSlice<u32>,
3691 optimistic_pending: u32,
3692 valid: &mut CudaSlice<u32>,
3693 ) -> Result<(), Box<dyn std::error::Error>> {
3694 let f = self.func("spec_fork_valid");
3695 let cfg = LaunchConfig {
3696 grid_dim: (1, 1, 1),
3697 block_dim: (1, 1, 1),
3698 shared_mem_bytes: 0,
3699 };
3700 let __s_bl = self.gpu.stream();
3701 let mut bl = __s_bl.launch_builder(&f);
3702 bl.arg(acc).arg(&optimistic_pending).arg(valid);
3703 unsafe {
3704 bl.launch(cfg)?;
3705 }
3706 Ok(())
3707 }
3708
3709 pub fn spec_fork_reconcile_kv(
3711 &self,
3712 len_ptrs: &CudaSlice<u64>,
3713 saved: &CudaSlice<i32>,
3714 acc: &CudaSlice<u32>,
3715 valid: &CudaSlice<u32>,
3716 base: usize,
3717 n_layer: usize,
3718 ) -> Result<(), Box<dyn std::error::Error>> {
3719 let f = self.func("spec_fork_reconcile_kv");
3720 let (b, nl) = (base as i32, n_layer as i32);
3721 let cfg = LaunchConfig {
3722 grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
3723 block_dim: (64, 1, 1),
3724 shared_mem_bytes: 0,
3725 };
3726 let __s_bl = self.gpu.stream();
3727 let mut bl = __s_bl.launch_builder(&f);
3728 bl.arg(len_ptrs)
3729 .arg(saved)
3730 .arg(acc)
3731 .arg(valid)
3732 .arg(&b)
3733 .arg(&nl);
3734 unsafe {
3735 bl.launch(cfg)?;
3736 }
3737 Ok(())
3738 }
3739
3740 pub fn spec_fork_restore_f32(
3742 &self,
3743 snapshot: &CudaSlice<f32>,
3744 state: &mut CudaSlice<f32>,
3745 valid: &CudaSlice<u32>,
3746 ) -> Result<(), Box<dyn std::error::Error>> {
3747 assert_eq!(
3748 snapshot.len(),
3749 state.len(),
3750 "fork recurrent snapshot shape mismatch"
3751 );
3752 let f = self.func("spec_fork_restore_f32");
3753 let n = state.len() as i32;
3754 let blocks = state.len().div_ceil(256).min(65535).max(1) as u32;
3755 let cfg = LaunchConfig {
3756 grid_dim: (blocks, 1, 1),
3757 block_dim: (256, 1, 1),
3758 shared_mem_bytes: 0,
3759 };
3760 let __s_bl = self.gpu.stream();
3761 let mut bl = __s_bl.launch_builder(&f);
3762 bl.arg(snapshot).arg(state).arg(valid).arg(&n);
3763 unsafe {
3764 bl.launch(cfg)?;
3765 }
3766 Ok(())
3767 }
3768
3769 pub fn spec_seed_gather(
3772 &self,
3773 vx: &CudaSlice<f32>,
3774 fill_prev: &CudaSlice<f32>,
3775 acc: &CudaSlice<u32>,
3776 h_seed: &mut CudaSlice<f32>,
3777 base: usize,
3778 n_embd: usize,
3779 ) -> Result<(), Box<dyn std::error::Error>> {
3780 let f = self.func("spec_seed_gather");
3781 let (b, ne) = (base as i32, n_embd as i32);
3782 let cfg = LaunchConfig {
3783 grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
3784 block_dim: (256, 1, 1),
3785 shared_mem_bytes: 0,
3786 };
3787 let __s_bl = self.gpu.stream();
3788 let mut bl = __s_bl.launch_builder(&f);
3789 bl.arg(vx)
3790 .arg(fill_prev)
3791 .arg(acc)
3792 .arg(h_seed)
3793 .arg(&b)
3794 .arg(&ne);
3795 unsafe {
3796 bl.launch(cfg)?;
3797 }
3798 Ok(())
3799 }
3800
3801 pub fn spec_accept_greedy(
3803 &self,
3804 preds: &CudaSlice<u32>,
3805 draft: &CudaSlice<u32>,
3806 last_pred: u32,
3807 base: usize,
3808 k_round: usize,
3809 out: &mut CudaSlice<u32>,
3810 ) -> Result<(), Box<dyn std::error::Error>> {
3811 let f = self.func("spec_accept_greedy");
3812 let (b, k) = (base as i32, k_round as i32);
3813 let cfg = LaunchConfig {
3814 grid_dim: (1, 1, 1),
3815 block_dim: (32, 1, 1),
3816 shared_mem_bytes: 0,
3817 };
3818 let __s_bl = self.gpu.stream();
3819 let mut bl = __s_bl.launch_builder(&f);
3820 bl.arg(preds)
3821 .arg(draft)
3822 .arg(&last_pred)
3823 .arg(&b)
3824 .arg(&k)
3825 .arg(out);
3826 unsafe {
3827 bl.launch(cfg)?;
3828 }
3829 Ok(())
3830 }
3831
3832 pub fn gumbel_perturb(
3839 &self,
3840 x: &CudaSlice<f32>,
3841 y: &mut CudaSlice<f32>,
3842 n: usize,
3843 seed: u64,
3844 stream_pos: u32,
3845 temp: f32,
3846 ) -> Result<(), Box<dyn std::error::Error>> {
3847 let f = self.func("gumbel_perturb_f32");
3848 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3849 let cfg = LaunchConfig {
3850 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3851 block_dim: (256, 1, 1),
3852 shared_mem_bytes: 0,
3853 };
3854 let __s_b = self.gpu.stream();
3855 let mut b = __s_b.launch_builder(&f);
3856 b.arg(x)
3857 .arg(&mut *y)
3858 .arg(&ni)
3859 .arg(&slo)
3860 .arg(&shi)
3861 .arg(&stream_pos)
3862 .arg(&temp);
3863 unsafe {
3864 b.launch(cfg)?;
3865 }
3866 Ok(())
3867 }
3868
3869 pub fn mask_logits_col(
3877 &self,
3878 logits: &mut CudaSlice<f32>,
3879 mask: &CudaSlice<u32>,
3880 col: usize,
3881 n: usize,
3882 mask_words: usize,
3883 ) -> Result<(), Box<dyn std::error::Error>> {
3884 let f = self.func("mask_logits_f32");
3885 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
3886 let cfg = LaunchConfig {
3887 grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
3888 block_dim: (256, 1, 1),
3889 shared_mem_bytes: 0,
3890 };
3891 let __s_b = self.gpu.stream();
3892 let mut b = __s_b.launch_builder(&f);
3893 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
3894 unsafe {
3895 b.launch(cfg)?;
3896 }
3897 Ok(())
3898 }
3899
3900 pub fn gumbel_perturb_col(
3907 &self,
3908 x: &CudaSlice<f32>,
3909 col: usize,
3910 y: &mut CudaSlice<f32>,
3911 n: usize,
3912 seed: u64,
3913 stream_pos: u32,
3914 temp: f32,
3915 ) -> Result<(), Box<dyn std::error::Error>> {
3916 let f = self.func("gumbel_perturb_f32");
3917 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3918 let col_view = x.slice(col * n..(col + 1) * n);
3919 let cfg = LaunchConfig {
3920 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3921 block_dim: (256, 1, 1),
3922 shared_mem_bytes: 0,
3923 };
3924 let __s_b = self.gpu.stream();
3925 let mut b = __s_b.launch_builder(&f);
3926 b.arg(&col_view)
3927 .arg(&mut *y)
3928 .arg(&ni)
3929 .arg(&slo)
3930 .arg(&shi)
3931 .arg(&stream_pos)
3932 .arg(&temp);
3933 unsafe {
3934 b.launch(cfg)?;
3935 }
3936 Ok(())
3937 }
3938
3939 #[allow(clippy::too_many_arguments)]
3945 pub fn gumbel_perturb_filtered_col(
3946 &self,
3947 x: &CudaSlice<f32>,
3948 col: usize,
3949 y: &mut CudaSlice<f32>,
3950 n: usize,
3951 seed: u64,
3952 stream_pos: u32,
3953 temp: f32,
3954 stat_max: &CudaSlice<f32>,
3955 stat_th: &CudaSlice<f32>,
3956 stat_idx: usize,
3957 ) -> Result<(), Box<dyn std::error::Error>> {
3958 let f = self.func("gumbel_perturb_filtered_col_f32");
3959 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
3960 let (ci, si) = (col as i32, stat_idx as i32);
3961 let cfg = LaunchConfig {
3962 grid_dim: (n.div_ceil(256) as u32, 1, 1),
3963 block_dim: (256, 1, 1),
3964 shared_mem_bytes: 0,
3965 };
3966 let __s_b = self.gpu.stream();
3967 let mut b = __s_b.launch_builder(&f);
3968 b.arg(x)
3969 .arg(&ci)
3970 .arg(&mut *y)
3971 .arg(&ni)
3972 .arg(&slo)
3973 .arg(&shi)
3974 .arg(&stream_pos)
3975 .arg(&temp)
3976 .arg(stat_max)
3977 .arg(stat_th)
3978 .arg(&si);
3979 unsafe {
3980 b.launch(cfg)?;
3981 }
3982 Ok(())
3983 }
3984
3985 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
3990 let f = self.func("memra_sctr_inc");
3991 let cfg = LaunchConfig {
3992 grid_dim: (1, 1, 1),
3993 block_dim: (1, 1, 1),
3994 shared_mem_bytes: 0,
3995 };
3996 let __s_b = self.gpu.stream();
3997 let mut b = __s_b.launch_builder(&f);
3998 b.arg(&mut *ctr);
3999 unsafe {
4000 b.launch(cfg)?;
4001 }
4002 Ok(())
4003 }
4004
4005 pub fn gumbel_perturb_ctr(
4010 &self,
4011 x: &CudaSlice<f32>,
4012 y: &mut CudaSlice<f32>,
4013 n: usize,
4014 seed: u64,
4015 ctr: &CudaSlice<u32>,
4016 temp: f32,
4017 ) -> Result<(), Box<dyn std::error::Error>> {
4018 let f = self.func("gumbel_perturb_ctr_f32");
4019 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4020 let cfg = LaunchConfig {
4021 grid_dim: (n.div_ceil(256) as u32, 1, 1),
4022 block_dim: (256, 1, 1),
4023 shared_mem_bytes: 0,
4024 };
4025 let __s_b = self.gpu.stream();
4026 let mut b = __s_b.launch_builder(&f);
4027 b.arg(x)
4028 .arg(&mut *y)
4029 .arg(&ni)
4030 .arg(&slo)
4031 .arg(&shi)
4032 .arg(ctr)
4033 .arg(&temp);
4034 unsafe {
4035 b.launch(cfg)?;
4036 }
4037 Ok(())
4038 }
4039
4040 pub fn softmax_gather(
4044 &self,
4045 x: &CudaSlice<f32>,
4046 row_stride: usize,
4047 ids: &CudaSlice<u32>,
4048 rows: &CudaSlice<i32>,
4049 out: &mut CudaSlice<f32>,
4050 n: usize,
4051 npair: usize,
4052 temp: f32,
4053 ) -> Result<(), Box<dyn std::error::Error>> {
4054 let f = self.func("softmax_gather_f32");
4055 let (ni, rs) = (n as i32, row_stride as i64);
4056 let np = npair as i32;
4057 let cfg = LaunchConfig {
4058 grid_dim: (npair as u32, 1, 1),
4059 block_dim: (256, 1, 1),
4060 shared_mem_bytes: 0,
4061 };
4062 let __s_b = self.gpu.stream();
4063 let mut b = __s_b.launch_builder(&f);
4064 b.arg(x)
4065 .arg(&rs)
4066 .arg(ids)
4067 .arg(rows)
4068 .arg(&mut *out)
4069 .arg(&ni)
4070 .arg(&np)
4071 .arg(&temp);
4072 unsafe {
4073 b.launch(cfg)?;
4074 }
4075 Ok(())
4076 }
4077
4078 pub fn residual_sample(
4082 &self,
4083 p: &CudaSlice<f32>,
4084 q: Option<&CudaSlice<f32>>,
4085 n: usize,
4086 temp: f32,
4087 seed: u64,
4088 stream_pos: u32,
4089 out_tok: &mut CudaSlice<u32>,
4090 ) -> Result<(), Box<dyn std::error::Error>> {
4091 let f = self.func("residual_sample_f32");
4092 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
4093 let nth = 1024u32;
4094 let cfg = LaunchConfig {
4095 grid_dim: (1, 1, 1),
4096 block_dim: (nth, 1, 1),
4097 shared_mem_bytes: 0,
4098 };
4099 let has_q: i32 = q.is_some() as i32;
4100 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
4102 let mut b = __s_b.launch_builder(&f);
4103 b.arg(p)
4104 .arg(qbuf)
4105 .arg(&has_q)
4106 .arg(&ni)
4107 .arg(&temp)
4108 .arg(&slo)
4109 .arg(&shi)
4110 .arg(&stream_pos)
4111 .arg(&mut *out_tok);
4112 unsafe {
4113 b.launch(cfg)?;
4114 }
4115 Ok(())
4116 }
4117
4118 pub fn with_moe_cache<R>(
4123 &self,
4124 max_block_bytes: usize,
4125 f: impl FnOnce(
4126 &mut crate::moe_cache::MoeSlotCache,
4127 &Engine,
4128 ) -> Result<R, Box<dyn std::error::Error>>,
4129 ) -> Result<R, Box<dyn std::error::Error>> {
4130 let mut guard = self.moe_cache.lock().unwrap();
4131 if guard.is_none() {
4132 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
4133 }
4134 let cache = guard.as_mut().unwrap();
4135 f(cache, self)
4136 }
4137
4138 pub fn freeze_moe_cache(&self) {
4141 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
4142 cache.freeze();
4143 }
4144 }
4145
4146 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
4149 self.moe_cache
4150 .lock()
4151 .unwrap()
4152 .as_ref()
4153 .map(crate::moe_cache::MoeSlotCache::export_residency)
4154 }
4155
4156 pub(crate) fn moe_cache_frozen(&self) -> bool {
4157 self.moe_cache
4158 .lock()
4159 .unwrap()
4160 .as_ref()
4161 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
4162 }
4163
4164 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
4171 crate::cpu_experts::configured()
4172 && self.moe_cache_frozen()
4173 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
4174 }
4175
4176 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
4178 assert!(
4179 self.moe_cache.lock().unwrap().is_none(),
4180 "MoE cache layout configured after cache construction"
4181 );
4182 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
4183 }
4184
4185 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
4186 self.moe_cache_layout.lock().unwrap().clone()
4187 }
4188
4189 pub fn moe_cache_enabled() -> bool {
4191 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
4192 }
4193
4194 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
4197 let guard = self.moe_cache.lock().unwrap();
4198 guard
4199 .as_ref()
4200 .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
4201 }
4202
4203 pub fn cpu_expert_stats(
4207 &self,
4208 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
4209 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
4210 }
4211
4212 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
4215 crate::cpu_experts::predictor_stats()
4216 }
4217
4218 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
4219 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
4220 }
4221
4222 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
4225 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
4226 }
4227
4228 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
4231 let guard = self.moe_cache.lock().unwrap();
4232 guard
4233 .as_ref()
4234 .and_then(|cache| cache.pread_stats())
4235 .map(|stats| {
4236 (
4237 stats.reads,
4238 stats.bytes,
4239 stats.read_errors,
4240 stats.short_reads,
4241 stats.fallbacks,
4242 stats.buffer_waits,
4243 stats.ring_full,
4244 )
4245 })
4246 }
4247
4248 pub fn spill_config_fallbacks(&self) -> u64 {
4250 crate::spill_pread::config_fallbacks()
4251 }
4252
4253 pub fn moe_cache_reset_counters(&self) {
4255 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() {
4256 c.reset_counters();
4257 }
4258 }
4259
4260 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4261 Ok(self.gpu.stream().clone_htod(v)?)
4262 }
4263
4264 pub fn htod_bytes_padded(
4268 &self,
4269 v: &[u8],
4270 pad: usize,
4271 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4272 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
4273 {
4274 let mut view = d.slice_mut(0..v.len());
4275 self.gpu.stream().memcpy_htod(v, &mut view)?;
4276 }
4277 Ok(d)
4278 }
4279
4280 pub fn copy_into(
4282 &self,
4283 dst: &mut CudaSlice<f32>,
4284 off: usize,
4285 src: &CudaSlice<f32>,
4286 len: usize,
4287 ) -> Result<(), Box<dyn std::error::Error>> {
4288 let mut view = dst.slice_mut(off..off + len);
4289 self.gpu
4290 .stream()
4291 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4292 Ok(())
4293 }
4294
4295 pub fn copy_range_into(
4299 &self,
4300 dst: &mut CudaSlice<f32>,
4301 dst_off: usize,
4302 src: &CudaSlice<f32>,
4303 src_off: usize,
4304 len: usize,
4305 ) -> Result<(), Box<dyn std::error::Error>> {
4306 let mut view = dst.slice_mut(dst_off..dst_off + len);
4307 self.gpu
4308 .stream()
4309 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut view)?;
4310 Ok(())
4311 }
4312
4313 pub fn copy_u8_into(
4316 &self,
4317 dst: &mut CudaSlice<u8>,
4318 off: usize,
4319 src: &CudaSlice<u8>,
4320 len: usize,
4321 ) -> Result<(), Box<dyn std::error::Error>> {
4322 let cap = dst.len();
4326 let mut view = dst.try_slice_mut(off..off + len).ok_or_else(|| {
4327 format!(
4328 "copy_u8_into dst range [{off},{}) exceeds capacity {cap}",
4329 off + len,
4330 )
4331 })?;
4332 self.gpu
4333 .stream()
4334 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4335 Ok(())
4336 }
4337
4338 pub fn copy_u8_range_into(
4340 &self,
4341 dst: &mut CudaSlice<u8>,
4342 dst_off: usize,
4343 src: &CudaSlice<u8>,
4344 src_off: usize,
4345 len: usize,
4346 ) -> Result<(), Box<dyn std::error::Error>> {
4347 let cap = dst.len();
4350 let mut dst_view = dst.try_slice_mut(dst_off..dst_off + len).ok_or_else(|| {
4351 format!(
4352 "copy_u8_range_into dst range [{dst_off},{}) exceeds capacity {cap}",
4353 dst_off + len,
4354 )
4355 })?;
4356 self.gpu
4357 .stream()
4358 .memcpy_dtod(&src.slice(src_off..src_off + len), &mut dst_view)?;
4359 Ok(())
4360 }
4361
4362 #[track_caller]
4372 pub fn prepare_kv_append(
4373 &self,
4374 kv: &mut crate::cache::KvLayer,
4375 retain_from: usize,
4376 append_rows: usize,
4377 ) -> Result<usize, Box<dyn std::error::Error>> {
4378 let caller = std::panic::Location::caller();
4379 let base_before = kv.ring.as_ref().map(|r| r.base());
4380 let Some(plan) = kv
4381 .ring
4382 .as_ref()
4383 .map(|ring| ring.append_plan(kv.len, retain_from, append_rows))
4384 .transpose()
4385 .map_err(|err| -> Box<dyn std::error::Error> {
4386 format!(
4387 "{err} [append len={} retain_from={retain_from} append_rows={append_rows} base={base_before:?} called from {caller}]",
4388 kv.len
4389 )
4390 .into()
4391 })?
4392 else {
4393 return Ok(kv.len);
4394 };
4395 match plan {
4396 crate::cache::KvRingAppend::Contiguous { write_row } => Ok(write_row),
4397 crate::cache::KvRingAppend::Rebase {
4398 src_row,
4399 keep_rows,
4400 new_base,
4401 write_row,
4402 } => {
4403 if keep_rows > 0 {
4404 let k_len = keep_rows * kv.k_tok_bytes;
4405 let v_len = keep_rows * kv.v_tok_bytes;
4406 let mut k_tmp = self.alloc_u8_uninit(k_len)?;
4407 let mut v_tmp = self.alloc_u8_uninit(v_len)?;
4408 self.copy_u8_range_into(&mut k_tmp, 0, &kv.k, src_row * kv.k_tok_bytes, k_len)?;
4409 self.copy_u8_range_into(&mut v_tmp, 0, &kv.v, src_row * kv.v_tok_bytes, v_len)?;
4410 self.copy_u8_into(&mut kv.k, 0, &k_tmp, k_len)?;
4411 self.copy_u8_into(&mut kv.v, 0, &v_tmp, v_len)?;
4412 }
4413 if std::env::var("MEMRA_KV_REBASE_TRACE").as_deref() == Ok("1") {
4416 eprintln!(
4417 "[kv-rebase] new_base={new_base} keep_rows={keep_rows} len={} \
4418 retain_from={retain_from} called from {caller}",
4419 kv.len
4420 );
4421 }
4422 kv.ring.as_mut().unwrap().apply_rebase(new_base);
4423 if let Some(base_d) = kv.base_d.as_mut() {
4427 self.set_i32_one(base_d, new_base as i32)?;
4428 }
4429 Ok(write_row)
4430 }
4431 }
4432 }
4433
4434 pub fn htod_u8_into(
4437 &self,
4438 dst: &mut CudaSlice<u8>,
4439 off: usize,
4440 src: &[u8],
4441 ) -> Result<(), Box<dyn std::error::Error>> {
4442 let mut view = dst.slice_mut(off..off + src.len());
4443 self.gpu.stream().memcpy_htod(src, &mut view)?;
4444 Ok(())
4445 }
4446
4447 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
4448 b.slice(0..len)
4449 }
4450
4451 pub fn view_u8_range<'a>(
4454 &self,
4455 b: &'a CudaSlice<u8>,
4456 start: usize,
4457 end: usize,
4458 ) -> cudarc::driver::CudaView<'a, u8> {
4459 b.slice(start..end)
4460 }
4461 pub fn view_u8<'a>(
4462 &self,
4463 b: &'a CudaSlice<u8>,
4464 len: usize,
4465 ) -> cudarc::driver::CudaView<'a, u8> {
4466 b.slice(0..len)
4467 }
4468
4469 pub fn append_kv_quantized(
4473 &self,
4474 k_row: &CudaSlice<f32>,
4475 v_row: &CudaSlice<f32>,
4476 kc: &mut CudaSlice<u8>,
4477 vc: &mut CudaSlice<u8>,
4478 t: usize,
4479 kv_dim_k: usize,
4480 kv_dim_v: usize,
4481 k_tok_bytes: usize,
4482 v_tok_bytes: usize,
4483 g: bool,
4484 ) -> Result<(), Box<dyn std::error::Error>> {
4485 let f = if g {
4486 self.func_g("append_quantize_kv_q8_0_q5_1")
4487 } else {
4488 self.func("append_quantize_kv_q8_0_q5_1")
4489 };
4490 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4491 let cfg = LaunchConfig {
4492 grid_dim: (nblk, 1, 1),
4493 block_dim: (32, 1, 1),
4494 shared_mem_bytes: 0,
4495 };
4496 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4497 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4498 let __s_b = self.gpu.stream();
4499 let mut b = __s_b.launch_builder(&f);
4500 b.arg(k_row)
4501 .arg(v_row)
4502 .arg(kc)
4503 .arg(vc)
4504 .arg(&ti)
4505 .arg(&kdk)
4506 .arg(&kdv)
4507 .arg(&ktb)
4508 .arg(&vtb);
4509 unsafe {
4510 b.launch(cfg)?;
4511 }
4512 Ok(())
4513 }
4514
4515 pub fn append_kv_quantized_dc(
4519 &self,
4520 k_row: &CudaSlice<f32>,
4521 v_row: &CudaSlice<f32>,
4522 kc: &mut CudaSlice<u8>,
4523 vc: &mut CudaSlice<u8>,
4524 t_dev: &CudaSlice<i32>,
4525 kv_dim_k: usize,
4526 kv_dim_v: usize,
4527 k_tok_bytes: usize,
4528 v_tok_bytes: usize,
4529 g: bool,
4530 ) -> Result<(), Box<dyn std::error::Error>> {
4531 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4532 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
4533 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4534 if Self::pdl_on() && Self::pdl_wb_on() {
4536 use cudarc::driver::{DevicePtr, DevicePtrMut};
4537 let s = &self.gpu.stream();
4538 let (pk, _g0) = k_row.device_ptr(s);
4539 let (pv, _g1) = v_row.device_ptr(s);
4540 let (pkc, _g2) = kc.device_ptr_mut(s);
4541 let (pvc, _g3) = vc.device_ptr_mut(s);
4542 let (pt, _g4) = t_dev.device_ptr(s);
4543 let mut ps = [
4544 &pk as *const _ as *mut std::ffi::c_void,
4545 &pv as *const _ as *mut _,
4546 &pkc as *const _ as *mut _,
4547 &pvc as *const _ as *mut _,
4548 &pt as *const _ as *mut _,
4549 &kdk as *const _ as *mut _,
4550 &kdv as *const _ as *mut _,
4551 &ktb as *const _ as *mut _,
4552 &vtb as *const _ as *mut _,
4553 ];
4554 unsafe {
4555 self.launch_pdl_flash(
4556 g,
4557 "append_quantize_kv_q8_0_q5_1_dc",
4558 (nblk, 1, 1),
4559 (32, 1, 1),
4560 0,
4561 &mut ps,
4562 )?;
4563 }
4564 return Ok(());
4565 }
4566 let f = if g {
4567 self.func_g("append_quantize_kv_q8_0_q5_1_dc")
4568 } else {
4569 self.func("append_quantize_kv_q8_0_q5_1_dc")
4570 };
4571 let cfg = LaunchConfig {
4572 grid_dim: (nblk, 1, 1),
4573 block_dim: (32, 1, 1),
4574 shared_mem_bytes: 0,
4575 };
4576 let __s_b = self.gpu.stream();
4577 let mut b = __s_b.launch_builder(&f);
4578 b.arg(k_row)
4579 .arg(v_row)
4580 .arg(kc)
4581 .arg(vc)
4582 .arg(t_dev)
4583 .arg(&kdk)
4584 .arg(&kdv)
4585 .arg(&ktb)
4586 .arg(&vtb);
4587 unsafe {
4588 b.launch(cfg)?;
4589 }
4590 Ok(())
4591 }
4592
4593 #[allow(clippy::too_many_arguments)]
4600 pub fn append_kv_quantized_rows(
4601 &self,
4602 k_rows: &CudaSlice<f32>,
4603 v_rows: &CudaSlice<f32>,
4604 kc: &mut CudaSlice<u8>,
4605 vc: &mut CudaSlice<u8>,
4606 t0: usize,
4607 t: usize,
4608 kv_dim_k: usize,
4609 kv_dim_v: usize,
4610 k_tok_bytes: usize,
4611 v_tok_bytes: usize,
4612 g: bool,
4613 ) -> Result<(), Box<dyn std::error::Error>> {
4614 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
4615 for i in 0..t {
4616 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
4617 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
4618 self.append_kv_quantized_view(
4619 &k_row,
4620 &v_row,
4621 kc,
4622 vc,
4623 t0 + i,
4624 kv_dim_k,
4625 kv_dim_v,
4626 k_tok_bytes,
4627 v_tok_bytes,
4628 g,
4629 )?;
4630 }
4631 return Ok(());
4632 }
4633 let f = if g {
4634 self.func_g("append_quantize_kv_q8_0_q5_1_rows")
4635 } else {
4636 self.func("append_quantize_kv_q8_0_q5_1_rows")
4637 };
4638 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4639 let cfg = LaunchConfig {
4640 grid_dim: (nblk, t as u32, 1),
4641 block_dim: (32, 1, 1),
4642 shared_mem_bytes: 0,
4643 };
4644 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
4645 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4646 let __s_b = self.gpu.stream();
4647 let mut b = __s_b.launch_builder(&f);
4648 b.arg(k_rows)
4649 .arg(v_rows)
4650 .arg(kc)
4651 .arg(vc)
4652 .arg(&t0i)
4653 .arg(&kdk)
4654 .arg(&kdv)
4655 .arg(&ktb)
4656 .arg(&vtb);
4657 unsafe {
4658 b.launch(cfg)?;
4659 }
4660 Ok(())
4661 }
4662
4663 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
4667 let f = self.func("inc_i32");
4668 let cfg = LaunchConfig {
4669 grid_dim: (1, 1, 1),
4670 block_dim: (1, 1, 1),
4671 shared_mem_bytes: 0,
4672 };
4673 let __s_b = self.gpu.stream();
4674 let mut b = __s_b.launch_builder(&f);
4675 b.arg(p);
4676 unsafe {
4677 b.launch(cfg)?;
4678 }
4679 Ok(())
4680 }
4681
4682 pub fn append_kv_quantized_view(
4685 &self,
4686 k_row: &cudarc::driver::CudaView<f32>,
4687 v_row: &cudarc::driver::CudaView<f32>,
4688 kc: &mut CudaSlice<u8>,
4689 vc: &mut CudaSlice<u8>,
4690 t: usize,
4691 kv_dim_k: usize,
4692 kv_dim_v: usize,
4693 k_tok_bytes: usize,
4694 v_tok_bytes: usize,
4695 g: bool,
4696 ) -> Result<(), Box<dyn std::error::Error>> {
4697 let stream = self.gpu.stream();
4698 ensure_tensor_stream_device(k_row, &stream, "append_kv_quantized_view.k_row")?;
4699 ensure_tensor_stream_device(v_row, &stream, "append_kv_quantized_view.v_row")?;
4700 ensure_tensor_stream_device(kc, &stream, "append_kv_quantized_view.k_cache")?;
4701 ensure_tensor_stream_device(vc, &stream, "append_kv_quantized_view.v_cache")?;
4702 let f = if g {
4703 self.func_g("append_quantize_kv_q8_0_q5_1")
4704 } else {
4705 self.func("append_quantize_kv_q8_0_q5_1")
4706 };
4707 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
4708 let cfg = LaunchConfig {
4709 grid_dim: (nblk, 1, 1),
4710 block_dim: (32, 1, 1),
4711 shared_mem_bytes: 0,
4712 };
4713 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
4714 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
4715 let mut b = stream.launch_builder(&f);
4716 b.arg(k_row)
4717 .arg(v_row)
4718 .arg(kc)
4719 .arg(vc)
4720 .arg(&ti)
4721 .arg(&kdk)
4722 .arg(&kdv)
4723 .arg(&ktb)
4724 .arg(&vtb);
4725 unsafe {
4726 b.launch(cfg)?;
4727 }
4728 Ok(())
4729 }
4730
4731 pub fn copy_view_into(
4734 &self,
4735 dst: &mut CudaSlice<f32>,
4736 off: usize,
4737 src: &cudarc::driver::CudaView<f32>,
4738 len: usize,
4739 ) -> Result<(), Box<dyn std::error::Error>> {
4740 let mut view = dst.slice_mut(off..off + len);
4741 self.gpu
4742 .stream()
4743 .memcpy_dtod(&src.slice(0..len), &mut view)?;
4744 Ok(())
4745 }
4746
4747 pub fn clone_dtod(
4751 &self,
4752 src: &CudaSlice<f32>,
4753 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4754 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
4755 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
4756 Ok(dst)
4757 }
4758
4759 pub fn dtod_copy_view(
4762 &self,
4763 src: &cudarc::driver::CudaView<f32>,
4764 dst: &mut CudaSlice<f32>,
4765 ) -> Result<(), Box<dyn std::error::Error>> {
4766 self.gpu.stream().memcpy_dtod(src, dst)?;
4767 Ok(())
4768 }
4769
4770 pub fn dtod_copy_view_i8(
4772 &self,
4773 src: &cudarc::driver::CudaView<i8>,
4774 dst: &mut CudaSlice<i8>,
4775 ) -> Result<(), Box<dyn std::error::Error>> {
4776 self.gpu.stream().memcpy_dtod(src, dst)?;
4777 Ok(())
4778 }
4779
4780 pub fn dtod_copy_into(
4782 &self,
4783 src: &CudaSlice<f32>,
4784 dst: &mut CudaSlice<f32>,
4785 offset: usize,
4786 ) -> Result<(), Box<dyn std::error::Error>> {
4787 let n = src.len();
4788 let mut dv = dst.slice_mut(offset..offset + n);
4789 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
4790 Ok(())
4791 }
4792
4793 pub fn copy_batch_uniform_f32(
4799 &self,
4800 table: &CudaSlice<u64>,
4801 n: usize,
4802 words: usize,
4803 ) -> Result<(), Box<dyn std::error::Error>> {
4804 if n == 0 || words == 0 {
4805 return Ok(());
4806 }
4807 debug_assert!(
4808 table.len() >= 2 * n,
4809 "pointer table must hold n srcs + n dsts"
4810 );
4811 let f = self.func("copy_batch_uniform_f32");
4812 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4815 let (ni, wi) = (n as i32, words as i32);
4816 let cfg = LaunchConfig {
4817 grid_dim: (chunks, n as u32, 1),
4818 block_dim: (256, 1, 1),
4819 shared_mem_bytes: 0,
4820 };
4821 let __s = self.gpu.stream();
4822 let mut b = __s.launch_builder(&f);
4823 b.arg(table).arg(&ni).arg(&wi);
4824 unsafe {
4825 b.launch(cfg)?;
4826 }
4827 Ok(())
4828 }
4829
4830 pub fn htod_u64_into(
4833 &self,
4834 v: &[u64],
4835 dst: &mut CudaSlice<u64>,
4836 ) -> Result<(), Box<dyn std::error::Error>> {
4837 let mut view = dst.slice_mut(0..v.len());
4838 self.gpu.stream().memcpy_htod(v, &mut view)?;
4839 Ok(())
4840 }
4841
4842 pub fn copy_indirect_src_f32(
4847 &self,
4848 src_entry: &cudarc::driver::CudaView<u64>,
4849 dst: &mut CudaSlice<f32>,
4850 dst_off: usize,
4851 words: usize,
4852 ) -> Result<(), Box<dyn std::error::Error>> {
4853 let f = self.func("copy_indirect_src_f32");
4854 let chunks = (words / 4).max(1).div_ceil(256).min(48) as u32;
4855 let wi = words as i32;
4856 let cfg = LaunchConfig {
4857 grid_dim: (chunks, 1, 1),
4858 block_dim: (256, 1, 1),
4859 shared_mem_bytes: 0,
4860 };
4861 let mut dv = dst.slice_mut(dst_off..dst_off + words);
4862 let __s = self.gpu.stream();
4863 let mut b = __s.launch_builder(&f);
4864 b.arg(src_entry).arg(&mut dv).arg(&wi);
4865 unsafe {
4866 b.launch(cfg)?;
4867 }
4868 Ok(())
4869 }
4870
4871 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4873 self.alloc_uninit::<i8>(n)
4874 }
4875
4876 pub fn qmatvec(
4878 &self,
4879 w: &CudaSlice<u8>,
4880 x: &CudaSlice<f32>,
4881 m: usize,
4882 in_f: usize,
4883 out_f: usize,
4884 qtype: i32,
4885 row_bytes: usize,
4886 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4887 let f = self.func("qmatvec_f32");
4888 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
4890 grid_dim: (out_f as u32, m as u32, 1),
4891 block_dim: (256, 1, 1),
4892 shared_mem_bytes: 0,
4893 };
4894 let (inf, outf, mi, qt, rb) =
4895 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
4896 let __s_b = self.gpu.stream();
4897 let mut b = __s_b.launch_builder(&f);
4898 b.arg(w)
4899 .arg(x)
4900 .arg(&mut y)
4901 .arg(&inf)
4902 .arg(&outf)
4903 .arg(&mi)
4904 .arg(&qt)
4905 .arg(&rb);
4906 unsafe {
4907 b.launch(cfg)?;
4908 }
4909 Ok(y)
4910 }
4911
4912 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4914 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
4915 self.keep_if_capturing(&s);
4916 Ok(s)
4917 }
4918
4919 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4923 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
4924 self.keep_if_capturing(&s);
4925 Ok(s)
4926 }
4927
4928 pub fn memset_zeros_view(
4931 &self,
4932 dst: &mut cudarc::driver::CudaViewMut<f32>,
4933 ) -> Result<(), Box<dyn std::error::Error>> {
4934 self.gpu.stream().memset_zeros(dst)?;
4935 Ok(())
4936 }
4937
4938 pub fn stage_expert(
4944 &self,
4945 host_bytes: &[u8],
4946 scratch: &mut CudaSlice<u8>,
4947 off: usize,
4948 ) -> Result<(), Box<dyn std::error::Error>> {
4949 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
4952 }
4953
4954 pub fn moe_router_topk(
4960 &self,
4961 logits: &CudaSlice<f32>,
4962 t: usize,
4963 n_expert: usize,
4964 n_used: usize,
4965 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4966 let f = self.func("moe_router_topk_f32");
4967 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 {
4970 grid_dim: (t as u32, 1, 1),
4971 block_dim: (n_expert as u32, 1, 1),
4972 shared_mem_bytes: 0,
4973 };
4974 let (ne, nu) = (n_expert as i32, n_used as i32);
4975 let __s_b = self.gpu.stream();
4976 let mut b = __s_b.launch_builder(&f);
4977 b.arg(logits)
4978 .arg(&mut sel_idx)
4979 .arg(&mut sel_w)
4980 .arg(&ne)
4981 .arg(&nu);
4982 unsafe {
4983 b.launch(cfg)?;
4984 }
4985 Ok((sel_idx, sel_w))
4986 }
4987
4988 pub fn moe_router_topk_scaled(
4991 &self,
4992 logits: &CudaSlice<f32>,
4993 t: usize,
4994 n_expert: usize,
4995 n_used: usize,
4996 ex_scale: &CudaSlice<f32>,
4997 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4998 let f = self.func("moe_router_topk_scaled_f32");
5003 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
5004 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
5005 let cfg = LaunchConfig {
5006 grid_dim: (t as u32, 1, 1),
5007 block_dim: (n_expert as u32, 1, 1),
5008 shared_mem_bytes: 0,
5009 };
5010 let (ne, nu) = (n_expert as i32, n_used as i32);
5011 let __s_b = self.gpu.stream();
5012 let mut b = __s_b.launch_builder(&f);
5013 b.arg(logits)
5014 .arg(&mut sel_idx)
5015 .arg(&mut sel_w)
5016 .arg(&ne)
5017 .arg(&nu)
5018 .arg(ex_scale);
5019 unsafe {
5020 b.launch(cfg)?;
5021 }
5022 Ok((sel_idx, sel_w))
5023 }
5024
5025 pub fn moe_router_topk_host(
5033 &self,
5034 logits: &CudaSlice<f32>,
5035 t: usize,
5036 n_expert: usize,
5037 n_used: usize,
5038 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5039 let f = self.func("moe_router_topk_f32");
5040 let n = t * n_used;
5041 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
5042 let mut sel_w = self.alloc_uninit::<f32>(n)?;
5043 let cfg = LaunchConfig {
5044 grid_dim: (t as u32, 1, 1),
5045 block_dim: (n_expert as u32, 1, 1),
5046 shared_mem_bytes: 0,
5047 };
5048 let (ne, nu) = (n_expert as i32, n_used as i32);
5049 let __s_b = self.gpu.stream();
5050 let mut b = __s_b.launch_builder(&f);
5051 b.arg(logits)
5052 .arg(&mut sel_idx)
5053 .arg(&mut sel_w)
5054 .arg(&ne)
5055 .arg(&nu);
5056 unsafe {
5057 b.launch(cfg)?;
5058 }
5059 let bytes = n * 8;
5061 let mut guard = self.router_stage.lock().unwrap();
5062 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5063 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5064 }
5065 let stage = guard.as_mut().unwrap();
5066 let (si, sw) = unsafe {
5067 (
5068 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5069 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5070 )
5071 };
5072 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()))
5076 }
5077
5078 #[allow(clippy::too_many_arguments)]
5082 pub fn moe_router_sigmoid_topk(
5083 &self,
5084 logits: &CudaSlice<f32>,
5085 t: usize,
5086 n_expert: usize,
5087 n_used: usize,
5088 active_count: usize,
5089 correction_bias: &CudaSlice<f32>,
5090 active: &CudaSlice<u8>,
5091 scaling_factor: f32,
5092 route_norm: bool,
5093 ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5094 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5095 if n_expert == 0 || n_expert > 1024 || n_used == 0 || n_used > n_expert {
5096 return Err(format!(
5097 "sigmoid router shape unsupported: n_expert={n_expert}, n_used={n_used}",
5098 )
5099 .into());
5100 }
5101 if logits.len() < t * n_expert
5102 || correction_bias.len() != n_expert
5103 || active.len() != n_expert
5104 {
5105 return Err(format!(
5106 "sigmoid router buffer mismatch: logits={} bias={} active={} expected logits>={} row={}",
5107 logits.len(), correction_bias.len(), active.len(), t * n_expert, n_expert,
5108 ).into());
5109 }
5110 let f = self.func(crate::sigmoid_topk_kernel(
5111 crate::sig_expf_dev_on(),
5112 crate::topk_fast_on(),
5113 n_used,
5114 ));
5115 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
5116 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
5117 let threads = n_expert.div_ceil(32) * 32;
5118 let cfg = LaunchConfig {
5119 grid_dim: (t as u32, 1, 1),
5120 block_dim: (threads as u32, 1, 1),
5121 shared_mem_bytes: 0,
5122 };
5123 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5124 let __s_b = self.gpu.stream();
5125 let mut b = __s_b.launch_builder(&f);
5126 b.arg(logits)
5127 .arg(correction_bias)
5128 .arg(active)
5129 .arg(&mut sel_idx)
5130 .arg(&mut sel_w)
5131 .arg(&ne)
5132 .arg(&nu)
5133 .arg(&scaling_factor)
5134 .arg(&rn);
5135 unsafe {
5136 b.launch(cfg)?;
5137 }
5138 Ok((sel_idx, sel_w))
5139 }
5140
5141 #[allow(clippy::too_many_arguments)]
5144 pub fn ring_flag_raw(&self, ptr: u64, value: u32) -> Result<(), Box<dyn std::error::Error>> {
5148 if ptr == 0 {
5149 return Err("ring_flag_raw: unarmed flag".into());
5150 }
5151 let f = self.func("memra_ring_flag");
5152 let cfg = LaunchConfig {
5153 grid_dim: (1, 1, 1),
5154 block_dim: (32, 1, 1),
5155 shared_mem_bytes: 0,
5156 };
5157 let __s_b = self.gpu.stream();
5158 let mut b = __s_b.launch_builder(&f);
5159 b.arg(&ptr).arg(&value);
5160 unsafe {
5161 b.launch(cfg)?;
5162 }
5163 Ok(())
5164 }
5165
5166 pub fn moe_sel_w_mirror(
5169 &self,
5170 sel_src: &CudaSlice<i32>,
5171 w_src: &CudaSlice<f32>,
5172 sel_dst: &mut CudaSlice<i32>,
5173 w_dst: &mut CudaSlice<f32>,
5174 n: usize,
5175 ) -> Result<(), Box<dyn std::error::Error>> {
5176 if n == 0
5177 || n > 32
5178 || sel_src.len() < n
5179 || w_src.len() < n
5180 || sel_dst.len() < n
5181 || w_dst.len() < n
5182 {
5183 return Err(format!("moe_sel_w_mirror geometry n={n}").into());
5184 }
5185 let f = self.func("moe_sel_w_mirror");
5186 let cfg = LaunchConfig {
5187 grid_dim: (1, 1, 1),
5188 block_dim: (32, 1, 1),
5189 shared_mem_bytes: 0,
5190 };
5191 let ni = n as i32;
5192 let __s_b = self.gpu.stream();
5193 let mut b = __s_b.launch_builder(&f);
5194 b.arg(sel_src).arg(w_src).arg(sel_dst).arg(w_dst).arg(&ni);
5195 unsafe {
5196 b.launch(cfg)?;
5197 }
5198 Ok(())
5199 }
5200
5201 pub fn moe_router_sigmoid_topk_into(
5202 &self,
5203 logits: &CudaSlice<f32>,
5204 t: usize,
5205 n_expert: usize,
5206 n_used: usize,
5207 active_count: usize,
5208 correction_bias: &CudaSlice<f32>,
5209 active: &CudaSlice<u8>,
5210 scaling_factor: f32,
5211 route_norm: bool,
5212 sel_idx: &mut CudaSlice<i32>,
5213 sel_w: &mut CudaSlice<f32>,
5214 ) -> Result<(), Box<dyn std::error::Error>> {
5215 crate::sigrouter_contract::validate_active_count(n_used, active_count)?;
5216 if n_expert == 0
5217 || n_expert > 1024
5218 || n_used == 0
5219 || n_used > 32 || n_used > n_expert
5221 || logits.len() < t * n_expert
5222 || correction_bias.len() != n_expert
5223 || active.len() != n_expert
5224 || sel_idx.len() < t * n_used
5225 || sel_w.len() < t * n_used
5226 {
5227 return Err("sigmoid router _into geometry mismatch".into());
5228 }
5229 let f = self.func(crate::sigmoid_topk_kernel(
5230 crate::sig_expf_dev_on(),
5231 crate::topk_fast_on(),
5232 n_used,
5233 ));
5234 let threads = n_expert.div_ceil(32) * 32;
5235 let cfg = LaunchConfig {
5236 grid_dim: (t as u32, 1, 1),
5237 block_dim: (threads as u32, 1, 1),
5238 shared_mem_bytes: 0,
5239 };
5240 let (ne, nu, rn) = (n_expert as i32, n_used as i32, i32::from(route_norm));
5241 let __s_b = self.gpu.stream();
5242 let mut b = __s_b.launch_builder(&f);
5243 b.arg(logits)
5244 .arg(correction_bias)
5245 .arg(active)
5246 .arg(&mut *sel_idx)
5247 .arg(&mut *sel_w)
5248 .arg(&ne)
5249 .arg(&nu)
5250 .arg(&scaling_factor)
5251 .arg(&rn);
5252 unsafe {
5253 b.launch(cfg)?;
5254 }
5255 Ok(())
5256 }
5257
5258 #[allow(clippy::too_many_arguments)]
5261 pub fn moe_router_sigmoid_topk_host(
5262 &self,
5263 logits: &CudaSlice<f32>,
5264 t: usize,
5265 n_expert: usize,
5266 n_used: usize,
5267 active_count: usize,
5268 correction_bias: &CudaSlice<f32>,
5269 active: &CudaSlice<u8>,
5270 scaling_factor: f32,
5271 route_norm: bool,
5272 ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
5273 let (sel_idx, sel_w) = self.moe_router_sigmoid_topk(
5274 logits,
5275 t,
5276 n_expert,
5277 n_used,
5278 active_count,
5279 correction_bias,
5280 active,
5281 scaling_factor,
5282 route_norm,
5283 )?;
5284 let n = t * n_used;
5285 let bytes = n * 8;
5286 let mut guard = self.router_stage.lock().unwrap();
5287 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
5288 *guard = Some(PinnedStage::new(bytes.max(4096))?);
5289 }
5290 let stage = guard.as_mut().unwrap();
5291 let (si, sw) = unsafe {
5292 (
5293 std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
5294 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n),
5295 )
5296 };
5297 self.gpu.stream().memcpy_dtoh(&sel_idx, si)?;
5298 self.gpu.stream().memcpy_dtoh(&sel_w, sw)?;
5299 self.gpu.stream().synchronize()?;
5300 Ok((si.iter().map(|&i| i as u32).collect(), sw.to_vec()))
5301 }
5302
5303 pub fn stage_expert_async(
5307 &self,
5308 host_bytes: &[u8],
5309 scratch: &mut CudaSlice<u8>,
5310 off: usize,
5311 ) -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
5312 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
5313 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
5314 Ok(self.copy_stream.record_event(None)?)
5315 }
5316
5317 pub fn compute_wait(
5319 &self,
5320 ev: &cudarc::driver::CudaEvent,
5321 ) -> Result<(), Box<dyn std::error::Error>> {
5322 self.gpu.stream().wait(ev)?;
5323 Ok(())
5324 }
5325
5326 pub fn qmatvec_view(
5331 &self,
5332 w: &CudaSlice<u8>,
5333 range: std::ops::Range<usize>,
5334 x: &cudarc::driver::CudaView<f32>,
5335 m: usize,
5336 in_f: usize,
5337 out_f: usize,
5338 qtype: i32,
5339 row_bytes: usize,
5340 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5341 let f = self.func("qmatvec_f32");
5342 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
5345 grid_dim: (out_f as u32, m as u32, 1),
5346 block_dim: (256, 1, 1),
5347 shared_mem_bytes: 0,
5348 };
5349 let (inf, outf, mi, qt, rb) =
5350 (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
5351 let __s_b = self.gpu.stream();
5352 let mut b = __s_b.launch_builder(&f);
5353 b.arg(&wv)
5354 .arg(x)
5355 .arg(&mut y)
5356 .arg(&inf)
5357 .arg(&outf)
5358 .arg(&mi)
5359 .arg(&qt)
5360 .arg(&rb);
5361 unsafe {
5362 b.launch(cfg)?;
5363 }
5364 Ok(y)
5365 }
5366
5367 #[allow(clippy::too_many_arguments)]
5374 pub fn moe_gate_up_silu8_q8(
5378 &self,
5379 gp: WPtr8,
5380 up: WPtr8,
5381 aq: &CudaSlice<i8>,
5382 ad: &CudaSlice<f32>,
5383 in_f: usize,
5384 n_ff: usize,
5385 n_used: usize,
5386 qt_g: i32,
5387 qt_u: i32,
5388 rb_g: usize,
5389 rb_u: usize,
5390 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5391 let f = self.func("moe_gate_up_silu8_q8");
5392 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5393 let cfg = LaunchConfig {
5394 grid_dim: (n_ff as u32, n_used as u32, 1),
5395 block_dim: (32, 1, 1),
5396 shared_mem_bytes: 0,
5397 };
5398 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5399 let __s_b = self.gpu.stream();
5400 let mut b = __s_b.launch_builder(&f);
5401 b.arg(&gp)
5402 .arg(&up)
5403 .arg(aq)
5404 .arg(ad)
5405 .arg(&mut act)
5406 .arg(&inf)
5407 .arg(&nff)
5408 .arg(&qt_g)
5409 .arg(&qt_u)
5410 .arg(&rbg)
5411 .arg(&rbu);
5412 unsafe {
5413 b.launch(cfg)?;
5414 }
5415 Ok(act)
5416 }
5417
5418 #[allow(clippy::too_many_arguments)]
5419 pub fn moe_down8_fma_q8(
5420 &self,
5421 dp: WPtr8,
5422 w: F32x8,
5423 aq2: &CudaSlice<i8>,
5424 ad2: &CudaSlice<f32>,
5425 dst: &mut cudarc::driver::CudaViewMut<f32>,
5426 in_f: usize,
5427 out_f: usize,
5428 n_used: usize,
5429 qt: i32,
5430 rb: usize,
5431 ) -> Result<(), Box<dyn std::error::Error>> {
5432 let f = self.func("moe_down8_fma_q8");
5433 let cfg = LaunchConfig {
5434 grid_dim: (out_f as u32, 1, 1),
5435 block_dim: (32, 1, 1),
5436 shared_mem_bytes: 0,
5437 };
5438 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5439 let __s_b = self.gpu.stream();
5440 let mut b = __s_b.launch_builder(&f);
5441 b.arg(&dp)
5442 .arg(&w)
5443 .arg(aq2)
5444 .arg(ad2)
5445 .arg(dst)
5446 .arg(&inf)
5447 .arg(&outf)
5448 .arg(&nu)
5449 .arg(&qt)
5450 .arg(&rbi);
5451 unsafe {
5452 b.launch(cfg)?;
5453 }
5454 Ok(())
5455 }
5456
5457 pub fn qmatvec_expert_q8(
5459 &self,
5460 w: &CudaSlice<u8>,
5461 range: std::ops::Range<usize>,
5462 aq: &CudaSlice<i8>,
5463 ad: &CudaSlice<f32>,
5464 m: usize,
5465 in_f: usize,
5466 out_f: usize,
5467 qtype: i32,
5468 row_bytes: usize,
5469 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5470 let f = self.func("qmatvec_expert_q8");
5471 let wv = w.slice(range);
5472 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
5473 const ROWS: u32 = 4; let cfg = LaunchConfig {
5475 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
5476 block_dim: (32, ROWS, 1),
5477 shared_mem_bytes: 0,
5478 };
5479 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5480 let __s_b = self.gpu.stream();
5481 let mut b = __s_b.launch_builder(&f);
5482 b.arg(&wv)
5483 .arg(aq)
5484 .arg(ad)
5485 .arg(&mut y)
5486 .arg(&inf)
5487 .arg(&outf)
5488 .arg(&mi)
5489 .arg(&qtype)
5490 .arg(&rbi);
5491 unsafe {
5492 b.launch(cfg)?;
5493 }
5494 Ok(y)
5495 }
5496
5497 pub fn moe_gate_up_silu8(
5498 &self,
5499 gp: WPtr8,
5500 up: WPtr8,
5501 x: &cudarc::driver::CudaView<f32>,
5502 in_f: usize,
5503 n_ff: usize,
5504 n_used: usize,
5505 qt_g: i32,
5506 qt_u: i32,
5507 rb_g: usize,
5508 rb_u: usize,
5509 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5510 let f = self.func("moe_gate_up_silu8_f32");
5511 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
5513 grid_dim: (n_ff as u32, n_used as u32, 1),
5514 block_dim: (256, 1, 1),
5515 shared_mem_bytes: 0,
5516 };
5517 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
5518 let __s_b = self.gpu.stream();
5519 let mut b = __s_b.launch_builder(&f);
5520 b.arg(&gp)
5521 .arg(&up)
5522 .arg(x)
5523 .arg(&mut act)
5524 .arg(&inf)
5525 .arg(&nff)
5526 .arg(&qt_g)
5527 .arg(&qt_u)
5528 .arg(&rbg)
5529 .arg(&rbu);
5530 unsafe {
5531 b.launch(cfg)?;
5532 }
5533 Ok(act)
5534 }
5535
5536 #[allow(clippy::too_many_arguments)]
5542 pub fn moe_down8_fma_into(
5543 &self,
5544 dp: WPtr8,
5545 w: F32x8,
5546 act: &CudaSlice<f32>,
5547 dst: &mut cudarc::driver::CudaViewMut<f32>,
5548 in_f: usize,
5549 out_f: usize,
5550 n_used: usize,
5551 qt: i32,
5552 rb: usize,
5553 ) -> Result<(), Box<dyn std::error::Error>> {
5554 let f = self.func("moe_down8_fma_f32");
5555 let cfg = LaunchConfig {
5556 grid_dim: (out_f as u32, 1, 1),
5557 block_dim: (256, 1, 1),
5558 shared_mem_bytes: 0,
5559 };
5560 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
5561 let __s_b = self.gpu.stream();
5562 let mut b = __s_b.launch_builder(&f);
5563 b.arg(&dp)
5564 .arg(&w)
5565 .arg(act)
5566 .arg(dst)
5567 .arg(&inf)
5568 .arg(&outf)
5569 .arg(&nu)
5570 .arg(&qt)
5571 .arg(&rbv);
5572 unsafe {
5573 b.launch(cfg)?;
5574 }
5575 Ok(())
5576 }
5577
5578 #[allow(clippy::too_many_arguments)]
5583 #[allow(clippy::too_many_arguments)]
5598 #[allow(clippy::too_many_arguments)]
5600 pub fn moe_pairs_matvec_q8(
5601 &self,
5602 table: &CudaSlice<u64>,
5603 proj: i32,
5604 pair_tok: &CudaSlice<i32>,
5605 pair_ex: &CudaSlice<i32>,
5606 aq: &CudaSlice<i8>,
5607 ad: &CudaSlice<f32>,
5608 in_f: usize,
5609 out_f: usize,
5610 n_expert: usize,
5611 n_pairs: usize,
5612 qtype: i32,
5613 row_bytes: usize,
5614 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5615 let f = self.func("moe_pairs_matvec_q8");
5616 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5617 const ROWS: u32 = 4;
5618 let cfg = LaunchConfig {
5619 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
5620 block_dim: (32, ROWS, 1),
5621 shared_mem_bytes: 0,
5622 };
5623 let (inf, outf, ne, np, rbi) = (
5624 in_f as i32,
5625 out_f as i32,
5626 n_expert as i32,
5627 n_pairs as i32,
5628 row_bytes as i64,
5629 );
5630 let __s_b = self.gpu.stream();
5631 let mut b = __s_b.launch_builder(&f);
5632 b.arg(table)
5633 .arg(&proj)
5634 .arg(pair_tok)
5635 .arg(pair_ex)
5636 .arg(aq)
5637 .arg(ad)
5638 .arg(&mut y)
5639 .arg(&inf)
5640 .arg(&outf)
5641 .arg(&ne)
5642 .arg(&np)
5643 .arg(&qtype)
5644 .arg(&rbi);
5645 unsafe {
5646 b.launch(cfg)?;
5647 }
5648 Ok(y)
5649 }
5650
5651 #[allow(clippy::too_many_arguments)]
5653 pub fn moe_pairs_matvec_q8_em(
5654 &self,
5655 table: &CudaSlice<u64>,
5656 proj: i32,
5657 ex_ids: &CudaSlice<i32>,
5658 ex_off: &CudaSlice<i32>,
5659 ex_pairs: &CudaSlice<i32>,
5660 pair_tok: &CudaSlice<i32>,
5661 aq: &CudaSlice<i8>,
5662 ad: &CudaSlice<f32>,
5663 in_f: usize,
5664 out_f: usize,
5665 n_expert: usize,
5666 n_active: usize,
5667 n_pairs: usize,
5668 qtype: i32,
5669 row_bytes: usize,
5670 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5671 let f = self.func("moe_pairs_matvec_q8_em");
5672 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5673 const ROWS: u32 = 4;
5674 let cfg = LaunchConfig {
5675 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5676 block_dim: (32, ROWS, 1),
5677 shared_mem_bytes: 0,
5678 };
5679 let (inf, outf, ne, na, rbi) = (
5680 in_f as i32,
5681 out_f as i32,
5682 n_expert as i32,
5683 n_active as i32,
5684 row_bytes as i64,
5685 );
5686 let __s_b = self.gpu.stream();
5687 let mut b = __s_b.launch_builder(&f);
5688 b.arg(table)
5689 .arg(&proj)
5690 .arg(ex_ids)
5691 .arg(ex_off)
5692 .arg(ex_pairs)
5693 .arg(pair_tok)
5694 .arg(aq)
5695 .arg(ad)
5696 .arg(&mut y)
5697 .arg(&inf)
5698 .arg(&outf)
5699 .arg(&ne)
5700 .arg(&na)
5701 .arg(&qtype)
5702 .arg(&rbi);
5703 unsafe {
5704 b.launch(cfg)?;
5705 }
5706 Ok(y)
5707 }
5708
5709 #[allow(clippy::too_many_arguments)]
5712 pub fn moe_pairs_matvec_q8_dec(
5713 &self,
5714 table: &CudaSlice<u64>,
5715 proj: i32,
5716 ex_ids: &CudaSlice<i32>,
5717 ex_off: &CudaSlice<i32>,
5718 ex_pairs: &CudaSlice<i32>,
5719 pair_tok: &CudaSlice<i32>,
5720 aq: &CudaSlice<i8>,
5721 ad: &CudaSlice<f32>,
5722 in_f: usize,
5723 out_f: usize,
5724 n_expert: usize,
5725 n_active: usize,
5726 n_pairs: usize,
5727 qtype: i32,
5728 row_bytes: usize,
5729 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5730 let f = self.func("moe_pairs_matvec_q8_dec");
5731 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
5732 const ROWS: u32 = 4;
5733 let cfg = LaunchConfig {
5734 grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
5735 block_dim: (32, ROWS, 1),
5736 shared_mem_bytes: 0,
5737 };
5738 let (inf, outf, ne, na, rbi) = (
5739 in_f as i32,
5740 out_f as i32,
5741 n_expert as i32,
5742 n_active as i32,
5743 row_bytes as i64,
5744 );
5745 let __s_b = self.gpu.stream();
5746 let mut b = __s_b.launch_builder(&f);
5747 b.arg(table)
5748 .arg(&proj)
5749 .arg(ex_ids)
5750 .arg(ex_off)
5751 .arg(ex_pairs)
5752 .arg(pair_tok)
5753 .arg(aq)
5754 .arg(ad)
5755 .arg(&mut y)
5756 .arg(&inf)
5757 .arg(&outf)
5758 .arg(&ne)
5759 .arg(&na)
5760 .arg(&qtype)
5761 .arg(&rbi);
5762 unsafe {
5763 b.launch(cfg)?;
5764 }
5765 Ok(y)
5766 }
5767
5768 pub fn moe_pairs_gelu_mul(
5769 &self,
5770 gate: &CudaSlice<f32>,
5771 up: &CudaSlice<f32>,
5772 n: usize,
5773 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5774 let f = self.func("moe_pairs_gelu_mul");
5775 let mut act = self.alloc_uninit::<f32>(n)?;
5776 let cfg = LaunchConfig::for_num_elems(n as u32);
5777 let nl = n as i64;
5778 let __s_b = self.gpu.stream();
5779 let mut b = __s_b.launch_builder(&f);
5780 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5781 unsafe {
5782 b.launch(cfg)?;
5783 }
5784 Ok(act)
5785 }
5786
5787 pub fn moe_pairs_silu_mul(
5788 &self,
5789 gate: &CudaSlice<f32>,
5790 up: &CudaSlice<f32>,
5791 n: usize,
5792 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5793 let f = self.func("moe_pairs_silu_mul");
5794 let mut act = self.alloc_uninit::<f32>(n)?;
5795 let cfg = LaunchConfig::for_num_elems(n as u32);
5796 let nl = n as i64;
5797 let __s_b = self.gpu.stream();
5798 let mut b = __s_b.launch_builder(&f);
5799 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
5800 unsafe {
5801 b.launch(cfg)?;
5802 }
5803 Ok(act)
5804 }
5805
5806 #[allow(clippy::too_many_arguments)]
5807 pub fn moe_pairs_scatter(
5808 &self,
5809 y_down: &CudaSlice<f32>,
5810 pair_w: &CudaSlice<f32>,
5811 tok_pair_off: &CudaSlice<i32>,
5812 tok_pair_ids: &CudaSlice<i32>,
5813 moe_out: &mut CudaSlice<f32>,
5814 t: usize,
5815 n_embd: usize,
5816 ) -> Result<(), Box<dyn std::error::Error>> {
5817 let f = self.func("moe_pairs_scatter");
5818 let cfg = LaunchConfig {
5819 grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
5820 block_dim: (256, 1, 1),
5821 shared_mem_bytes: 0,
5822 };
5823 let ne = n_embd as i32;
5824 let __s_b = self.gpu.stream();
5825 let mut b = __s_b.launch_builder(&f);
5826 b.arg(y_down)
5827 .arg(pair_w)
5828 .arg(tok_pair_off)
5829 .arg(tok_pair_ids)
5830 .arg(moe_out)
5831 .arg(&ne);
5832 unsafe {
5833 b.launch(cfg)?;
5834 }
5835 Ok(())
5836 }
5837
5838 #[allow(clippy::too_many_arguments)]
5842 pub fn moe_gate_up_gelu8_dev_q8(
5843 &self,
5844 table: &CudaSlice<u64>,
5845 sel: &cudarc::driver::CudaView<i32>,
5846 aq: &CudaSlice<i8>,
5847 ad: &CudaSlice<f32>,
5848 in_f: usize,
5849 n_ff: usize,
5850 n_used: usize,
5851 n_expert: usize,
5852 qt_g: i32,
5853 qt_u: i32,
5854 rb_g: usize,
5855 rb_u: usize,
5856 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5857 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
5858 let (inf, nff, ne, rbg, rbu) = (
5859 in_f as i32,
5860 n_ff as i32,
5861 n_expert as i32,
5862 rb_g as i64,
5863 rb_u as i64,
5864 );
5865 let f = self.func("moe_gate_up_gelu8_dev_q8");
5866 let cfg = LaunchConfig {
5867 grid_dim: (n_ff as u32, n_used as u32, 1),
5868 block_dim: (32, 1, 1),
5869 shared_mem_bytes: 0,
5870 };
5871 let __s_b = self.gpu.stream();
5872 let mut b = __s_b.launch_builder(&f);
5873 b.arg(table)
5874 .arg(sel)
5875 .arg(aq)
5876 .arg(ad)
5877 .arg(&mut act)
5878 .arg(&inf)
5879 .arg(&nff)
5880 .arg(&ne)
5881 .arg(&qt_g)
5882 .arg(&qt_u)
5883 .arg(&rbg)
5884 .arg(&rbu);
5885 unsafe {
5886 b.launch(cfg)?;
5887 }
5888 Ok(act)
5889 }
5890
5891 #[allow(clippy::too_many_arguments)]
5893 pub fn moe_gate_up_gelu8_dev_q8_rows(
5894 &self,
5895 table: &CudaSlice<u64>,
5896 sel: &CudaSlice<i32>,
5897 aq: &CudaSlice<i8>,
5898 ad: &CudaSlice<f32>,
5899 t: usize,
5900 in_f: usize,
5901 n_ff: usize,
5902 n_used: usize,
5903 n_expert: usize,
5904 qt_g: i32,
5905 qt_u: i32,
5906 rb_g: usize,
5907 rb_u: usize,
5908 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5909 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
5910 let (inf, nff, ne, rbg, rbu, nu) = (
5911 in_f as i32,
5912 n_ff as i32,
5913 n_expert as i32,
5914 rb_g as i64,
5915 rb_u as i64,
5916 n_used as i32,
5917 );
5918 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
5919 let cfg = LaunchConfig {
5920 grid_dim: (n_ff as u32, n_used as u32, t as u32),
5921 block_dim: (32, 1, 1),
5922 shared_mem_bytes: 0,
5923 };
5924 let __s_b = self.gpu.stream();
5925 let mut b = __s_b.launch_builder(&f);
5926 b.arg(table)
5927 .arg(sel)
5928 .arg(aq)
5929 .arg(ad)
5930 .arg(&mut act)
5931 .arg(&inf)
5932 .arg(&nff)
5933 .arg(&ne)
5934 .arg(&qt_g)
5935 .arg(&qt_u)
5936 .arg(&rbg)
5937 .arg(&rbu)
5938 .arg(&nu);
5939 unsafe {
5940 b.launch(cfg)?;
5941 }
5942 Ok(act)
5943 }
5944
5945 #[allow(clippy::too_many_arguments)]
5947 pub fn moe_gate_up_gelu8_dev_q8_csr(
5948 &self,
5949 table: &CudaSlice<u64>,
5950 sel: &CudaSlice<i32>,
5951 aq: &CudaSlice<i8>,
5952 ad: &CudaSlice<f32>,
5953 n_pairs: usize,
5954 in_f: usize,
5955 n_ff: usize,
5956 n_used: usize,
5957 n_expert: usize,
5958 qt_g: i32,
5959 qt_u: i32,
5960 rb_g: usize,
5961 rb_u: usize,
5962 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5963 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
5964 let (inf, nff, ne, rbg, rbu, nu, npi) = (
5965 in_f as i32,
5966 n_ff as i32,
5967 n_expert as i32,
5968 rb_g as i64,
5969 rb_u as i64,
5970 n_used as i32,
5971 n_pairs as i32,
5972 );
5973 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
5974 let cfg = LaunchConfig {
5975 grid_dim: (n_ff as u32, n_pairs as u32, 1),
5976 block_dim: (32, 1, 1),
5977 shared_mem_bytes: 0,
5978 };
5979 let __s_b = self.gpu.stream();
5980 let mut b = __s_b.launch_builder(&f);
5981 b.arg(table)
5982 .arg(sel)
5983 .arg(aq)
5984 .arg(ad)
5985 .arg(&mut act)
5986 .arg(&inf)
5987 .arg(&nff)
5988 .arg(&ne)
5989 .arg(&qt_g)
5990 .arg(&qt_u)
5991 .arg(&rbg)
5992 .arg(&rbu)
5993 .arg(&nu)
5994 .arg(&npi);
5995 unsafe {
5996 b.launch(cfg)?;
5997 }
5998 Ok(act)
5999 }
6000
6001 #[allow(clippy::too_many_arguments)]
6003 pub fn moe_down8_fma_dev_q8_rows_g(
6004 &self,
6005 table: &CudaSlice<u64>,
6006 sel: &CudaSlice<i32>,
6007 w: &CudaSlice<f32>,
6008 aq2: &CudaSlice<i8>,
6009 ad2: &CudaSlice<f32>,
6010 dst: &mut CudaSlice<f32>,
6011 t: usize,
6012 in_f: usize,
6013 out_f: usize,
6014 n_used: usize,
6015 n_expert: usize,
6016 qt: i32,
6017 rb: usize,
6018 ) -> Result<(), Box<dyn std::error::Error>> {
6019 let (inf, outf, nu, ne, rbi) = (
6020 in_f as i32,
6021 out_f as i32,
6022 n_used as i32,
6023 n_expert as i32,
6024 rb as i64,
6025 );
6026 let step_b1_w8 = t == 1 && in_f == 1280 && out_f == 4096 && n_used == 8 && qt == QT_IQ4_XS;
6030 let f = self.func(if step_b1_w8 {
6031 "moe_down8_fma_dev_q8_rows_w8"
6032 } else {
6033 "moe_down8_fma_dev_q8_rows_g"
6034 });
6035 let cfg = LaunchConfig {
6036 grid_dim: (out_f as u32, 1, t as u32),
6037 block_dim: (32, if step_b1_w8 { 8 } else { 1 }, 1),
6038 shared_mem_bytes: 0,
6039 };
6040 let __s_b = self.gpu.stream();
6041 let mut b = __s_b.launch_builder(&f);
6042 b.arg(table)
6043 .arg(sel)
6044 .arg(w)
6045 .arg(aq2)
6046 .arg(ad2)
6047 .arg(dst)
6048 .arg(&inf)
6049 .arg(&outf)
6050 .arg(&nu)
6051 .arg(&ne)
6052 .arg(&qt)
6053 .arg(&rbi);
6054 unsafe {
6055 b.launch(cfg)?;
6056 }
6057 Ok(())
6058 }
6059
6060 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
6064 let (out_f, in_f) = (2048usize, 2816usize);
6065 let nblk = in_f / 32;
6066 let mut seed = 0x9E3779B97F4A7C15u64;
6067 let mut rng = move || {
6068 seed = seed
6069 .wrapping_mul(6364136223846793005)
6070 .wrapping_add(1442695040888963407);
6071 (seed >> 33) as u8
6072 };
6073 let mut w = vec![0u8; out_f * nblk * 18];
6074 for b in w.iter_mut() {
6075 *b = rng();
6076 }
6077 for r in 0..out_f {
6078 for g in 0..nblk {
6079 let off = (r * nblk + g) * 18;
6080 w[off] = 0x00;
6081 w[off + 1] = 0x2C; }
6083 }
6084 let qplane = out_f * nblk * 16;
6085 let mut wrp = vec![0u8; w.len()];
6086 for r in 0..out_f {
6087 for g in 0..nblk {
6088 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
6089 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
6090 .copy_from_slice(&src[0..2]);
6091 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
6092 }
6093 }
6094 let w_d = self.htod_bytes(&w)?;
6095 let wrp_d = self.htod_bytes(&wrp)?;
6096 let mut aq = vec![0i8; m * in_f];
6097 for v in aq.iter_mut() {
6098 *v = rng() as i8;
6099 }
6100 let aq_d = self.htod_i8(&aq)?;
6101 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
6102 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
6103 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
6104 const RPB: u32 = 4;
6105 let cfg = LaunchConfig {
6106 grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
6107 block_dim: (32, RPB, 1),
6108 shared_mem_bytes: 0,
6109 };
6110 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
6111 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
6112 let fb = self.func("qmatvec_q4_0_mmvq_b4");
6113 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
6114 {
6115 let __s_b = self.gpu.stream();
6116 let mut b = __s_b.launch_builder(&fb);
6117 b.arg(&w_d)
6118 .arg(&aq_d)
6119 .arg(&ad_d)
6120 .arg(&mut y0)
6121 .arg(&inf)
6122 .arg(&outf)
6123 .arg(&mi)
6124 .arg(&rb);
6125 unsafe {
6126 b.launch(cfg)?;
6127 }
6128 let __s_b = self.gpu.stream();
6129 let mut b = __s_b.launch_builder(&fr);
6130 b.arg(&wrp_d)
6131 .arg(&aq_d)
6132 .arg(&ad_d)
6133 .arg(&mut y1)
6134 .arg(&inf)
6135 .arg(&outf)
6136 .arg(&mi)
6137 .arg(&qp);
6138 unsafe {
6139 b.launch(cfg)?;
6140 }
6141 }
6142 self.gpu.stream().synchronize()?;
6143 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
6144 let nd = h0
6145 .iter()
6146 .zip(&h1)
6147 .filter(|(a, b)| a.to_bits() != b.to_bits())
6148 .count();
6149 if nd != 0 {
6150 return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into());
6151 }
6152 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
6153 self.gpu.stream().synchronize()?;
6154 let t0 = std::time::Instant::now();
6155 for _ in 0..500 {
6156 if rp {
6157 let __s_b = self.gpu.stream();
6158 let mut b = __s_b.launch_builder(&fr);
6159 b.arg(&wrp_d)
6160 .arg(&aq_d)
6161 .arg(&ad_d)
6162 .arg(&mut y1)
6163 .arg(&inf)
6164 .arg(&outf)
6165 .arg(&mi)
6166 .arg(&qp);
6167 unsafe {
6168 b.launch(cfg)?;
6169 }
6170 } else {
6171 let __s_b = self.gpu.stream();
6172 let mut b = __s_b.launch_builder(&fb);
6173 b.arg(&w_d)
6174 .arg(&aq_d)
6175 .arg(&ad_d)
6176 .arg(&mut y0)
6177 .arg(&inf)
6178 .arg(&outf)
6179 .arg(&mi)
6180 .arg(&rb);
6181 unsafe {
6182 b.launch(cfg)?;
6183 }
6184 }
6185 }
6186 self.gpu.stream().synchronize()?;
6187 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
6188 };
6189 let _ = time(false)?;
6190 let _ = time(true)?; Ok((time(false)?, time(true)?))
6192 }
6193
6194 pub fn build_q4_rp4(
6199 &self,
6200 t: &mut crate::model::GpuTensor,
6201 ) -> Result<(), Box<dyn std::error::Error>> {
6202 use crate::model::GpuTensor;
6203 let GpuTensor::Quant {
6204 bytes,
6205 qtype,
6206 row_bytes,
6207 ne,
6208 rp4,
6209 ..
6210 } = t
6211 else {
6212 return Ok(());
6213 };
6214 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 {
6215 return Ok(());
6216 }
6217 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6218 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 {
6219 return Ok(());
6220 }
6221 let nblk = in_f / 32;
6222 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
6223 let f = self.func("q4_0_split_rp_build");
6224 let n = (out_f * nblk) as i32;
6225 let cfg = LaunchConfig {
6226 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6227 block_dim: (256, 1, 1),
6228 shared_mem_bytes: 0,
6229 };
6230 let (of, nb) = (out_f as i32, nblk as i32);
6231 let _ = n;
6232 let __s_b = self.gpu.stream();
6233 let mut b = __s_b.launch_builder(&f);
6234 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6235 unsafe {
6236 b.launch(cfg)?;
6237 }
6238 *rp4 = Some(dst);
6239 Ok(())
6240 }
6241
6242 pub fn build_q8_rp4(
6247 &self,
6248 t: &mut crate::model::GpuTensor,
6249 ) -> Result<(), Box<dyn std::error::Error>> {
6250 use crate::model::GpuTensor;
6251 let GpuTensor::Quant {
6252 bytes,
6253 qtype,
6254 row_bytes,
6255 ne,
6256 rp4,
6257 ..
6258 } = t
6259 else {
6260 return Ok(());
6261 };
6262 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 {
6263 return Ok(());
6264 }
6265 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6266 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 {
6267 return Ok(());
6268 }
6269 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
6270 Ok(())
6271 }
6272
6273 pub fn build_q8_rp4_raw(
6276 &self,
6277 bytes: &CudaSlice<u8>,
6278 in_f: usize,
6279 out_f: usize,
6280 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6281 assert!(in_f % 32 == 0);
6282 let nblk = in_f / 32;
6283 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
6284 let f = self.func("q8_0_split_rp_build");
6285 let cfg = LaunchConfig {
6286 grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
6287 block_dim: (256, 1, 1),
6288 shared_mem_bytes: 0,
6289 };
6290 let (of, nb) = (out_f as i32, nblk as i32);
6291 let __s_b = self.gpu.stream();
6292 let mut b = __s_b.launch_builder(&f);
6293 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6294 unsafe {
6295 b.launch(cfg)?;
6296 }
6297 Ok(dst)
6298 }
6299
6300 pub fn build_q4k_rp4(
6308 &self,
6309 t: &mut crate::model::GpuTensor,
6310 ) -> Result<(), Box<dyn std::error::Error>> {
6311 use crate::model::GpuTensor;
6312 let GpuTensor::Quant {
6313 bytes,
6314 qtype,
6315 row_bytes,
6316 ne,
6317 rp4,
6318 ..
6319 } = t
6320 else {
6321 return Ok(());
6322 };
6323 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 {
6324 return Ok(());
6325 }
6326 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6327 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 {
6328 return Ok(());
6329 }
6330 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
6331 Ok(())
6332 }
6333
6334 pub fn build_q6k_rp4(
6335 &self,
6336 t: &mut crate::model::GpuTensor,
6337 ) -> Result<(), Box<dyn std::error::Error>> {
6338 use crate::model::GpuTensor;
6339 let GpuTensor::Quant {
6340 bytes,
6341 qtype,
6342 row_bytes,
6343 ne,
6344 rp4,
6345 ..
6346 } = t
6347 else {
6348 return Ok(());
6349 };
6350 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 {
6351 return Ok(());
6352 }
6353 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
6354 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 {
6355 return Ok(());
6356 }
6357 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
6358 Ok(())
6359 }
6360
6361 pub fn build_kq_rp4_raw(
6363 &self,
6364 bytes: &CudaSlice<u8>,
6365 in_f: usize,
6366 out_f: usize,
6367 qtype: i32,
6368 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
6369 assert!(in_f % 256 == 0);
6370 let nsbk = in_f / 256;
6371 let (sb_bytes, kname) = match qtype {
6372 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
6373 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
6374 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
6375 };
6376 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
6377 let f = self.func(kname);
6378 let cfg = LaunchConfig {
6379 grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
6380 block_dim: (256, 1, 1),
6381 shared_mem_bytes: 0,
6382 };
6383 let (of, nb) = (out_f as i32, nsbk as i32);
6384 let __s_b = self.gpu.stream();
6385 let mut b = __s_b.launch_builder(&f);
6386 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
6387 unsafe {
6388 b.launch(cfg)?;
6389 }
6390 Ok(dst)
6391 }
6392
6393 pub fn kqrp_enabled() -> bool {
6397 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6398 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
6399 Ok("0") => false,
6400 Ok(_) => true,
6401 Err(_) => cfg!(memra_hopper_mma),
6402 })
6403 }
6404
6405 pub fn build_q4_rp_swap(
6411 &self,
6412 t: &mut crate::model::GpuTensor,
6413 ) -> Result<bool, Box<dyn std::error::Error>> {
6414 use crate::model::GpuTensor;
6415 if !matches!(t, GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0) {
6425 return Ok(false);
6426 }
6427 self.build_q4_rp4(t)?;
6428 self.gpu.stream().synchronize()?; let GpuTensor::Quant { bytes, rp4, rp, .. } = t else {
6430 return Ok(false);
6431 };
6432 match rp4.take() {
6433 Some(split) => {
6434 *bytes = split; *rp = true;
6436 Ok(true)
6437 }
6438 None => Ok(false),
6439 }
6440 }
6441
6442 pub fn q4rp_enabled() -> bool {
6444 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6445 *ON.get_or_init(|| {
6446 std::env::var("MEMRA_Q4RP")
6447 .map(|v| v != "0")
6448 .unwrap_or(true)
6449 })
6450 }
6451
6452 pub fn copy_rows_strided(
6455 &self,
6456 src: &CudaSlice<f32>,
6457 dst: &mut CudaSlice<f32>,
6458 row_elems: usize,
6459 n_rows: usize,
6460 src_stride: usize,
6461 src_off: usize,
6462 ) -> Result<(), Box<dyn std::error::Error>> {
6463 let f = self.func("copy_rows_strided_f32");
6464 let cfg = LaunchConfig {
6465 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6466 block_dim: (256, 1, 1),
6467 shared_mem_bytes: 0,
6468 };
6469 let (re, nr) = (row_elems as i32, n_rows as i32);
6470 let (st, off) = (src_stride as i64, src_off as i64);
6471 let __s_b = self.gpu.stream();
6472 let mut b = __s_b.launch_builder(&f);
6473 b.arg(src)
6474 .arg(&mut *dst)
6475 .arg(&re)
6476 .arg(&nr)
6477 .arg(&st)
6478 .arg(&off);
6479 unsafe {
6480 b.launch(cfg)?;
6481 }
6482 Ok(())
6483 }
6484
6485 pub fn place_rows_strided(
6491 &self,
6492 src: &CudaSlice<f32>,
6493 dst: &mut CudaSlice<f32>,
6494 row_elems: usize,
6495 n_rows: usize,
6496 dst_stride: usize,
6497 dst_off: usize,
6498 ) -> Result<(), Box<dyn std::error::Error>> {
6499 if row_elems == 0 || n_rows == 0 {
6500 return Err("strided row placement requires nonzero rows and row width".into());
6501 }
6502 let src_len = n_rows
6503 .checked_mul(row_elems)
6504 .ok_or("strided row placement source size overflow")?;
6505 let dst_len = n_rows
6506 .checked_sub(1)
6507 .and_then(|rows| rows.checked_mul(dst_stride))
6508 .and_then(|base| base.checked_add(dst_off))
6509 .and_then(|base| base.checked_add(row_elems))
6510 .ok_or("strided row placement destination size overflow")?;
6511 let row_end = dst_off
6512 .checked_add(row_elems)
6513 .ok_or("strided row placement row size overflow")?;
6514 if src.len() < src_len || dst.len() < dst_len || row_end > dst_stride {
6515 return Err(format!(
6516 "strided row placement geometry mismatch: src={} need_src={src_len} \
6517 dst={} need_dst={dst_len} row_elems={row_elems} rows={n_rows} \
6518 dst_stride={dst_stride} dst_off={dst_off}",
6519 src.len(),
6520 dst.len(),
6521 )
6522 .into());
6523 }
6524 if row_elems > i32::MAX as usize || n_rows > i32::MAX as usize {
6525 return Err("strided row placement exceeds CUDA kernel geometry".into());
6526 }
6527 let f = self.func("place_rows_strided_f32");
6528 let cfg = LaunchConfig {
6529 grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
6530 block_dim: (256, 1, 1),
6531 shared_mem_bytes: 0,
6532 };
6533 let (re, nr) = (row_elems as i32, n_rows as i32);
6534 let (st, off) = (dst_stride as i64, dst_off as i64);
6535 let __s_b = self.gpu.stream();
6536 let mut b = __s_b.launch_builder(&f);
6537 b.arg(src)
6538 .arg(&mut *dst)
6539 .arg(&re)
6540 .arg(&nr)
6541 .arg(&st)
6542 .arg(&off);
6543 unsafe {
6544 b.launch(cfg)?;
6545 }
6546 Ok(())
6547 }
6548
6549 pub fn u32_set_k(
6551 &self,
6552 dst: &mut CudaSlice<u32>,
6553 v: u32,
6554 idx: usize,
6555 ) -> Result<(), Box<dyn std::error::Error>> {
6556 let f = self.func("u32_set_k");
6557 let cfg = LaunchConfig {
6558 grid_dim: (1, 1, 1),
6559 block_dim: (1, 1, 1),
6560 shared_mem_bytes: 0,
6561 };
6562 let ii = idx as i32;
6563 let __s_b = self.gpu.stream();
6564 let mut b = __s_b.launch_builder(&f);
6565 b.arg(dst).arg(&v).arg(&ii);
6566 unsafe {
6567 b.launch(cfg)?;
6568 }
6569 Ok(())
6570 }
6571
6572 pub fn i32_add_k(
6574 &self,
6575 d: &mut CudaSlice<i32>,
6576 v: i32,
6577 ) -> Result<(), Box<dyn std::error::Error>> {
6578 let f = self.func("i32_add_k");
6579 let cfg = LaunchConfig {
6580 grid_dim: (1, 1, 1),
6581 block_dim: (32, 1, 1),
6582 shared_mem_bytes: 0,
6583 };
6584 let __s_b = self.gpu.stream();
6585 let mut b = __s_b.launch_builder(&f);
6586 b.arg(d).arg(&v);
6587 unsafe {
6588 b.launch(cfg)?;
6589 }
6590 Ok(())
6591 }
6592
6593 pub fn i32_iota_from(
6595 &self,
6596 ctr: &CudaSlice<i32>,
6597 dst: &mut CudaSlice<i32>,
6598 n: usize,
6599 ) -> Result<(), Box<dyn std::error::Error>> {
6600 let f = self.func("i32_iota_from");
6601 let cfg = LaunchConfig::for_num_elems(n as u32);
6602 let ni = n as i32;
6603 let __s_b = self.gpu.stream();
6604 let mut b = __s_b.launch_builder(&f);
6605 b.arg(ctr).arg(dst).arg(&ni);
6606 unsafe {
6607 b.launch(cfg)?;
6608 }
6609 Ok(())
6610 }
6611
6612 pub fn u32_map_k(
6614 &self,
6615 buf: &mut CudaSlice<u32>,
6616 map: &CudaSlice<u32>,
6617 idx: usize,
6618 ) -> Result<(), Box<dyn std::error::Error>> {
6619 let f = self.func("u32_map_k");
6620 let cfg = LaunchConfig {
6621 grid_dim: (1, 1, 1),
6622 block_dim: (1, 1, 1),
6623 shared_mem_bytes: 0,
6624 };
6625 let ii = idx as i32;
6626 let __s_b = self.gpu.stream();
6627 let mut b = __s_b.launch_builder(&f);
6628 b.arg(buf).arg(map).arg(&ii);
6629 unsafe {
6630 b.launch(cfg)?;
6631 }
6632 Ok(())
6633 }
6634
6635 #[allow(clippy::too_many_arguments)]
6637 pub fn u32_pack2(
6638 &self,
6639 a: &CudaSlice<u32>,
6640 off_a: usize,
6641 n1: usize,
6642 b_in: &CudaSlice<u32>,
6643 n2: usize,
6644 out: &mut CudaSlice<u32>,
6645 ) -> Result<(), Box<dyn std::error::Error>> {
6646 let f = self.func("u32_pack2");
6647 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
6648 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
6649 let __s_b = self.gpu.stream();
6650 let mut b = __s_b.launch_builder(&f);
6651 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
6652 unsafe {
6653 b.launch(cfg)?;
6654 }
6655 Ok(())
6656 }
6657
6658 pub fn moe_w_exscale(
6660 &self,
6661 w: &mut CudaSlice<f32>,
6662 sel: &CudaSlice<i32>,
6663 s: &CudaSlice<f32>,
6664 n: usize,
6665 ) -> Result<(), Box<dyn std::error::Error>> {
6666 let f = self.func("moe_w_exscale");
6667 let cfg = LaunchConfig::for_num_elems(n as u32);
6668 let ni = n as i32;
6669 let __s_b = self.gpu.stream();
6670 let mut b = __s_b.launch_builder(&f);
6671 b.arg(w).arg(sel).arg(s).arg(&ni);
6672 unsafe {
6673 b.launch(cfg)?;
6674 }
6675 Ok(())
6676 }
6677
6678 pub fn moe_w_scale_by_expert(
6681 &self,
6682 w: &mut CudaSlice<f32>,
6683 sel: &CudaSlice<i32>,
6684 macros: &CudaSlice<f32>,
6685 n_expert: usize,
6686 n: usize,
6687 ) -> Result<(), Box<dyn std::error::Error>> {
6688 let f = self.func("moe_w_scale_by_expert");
6689 let cfg = LaunchConfig {
6690 grid_dim: (n.div_ceil(64) as u32, 1, 1),
6691 block_dim: (64, 1, 1),
6692 shared_mem_bytes: 0,
6693 };
6694 let (ne, nn) = (n_expert as i32, n as i32);
6695 let __s_b = self.gpu.stream();
6696 let mut b = __s_b.launch_builder(&f);
6697 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
6698 unsafe {
6699 b.launch(cfg)?;
6700 }
6701 Ok(())
6702 }
6703
6704 pub fn moe_gate_up_silu8_dev_q8(
6705 &self,
6706 table: &CudaSlice<u64>,
6707 sel: &cudarc::driver::CudaView<i32>,
6708 aq: &CudaSlice<i8>,
6709 ad: &CudaSlice<f32>,
6710 in_f: usize,
6711 n_ff: usize,
6712 n_used: usize,
6713 n_expert: usize,
6714 qt_g: i32,
6715 qt_u: i32,
6716 rb_g: usize,
6717 rb_u: usize,
6718 macros: &CudaSlice<f32>,
6719 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6720 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
6721 let (mode, wpb) = GU.get_or_init(|| {
6722 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
6723 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB")
6724 .ok()
6725 .and_then(|v| v.parse().ok())
6726 .unwrap_or(4u32)
6727 .clamp(1, 16);
6728 (mode, wpb)
6729 });
6730 let (mode, wpb) = (mode.as_str(), *wpb);
6731 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
6732 let (inf, nff, ne, rbg, rbu) = (
6733 in_f as i32,
6734 n_ff as i32,
6735 n_expert as i32,
6736 rb_g as i64,
6737 rb_u as i64,
6738 );
6739 let (f, cfg) = match mode {
6740 "1" | "2" | "4" => {
6741 let rpw: u32 = mode.parse().unwrap();
6742 let f = self.func(match rpw {
6743 1 => "moe_gate_up_silu8_dev_q8_r1",
6744 2 => "moe_gate_up_silu8_dev_q8_r2",
6745 _ => "moe_gate_up_silu8_dev_q8_r4",
6746 });
6747 let rows_per_block = (rpw * wpb) as usize;
6748 let gx = n_ff.div_ceil(rows_per_block) as u32;
6749 (
6750 f,
6751 LaunchConfig {
6752 grid_dim: (gx, n_used as u32, 1),
6753 block_dim: (32, wpb, 1),
6754 shared_mem_bytes: 0,
6755 },
6756 )
6757 }
6758 "j8" if n_used <= 32 => (
6759 self.func("moe_gate_up_silu8_dev_q8_j8"),
6760 LaunchConfig {
6761 grid_dim: (n_ff as u32, 1, 1),
6762 block_dim: (32, n_used as u32, 1),
6763 shared_mem_bytes: 0,
6764 },
6765 ),
6766 "vsm2" => {
6768 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
6769 let sh = (rb_g + rb_u) as u32;
6770 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6771 f.set_attribute(
6772 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6773 sh as i32,
6774 )?;
6775 (
6776 f,
6777 LaunchConfig {
6778 grid_dim: (n_ff as u32, n_used as u32, 1),
6779 block_dim: (32, 1, 1),
6780 shared_mem_bytes: sh,
6781 },
6782 )
6783 }
6784 "vsm" => {
6785 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
6786 let sh = (rb_g + rb_u) as u32;
6787 use cudarc::driver::sys::CUfunction_attribute_enum as A;
6788 f.set_attribute(
6789 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
6790 sh as i32,
6791 )?;
6792 (
6793 f,
6794 LaunchConfig {
6795 grid_dim: (n_ff as u32, n_used as u32, 1),
6796 block_dim: (32, 1, 1),
6797 shared_mem_bytes: sh,
6798 },
6799 )
6800 }
6801 "sg" => (
6802 self.func("moe_gate_up_silu8_dev_q8_sg"),
6803 LaunchConfig {
6804 grid_dim: (n_ff as u32, n_used as u32, 1),
6805 block_dim: (32, 1, 1),
6806 shared_mem_bytes: 0,
6807 },
6808 ),
6809 "j8sg" if n_used <= 32 => (
6810 self.func("moe_gate_up_silu8_dev_q8_j8sg"),
6811 LaunchConfig {
6812 grid_dim: (n_ff as u32, 1, 1),
6813 block_dim: (32, n_used as u32, 1),
6814 shared_mem_bytes: 0,
6815 },
6816 ),
6817 "u64" if in_f == 2048 => (
6818 self.func("moe_gate_up_silu8_dev_q8_u64"),
6819 LaunchConfig {
6820 grid_dim: (n_ff as u32, n_used as u32, 1),
6821 block_dim: (32, 1, 1),
6822 shared_mem_bytes: 0,
6823 },
6824 ),
6825 "gs4" if in_f == 2048 => (
6826 self.func("moe_gate_up_silu8_dev_q8_gs4"),
6827 LaunchConfig {
6828 grid_dim: (n_ff as u32, n_used as u32, 1),
6829 block_dim: (32, 4, 1),
6830 shared_mem_bytes: 0,
6831 },
6832 ),
6833 "v" | "" => (
6835 self.func("moe_gate_up_silu8_dev_q8_v"),
6836 LaunchConfig {
6837 grid_dim: (n_ff as u32, n_used as u32, 1),
6838 block_dim: (32, 1, 1),
6839 shared_mem_bytes: 0,
6840 },
6841 ),
6842 "s2" => (
6843 self.func("moe_gate_up_silu8_dev_q8_s2"),
6844 LaunchConfig {
6845 grid_dim: (n_ff as u32, n_used as u32, 1),
6846 block_dim: (32, 2, 1),
6847 shared_mem_bytes: 0,
6848 },
6849 ),
6850 "s2z" => {
6851 let rz = wpb.min(16); (
6853 self.func("moe_gate_up_silu8_dev_q8_s2z"),
6854 LaunchConfig {
6855 grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
6856 block_dim: (32, 2, rz),
6857 shared_mem_bytes: 0,
6858 },
6859 )
6860 }
6861 _ => (
6862 self.func("moe_gate_up_silu8_dev_q8"),
6863 LaunchConfig {
6864 grid_dim: (n_ff as u32, n_used as u32, 1),
6865 block_dim: (32, 1, 1),
6866 shared_mem_bytes: 0,
6867 },
6868 ),
6869 };
6870 let __s_b = self.gpu.stream();
6871 let mut b = __s_b.launch_builder(&f);
6872 b.arg(table)
6873 .arg(sel)
6874 .arg(aq)
6875 .arg(ad)
6876 .arg(&mut act)
6877 .arg(&inf)
6878 .arg(&nff)
6879 .arg(&ne)
6880 .arg(&qt_g)
6881 .arg(&qt_u)
6882 .arg(&rbg)
6883 .arg(&rbu)
6884 .arg(macros);
6885 unsafe {
6886 b.launch(cfg)?;
6887 }
6888 Ok(act)
6889 }
6890
6891 #[allow(clippy::too_many_arguments)]
6892 pub fn moe_down8_fma_dev_q8(
6893 &self,
6894 table: &CudaSlice<u64>,
6895 sel: &cudarc::driver::CudaView<i32>,
6896 w: &cudarc::driver::CudaView<f32>,
6897 aq2: &CudaSlice<i8>,
6898 ad2: &CudaSlice<f32>,
6899 dst: &mut cudarc::driver::CudaViewMut<f32>,
6900 in_f: usize,
6901 out_f: usize,
6902 n_used: usize,
6903 n_expert: usize,
6904 qt: i32,
6905 rb: usize,
6906 ) -> Result<(), Box<dyn std::error::Error>> {
6907 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
6908 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
6909 let (inf, outf, nu, ne, rbi) = (
6910 in_f as i32,
6911 out_f as i32,
6912 n_used as i32,
6913 n_expert as i32,
6914 rb as i64,
6915 );
6916 let (f, cfg) = match mode.as_str() {
6919 m @ ("1" | "2" | "4") if n_used <= 8 => {
6920 let rpw: usize = m.parse().unwrap();
6921 let f = self.func(match rpw {
6922 1 => "moe_down8_fma_dev_q8_w8r1",
6923 2 => "moe_down8_fma_dev_q8_w8r2",
6924 _ => "moe_down8_fma_dev_q8_w8r4",
6925 });
6926 (
6927 f,
6928 LaunchConfig {
6929 grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
6930 block_dim: (32, n_used as u32, 1),
6931 shared_mem_bytes: 0,
6932 },
6933 )
6934 }
6935 "h2" if in_f == 512 => (
6936 self.func("moe_down8_fma_dev_q8_h2"),
6937 LaunchConfig {
6938 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6939 block_dim: (32, 1, 1),
6940 shared_mem_bytes: 0,
6941 },
6942 ),
6943 "" if in_f == 704 && n_used <= 8 => (
6946 self.func("moe_down8_fma_dev_q8_w8r2"),
6947 LaunchConfig {
6948 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6949 block_dim: (32, n_used as u32, 1),
6950 shared_mem_bytes: 0,
6951 },
6952 ),
6953 "w8h2v" | "" if in_f == 512 && n_used <= 8 => (
6957 self.func("moe_down8_fma_dev_q8_w8h2v"),
6958 LaunchConfig {
6959 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6960 block_dim: (32, n_used as u32, 1),
6961 shared_mem_bytes: 0,
6962 },
6963 ),
6964 "w8h2r2v" if in_f == 512 && n_used <= 8 => (
6965 self.func("moe_down8_fma_dev_q8_w8h2r2v"),
6966 LaunchConfig {
6967 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6968 block_dim: (32, n_used as u32, 1),
6969 shared_mem_bytes: 0,
6970 },
6971 ),
6972 "w8h2r2" if in_f == 512 && n_used <= 8 => (
6973 self.func("moe_down8_fma_dev_q8_w8h2r2"),
6974 LaunchConfig {
6975 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
6976 block_dim: (32, n_used as u32, 1),
6977 shared_mem_bytes: 0,
6978 },
6979 ),
6980 "w8h2" if in_f == 512 && n_used <= 8 => (
6981 self.func("moe_down8_fma_dev_q8_w8h2"),
6982 LaunchConfig {
6983 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
6984 block_dim: (32, n_used as u32, 1),
6985 shared_mem_bytes: 0,
6986 },
6987 ),
6988 _ => (
6989 self.func("moe_down8_fma_dev_q8"),
6990 LaunchConfig {
6991 grid_dim: (out_f as u32, 1, 1),
6992 block_dim: (32, 1, 1),
6993 shared_mem_bytes: 0,
6994 },
6995 ),
6996 };
6997 let __s_b = self.gpu.stream();
6998 let mut b = __s_b.launch_builder(&f);
6999 b.arg(table)
7000 .arg(sel)
7001 .arg(w)
7002 .arg(aq2)
7003 .arg(ad2)
7004 .arg(dst)
7005 .arg(&inf)
7006 .arg(&outf)
7007 .arg(&nu)
7008 .arg(&ne)
7009 .arg(&qt)
7010 .arg(&rbi);
7011 unsafe {
7012 b.launch(cfg)?;
7013 }
7014 Ok(())
7015 }
7016
7017 #[allow(clippy::too_many_arguments)]
7024 pub fn moe_gate_up_silu8_dev_q8_rows(
7025 &self,
7026 table: &CudaSlice<u64>,
7027 sel: &CudaSlice<i32>,
7028 aq: &CudaSlice<i8>,
7029 ad: &CudaSlice<f32>,
7030 t: usize,
7031 in_f: usize,
7032 n_ff: usize,
7033 n_used: usize,
7034 n_expert: usize,
7035 qt_g: i32,
7036 qt_u: i32,
7037 rb_g: usize,
7038 rb_u: usize,
7039 macros: &CudaSlice<f32>,
7040 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7041 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
7042 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
7043 let cfg = LaunchConfig {
7044 grid_dim: (n_ff as u32, n_used as u32, t as u32),
7045 block_dim: (32, 1, 1),
7046 shared_mem_bytes: 0,
7047 };
7048 let (inf, nff, ne, nu, rbg, rbu) = (
7049 in_f as i32,
7050 n_ff as i32,
7051 n_expert as i32,
7052 n_used as i32,
7053 rb_g as i64,
7054 rb_u as i64,
7055 );
7056 let __s_b = self.gpu.stream();
7057 let mut b = __s_b.launch_builder(&f);
7058 b.arg(table)
7059 .arg(sel)
7060 .arg(aq)
7061 .arg(ad)
7062 .arg(&mut act)
7063 .arg(&inf)
7064 .arg(&nff)
7065 .arg(&ne)
7066 .arg(&qt_g)
7067 .arg(&qt_u)
7068 .arg(&rbg)
7069 .arg(&rbu)
7070 .arg(&nu)
7071 .arg(macros);
7072 unsafe {
7073 b.launch(cfg)?;
7074 }
7075 Ok(act)
7076 }
7077
7078 #[allow(clippy::too_many_arguments)]
7083 pub fn moe_down8_fma_dev_q8_rows(
7084 &self,
7085 table: &CudaSlice<u64>,
7086 sel: &CudaSlice<i32>,
7087 w: &CudaSlice<f32>,
7088 aq2: &CudaSlice<i8>,
7089 ad2: &CudaSlice<f32>,
7090 dst: &mut CudaSlice<f32>,
7091 t: usize,
7092 in_f: usize,
7093 out_f: usize,
7094 n_used: usize,
7095 n_expert: usize,
7096 qt: i32,
7097 rb: usize,
7098 ) -> Result<(), Box<dyn std::error::Error>> {
7099 assert!(
7100 in_f == 512 && n_used <= 8,
7101 "down rows twin is w8h2v shape-gated"
7102 );
7103 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
7104 let cfg = LaunchConfig {
7105 grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
7106 block_dim: (32, n_used as u32, 1),
7107 shared_mem_bytes: 0,
7108 };
7109 let (inf, outf, nu, ne, rbi) = (
7110 in_f as i32,
7111 out_f as i32,
7112 n_used as i32,
7113 n_expert as i32,
7114 rb as i64,
7115 );
7116 let __s_b = self.gpu.stream();
7117 let mut b = __s_b.launch_builder(&f);
7118 b.arg(table)
7119 .arg(sel)
7120 .arg(w)
7121 .arg(aq2)
7122 .arg(ad2)
7123 .arg(dst)
7124 .arg(&inf)
7125 .arg(&outf)
7126 .arg(&nu)
7127 .arg(&ne)
7128 .arg(&qt)
7129 .arg(&rbi);
7130 unsafe {
7131 b.launch(cfg)?;
7132 }
7133 Ok(())
7134 }
7135
7136 #[allow(clippy::too_many_arguments)]
7140 pub fn moe_gate_up_silu8_dev_q8_csr(
7141 &self,
7142 table: &CudaSlice<u64>,
7143 sel: &CudaSlice<i32>,
7144 aq: &CudaSlice<i8>,
7145 ad: &CudaSlice<f32>,
7146 n_pairs: usize,
7147 in_f: usize,
7148 n_ff: usize,
7149 n_used: usize,
7150 n_expert: usize,
7151 qt_g: i32,
7152 qt_u: i32,
7153 rb_g: usize,
7154 rb_u: usize,
7155 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7156 let f = if qt_g == crate::QT_NVFP4 {
7159 self.func("moe_gate_up_silu8_dev_q8_csr_nvfp4")
7160 } else {
7161 self.func("moe_gate_up_silu8_dev_q8_csr_iq4")
7162 };
7163 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
7164 let cfg = LaunchConfig {
7165 grid_dim: (n_ff as u32, n_pairs as u32, 1),
7166 block_dim: (32, 1, 1),
7167 shared_mem_bytes: 0,
7168 };
7169 let (inf, nff, ne, nu, npi, rbg, rbu) = (
7170 in_f as i32,
7171 n_ff as i32,
7172 n_expert as i32,
7173 n_used as i32,
7174 n_pairs as i32,
7175 rb_g as i64,
7176 rb_u as i64,
7177 );
7178 let __s_b = self.gpu.stream();
7179 let mut b = __s_b.launch_builder(&f);
7180 b.arg(table)
7181 .arg(sel)
7182 .arg(aq)
7183 .arg(ad)
7184 .arg(&mut act)
7185 .arg(&inf)
7186 .arg(&nff)
7187 .arg(&ne)
7188 .arg(&qt_g)
7189 .arg(&qt_u)
7190 .arg(&rbg)
7191 .arg(&rbu)
7192 .arg(&nu)
7193 .arg(&npi);
7194 unsafe {
7195 b.launch(cfg)?;
7196 }
7197 Ok(act)
7198 }
7199
7200 #[allow(clippy::too_many_arguments)]
7204 pub fn moe_down8_fma_dev_q8_variant(
7205 &self,
7206 variant: &str,
7207 table: &CudaSlice<u64>,
7208 sel: &cudarc::driver::CudaView<i32>,
7209 w: &cudarc::driver::CudaView<f32>,
7210 aq2: &CudaSlice<i8>,
7211 ad2: &CudaSlice<f32>,
7212 dst: &mut cudarc::driver::CudaViewMut<f32>,
7213 in_f: usize,
7214 out_f: usize,
7215 n_used: usize,
7216 n_expert: usize,
7217 qt: i32,
7218 rb: usize,
7219 ) -> Result<(), Box<dyn std::error::Error>> {
7220 let (inf, outf, nu, ne, rbi) = (
7221 in_f as i32,
7222 out_f as i32,
7223 n_used as i32,
7224 n_expert as i32,
7225 rb as i64,
7226 );
7227 let (f, cfg) = match variant {
7228 "w8h2" | "w8h2v" => (
7229 self.func(if variant == "w8h2" {
7230 "moe_down8_fma_dev_q8_w8h2"
7231 } else {
7232 "moe_down8_fma_dev_q8_w8h2v"
7233 }),
7234 LaunchConfig {
7235 grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
7236 block_dim: (32, n_used as u32, 1),
7237 shared_mem_bytes: 0,
7238 },
7239 ),
7240 "w8h2r2" | "w8h2r2v" => (
7241 self.func(if variant == "w8h2r2" {
7242 "moe_down8_fma_dev_q8_w8h2r2"
7243 } else {
7244 "moe_down8_fma_dev_q8_w8h2r2v"
7245 }),
7246 LaunchConfig {
7247 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
7248 block_dim: (32, n_used as u32, 1),
7249 shared_mem_bytes: 0,
7250 },
7251 ),
7252 _ => (
7253 self.func("moe_down8_fma_dev_q8"),
7254 LaunchConfig {
7255 grid_dim: (out_f as u32, 1, 1),
7256 block_dim: (32, 1, 1),
7257 shared_mem_bytes: 0,
7258 },
7259 ),
7260 };
7261 let __s_b = self.gpu.stream();
7262 let mut b = __s_b.launch_builder(&f);
7263 b.arg(table)
7264 .arg(sel)
7265 .arg(w)
7266 .arg(aq2)
7267 .arg(ad2)
7268 .arg(dst)
7269 .arg(&inf)
7270 .arg(&outf)
7271 .arg(&nu)
7272 .arg(&ne)
7273 .arg(&qt)
7274 .arg(&rbi);
7275 unsafe {
7276 b.launch(cfg)?;
7277 }
7278 Ok(())
7279 }
7280
7281 #[allow(clippy::too_many_arguments)]
7283 pub fn moe_gate_up_silu8_dev_q8_variant(
7284 &self,
7285 variant: &str,
7286 table: &CudaSlice<u64>,
7287 sel: &cudarc::driver::CudaView<i32>,
7288 aq: &CudaSlice<i8>,
7289 ad: &CudaSlice<f32>,
7290 in_f: usize,
7291 n_ff: usize,
7292 n_used: usize,
7293 n_expert: usize,
7294 qt_g: i32,
7295 qt_u: i32,
7296 rb_g: usize,
7297 rb_u: usize,
7298 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7299 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
7300 let (inf, nff, ne, rbg, rbu) = (
7301 in_f as i32,
7302 n_ff as i32,
7303 n_expert as i32,
7304 rb_g as i64,
7305 rb_u as i64,
7306 );
7307 let f = self.func(if variant == "v" {
7308 "moe_gate_up_silu8_dev_q8_v"
7309 } else {
7310 "moe_gate_up_silu8_dev_q8"
7311 });
7312 let cfg = LaunchConfig {
7313 grid_dim: (n_ff as u32, n_used as u32, 1),
7314 block_dim: (32, 1, 1),
7315 shared_mem_bytes: 0,
7316 };
7317 let __s_b = self.gpu.stream();
7318 let mut b = __s_b.launch_builder(&f);
7319 b.arg(table)
7320 .arg(sel)
7321 .arg(aq)
7322 .arg(ad)
7323 .arg(&mut act)
7324 .arg(&inf)
7325 .arg(&nff)
7326 .arg(&ne)
7327 .arg(&qt_g)
7328 .arg(&qt_u)
7329 .arg(&rbg)
7330 .arg(&rbu);
7331 unsafe {
7332 b.launch(cfg)?;
7333 }
7334 Ok(act)
7335 }
7336
7337 pub fn moe_gate_up_silu8_dev(
7338 &self,
7339 table: &CudaSlice<u64>,
7340 sel: &cudarc::driver::CudaView<i32>,
7341 x: &cudarc::driver::CudaView<f32>,
7342 in_f: usize,
7343 n_ff: usize,
7344 n_used: usize,
7345 n_expert: usize,
7346 qt_g: i32,
7347 qt_u: i32,
7348 rb_g: usize,
7349 rb_u: usize,
7350 macros: &CudaSlice<f32>,
7351 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7352 let f = self.func("moe_gate_up_silu8_dev");
7353 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig {
7355 grid_dim: (n_ff as u32, n_used as u32, 1),
7356 block_dim: (256, 1, 1),
7357 shared_mem_bytes: 0,
7358 };
7359 let (inf, nff, ne, rbg, rbu) = (
7360 in_f as i32,
7361 n_ff as i32,
7362 n_expert as i32,
7363 rb_g as i64,
7364 rb_u as i64,
7365 );
7366 let __s_b = self.gpu.stream();
7367 let mut b = __s_b.launch_builder(&f);
7368 b.arg(table)
7369 .arg(sel)
7370 .arg(x)
7371 .arg(&mut act)
7372 .arg(&inf)
7373 .arg(&nff)
7374 .arg(&ne)
7375 .arg(&qt_g)
7376 .arg(&qt_u)
7377 .arg(&rbg)
7378 .arg(&rbu)
7379 .arg(macros);
7380 unsafe {
7381 b.launch(cfg)?;
7382 }
7383 Ok(act)
7384 }
7385
7386 #[allow(clippy::too_many_arguments)]
7389 pub fn moe_down8_fma_dev(
7390 &self,
7391 table: &CudaSlice<u64>,
7392 sel: &cudarc::driver::CudaView<i32>,
7393 w: &cudarc::driver::CudaView<f32>,
7394 act: &CudaSlice<f32>,
7395 dst: &mut cudarc::driver::CudaViewMut<f32>,
7396 in_f: usize,
7397 out_f: usize,
7398 n_used: usize,
7399 n_expert: usize,
7400 qt: i32,
7401 rb: usize,
7402 ) -> Result<(), Box<dyn std::error::Error>> {
7403 let f = self.func("moe_down8_fma_dev");
7404 let cfg = LaunchConfig {
7405 grid_dim: (out_f as u32, 1, 1),
7406 block_dim: (256, 1, 1),
7407 shared_mem_bytes: 0,
7408 };
7409 let (inf, outf, nu, ne, rbv) = (
7410 in_f as i32,
7411 out_f as i32,
7412 n_used as i32,
7413 n_expert as i32,
7414 rb as i64,
7415 );
7416 let __s_b = self.gpu.stream();
7417 let mut b = __s_b.launch_builder(&f);
7418 b.arg(table)
7419 .arg(sel)
7420 .arg(w)
7421 .arg(act)
7422 .arg(dst)
7423 .arg(&inf)
7424 .arg(&outf)
7425 .arg(&nu)
7426 .arg(&ne)
7427 .arg(&qt)
7428 .arg(&rbv);
7429 unsafe {
7430 b.launch(cfg)?;
7431 }
7432 Ok(())
7433 }
7434
7435 pub fn axpy_into(
7437 &self,
7438 src: &CudaSlice<f32>,
7439 alpha: f32,
7440 dst: &mut cudarc::driver::CudaViewMut<f32>,
7441 n: usize,
7442 ) -> Result<(), Box<dyn std::error::Error>> {
7443 let f = self.func("axpy_f32");
7444 let cfg = LaunchConfig::for_num_elems(n as u32);
7445 let (a, ni) = (alpha, n as i32);
7446 let __s_b = self.gpu.stream();
7447 let mut b = __s_b.launch_builder(&f);
7448 b.arg(src).arg(dst).arg(&a).arg(&ni);
7449 unsafe {
7450 b.launch(cfg)?;
7451 }
7452 Ok(())
7453 }
7454
7455 pub fn axpy_host_into(
7457 &self,
7458 src: &cudarc::driver::CudaView<'_, f32>,
7459 alpha: f32,
7460 dst: &mut cudarc::driver::CudaViewMut<f32>,
7461 n: usize,
7462 ) -> Result<(), Box<dyn std::error::Error>> {
7463 let f = self.func("axpy_host_f32");
7464 let cfg = LaunchConfig::for_num_elems(n as u32);
7465 let (a, ni) = (alpha, n as i32);
7466 let __s_b = self.gpu.stream();
7467 let mut b = __s_b.launch_builder(&f);
7468 b.arg(src).arg(dst).arg(&a).arg(&ni);
7469 unsafe {
7470 b.launch(cfg)?;
7471 }
7472 Ok(())
7473 }
7474
7475 pub fn add_scaled_rows(
7477 &self,
7478 src: &CudaSlice<f32>,
7479 scale: &CudaSlice<f32>,
7480 dst: &mut CudaSlice<f32>,
7481 ncols: usize,
7482 nrows: usize,
7483 ) -> Result<(), Box<dyn std::error::Error>> {
7484 let f = self.func("add_scaled_rows_f32");
7485 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7486 let (nc, nr) = (ncols as i32, nrows as i32);
7487 let __s_b = self.gpu.stream();
7488 let mut b = __s_b.launch_builder(&f);
7489 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
7490 unsafe {
7491 b.launch(cfg)?;
7492 }
7493 Ok(())
7494 }
7495
7496 pub fn scale_rows(
7499 &self,
7500 y: &mut CudaSlice<f32>,
7501 s: &CudaSlice<f32>,
7502 ncols: usize,
7503 nrows: usize,
7504 ) -> Result<(), Box<dyn std::error::Error>> {
7505 let f = self.func("scale_rows_f32");
7506 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
7507 let (nc, nr) = (ncols as i32, nrows as i32);
7508 let __s_b = self.gpu.stream();
7509 let mut b = __s_b.launch_builder(&f);
7510 b.arg(&mut *y).arg(s).arg(&nc).arg(&nr);
7511 unsafe {
7512 b.launch(cfg)?;
7513 }
7514 Ok(())
7515 }
7516
7517 #[allow(clippy::too_many_arguments)]
7521 pub fn moe_prime_join_scatter(
7522 &self,
7523 y0: &CudaSlice<f32>,
7524 y1: &CudaSlice<f32>,
7525 inv: &CudaSlice<i32>,
7526 w: &CudaSlice<f32>,
7527 out: &mut CudaSlice<f32>,
7528 ncols: usize,
7529 n_used: usize,
7530 t: usize,
7531 ) -> Result<(), Box<dyn std::error::Error>> {
7532 let f = self.func("moe_prime_join_scatter_f32");
7533 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7534 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7535 let __s_b = self.gpu.stream();
7536 let mut b = __s_b.launch_builder(&f);
7537 b.arg(y0)
7538 .arg(y1)
7539 .arg(inv)
7540 .arg(w)
7541 .arg(&mut *out)
7542 .arg(&nc)
7543 .arg(&nu)
7544 .arg(&ti);
7545 unsafe {
7546 b.launch(cfg)?;
7547 }
7548 Ok(())
7549 }
7550
7551 pub fn moe_pairs_weighted_scatter(
7554 &self,
7555 y: &CudaSlice<f32>,
7556 w: &CudaSlice<f32>,
7557 out: &mut CudaSlice<f32>,
7558 ncols: usize,
7559 n_used: usize,
7560 t: usize,
7561 ) -> Result<(), Box<dyn std::error::Error>> {
7562 let f = self.func("moe_pairs_weighted_scatter_f32");
7563 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7564 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7565 let __s_b = self.gpu.stream();
7566 let mut b = __s_b.launch_builder(&f);
7567 b.arg(y).arg(w).arg(&mut *out).arg(&nc).arg(&nu).arg(&ti);
7568 unsafe {
7569 b.launch(cfg)?;
7570 }
7571 Ok(())
7572 }
7573
7574 pub fn gather_rows(
7578 &self,
7579 src: &CudaSlice<f32>,
7580 idx: &CudaSlice<i32>,
7581 dst: &mut CudaSlice<f32>,
7582 ncols: usize,
7583 m_e: usize,
7584 ) -> Result<(), Box<dyn std::error::Error>> {
7585 let f = self.func("gather_rows_f32");
7586 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7587 let (nc, me) = (ncols as i32, m_e as i32);
7588 let __s_b = self.gpu.stream();
7589 let mut b = __s_b.launch_builder(&f);
7590 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
7591 unsafe {
7592 b.launch(cfg)?;
7593 }
7594 Ok(())
7595 }
7596
7597 pub fn scatter_slot(
7602 &self,
7603 src: &CudaSlice<f32>,
7604 tok_idx: &CudaSlice<i32>,
7605 slot_idx: &CudaSlice<i32>,
7606 weight: &CudaSlice<f32>,
7607 dst: &mut CudaSlice<f32>,
7608 wbuf: &mut CudaSlice<f32>,
7609 ncols: usize,
7610 n_used: usize,
7611 m_e: usize,
7612 ) -> Result<(), Box<dyn std::error::Error>> {
7613 let f = self.func("scatter_add_slot_f32");
7614 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
7615 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
7616 let __s_b = self.gpu.stream();
7617 let mut b = __s_b.launch_builder(&f);
7618 b.arg(src)
7619 .arg(tok_idx)
7620 .arg(slot_idx)
7621 .arg(weight)
7622 .arg(dst)
7623 .arg(wbuf)
7624 .arg(&nc)
7625 .arg(&nu)
7626 .arg(&me);
7627 unsafe {
7628 b.launch(cfg)?;
7629 }
7630 Ok(())
7631 }
7632
7633 pub fn reduce_slots(
7637 &self,
7638 slots: &CudaSlice<f32>,
7639 wbuf: &CudaSlice<f32>,
7640 dst: &mut CudaSlice<f32>,
7641 ncols: usize,
7642 n_used: usize,
7643 t: usize,
7644 ) -> Result<(), Box<dyn std::error::Error>> {
7645 let f = self.func("reduce_slots_f32");
7646 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7647 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7648 let __s_b = self.gpu.stream();
7649 let mut b = __s_b.launch_builder(&f);
7650 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7651 unsafe {
7652 b.launch(cfg)?;
7653 }
7654 Ok(())
7655 }
7656
7657 pub fn reduce_slots_host(
7662 &self,
7663 slots: &CudaSlice<f32>,
7664 wbuf: &CudaSlice<f32>,
7665 dst: &mut CudaSlice<f32>,
7666 ncols: usize,
7667 n_used: usize,
7668 t: usize,
7669 ) -> Result<(), Box<dyn std::error::Error>> {
7670 let f = self.func("reduce_slots_host_f32");
7671 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
7672 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
7673 let __s_b = self.gpu.stream();
7674 let mut b = __s_b.launch_builder(&f);
7675 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
7676 unsafe {
7677 b.launch(cfg)?;
7678 }
7679 Ok(())
7680 }
7681
7682 pub fn quantize_q8_1_view(
7689 &self,
7690 x: &cudarc::driver::CudaView<f32>,
7691 m: usize,
7692 in_f: usize,
7693 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7694 let f = self.func("quantize_q8_1");
7695 let nblk = in_f / 32;
7696 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
7697 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
7698 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
7699 let (inf, mi) = (in_f as i32, m as i32);
7700 let __s_b = self.gpu.stream();
7701 let mut b = __s_b.launch_builder(&f);
7702 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7703 unsafe {
7704 b.launch(cfg)?;
7705 }
7706 Ok((q, d))
7707 }
7708
7709 pub fn quantize_q8_1(
7710 &self,
7711 x: &CudaSlice<f32>,
7712 m: usize,
7713 in_f: usize,
7714 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7715 let nblk = in_f / 32;
7716 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);
7720 let (inf, mi) = (in_f as i32, m as i32);
7721 if Self::pdl_on() && Self::pdl_wb_on() {
7722 {
7723 use cudarc::driver::{DevicePtr, DevicePtrMut};
7724 let s = &self.gpu.stream();
7725 let (px, _g0) = x.device_ptr(s);
7726 let (pq, _g1) = q.device_ptr_mut(s);
7727 let (pd, _g2) = d.device_ptr_mut(s);
7728 let mut ps = [
7729 &px as *const _ as *mut std::ffi::c_void,
7730 &pq as *const _ as *mut _,
7731 &pd as *const _ as *mut _,
7732 &inf as *const _ as *mut _,
7733 &mi as *const _ as *mut _,
7734 ];
7735 unsafe {
7736 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
7737 }
7738 }
7739 return Ok((q, d));
7740 }
7741 let f = self.func("quantize_q8_1");
7742 let __s_b = self.gpu.stream();
7743 let mut b = __s_b.launch_builder(&f);
7744 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
7745 unsafe {
7746 b.launch(cfg)?;
7747 }
7748 Ok((q, d))
7749 }
7750
7751 pub fn quantize_fp4_act(
7755 &self,
7756 x: &CudaSlice<f32>,
7757 m: usize,
7758 in_f: usize,
7759 ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
7760 let f = self.func("quantize_fp4_act");
7761 let nb16 = in_f / 16;
7762 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);
7765 let (inf, mi) = (in_f as i32, m as i32);
7766 let __s_b = self.gpu.stream();
7767 let mut b = __s_b.launch_builder(&f);
7768 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
7769 unsafe {
7770 b.launch(cfg)?;
7771 }
7772 Ok((aq4, ad4))
7773 }
7774
7775 pub fn qmatvec_gemm_nvfp4_fp4(
7780 &self,
7781 bytes: &CudaSlice<u8>,
7782 x: &CudaSlice<f32>,
7783 m: usize,
7784 in_f: usize,
7785 out_f: usize,
7786 row_bytes: usize,
7787 scale: f32,
7788 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7789 assert!(
7790 in_f % 64 == 0,
7791 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7792 );
7793 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7794 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
7795 if scale != 1.0 {
7796 self.scale_inplace(&mut y, scale, m * out_f)?;
7797 }
7798 Ok(y)
7799 }
7800
7801 fn fp4_gemm_launch(
7804 &self,
7805 bytes: &CudaSlice<u8>,
7806 aq4: &CudaSlice<u32>,
7807 ad4: &CudaSlice<u8>,
7808 m: usize,
7809 in_f: usize,
7810 out_f: usize,
7811 row_bytes: usize,
7812 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7813 let f = self.func("qmatvec_gemm_nvfp4_fp4");
7814 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64;
7816 const BN: u32 = 256;
7817 let cfg = LaunchConfig {
7818 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
7819 block_dim: (32, 4, 1),
7820 shared_mem_bytes: 0,
7821 };
7822 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7823 let __s_b = self.gpu.stream();
7824 let mut b = __s_b.launch_builder(&f);
7825 b.arg(bytes)
7826 .arg(aq4)
7827 .arg(ad4)
7828 .arg(&mut y)
7829 .arg(&inf)
7830 .arg(&outf)
7831 .arg(&mi)
7832 .arg(&rb);
7833 unsafe {
7834 b.launch(cfg)?;
7835 }
7836 Ok(y)
7837 }
7838
7839 pub fn qmatvec_gemm_nvfp4_fp4_raw(
7841 &self,
7842 bytes: &CudaSlice<u8>,
7843 x: &CudaSlice<f32>,
7844 m: usize,
7845 in_f: usize,
7846 out_f: usize,
7847 row_bytes: usize,
7848 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7849 assert!(
7850 in_f % 64 == 0,
7851 "FP4 GEMM requires in_f % 64 == 0, got {in_f}"
7852 );
7853 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
7854 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
7855 }
7856
7857 pub fn qmatvec_q8_0_fast(
7859 &self,
7860 w: &CudaSlice<u8>,
7861 x: &CudaSlice<f32>,
7862 m: usize,
7863 in_f: usize,
7864 out_f: usize,
7865 row_bytes: usize,
7866 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7867 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7868 let f = self.func("qmatvec_q8_0_dp4a");
7869 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7871 grid_dim: (out_f as u32, m as u32, 1),
7872 block_dim: (128, 1, 1),
7873 shared_mem_bytes: 0,
7874 };
7875 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7876 let __s_b = self.gpu.stream();
7877 let mut b = __s_b.launch_builder(&f);
7878 b.arg(w)
7879 .arg(&aq)
7880 .arg(&ad)
7881 .arg(&mut y)
7882 .arg(&inf)
7883 .arg(&outf)
7884 .arg(&mi)
7885 .arg(&rb);
7886 unsafe {
7887 b.launch(cfg)?;
7888 }
7889 Ok(y)
7890 }
7891
7892 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(
7895 &self,
7896 w: &CudaSlice<u8>,
7897 x: &CudaSlice<f32>,
7898 m: usize,
7899 in_f: usize,
7900 out_f: usize,
7901 row_bytes: usize,
7902 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7903 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7904 let f = self.func("qmatvec_q4_K_dp4a");
7905 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7907 grid_dim: (out_f as u32, m as u32, 1),
7908 block_dim: (128, 1, 1),
7909 shared_mem_bytes: 0,
7910 };
7911 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7912 let __s_b = self.gpu.stream();
7913 let mut b = __s_b.launch_builder(&f);
7914 b.arg(w)
7915 .arg(&aq)
7916 .arg(&ad)
7917 .arg(&mut y)
7918 .arg(&inf)
7919 .arg(&outf)
7920 .arg(&mi)
7921 .arg(&rb);
7922 unsafe {
7923 b.launch(cfg)?;
7924 }
7925 Ok(y)
7926 }
7927
7928 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(
7931 &self,
7932 w: &CudaSlice<u8>,
7933 x: &CudaSlice<f32>,
7934 m: usize,
7935 in_f: usize,
7936 out_f: usize,
7937 row_bytes: usize,
7938 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7939 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7940 let f = self.func("qmatvec_q6_K_dp4a");
7941 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7943 grid_dim: (out_f as u32, m as u32, 1),
7944 block_dim: (128, 1, 1),
7945 shared_mem_bytes: 0,
7946 };
7947 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7948 let __s_b = self.gpu.stream();
7949 let mut b = __s_b.launch_builder(&f);
7950 b.arg(w)
7951 .arg(&aq)
7952 .arg(&ad)
7953 .arg(&mut y)
7954 .arg(&inf)
7955 .arg(&outf)
7956 .arg(&mi)
7957 .arg(&rb);
7958 unsafe {
7959 b.launch(cfg)?;
7960 }
7961 Ok(y)
7962 }
7963
7964 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(
7967 &self,
7968 w: &CudaSlice<u8>,
7969 x: &CudaSlice<f32>,
7970 m: usize,
7971 in_f: usize,
7972 out_f: usize,
7973 row_bytes: usize,
7974 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7975 self.qmatvec_dp4a_named(
7976 "qmatvec_q5_K_dp4a",
7977 &w.slice(0..w.len()),
7978 x,
7979 m,
7980 in_f,
7981 out_f,
7982 row_bytes,
7983 )
7984 }
7985 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(
7988 &self,
7989 w: &CudaSlice<u8>,
7990 x: &CudaSlice<f32>,
7991 m: usize,
7992 in_f: usize,
7993 out_f: usize,
7994 row_bytes: usize,
7995 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7996 self.qmatvec_dp4a_named(
7997 "qmatvec_q3_K_dp4a",
7998 &w.slice(0..w.len()),
7999 x,
8000 m,
8001 in_f,
8002 out_f,
8003 row_bytes,
8004 )
8005 }
8006 pub fn qmatvec_nvfp4_fast_rp(
8008 &self,
8009 w: &CudaSlice<u8>,
8010 x: &CudaSlice<f32>,
8011 m: usize,
8012 in_f: usize,
8013 out_f: usize,
8014 row_bytes: usize,
8015 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8016 assert!(
8017 in_f % 64 == 0,
8018 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8019 );
8020 self.qmatvec_dp4a_named(
8021 "qmatvec_nvfp4_dp4a_rp",
8022 &w.slice(0..w.len()),
8023 x,
8024 m,
8025 in_f,
8026 out_f,
8027 row_bytes,
8028 )
8029 }
8030 pub fn qmatvec_nvfp4_fast(
8032 &self,
8033 w: &cudarc::driver::CudaView<'_, u8>,
8034 x: &CudaSlice<f32>,
8035 m: usize,
8036 in_f: usize,
8037 out_f: usize,
8038 row_bytes: usize,
8039 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8040 assert!(
8043 in_f % 64 == 0,
8044 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8045 );
8046 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
8047 }
8048 pub fn qmatvec_nvfp4_fast_v2(
8051 &self,
8052 w: &cudarc::driver::CudaView<'_, u8>,
8053 x: &CudaSlice<f32>,
8054 m: usize,
8055 in_f: usize,
8056 out_f: usize,
8057 row_bytes: usize,
8058 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8059 assert!(
8060 in_f % 64 == 0,
8061 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8062 );
8063 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_v2", w, x, m, in_f, out_f, row_bytes)
8064 }
8065 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(
8068 &self,
8069 w: &CudaSlice<u8>,
8070 x: &CudaSlice<f32>,
8071 m: usize,
8072 in_f: usize,
8073 out_f: usize,
8074 row_bytes: usize,
8075 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8076 self.qmatvec_dp4a_named(
8077 "qmatvec_iq4_XS_dp4a",
8078 &w.slice(0..w.len()),
8079 x,
8080 m,
8081 in_f,
8082 out_f,
8083 row_bytes,
8084 )
8085 }
8086
8087 fn qmatvec_dp4a_named(
8089 &self,
8090 name: &str,
8091 w: &cudarc::driver::CudaView<'_, u8>,
8092 x: &CudaSlice<f32>,
8093 m: usize,
8094 in_f: usize,
8095 out_f: usize,
8096 row_bytes: usize,
8097 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8098 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
8099 let f = self.func(name);
8100 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
8102 grid_dim: (out_f as u32, m as u32, 1),
8103 block_dim: (128, 1, 1),
8104 shared_mem_bytes: 0,
8105 };
8106 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8107 let __s_b = self.gpu.stream();
8108 let mut b = __s_b.launch_builder(&f);
8109 b.arg(w)
8110 .arg(&aq)
8111 .arg(&ad)
8112 .arg(&mut y)
8113 .arg(&inf)
8114 .arg(&outf)
8115 .arg(&mi)
8116 .arg(&rb);
8117 unsafe {
8118 b.launch(cfg)?;
8119 }
8120 Ok(y)
8121 }
8122
8123 #[allow(clippy::too_many_arguments)]
8129 pub fn qmatvec_nvfp4_fast_prequant_into(
8130 &self,
8131 w: &CudaSlice<u8>,
8132 aq: &CudaSlice<i8>,
8133 ad: &CudaSlice<f32>,
8134 y: &mut CudaSlice<f32>,
8135 m: usize,
8136 in_f: usize,
8137 out_f: usize,
8138 row_bytes: usize,
8139 ) -> Result<(), Box<dyn std::error::Error>> {
8140 assert!(
8141 in_f % 64 == 0,
8142 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8143 );
8144 if y.len() < m * out_f {
8145 return Err(format!(
8146 "NVFP4 prequant output {} is shorter than {m}x{out_f}",
8147 y.len()
8148 )
8149 .into());
8150 }
8151 let f = self.func("qmatvec_nvfp4_dp4a");
8152 let cfg = LaunchConfig {
8153 grid_dim: (out_f as u32, m as u32, 1),
8154 block_dim: (128, 1, 1),
8155 shared_mem_bytes: 0,
8156 };
8157 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
8158 let __s_b = self.gpu.stream();
8159 let mut b = __s_b.launch_builder(&f);
8160 b.arg(w)
8161 .arg(aq)
8162 .arg(ad)
8163 .arg(y)
8164 .arg(&inf)
8165 .arg(&outf)
8166 .arg(&mi)
8167 .arg(&rb);
8168 unsafe {
8169 b.launch(cfg)?;
8170 }
8171 Ok(())
8172 }
8173
8174 #[allow(clippy::too_many_arguments)]
8177 pub fn matvec_f32_qkv_into(
8178 &self,
8179 wq: &CudaSlice<f32>,
8180 wk: &CudaSlice<f32>,
8181 wv: &CudaSlice<f32>,
8182 wg: &CudaSlice<f32>,
8183 x: &CudaSlice<f32>,
8184 yq: &mut CudaSlice<f32>,
8185 yk: &mut CudaSlice<f32>,
8186 yv: &mut CudaSlice<f32>,
8187 yg: &mut CudaSlice<f32>,
8188 in_f: usize,
8189 out_q: usize,
8190 out_kv: usize,
8191 out_g: usize,
8192 ) -> Result<(), Box<dyn std::error::Error>> {
8193 if in_f % 4 != 0
8194 || wq.len() != out_q * in_f
8195 || wk.len() != out_kv * in_f
8196 || wv.len() != out_kv * in_f
8197 || wg.len() < out_g * in_f
8198 || x.len() < in_f
8199 || yq.len() < out_q
8200 || yk.len() < out_kv
8201 || yv.len() < out_kv
8202 || (out_g > 0 && yg.len() < out_g)
8203 {
8204 return Err(format!(
8205 "fused QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g} \
8206 wq={} wk={} wv={} wg={}",
8207 wq.len(),
8208 wk.len(),
8209 wv.len(),
8210 wg.len()
8211 )
8212 .into());
8213 }
8214 let f = self.func("matvec_f32_qkv");
8215 let cfg = LaunchConfig {
8216 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
8217 block_dim: (128, 1, 1),
8218 shared_mem_bytes: 0,
8219 };
8220 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
8221 let __s_b = self.gpu.stream();
8222 let mut b = __s_b.launch_builder(&f);
8223 b.arg(wq)
8224 .arg(wk)
8225 .arg(wv)
8226 .arg(wg)
8227 .arg(x)
8228 .arg(yq)
8229 .arg(yk)
8230 .arg(yv)
8231 .arg(yg)
8232 .arg(&inf)
8233 .arg(&oq)
8234 .arg(&okv)
8235 .arg(&og);
8236 unsafe {
8237 b.launch(cfg)?;
8238 }
8239 Ok(())
8240 }
8241
8242 #[allow(clippy::too_many_arguments)]
8248 #[allow(clippy::too_many_arguments)]
8251 pub fn qmatvec_nvfp4_sel_gu_into(
8252 &self,
8253 gate_bank: &CudaSlice<u8>,
8254 up_bank: &CudaSlice<u8>,
8255 sel: &CudaSlice<i32>,
8256 aq: &CudaSlice<i8>,
8257 ad: &CudaSlice<f32>,
8258 yg: &mut CudaSlice<f32>,
8259 yu: &mut CudaSlice<f32>,
8260 n_sel: usize,
8261 in_f: usize,
8262 out_f: usize,
8263 row_bytes: usize,
8264 expert_stride: usize,
8265 ) -> Result<(), Box<dyn std::error::Error>> {
8266 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8267 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8268 return Err("NVFP4 gu sel geometry".into());
8269 }
8270 static RPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
8273 let rpw = *RPW.get_or_init(|| {
8274 std::env::var("MEMRA_SEL_GU_RPW")
8275 .ok()
8276 .and_then(|v| v.parse().ok())
8277 .filter(|r| *r == 2 || *r == 4)
8278 .unwrap_or(1)
8279 });
8280 let rpw = if out_f % rpw == 0 { rpw } else { 1 };
8281 static WPR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8284 let wpr = *WPR.get_or_init(|| std::env::var("MEMRA_SEL_GU_WPR").as_deref() == Ok("1"));
8285 let f = self.func(match (wpr, rpw) {
8286 (true, _) => "qmatvec_nvfp4_dp4a_sel_v2_gu_wpr",
8287 (_, 4) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r4",
8288 (_, 2) => "qmatvec_nvfp4_dp4a_sel_v2_gu_r2",
8289 _ => "qmatvec_nvfp4_dp4a_sel_v2_gu",
8290 });
8291 let cfg = LaunchConfig {
8292 grid_dim: if wpr {
8293 (((2 * out_f) as u32).div_ceil(4), n_sel as u32, 1)
8294 } else if rpw == 1 {
8295 ((2 * out_f) as u32, n_sel as u32, 1)
8296 } else {
8297 ((out_f / rpw) as u32, n_sel as u32, 1)
8298 },
8299 block_dim: if wpr { (32, 4, 1) } else { (128, 1, 1) },
8300 shared_mem_bytes: 0,
8301 };
8302 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8303 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8304 let (ars, adrs) = (0i64, 0i64);
8305 let __s_b = self.gpu.stream();
8306 let mut b = __s_b.launch_builder(&f);
8307 b.arg(gate_bank)
8308 .arg(up_bank)
8309 .arg(sel)
8310 .arg(aq)
8311 .arg(ad)
8312 .arg(yg)
8313 .arg(yu)
8314 .arg(&inf)
8315 .arg(&outf)
8316 .arg(&ns)
8317 .arg(&rb)
8318 .arg(&es)
8319 .arg(&ars)
8320 .arg(&adrs);
8321 unsafe {
8322 b.launch(cfg)?;
8323 }
8324 Ok(())
8325 }
8326
8327 #[allow(clippy::too_many_arguments)]
8335 pub fn qmatvec_nvfp4_sel_down8_into(
8336 &self,
8337 bank: &CudaSlice<u8>,
8338 sel: &CudaSlice<i32>,
8339 aq: &CudaSlice<i8>,
8340 ad: &CudaSlice<f32>,
8341 route_w: &CudaSlice<f32>,
8342 md: &CudaSlice<f32>,
8343 dst: &mut CudaSlice<f32>,
8344 n_sel: usize,
8345 in_f: usize,
8346 out_f: usize,
8347 row_bytes: usize,
8348 expert_stride: usize,
8349 act_row_stride: usize,
8350 ad_row_stride: usize,
8351 ) -> Result<(), Box<dyn std::error::Error>> {
8352 if in_f % 64 != 0
8353 || n_sel == 0
8354 || n_sel > 8
8355 || (in_f >> 5) > 32
8356 || dst.len() < out_f
8357 || sel.len() < n_sel
8358 || route_w.len() < n_sel
8359 {
8360 return Err(format!(
8361 "NVFP4 sel down8 geometry in_f={in_f} out_f={out_f} n_sel={n_sel} dst={}",
8362 dst.len()
8363 )
8364 .into());
8365 }
8366 if !crate::tp::nvfp4_bank_v2_on() {
8367 return Err("NVFP4 sel down8 requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into());
8368 }
8369 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8");
8370 let cfg = LaunchConfig {
8371 grid_dim: (out_f as u32, 1, 1),
8372 block_dim: (32, n_sel as u32, 1),
8373 shared_mem_bytes: 0,
8374 };
8375 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8376 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8377 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8378 let __s_b = self.gpu.stream();
8379 let mut b = __s_b.launch_builder(&f);
8380 b.arg(bank)
8381 .arg(sel)
8382 .arg(aq)
8383 .arg(ad)
8384 .arg(route_w)
8385 .arg(md)
8386 .arg(dst)
8387 .arg(&inf)
8388 .arg(&outf)
8389 .arg(&ns)
8390 .arg(&rb)
8391 .arg(&es)
8392 .arg(&ars)
8393 .arg(&adrs);
8394 unsafe {
8395 b.launch(cfg)?;
8396 }
8397 Ok(())
8398 }
8399
8400 #[allow(clippy::too_many_arguments)]
8404 pub fn qmatvec_nvfp4_sel_down8_rows_into(
8405 &self,
8406 bank: &CudaSlice<u8>,
8407 sel: &CudaSlice<i32>,
8408 aq: &CudaSlice<i8>,
8409 ad: &CudaSlice<f32>,
8410 route_w: &CudaSlice<f32>,
8411 md: &CudaSlice<f32>,
8412 dst: &mut CudaSlice<f32>,
8413 t: usize,
8414 n_sel_col: usize,
8415 in_f: usize,
8416 out_f: usize,
8417 row_bytes: usize,
8418 expert_stride: usize,
8419 act_row_stride: usize,
8420 ad_row_stride: usize,
8421 ) -> Result<(), Box<dyn std::error::Error>> {
8422 let n_sel = t * n_sel_col;
8423 if in_f % 64 != 0
8424 || n_sel_col == 0
8425 || n_sel_col > 8
8426 || t == 0
8427 || t > 64
8428 || (in_f >> 5) > 32
8429 || dst.len() < t * out_f
8430 || sel.len() < n_sel
8431 || route_w.len() < n_sel
8432 {
8433 return Err(format!(
8434 "NVFP4 sel down8 rows geometry in_f={in_f} out_f={out_f} t={t} dst={}",
8435 dst.len()
8436 )
8437 .into());
8438 }
8439 if !crate::tp::nvfp4_bank_v2_on() {
8440 return Err(
8441 "NVFP4 sel down8 rows requires the v2 banks (MEMRA_NVFP4_BANK_V2=1)".into(),
8442 );
8443 }
8444 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_rows");
8445 let cfg = LaunchConfig {
8446 grid_dim: (out_f as u32, t as u32, 1),
8447 block_dim: (32, n_sel_col as u32, 1),
8448 shared_mem_bytes: 0,
8449 };
8450 let (inf, outf, nsc) = (in_f as i32, out_f as i32, n_sel_col as i32);
8451 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8452 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8453 let __s_b = self.gpu.stream();
8454 let mut b = __s_b.launch_builder(&f);
8455 b.arg(bank)
8456 .arg(sel)
8457 .arg(aq)
8458 .arg(ad)
8459 .arg(route_w)
8460 .arg(md)
8461 .arg(dst)
8462 .arg(&inf)
8463 .arg(&outf)
8464 .arg(&nsc)
8465 .arg(&rb)
8466 .arg(&es)
8467 .arg(&ars)
8468 .arg(&adrs);
8469 unsafe {
8470 b.launch(cfg)?;
8471 }
8472 Ok(())
8473 }
8474
8475 #[allow(clippy::too_many_arguments)]
8478 pub fn qmatvec_nvfp4_sel_gu_ep_into(
8479 &self,
8480 gate_bank: &CudaSlice<u8>,
8481 up_bank: &CudaSlice<u8>,
8482 sel: &CudaSlice<i32>,
8483 aq: &CudaSlice<i8>,
8484 ad: &CudaSlice<f32>,
8485 yg: &mut CudaSlice<f32>,
8486 yu: &mut CudaSlice<f32>,
8487 n_sel: usize,
8488 in_f: usize,
8489 out_f: usize,
8490 row_bytes: usize,
8491 expert_stride: usize,
8492 owner: usize,
8493 ) -> Result<(), Box<dyn std::error::Error>> {
8494 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
8495 if yg.len() < n_sel * out_f || yu.len() < n_sel * out_f || sel.len() < n_sel {
8496 return Err("NVFP4 gu ep geometry".into());
8497 }
8498 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_ep");
8499 let cfg = LaunchConfig {
8500 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
8501 block_dim: (128, 1, 1),
8502 shared_mem_bytes: 0,
8503 };
8504 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8505 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8506 let (ars, adrs) = (0i64, 0i64);
8507 let __s_b = self.gpu.stream();
8508 let mut b = __s_b.launch_builder(&f);
8509 b.arg(gate_bank)
8510 .arg(up_bank)
8511 .arg(sel)
8512 .arg(aq)
8513 .arg(ad)
8514 .arg(yg)
8515 .arg(yu)
8516 .arg(&inf)
8517 .arg(&outf)
8518 .arg(&ns)
8519 .arg(&rb)
8520 .arg(&es)
8521 .arg(&ars)
8522 .arg(&adrs)
8523 .arg(&own);
8524 unsafe {
8525 b.launch(cfg)?;
8526 }
8527 Ok(())
8528 }
8529
8530 #[allow(clippy::too_many_arguments)]
8532 pub fn silu_mul_scaled_q8_1_sel_ep_into(
8533 &self,
8534 gate: &CudaSlice<f32>,
8535 up: &CudaSlice<f32>,
8536 gmac: &CudaSlice<f32>,
8537 umac: &CudaSlice<f32>,
8538 sel: &CudaSlice<i32>,
8539 limit: Option<f32>,
8540 out_q: &mut CudaSlice<i8>,
8541 out_d: &mut CudaSlice<f32>,
8542 n_per: usize,
8543 n_sel: usize,
8544 owner: usize,
8545 ) -> Result<(), Box<dyn std::error::Error>> {
8546 if n_per % 32 != 0 || out_q.len() < n_sel * n_per || out_d.len() < n_sel * n_per / 32 {
8547 return Err("NVFP4 silu ep geometry".into());
8548 }
8549 let f = self.func("silu_mul_scaled_q8_1_sel_ep");
8550 let warps = n_sel * n_per / 32;
8551 let cfg = LaunchConfig {
8552 grid_dim: ((warps as u32).div_ceil(4), 1, 1),
8553 block_dim: (128, 1, 1),
8554 shared_mem_bytes: 0,
8555 };
8556 let (np, ns, own) = (n_per as i32, n_sel as i32, owner as i32);
8557 let (lim, has) = match limit {
8558 Some(l) => (l, 1i32),
8559 None => (0.0f32, 0i32),
8560 };
8561 let __s_b = self.gpu.stream();
8562 let mut b = __s_b.launch_builder(&f);
8563 b.arg(gate)
8564 .arg(up)
8565 .arg(gmac)
8566 .arg(umac)
8567 .arg(sel)
8568 .arg(&lim)
8569 .arg(&has)
8570 .arg(out_q)
8571 .arg(out_d)
8572 .arg(&np)
8573 .arg(&ns)
8574 .arg(&own);
8575 unsafe {
8576 b.launch(cfg)?;
8577 }
8578 Ok(())
8579 }
8580
8581 #[allow(clippy::too_many_arguments)]
8583 pub fn qmatvec_nvfp4_sel_down8_ep_into(
8584 &self,
8585 bank: &CudaSlice<u8>,
8586 sel: &CudaSlice<i32>,
8587 aq: &CudaSlice<i8>,
8588 ad: &CudaSlice<f32>,
8589 route_w: &CudaSlice<f32>,
8590 md: &CudaSlice<f32>,
8591 dst: &mut CudaSlice<f32>,
8592 n_sel: usize,
8593 in_f: usize,
8594 out_f: usize,
8595 row_bytes: usize,
8596 expert_stride: usize,
8597 act_row_stride: usize,
8598 ad_row_stride: usize,
8599 owner: usize,
8600 ) -> Result<(), Box<dyn std::error::Error>> {
8601 if in_f % 64 != 0 || n_sel == 0 || n_sel > 8 || (in_f >> 5) > 64 || dst.len() < out_f {
8602 return Err("NVFP4 down8 ep geometry".into());
8603 }
8604 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_down8_ep");
8605 let cfg = LaunchConfig {
8606 grid_dim: (out_f as u32, 1, 1),
8607 block_dim: (32, n_sel as u32, 1),
8608 shared_mem_bytes: 0,
8609 };
8610 let (inf, outf, ns, own) = (in_f as i32, out_f as i32, n_sel as i32, owner as i32);
8611 let (rb, es) = (row_bytes as i64, expert_stride as i64);
8612 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
8613 let __s_b = self.gpu.stream();
8614 let mut b = __s_b.launch_builder(&f);
8615 b.arg(bank)
8616 .arg(sel)
8617 .arg(aq)
8618 .arg(ad)
8619 .arg(route_w)
8620 .arg(md)
8621 .arg(dst)
8622 .arg(&inf)
8623 .arg(&outf)
8624 .arg(&ns)
8625 .arg(&rb)
8626 .arg(&es)
8627 .arg(&ars)
8628 .arg(&adrs)
8629 .arg(&own);
8630 unsafe {
8631 b.launch(cfg)?;
8632 }
8633 Ok(())
8634 }
8635
8636 pub fn qmatvec_nvfp4_sel_into(
8637 &self,
8638 bank: &CudaSlice<u8>,
8639 sel: &CudaSlice<i32>,
8640 aq: &CudaSlice<i8>,
8641 ad: &CudaSlice<f32>,
8642 y: &mut CudaSlice<f32>,
8643 n_sel: usize,
8644 in_f: usize,
8645 out_f: usize,
8646 row_bytes: usize,
8647 expert_stride: usize,
8648 act_row_stride: usize,
8649 ad_row_stride: usize,
8650 ) -> Result<(), Box<dyn std::error::Error>> {
8651 assert!(
8652 in_f % 64 == 0,
8653 "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}"
8654 );
8655 if y.len() < n_sel * out_f || sel.len() < n_sel {
8656 return Err(format!(
8657 "NVFP4 sel output {} / sel {} shorter than {n_sel}x{out_f}",
8658 y.len(),
8659 sel.len()
8660 )
8661 .into());
8662 }
8663 static MR: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
8670 let mode = *MR.get_or_init(|| {
8671 if crate::tp::nvfp4_bank_v2_on() {
8672 3
8673 } else if std::env::var("MEMRA_SEL_STREAM").as_deref() == Ok("1") {
8674 2
8675 } else if std::env::var("MEMRA_SEL_MR").as_deref() == Ok("1") {
8676 1
8677 } else {
8678 0
8679 }
8680 });
8681 let mode = if mode == 2 && in_f > 4096 { 0 } else { mode };
8682 static V2S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8686 let v2s = mode == 3
8687 && *V2S.get_or_init(|| std::env::var("MEMRA_SEL_V2S").as_deref() == Ok("1"))
8688 && row_bytes % 16 == 0
8689 && in_f <= 4096;
8690 let f = match (mode, v2s) {
8691 (3, true) => self.func("qmatvec_nvfp4_dp4a_sel_v2s"),
8692 (3, false) => self.func("qmatvec_nvfp4_dp4a_sel_v2"),
8693 (2, _) => self.func("qmatvec_nvfp4_dp4a_sel_stream"),
8694 (1, _) => self.func("qmatvec_nvfp4_dp4a_sel_mr4"),
8695 _ => self.func("qmatvec_nvfp4_dp4a_sel"),
8696 };
8697 let nsb = in_f >> 5;
8702 let fit_block: u32 = if (mode == 0 || mode == 3) && !v2s && nsb <= 32 {
8703 32
8704 } else if mode == 1 {
8705 512
8706 } else {
8707 128
8708 };
8709 let cfg = LaunchConfig {
8710 grid_dim: (
8711 if v2s {
8712 (out_f as u32).div_ceil(8)
8713 } else {
8714 match mode {
8715 2 => (out_f as u32).div_ceil(16),
8716 1 => (out_f as u32).div_ceil(4),
8717 _ => out_f as u32,
8718 }
8719 },
8720 n_sel as u32,
8721 1,
8722 ),
8723 block_dim: (fit_block, 1, 1),
8724 shared_mem_bytes: 0,
8725 };
8726 let (inf, outf, ns) = (in_f as i32, out_f as i32, n_sel as i32);
8727 let (rb, es, ars, adrs) = (
8728 row_bytes as i64,
8729 expert_stride as i64,
8730 act_row_stride as i64,
8731 ad_row_stride as i64,
8732 );
8733 let __s_b = self.gpu.stream();
8734 let mut b = __s_b.launch_builder(&f);
8735 b.arg(bank)
8736 .arg(sel)
8737 .arg(aq)
8738 .arg(ad)
8739 .arg(y)
8740 .arg(&inf)
8741 .arg(&outf)
8742 .arg(&ns)
8743 .arg(&rb)
8744 .arg(&es)
8745 .arg(&ars)
8746 .arg(&adrs);
8747 unsafe {
8748 b.launch(cfg)?;
8749 }
8750 Ok(())
8751 }
8752
8753 #[allow(clippy::too_many_arguments)]
8758 pub fn silu_mul_scaled_q8_1_sel_into(
8759 &self,
8760 gate: &CudaSlice<f32>,
8761 up: &CudaSlice<f32>,
8762 gmac: &CudaSlice<f32>,
8763 umac: &CudaSlice<f32>,
8764 sel: &CudaSlice<i32>,
8765 limit: Option<f32>,
8766 out_q: &mut CudaSlice<i8>,
8767 out_d: &mut CudaSlice<f32>,
8768 n_per: usize,
8769 n_sel: usize,
8770 ) -> Result<(), Box<dyn std::error::Error>> {
8771 let n = n_per * n_sel;
8772 if n_per % 32 != 0 || out_q.len() < n || out_d.len() < n / 32 {
8773 return Err(format!(
8774 "silu sel geometry n_per={n_per} n_sel={n_sel} q={} d={}",
8775 out_q.len(),
8776 out_d.len()
8777 )
8778 .into());
8779 }
8780 if let Some(limit) = limit {
8781 if limit <= 1e-6 {
8782 return Err(format!(
8783 "silu sel clamp limit {limit} is at or below the 1e-6 eps gate"
8784 )
8785 .into());
8786 }
8787 let f = self.func("silu_mul_scaled_q8_1_sel_clamp");
8788 let cfg = LaunchConfig::for_num_elems(n as u32);
8789 let (np, ns) = (n_per as i32, n_sel as i32);
8790 let __s_b = self.gpu.stream();
8791 let mut b = __s_b.launch_builder(&f);
8792 b.arg(gate)
8793 .arg(up)
8794 .arg(gmac)
8795 .arg(umac)
8796 .arg(sel)
8797 .arg(&limit)
8798 .arg(out_q)
8799 .arg(out_d)
8800 .arg(&np)
8801 .arg(&ns);
8802 unsafe {
8803 b.launch(cfg)?;
8804 }
8805 return Ok(());
8806 }
8807 let f = self.func("silu_mul_scaled_q8_1_sel");
8808 let cfg = LaunchConfig::for_num_elems(n as u32);
8809 let (np, ns) = (n_per as i32, n_sel as i32);
8810 let __s_b = self.gpu.stream();
8811 let mut b = __s_b.launch_builder(&f);
8812 b.arg(gate)
8813 .arg(up)
8814 .arg(gmac)
8815 .arg(umac)
8816 .arg(sel)
8817 .arg(out_q)
8818 .arg(out_d)
8819 .arg(&np)
8820 .arg(&ns);
8821 unsafe {
8822 b.launch(cfg)?;
8823 }
8824 Ok(())
8825 }
8826
8827 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8828 Ok(self.gpu.stream().clone_htod(v)?)
8829 }
8830 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
8831 Ok(self.gpu.stream().clone_htod(v)?)
8832 }
8833 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
8835 Ok(self.gpu.stream().clone_htod(v)?)
8836 }
8837 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
8838 Ok(self.gpu.stream().clone_htod(v)?)
8839 }
8840 pub fn dtoh_view(
8842 &self,
8843 d: &cudarc::driver::CudaView<f32>,
8844 ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8845 let v = self.gpu.stream().clone_dtoh(d)?;
8846 self.gpu.stream().synchronize()?;
8847 Ok(v)
8848 }
8849 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
8850 let v = self.gpu.stream().clone_dtoh(d)?;
8851 self.gpu.stream().synchronize()?;
8852 Ok(v)
8853 }
8854 pub fn dtoh_pair(
8858 &self,
8859 a: &CudaSlice<f32>,
8860 b: &CudaSlice<f32>,
8861 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8862 let av = self.gpu.stream().clone_dtoh(a)?;
8863 let bv = self.gpu.stream().clone_dtoh(b)?;
8864 self.gpu.stream().synchronize()?;
8865 Ok((av, bv))
8866 }
8867 pub fn dtoh_pair_views(
8870 &self,
8871 a: &cudarc::driver::CudaView<f32>,
8872 b: &cudarc::driver::CudaView<f32>,
8873 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
8874 let av = self.gpu.stream().clone_dtoh(a)?;
8875 let bv = self.gpu.stream().clone_dtoh(b)?;
8876 self.gpu.stream().synchronize()?;
8877 Ok((av, bv))
8878 }
8879 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
8881 let v = self.gpu.stream().clone_dtoh(d)?;
8882 self.gpu.stream().synchronize()?;
8883 Ok(v)
8884 }
8885 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8887 let v = self.gpu.stream().clone_dtoh(d)?;
8888 self.gpu.stream().synchronize()?;
8889 Ok(v)
8890 }
8891 pub fn dtoh_u8_view(
8892 &self,
8893 d: &cudarc::driver::CudaView<u8>,
8894 ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
8895 let v = self.gpu.stream().clone_dtoh(d)?;
8896 self.gpu.stream().synchronize()?;
8897 Ok(v)
8898 }
8899 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8900 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
8901 self.keep_if_capturing(&s);
8902 Ok(s)
8903 }
8904
8905 pub fn prob_of_token_device(
8914 &self,
8915 logits: &CudaSlice<f32>,
8916 tok: &CudaSlice<u32>,
8917 n_vocab: usize,
8918 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8919 let nb = ARGMAX_NB;
8920 let mut part = self.alloc_uninit::<f32>(nb)?;
8921 let mut p = self.alloc_uninit::<f32>(1)?;
8922 let f1 = self.func("prob_of_token_partial_f32");
8923 let cfg1 = LaunchConfig {
8924 grid_dim: (nb as u32, 1, 1),
8925 block_dim: (256, 1, 1),
8926 shared_mem_bytes: 0,
8927 };
8928 let nv = n_vocab as i32;
8929 let __s_b1 = self.gpu.stream();
8930 let mut b1 = __s_b1.launch_builder(&f1);
8931 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
8932 unsafe {
8933 b1.launch(cfg1)?;
8934 }
8935 let f2 = self.func("prob_of_token_final_f32");
8936 let cfg2 = LaunchConfig {
8937 grid_dim: (1, 1, 1),
8938 block_dim: (256, 1, 1),
8939 shared_mem_bytes: 0,
8940 };
8941 let nbi = nb as i32;
8942 let __s_b2 = self.gpu.stream();
8943 let mut b2 = __s_b2.launch_builder(&f2);
8944 b2.arg(&part).arg(&mut p).arg(&nbi);
8945 unsafe {
8946 b2.launch(cfg2)?;
8947 }
8948 Ok(p)
8949 }
8950
8951 pub fn prob_of_token_device_col(
8958 &self,
8959 logits: &CudaSlice<f32>,
8960 tok_all: &CudaSlice<u32>,
8961 tok_idx: usize,
8962 p_out: &mut CudaSlice<f32>,
8963 p_idx: usize,
8964 n_vocab: usize,
8965 ) -> Result<(), Box<dyn std::error::Error>> {
8966 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
8967 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
8968 let nb = ARGMAX_NB;
8969 let mut part = self.alloc_uninit::<f32>(nb)?;
8970 let f1 = self.func("prob_of_token_partial_f32");
8971 let cfg1 = LaunchConfig {
8972 grid_dim: (nb as u32, 1, 1),
8973 block_dim: (256, 1, 1),
8974 shared_mem_bytes: 0,
8975 };
8976 let nv = n_vocab as i32;
8977 let __s_b1 = self.gpu.stream();
8978 let mut b1 = __s_b1.launch_builder(&f1);
8979 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
8980 unsafe {
8981 b1.launch(cfg1)?;
8982 }
8983 let f2 = self.func("prob_of_token_final_f32");
8984 let cfg2 = LaunchConfig {
8985 grid_dim: (1, 1, 1),
8986 block_dim: (256, 1, 1),
8987 shared_mem_bytes: 0,
8988 };
8989 let nbi = nb as i32;
8990 let __s_b2 = self.gpu.stream();
8991 let mut b2 = __s_b2.launch_builder(&f2);
8992 b2.arg(&part).arg(&mut p_v).arg(&nbi);
8993 unsafe {
8994 b2.launch(cfg2)?;
8995 }
8996 Ok(())
8997 }
8998
8999 pub fn prob_of_token_device_into(
9000 &self,
9001 logits: &CudaSlice<f32>,
9002 tok: &CudaSlice<u32>,
9003 p_out: &mut CudaSlice<f32>,
9004 n_vocab: usize,
9005 ) -> Result<(), Box<dyn std::error::Error>> {
9006 let nb = ARGMAX_NB;
9007 let mut part = self.alloc_uninit::<f32>(nb)?;
9008 let f1 = self.func("prob_of_token_partial_f32");
9009 let cfg1 = LaunchConfig {
9010 grid_dim: (nb as u32, 1, 1),
9011 block_dim: (256, 1, 1),
9012 shared_mem_bytes: 0,
9013 };
9014 let nv = n_vocab as i32;
9015 let __s_b1 = self.gpu.stream();
9016 let mut b1 = __s_b1.launch_builder(&f1);
9017 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
9018 unsafe {
9019 b1.launch(cfg1)?;
9020 }
9021 let f2 = self.func("prob_of_token_final_f32");
9022 let cfg2 = LaunchConfig {
9023 grid_dim: (1, 1, 1),
9024 block_dim: (256, 1, 1),
9025 shared_mem_bytes: 0,
9026 };
9027 let nbi = nb as i32;
9028 let __s_b2 = self.gpu.stream();
9029 let mut b2 = __s_b2.launch_builder(&f2);
9030 b2.arg(&part).arg(p_out).arg(&nbi);
9031 unsafe {
9032 b2.launch(cfg2)?;
9033 }
9034 Ok(())
9035 }
9036
9037 pub fn u32_hist_append(
9040 &self,
9041 tok: &CudaSlice<u32>,
9042 hist: &mut CudaSlice<u32>,
9043 idx: &mut CudaSlice<i32>,
9044 ) -> Result<(), Box<dyn std::error::Error>> {
9045 let f = self.func("u32_hist_append");
9046 let cfg = LaunchConfig {
9047 grid_dim: (1, 1, 1),
9048 block_dim: (32, 1, 1),
9049 shared_mem_bytes: 0,
9050 };
9051 let __s_b = self.gpu.stream();
9052 let mut b = __s_b.launch_builder(&f);
9053 b.arg(tok).arg(&mut *hist).arg(&mut *idx);
9054 unsafe {
9055 b.launch(cfg)?;
9056 }
9057 Ok(())
9058 }
9059
9060 pub fn argmax_token_device(
9061 &self,
9062 logits: &CudaSlice<f32>,
9063 n_vocab: usize,
9064 ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9065 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
9066 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
9067 Ok(tok)
9068 }
9069 pub fn argmax_token_device_into(
9076 &self,
9077 logits: &CudaSlice<f32>,
9078 tok: &mut CudaSlice<u32>,
9079 n_vocab: usize,
9080 ) -> Result<(), Box<dyn std::error::Error>> {
9081 let nb = ARGMAX_NB;
9082 let f1 = self.func("argmax_partial_f32");
9083 let f2 = self.func("argmax_final_f32");
9084 let mut guard = self.argmax_partials.lock().unwrap();
9085 if guard.is_none() {
9086 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
9089 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
9090 *guard = Some((pv, pi));
9091 }
9092 let (part_v, part_i) = guard.as_mut().unwrap();
9093 let nv = n_vocab as i32;
9094 let nbi = nb as i32;
9095 let cfg1 = LaunchConfig {
9097 grid_dim: (nb as u32, 1, 1),
9098 block_dim: (256, 1, 1),
9099 shared_mem_bytes: 0,
9100 };
9101 let __s_b1 = self.gpu.stream();
9102 let mut b1 = __s_b1.launch_builder(&f1);
9103 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
9104 unsafe {
9105 b1.launch(cfg1)?;
9106 }
9107 let cfg2 = LaunchConfig {
9109 grid_dim: (1, 1, 1),
9110 block_dim: (256, 1, 1),
9111 shared_mem_bytes: 0,
9112 };
9113 let __s_b2 = self.gpu.stream();
9114 let mut b2 = __s_b2.launch_builder(&f2);
9115 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
9116 unsafe {
9117 b2.launch(cfg2)?;
9118 }
9119 Ok(())
9120 }
9121 pub fn argmax_token_device_col(
9127 &self,
9128 logits: &CudaSlice<f32>,
9129 col: usize,
9130 n_vocab: usize,
9131 toks: &mut CudaSlice<u32>,
9132 out_idx: usize,
9133 ) -> Result<(), Box<dyn std::error::Error>> {
9134 let nb = ARGMAX_NB;
9135 let f1 = self.func("argmax_partial_f32");
9136 let f2 = self.func("argmax_final_f32");
9137 let mut guard = self.argmax_partials.lock().unwrap();
9138 if guard.is_none() {
9139 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
9140 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
9141 *guard = Some((pv, pi));
9142 }
9143 let (part_v, part_i) = guard.as_mut().unwrap();
9144 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
9145 let nv = n_vocab as i32;
9146 let nbi = nb as i32;
9147 let cfg1 = LaunchConfig {
9148 grid_dim: (nb as u32, 1, 1),
9149 block_dim: (256, 1, 1),
9150 shared_mem_bytes: 0,
9151 };
9152 let __s_b1 = self.gpu.stream();
9153 let mut b1 = __s_b1.launch_builder(&f1);
9154 b1.arg(&col_view)
9155 .arg(&mut *part_v)
9156 .arg(&mut *part_i)
9157 .arg(&nv);
9158 unsafe {
9159 b1.launch(cfg1)?;
9160 }
9161 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
9162 let cfg2 = LaunchConfig {
9163 grid_dim: (1, 1, 1),
9164 block_dim: (256, 1, 1),
9165 shared_mem_bytes: 0,
9166 };
9167 let __s_b2 = self.gpu.stream();
9168 let mut b2 = __s_b2.launch_builder(&f2);
9169 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
9170 unsafe {
9171 b2.launch(cfg2)?;
9172 }
9173 Ok(())
9174 }
9175 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9177 Ok(self.gpu.stream().clone_htod(v)?)
9178 }
9179 pub fn dtoh_u64(&self, d: &CudaSlice<u64>) -> Result<Vec<u64>, Box<dyn std::error::Error>> {
9180 let v = self.gpu.stream().clone_dtoh(d)?;
9181 self.gpu.stream().synchronize()?;
9182 Ok(v)
9183 }
9184
9185 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
9186 let v = self.gpu.stream().clone_dtoh(d)?;
9187 self.gpu.stream().synchronize()?;
9188 Ok(v)
9189 }
9190 pub fn htod_u32_into(
9194 &self,
9195 dst: &mut CudaSlice<u32>,
9196 src: &[u32],
9197 ) -> Result<(), Box<dyn std::error::Error>> {
9198 let mut view = dst.slice_mut(0..src.len());
9199 self.gpu.stream().memcpy_htod(src, &mut view)?;
9200 Ok(())
9201 }
9202
9203 pub fn htod_i32_into(
9206 &self,
9207 dst: &mut CudaSlice<i32>,
9208 src: &[i32],
9209 ) -> Result<(), Box<dyn std::error::Error>> {
9210 let mut view = dst.slice_mut(0..src.len());
9211 self.gpu.stream().memcpy_htod(src, &mut view)?;
9212 Ok(())
9213 }
9214
9215 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
9216 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
9217 self.keep_if_capturing(&s);
9218 Ok(s)
9219 }
9220 pub fn embed_gather_device_into(
9223 &self,
9224 embd: &CudaSlice<u8>,
9225 token_d: &CudaSlice<u32>,
9226 x_out: &mut CudaSlice<f32>,
9227 n_embd: usize,
9228 qtype: i32,
9229 row_bytes: usize,
9230 ) -> Result<(), Box<dyn std::error::Error>> {
9231 let f = self.func("embed_gather_u32");
9232 let cfg = LaunchConfig {
9233 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9234 block_dim: (256, 1, 1),
9235 shared_mem_bytes: 0,
9236 };
9237 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9238 let __s_b = self.gpu.stream();
9239 let mut b = __s_b.launch_builder(&f);
9240 b.arg(embd)
9241 .arg(token_d)
9242 .arg(x_out)
9243 .arg(&ne)
9244 .arg(&qt)
9245 .arg(&rb);
9246 unsafe {
9247 b.launch(cfg)?;
9248 }
9249 Ok(())
9250 }
9251 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
9253 let v = self.gpu.stream().clone_dtoh(d)?;
9254 self.gpu.stream().synchronize()?;
9255 Ok(v[0])
9256 }
9257 pub fn i32_set_k(
9264 &self,
9265 dst: &mut CudaSlice<i32>,
9266 v: i32,
9267 ) -> Result<(), Box<dyn std::error::Error>> {
9268 let f = self.func("i32_set_k");
9269 let cfg = LaunchConfig {
9270 grid_dim: (1, 1, 1),
9271 block_dim: (1, 1, 1),
9272 shared_mem_bytes: 0,
9273 };
9274 let idx = 0i32;
9275 let __s_b = self.gpu.stream();
9276 let mut b = __s_b.launch_builder(&f);
9277 b.arg(dst).arg(&v).arg(&idx);
9278 unsafe {
9279 b.launch(cfg)?;
9280 }
9281 Ok(())
9282 }
9283
9284 pub fn set_i32_one(
9285 &self,
9286 d: &mut CudaSlice<i32>,
9287 v: i32,
9288 ) -> Result<(), Box<dyn std::error::Error>> {
9289 self.gpu.stream().memcpy_htod(&[v], d)?;
9290 Ok(())
9291 }
9292 pub fn set_u32_one(
9295 &self,
9296 d: &mut CudaSlice<u32>,
9297 v: u32,
9298 ) -> Result<(), Box<dyn std::error::Error>> {
9299 self.gpu.stream().memcpy_htod(&[v], d)?;
9300 Ok(())
9301 }
9302 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
9304 let v = self.gpu.stream().clone_dtoh(d)?;
9305 self.gpu.stream().synchronize()?;
9306 Ok(v[0])
9307 }
9308 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
9310 Ok(self.gpu.stream().clone_htod(bytes)?)
9311 }
9312 pub fn embed_gather_device(
9316 &self,
9317 embd: &CudaSlice<u8>,
9318 token_d: &CudaSlice<u32>,
9319 n_embd: usize,
9320 qtype: i32,
9321 row_bytes: usize,
9322 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9323 let f = self.func("embed_gather_u32");
9324 let mut x = self.alloc_uninit::<f32>(n_embd)?;
9325 let cfg = LaunchConfig {
9326 grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
9327 block_dim: (256, 1, 1),
9328 shared_mem_bytes: 0,
9329 };
9330 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
9331 let __s_b = self.gpu.stream();
9332 let mut b = __s_b.launch_builder(&f);
9333 b.arg(embd)
9334 .arg(token_d)
9335 .arg(&mut x)
9336 .arg(&ne)
9337 .arg(&qt)
9338 .arg(&rb);
9339 unsafe {
9340 b.launch(cfg)?;
9341 }
9342 Ok(x)
9343 }
9344
9345 pub fn embed_gather_device_t(
9349 &self,
9350 embd: &CudaSlice<u8>,
9351 tokens: &[u32],
9352 n_embd: usize,
9353 qtype: i32,
9354 row_bytes: usize,
9355 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9356 let t = tokens.len();
9357 let tok_d = self.gpu.stream().clone_htod(tokens)?;
9358 let f = self.func("embed_gather_u32_t");
9359 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9360 let cfg = LaunchConfig {
9361 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9362 block_dim: (256, 1, 1),
9363 shared_mem_bytes: 0,
9364 };
9365 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9366 let __s_b = self.gpu.stream();
9367 let mut b = __s_b.launch_builder(&f);
9368 b.arg(embd)
9369 .arg(&tok_d)
9370 .arg(&mut x)
9371 .arg(&ne)
9372 .arg(&qt)
9373 .arg(&rb)
9374 .arg(&ti);
9375 unsafe {
9376 b.launch(cfg)?;
9377 }
9378 Ok(x)
9379 }
9380
9381 pub fn embed_gather_device_tv(
9386 &self,
9387 embd: &CudaSlice<u8>,
9388 tok_v: &cudarc::driver::CudaView<u32>,
9389 t: usize,
9390 n_embd: usize,
9391 qtype: i32,
9392 row_bytes: usize,
9393 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9394 let f = self.func("embed_gather_u32_t");
9395 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9396 let cfg = LaunchConfig {
9397 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9398 block_dim: (256, 1, 1),
9399 shared_mem_bytes: 0,
9400 };
9401 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9402 let __s_b = self.gpu.stream();
9403 let mut b = __s_b.launch_builder(&f);
9404 b.arg(embd)
9405 .arg(tok_v)
9406 .arg(&mut x)
9407 .arg(&ne)
9408 .arg(&qt)
9409 .arg(&rb)
9410 .arg(&ti);
9411 unsafe {
9412 b.launch(cfg)?;
9413 }
9414 Ok(x)
9415 }
9416
9417 pub fn embed_gather_device_td(
9418 &self,
9419 embd: &CudaSlice<u8>,
9420 tok_d: &CudaSlice<u32>,
9421 t: usize,
9422 n_embd: usize,
9423 qtype: i32,
9424 row_bytes: usize,
9425 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9426 let f = self.func("embed_gather_u32_t");
9427 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
9428 let cfg = LaunchConfig {
9429 grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
9430 block_dim: (256, 1, 1),
9431 shared_mem_bytes: 0,
9432 };
9433 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
9434 let __s_b = self.gpu.stream();
9435 let mut b = __s_b.launch_builder(&f);
9436 b.arg(embd)
9437 .arg(tok_d)
9438 .arg(&mut x)
9439 .arg(&ne)
9440 .arg(&qt)
9441 .arg(&rb)
9442 .arg(&ti);
9443 unsafe {
9444 b.launch(cfg)?;
9445 }
9446 Ok(x)
9447 }
9448
9449 #[inline]
9455 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
9457 if self
9458 .capture_keep_on
9459 .load(std::sync::atomic::Ordering::Relaxed)
9460 {
9461 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
9462 }
9463 }
9464
9465 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(
9466 &self,
9467 n: usize,
9468 ) -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
9469 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
9470 {
9474 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9475 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
9476 use cudarc::driver::DevicePtrMut;
9478 let n_bytes = s.len() * std::mem::size_of::<T>();
9479 let stream = self.gpu.stream();
9480 let (p_, _g) = s.device_ptr_mut(&stream);
9481 unsafe {
9482 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
9483 .result()?;
9484 }
9485 }
9486 }
9487 self.keep_if_capturing(&s);
9488 Ok(s)
9489 }
9490
9491 pub fn uninit_q8_pair(
9496 &self,
9497 n: usize,
9498 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9499 Ok((
9500 self.alloc_uninit::<i8>(n)?,
9501 self.alloc_uninit::<f32>(n / 32)?,
9502 ))
9503 }
9504
9505 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9506 self.alloc_uninit::<f32>(n)
9507 }
9508
9509 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
9511 self.alloc_uninit::<i8>(n)
9512 }
9513
9514 #[allow(clippy::too_many_arguments)]
9518 pub fn rms_norm3(
9519 &self,
9520 x: &CudaSlice<f32>,
9521 w0: &CudaSlice<f32>,
9522 w1: &CudaSlice<f32>,
9523 w2: &CudaSlice<f32>,
9524 d0: &mut CudaSlice<f32>,
9525 d1: &mut CudaSlice<f32>,
9526 d2: &mut CudaSlice<f32>,
9527 ncols: usize,
9528 nrows: usize,
9529 eps: f32,
9530 ) -> Result<(), Box<dyn std::error::Error>> {
9531 let f = self.func("rms_norm3_f32");
9532 let cfg = LaunchConfig {
9533 grid_dim: (nrows as u32, 1, 1),
9534 block_dim: (rms_block(), 1, 1),
9535 shared_mem_bytes: 0,
9536 };
9537 let (nc, e) = (ncols as i32, eps);
9538 let __s_b = self.gpu.stream();
9539 let mut b = __s_b.launch_builder(&f);
9540 b.arg(x)
9541 .arg(w0)
9542 .arg(w1)
9543 .arg(w2)
9544 .arg(d0)
9545 .arg(d1)
9546 .arg(d2)
9547 .arg(&nc)
9548 .arg(&e);
9549 unsafe {
9550 b.launch(cfg)?;
9551 }
9552 Ok(())
9553 }
9554
9555 #[allow(clippy::too_many_arguments)]
9557 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
9560 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9561 *WARP_ON.get_or_init(|| {
9562 std::env::var("MEMRA_QKVNORM_W")
9563 .map(|v| v != "0")
9564 .unwrap_or(true)
9565 }) && ncols % 4 == 0
9566 && rows >= 64
9567 }
9568
9569 #[allow(clippy::too_many_arguments)]
9572 pub fn rms_norm_qkv_w4b(
9573 &self,
9574 q: &CudaSlice<f32>,
9575 k: &CudaSlice<f32>,
9576 v: &CudaSlice<f32>,
9577 wq: &CudaSlice<f32>,
9578 wk: &CudaSlice<f32>,
9579 wv: &CudaSlice<f32>,
9580 dq: &mut CudaSlice<f32>,
9581 dk: &mut CudaSlice<f32>,
9582 dv: &mut CudaSlice<f32>,
9583 dvb: &mut CudaSlice<u8>,
9584 ncols: usize,
9585 rq: usize,
9586 rk: usize,
9587 eps: f32,
9588 vf16: bool,
9589 ) -> Result<(), Box<dyn std::error::Error>> {
9590 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
9591 let f = self.func("rms_norm_qkv_w4b_f32");
9592 let rows = (rq + 2 * rk) as u32;
9593 let cfg = LaunchConfig {
9594 grid_dim: (rows.div_ceil(8), 1, 1),
9595 block_dim: (256, 1, 1),
9596 shared_mem_bytes: 0,
9597 };
9598 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9599 let vf = vf16 as i32;
9600 let __s_b = self.gpu.stream();
9601 let mut b = __s_b.launch_builder(&f);
9602 b.arg(q)
9603 .arg(k)
9604 .arg(v)
9605 .arg(wq)
9606 .arg(wk)
9607 .arg(wv)
9608 .arg(dq)
9609 .arg(dk)
9610 .arg(dv)
9611 .arg(&mut *dvb)
9612 .arg(&nc)
9613 .arg(&rqi)
9614 .arg(&rki)
9615 .arg(&rvi)
9616 .arg(&e)
9617 .arg(&vf);
9618 unsafe {
9619 b.launch(cfg)?;
9620 }
9621 Ok(())
9622 }
9623
9624 pub fn rms_norm_qkv(
9625 &self,
9626 q: &CudaSlice<f32>,
9627 k: &CudaSlice<f32>,
9628 v: &CudaSlice<f32>,
9629 wq: &CudaSlice<f32>,
9630 wk: &CudaSlice<f32>,
9631 wv: &CudaSlice<f32>,
9632 dq: &mut CudaSlice<f32>,
9633 dk: &mut CudaSlice<f32>,
9634 dv: &mut CudaSlice<f32>,
9635 ncols: usize,
9636 rq: usize,
9637 rk: usize,
9638 eps: f32,
9639 ) -> Result<(), Box<dyn std::error::Error>> {
9640 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9644 let warp_on = *WARP_ON.get_or_init(|| {
9645 std::env::var("MEMRA_QKVNORM_W")
9646 .map(|v| v != "0")
9647 .unwrap_or(true)
9648 });
9649 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
9652 let f = self.func("rms_norm_qkv_w4_f32");
9653 let rows = (rq + 2 * rk) as u32;
9654 let cfg = LaunchConfig {
9655 grid_dim: (rows.div_ceil(8), 1, 1),
9656 block_dim: (256, 1, 1),
9657 shared_mem_bytes: 0,
9658 };
9659 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
9660 let __s_b = self.gpu.stream();
9661 let mut b = __s_b.launch_builder(&f);
9662 b.arg(q)
9663 .arg(k)
9664 .arg(v)
9665 .arg(wq)
9666 .arg(wk)
9667 .arg(wv)
9668 .arg(dq)
9669 .arg(dk)
9670 .arg(dv)
9671 .arg(&nc)
9672 .arg(&rqi)
9673 .arg(&rki)
9674 .arg(&rvi)
9675 .arg(&e);
9676 unsafe {
9677 b.launch(cfg)?;
9678 }
9679 return Ok(());
9680 }
9681 let f = self.func("rms_norm_qkv_f32");
9682 let grid = (rq + 2 * rk) as u32;
9683 let cfg = LaunchConfig {
9684 grid_dim: (grid, 1, 1),
9685 block_dim: (rms_block(), 1, 1),
9686 shared_mem_bytes: 0,
9687 };
9688 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
9689 let __s_b = self.gpu.stream();
9690 let mut b = __s_b.launch_builder(&f);
9691 b.arg(q)
9692 .arg(k)
9693 .arg(v)
9694 .arg(wq)
9695 .arg(wk)
9696 .arg(wv)
9697 .arg(dq)
9698 .arg(dk)
9699 .arg(dv)
9700 .arg(&nc)
9701 .arg(&rqi)
9702 .arg(&rki)
9703 .arg(&e);
9704 unsafe {
9705 b.launch(cfg)?;
9706 }
9707 Ok(())
9708 }
9709
9710 #[allow(clippy::too_many_arguments)]
9712 pub fn rms_norm2x(
9713 &self,
9714 a: &CudaSlice<f32>,
9715 bb: &CudaSlice<f32>,
9716 wa: &CudaSlice<f32>,
9717 wb: &CudaSlice<f32>,
9718 da: &mut CudaSlice<f32>,
9719 db: &mut CudaSlice<f32>,
9720 ncols: usize,
9721 nrows: usize,
9722 eps: f32,
9723 ) -> Result<(), Box<dyn std::error::Error>> {
9724 let f = self.func("rms_norm2x_f32");
9725 let cfg = LaunchConfig {
9726 grid_dim: (2 * nrows as u32, 1, 1),
9727 block_dim: (rms_block(), 1, 1),
9728 shared_mem_bytes: 0,
9729 };
9730 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
9731 let __s_b = self.gpu.stream();
9732 let mut b = __s_b.launch_builder(&f);
9733 b.arg(a)
9734 .arg(bb)
9735 .arg(wa)
9736 .arg(wb)
9737 .arg(da)
9738 .arg(db)
9739 .arg(&nc)
9740 .arg(&nr)
9741 .arg(&e);
9742 unsafe {
9743 b.launch(cfg)?;
9744 }
9745 Ok(())
9746 }
9747
9748 pub fn softcap(
9750 &self,
9751 y: &mut CudaSlice<f32>,
9752 cap: f32,
9753 n: usize,
9754 ) -> Result<(), Box<dyn std::error::Error>> {
9755 let f = self.func("softcap_f32");
9756 let cfg = LaunchConfig::for_num_elems(n as u32);
9757 let ni = n as i32;
9758 let __s_b = self.gpu.stream();
9759 let mut b = __s_b.launch_builder(&f);
9760 b.arg(y).arg(&cap).arg(&ni);
9761 unsafe {
9762 b.launch(cfg)?;
9763 }
9764 Ok(())
9765 }
9766
9767 pub fn mask_ids_rows(
9770 &self,
9771 y: &mut CudaSlice<f32>,
9772 ids: &CudaSlice<i32>,
9773 n_ids: usize,
9774 n_vocab: usize,
9775 t: usize,
9776 ) -> Result<(), Box<dyn std::error::Error>> {
9777 let f = self.func("mask_ids_rows_f32");
9778 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
9779 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
9780 let __s_b = self.gpu.stream();
9781 let mut b = __s_b.launch_builder(&f);
9782 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
9783 unsafe {
9784 b.launch(cfg)?;
9785 }
9786 Ok(())
9787 }
9788
9789 #[allow(clippy::too_many_arguments)]
9791 pub fn add_scale_rms_norm(
9792 &self,
9793 a: &CudaSlice<f32>,
9794 b_in: &CudaSlice<f32>,
9795 c: f32,
9796 w: &CudaSlice<f32>,
9797 res: &mut CudaSlice<f32>,
9798 dst: &mut CudaSlice<f32>,
9799 ncols: usize,
9800 nrows: usize,
9801 eps: f32,
9802 ) -> Result<(), Box<dyn std::error::Error>> {
9803 let f = self.func("add_scale_rms_norm_f32");
9804 let cfg = LaunchConfig {
9805 grid_dim: (nrows as u32, 1, 1),
9806 block_dim: (rms_block(), 1, 1),
9807 shared_mem_bytes: 0,
9808 };
9809 let (nc, e2) = (ncols as i32, eps);
9810 let __s_b = self.gpu.stream();
9811 let mut b = __s_b.launch_builder(&f);
9812 b.arg(a)
9813 .arg(b_in)
9814 .arg(&c)
9815 .arg(w)
9816 .arg(res)
9817 .arg(dst)
9818 .arg(&nc)
9819 .arg(&e2);
9820 unsafe {
9821 b.launch(cfg)?;
9822 }
9823 Ok(())
9824 }
9825
9826 #[allow(clippy::too_many_arguments)]
9829 pub fn add_scale_rms_norm_q8_1(
9830 &self,
9831 a: &CudaSlice<f32>,
9832 b_in: &CudaSlice<f32>,
9833 c: f32,
9834 w: &CudaSlice<f32>,
9835 res: &mut CudaSlice<f32>,
9836 ncols: usize,
9837 nrows: usize,
9838 eps: f32,
9839 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9840 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9841 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9842 let (nc, e2) = (ncols as i32, eps);
9843 if Self::pdl_on() && Self::pdl_wb_on() {
9844 {
9845 use cudarc::driver::{DevicePtr, DevicePtrMut};
9846 let s = &self.gpu.stream();
9847 let (pa, _g0) = a.device_ptr(s);
9848 let (pb, _g1) = b_in.device_ptr(s);
9849 let (pw, _g2) = w.device_ptr(s);
9850 let (pr, _g3) = res.device_ptr_mut(s);
9851 let (pq, _g4) = out_q.device_ptr_mut(s);
9852 let (pd, _g5) = out_d.device_ptr_mut(s);
9853 let mut ps = [
9854 &pa as *const _ as *mut std::ffi::c_void,
9855 &pb as *const _ as *mut _,
9856 &c as *const _ as *mut _,
9857 &pw as *const _ as *mut _,
9858 &pr as *const _ as *mut _,
9859 &pq as *const _ as *mut _,
9860 &pd as *const _ as *mut _,
9861 &nc as *const _ as *mut _,
9862 &e2 as *const _ as *mut _,
9863 ];
9864 unsafe {
9865 self.launch_pdl(
9866 "add_scale_rms_norm_q8_1",
9867 (nrows as u32, 1, 1),
9868 (rms_block(), 1, 1),
9869 &mut ps,
9870 )?;
9871 }
9872 }
9873 return Ok((out_q, out_d));
9874 }
9875 let f = self.func("add_scale_rms_norm_q8_1");
9876 let cfg = LaunchConfig {
9877 grid_dim: (nrows as u32, 1, 1),
9878 block_dim: (rms_block(), 1, 1),
9879 shared_mem_bytes: 0,
9880 };
9881 let __s_b = self.gpu.stream();
9882 let mut b = __s_b.launch_builder(&f);
9883 b.arg(a)
9884 .arg(b_in)
9885 .arg(&c)
9886 .arg(w)
9887 .arg(res)
9888 .arg(&mut out_q)
9889 .arg(&mut out_d)
9890 .arg(&nc)
9891 .arg(&e2);
9892 unsafe {
9893 b.launch(cfg)?;
9894 }
9895 Ok((out_q, out_d))
9896 }
9897
9898 #[allow(clippy::too_many_arguments)]
9900 pub fn add_scale_rms_norm_q8_1_into(
9901 &self,
9902 a: &CudaSlice<f32>,
9903 b_in: &CudaSlice<f32>,
9904 c: f32,
9905 w: &CudaSlice<f32>,
9906 res: &mut CudaSlice<f32>,
9907 ncols: usize,
9908 nrows: usize,
9909 eps: f32,
9910 out_q: &mut CudaSlice<i8>,
9911 out_d: &mut CudaSlice<f32>,
9912 ) -> Result<(), Box<dyn std::error::Error>> {
9913 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
9914 let (nc, e2) = (ncols as i32, eps);
9915 if Self::pdl_on() && Self::pdl_wb_on() {
9916 use cudarc::driver::{DevicePtr, DevicePtrMut};
9917 let s = &self.gpu.stream();
9918 let (pa, _g0) = a.device_ptr(s);
9919 let (pb, _g1) = b_in.device_ptr(s);
9920 let (pw, _g2) = w.device_ptr(s);
9921 let (pr, _g3) = res.device_ptr_mut(s);
9922 let (pq, _g4) = out_q.device_ptr_mut(s);
9923 let (pd, _g5) = out_d.device_ptr_mut(s);
9924 let mut ps = [
9925 &pa as *const _ as *mut std::ffi::c_void,
9926 &pb as *const _ as *mut _,
9927 &c as *const _ as *mut _,
9928 &pw as *const _ as *mut _,
9929 &pr as *const _ as *mut _,
9930 &pq as *const _ as *mut _,
9931 &pd as *const _ as *mut _,
9932 &nc as *const _ as *mut _,
9933 &e2 as *const _ as *mut _,
9934 ];
9935 unsafe {
9936 self.launch_pdl(
9937 "add_scale_rms_norm_q8_1",
9938 (nrows as u32, 1, 1),
9939 (rms_block(), 1, 1),
9940 &mut ps,
9941 )?;
9942 }
9943 return Ok(());
9944 }
9945 let f = self.func("add_scale_rms_norm_q8_1");
9946 let cfg = LaunchConfig {
9947 grid_dim: (nrows as u32, 1, 1),
9948 block_dim: (rms_block(), 1, 1),
9949 shared_mem_bytes: 0,
9950 };
9951 let __s_b = self.gpu.stream();
9952 let mut b = __s_b.launch_builder(&f);
9953 b.arg(a)
9954 .arg(b_in)
9955 .arg(&c)
9956 .arg(w)
9957 .arg(res)
9958 .arg(&mut *out_q)
9959 .arg(&mut *out_d)
9960 .arg(&nc)
9961 .arg(&e2);
9962 unsafe {
9963 b.launch(cfg)?;
9964 }
9965 Ok(())
9966 }
9967
9968 #[allow(clippy::too_many_arguments)]
9971 pub fn rms_pre_add_scale_rms_norm_q8_1(
9972 &self,
9973 a: &CudaSlice<f32>,
9974 wa: &CudaSlice<f32>,
9975 b_in: &CudaSlice<f32>,
9976 c: f32,
9977 w: &CudaSlice<f32>,
9978 res: &mut CudaSlice<f32>,
9979 ncols: usize,
9980 nrows: usize,
9981 eps: f32,
9982 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9983 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
9984 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
9985 let (nc, e2) = (ncols as i32, eps);
9986 if Self::pdl_on() {
9987 {
9988 use cudarc::driver::{DevicePtr, DevicePtrMut};
9989 let s = &self.gpu.stream();
9990 let (pa, _g0) = a.device_ptr(s);
9991 let (pwa, _g1) = wa.device_ptr(s);
9992 let (pb, _g2) = b_in.device_ptr(s);
9993 let (pw, _g3) = w.device_ptr(s);
9994 let (pr, _g4) = res.device_ptr_mut(s);
9995 let (pq, _g5) = out_q.device_ptr_mut(s);
9996 let (pd, _g6) = out_d.device_ptr_mut(s);
9997 let mut ps = [
9998 &pa as *const _ as *mut std::ffi::c_void,
9999 &pwa as *const _ as *mut _,
10000 &pb as *const _ as *mut _,
10001 &c as *const _ as *mut _,
10002 &pw as *const _ as *mut _,
10003 &pr as *const _ as *mut _,
10004 &pq as *const _ as *mut _,
10005 &pd as *const _ as *mut _,
10006 &nc as *const _ as *mut _,
10007 &e2 as *const _ as *mut _,
10008 ];
10009 unsafe {
10010 self.launch_pdl(
10011 "rms_pre_add_scale_rms_norm_q8_1",
10012 (nrows as u32, 1, 1),
10013 (rms_block(), 1, 1),
10014 &mut ps,
10015 )?;
10016 }
10017 }
10018 return Ok((out_q, out_d));
10019 }
10020 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10021 let cfg = LaunchConfig {
10022 grid_dim: (nrows as u32, 1, 1),
10023 block_dim: (rms_block(), 1, 1),
10024 shared_mem_bytes: 0,
10025 };
10026 let __s_b = self.gpu.stream();
10027 let mut b = __s_b.launch_builder(&f);
10028 b.arg(a)
10029 .arg(wa)
10030 .arg(b_in)
10031 .arg(&c)
10032 .arg(w)
10033 .arg(res)
10034 .arg(&mut out_q)
10035 .arg(&mut out_d)
10036 .arg(&nc)
10037 .arg(&e2);
10038 unsafe {
10039 b.launch(cfg)?;
10040 }
10041 Ok((out_q, out_d))
10042 }
10043
10044 pub fn gelu_tanh_mul_q8_1(
10047 &self,
10048 gate: &CudaSlice<f32>,
10049 up: &cudarc::driver::CudaView<f32>,
10050 act: &mut CudaSlice<f32>,
10051 ncols: usize,
10052 nrows: usize,
10053 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10054 debug_assert!(ncols % 128 == 0);
10055 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10056 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10057 let nc = ncols as i32;
10058 if Self::pdl_on() {
10059 {
10060 use cudarc::driver::{DevicePtr, DevicePtrMut};
10061 let s = &self.gpu.stream();
10062 let (pg, _g0) = gate.device_ptr(s);
10063 let (pu, _g1) = up.device_ptr(s);
10064 let (pact, _g2) = act.device_ptr_mut(s);
10065 let (pq, _g3) = out_q.device_ptr_mut(s);
10066 let (pd, _g4) = out_d.device_ptr_mut(s);
10067 let mut ps = [
10068 &pg as *const _ as *mut std::ffi::c_void,
10069 &pu as *const _ as *mut _,
10070 &pact as *const _ as *mut _,
10071 &pq as *const _ as *mut _,
10072 &pd as *const _ as *mut _,
10073 &nc as *const _ as *mut _,
10074 ];
10075 unsafe {
10076 self.launch_pdl(
10077 "gelu_tanh_mul_q8_1",
10078 (nrows as u32, 1, 1),
10079 (rms_block(), 1, 1),
10080 &mut ps,
10081 )?;
10082 }
10083 }
10084 return Ok((out_q, out_d));
10085 }
10086 let f = self.func("gelu_tanh_mul_q8_1");
10087 let cfg = LaunchConfig {
10088 grid_dim: (nrows as u32, 1, 1),
10089 block_dim: (rms_block(), 1, 1),
10090 shared_mem_bytes: 0,
10091 };
10092 let __s_b = self.gpu.stream();
10093 let mut b = __s_b.launch_builder(&f);
10094 b.arg(gate)
10095 .arg(up)
10096 .arg(act)
10097 .arg(&mut out_q)
10098 .arg(&mut out_d)
10099 .arg(&nc);
10100 unsafe {
10101 b.launch(cfg)?;
10102 }
10103 Ok((out_q, out_d))
10104 }
10105
10106 #[allow(clippy::too_many_arguments)]
10108 pub fn gelu_tanh_mul_q8_1_into(
10109 &self,
10110 gate: &CudaSlice<f32>,
10111 up: &cudarc::driver::CudaView<f32>,
10112 act: &mut CudaSlice<f32>,
10113 ncols: usize,
10114 nrows: usize,
10115 out_q: &mut CudaSlice<i8>,
10116 out_d: &mut CudaSlice<f32>,
10117 ) -> Result<(), Box<dyn std::error::Error>> {
10118 debug_assert!(ncols % 128 == 0);
10119 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
10120 let nc = ncols as i32;
10121 if Self::pdl_on() {
10122 use cudarc::driver::{DevicePtr, DevicePtrMut};
10123 let s = &self.gpu.stream();
10124 let (pg, _g0) = gate.device_ptr(s);
10125 let (pu, _g1) = up.device_ptr(s);
10126 let (pact, _g2) = act.device_ptr_mut(s);
10127 let (pq, _g3) = out_q.device_ptr_mut(s);
10128 let (pd, _g4) = out_d.device_ptr_mut(s);
10129 let mut ps = [
10130 &pg as *const _ as *mut std::ffi::c_void,
10131 &pu as *const _ as *mut _,
10132 &pact as *const _ as *mut _,
10133 &pq as *const _ as *mut _,
10134 &pd as *const _ as *mut _,
10135 &nc as *const _ as *mut _,
10136 ];
10137 unsafe {
10138 self.launch_pdl(
10139 "gelu_tanh_mul_q8_1",
10140 (nrows as u32, 1, 1),
10141 (rms_block(), 1, 1),
10142 &mut ps,
10143 )?;
10144 }
10145 return Ok(());
10146 }
10147 let f = self.func("gelu_tanh_mul_q8_1");
10148 let cfg = LaunchConfig {
10149 grid_dim: (nrows as u32, 1, 1),
10150 block_dim: (rms_block(), 1, 1),
10151 shared_mem_bytes: 0,
10152 };
10153 let __s_b = self.gpu.stream();
10154 let mut b = __s_b.launch_builder(&f);
10155 b.arg(gate)
10156 .arg(up)
10157 .arg(&mut *act)
10158 .arg(&mut *out_q)
10159 .arg(&mut *out_d)
10160 .arg(&nc);
10161 unsafe {
10162 b.launch(cfg)?;
10163 }
10164 Ok(())
10165 }
10166
10167 #[allow(clippy::too_many_arguments)]
10169 pub fn add_rms_norm3_q8z(
10170 &self,
10171 a: &CudaSlice<f32>,
10172 b_in: &CudaSlice<f32>,
10173 w0: &CudaSlice<f32>,
10174 w1: &CudaSlice<f32>,
10175 w2: &CudaSlice<f32>,
10176 res: &mut CudaSlice<f32>,
10177 out1: &mut CudaSlice<f32>,
10178 ncols: usize,
10179 nrows: usize,
10180 eps: f32,
10181 ) -> Result<
10182 (
10183 (CudaSlice<i8>, CudaSlice<f32>),
10184 (CudaSlice<i8>, CudaSlice<f32>),
10185 ),
10186 Box<dyn std::error::Error>,
10187 > {
10188 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
10189 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10190 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
10191 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10192 let f = self.func("add_rms_norm3_q8z_f32");
10193 let cfg = LaunchConfig {
10194 grid_dim: (nrows as u32, 1, 1),
10195 block_dim: (rms_block(), 1, 1),
10196 shared_mem_bytes: 0,
10197 };
10198 let (nc, e2) = (ncols as i32, eps);
10199 let __s_b = self.gpu.stream();
10200 let mut b = __s_b.launch_builder(&f);
10201 b.arg(a)
10202 .arg(b_in)
10203 .arg(w0)
10204 .arg(w1)
10205 .arg(w2)
10206 .arg(res)
10207 .arg(&mut q0)
10208 .arg(&mut d0)
10209 .arg(out1)
10210 .arg(&mut q2)
10211 .arg(&mut d2)
10212 .arg(&nc)
10213 .arg(&e2);
10214 unsafe {
10215 b.launch(cfg)?;
10216 }
10217 Ok(((q0, d0), (q2, d2)))
10218 }
10219
10220 #[allow(clippy::too_many_arguments)]
10222 pub fn add_rms_norm3(
10223 &self,
10224 a: &CudaSlice<f32>,
10225 b_in: &CudaSlice<f32>,
10226 w0: &CudaSlice<f32>,
10227 w1: &CudaSlice<f32>,
10228 w2: &CudaSlice<f32>,
10229 res: &mut CudaSlice<f32>,
10230 d0: &mut CudaSlice<f32>,
10231 d1: &mut CudaSlice<f32>,
10232 d2: &mut CudaSlice<f32>,
10233 ncols: usize,
10234 nrows: usize,
10235 eps: f32,
10236 ) -> Result<(), Box<dyn std::error::Error>> {
10237 let f = self.func("add_rms_norm3_f32");
10238 let cfg = LaunchConfig {
10239 grid_dim: (nrows as u32, 1, 1),
10240 block_dim: (rms_block(), 1, 1),
10241 shared_mem_bytes: 0,
10242 };
10243 let (nc, e2) = (ncols as i32, eps);
10244 let __s_b = self.gpu.stream();
10245 let mut b = __s_b.launch_builder(&f);
10246 b.arg(a)
10247 .arg(b_in)
10248 .arg(w0)
10249 .arg(w1)
10250 .arg(w2)
10251 .arg(res)
10252 .arg(d0)
10253 .arg(d1)
10254 .arg(d2)
10255 .arg(&nc)
10256 .arg(&e2);
10257 unsafe {
10258 b.launch(cfg)?;
10259 }
10260 Ok(())
10261 }
10262
10263 pub fn add_scale(
10265 &self,
10266 a: &CudaSlice<f32>,
10267 b_in: &CudaSlice<f32>,
10268 c: f32,
10269 dst: &mut CudaSlice<f32>,
10270 n: usize,
10271 ) -> Result<(), Box<dyn std::error::Error>> {
10272 let f = self.func("add_scale_f32");
10273 let cfg = LaunchConfig::for_num_elems(n as u32);
10274 let ni = n as i32;
10275 let __s_b = self.gpu.stream();
10276 let mut b = __s_b.launch_builder(&f);
10277 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
10278 unsafe {
10279 b.launch(cfg)?;
10280 }
10281 Ok(())
10282 }
10283
10284 pub fn layer_norm_bias(
10286 &self,
10287 x: &CudaSlice<f32>,
10288 w: &CudaSlice<f32>,
10289 b: &CudaSlice<f32>,
10290 dst: &mut CudaSlice<f32>,
10291 ncols: usize,
10292 nrows: usize,
10293 eps: f32,
10294 ) -> Result<(), Box<dyn std::error::Error>> {
10295 let f = self.func("layer_norm_bias_f32");
10296 let (nc, e) = (ncols as i32, eps);
10297 let cfg = LaunchConfig {
10298 grid_dim: (nrows as u32, 1, 1),
10299 block_dim: (256, 1, 1),
10300 shared_mem_bytes: 0,
10301 };
10302 let __s_b = self.gpu.stream();
10303 let mut lb = __s_b.launch_builder(&f);
10304 lb.arg(x).arg(w).arg(b).arg(&mut *dst).arg(&nc).arg(&e);
10305 unsafe {
10306 lb.launch(cfg)?;
10307 }
10308 Ok(())
10309 }
10310
10311 pub fn gelu_tanh(
10313 &self,
10314 x: &CudaSlice<f32>,
10315 dst: &mut CudaSlice<f32>,
10316 n: usize,
10317 ) -> Result<(), Box<dyn std::error::Error>> {
10318 let f = self.func("gelu_tanh_f32");
10319 let ni = n as i64;
10320 let cfg = LaunchConfig {
10321 grid_dim: (n.div_ceil(256) as u32, 1, 1),
10322 block_dim: (256, 1, 1),
10323 shared_mem_bytes: 0,
10324 };
10325 let __s_b = self.gpu.stream();
10326 let mut lb = __s_b.launch_builder(&f);
10327 lb.arg(x).arg(&mut *dst).arg(&ni);
10328 unsafe {
10329 lb.launch(cfg)?;
10330 }
10331 Ok(())
10332 }
10333
10334 pub fn row_softmax(
10336 &self,
10337 x: &mut CudaSlice<f32>,
10338 ncols: usize,
10339 nrows: usize,
10340 ) -> Result<(), Box<dyn std::error::Error>> {
10341 let f = self.func("row_softmax_f32");
10342 let nc = ncols as i32;
10343 let cfg = LaunchConfig {
10344 grid_dim: (nrows as u32, 1, 1),
10345 block_dim: (256, 1, 1),
10346 shared_mem_bytes: 0,
10347 };
10348 let __s_b = self.gpu.stream();
10349 let mut lb = __s_b.launch_builder(&f);
10350 lb.arg(&mut *x).arg(&nc);
10351 unsafe {
10352 lb.launch(cfg)?;
10353 }
10354 Ok(())
10355 }
10356
10357 pub fn rms_norm(
10358 &self,
10359 x: &CudaSlice<f32>,
10360 w: &CudaSlice<f32>,
10361 dst: &mut CudaSlice<f32>,
10362 ncols: usize,
10363 nrows: usize,
10364 eps: f32,
10365 ) -> Result<(), Box<dyn std::error::Error>> {
10366 let (nc, e) = (ncols as i32, eps);
10367 let kname = if Self::norm_ilp_on() {
10368 "rms_norm_f32_v2"
10369 } else {
10370 "rms_norm_f32"
10371 };
10372 if Self::pdl_on() && Self::pdl_wb_on() {
10373 use cudarc::driver::{DevicePtr, DevicePtrMut};
10374 let s = &self.gpu.stream();
10375 let (px, _g0) = x.device_ptr(s);
10376 let (pw, _g1) = w.device_ptr(s);
10377 let (pd, _g2) = dst.device_ptr_mut(s);
10378 let mut ps = [
10379 &px as *const _ as *mut std::ffi::c_void,
10380 &pw as *const _ as *mut _,
10381 &pd as *const _ as *mut _,
10382 &nc as *const _ as *mut _,
10383 &e as *const _ as *mut _,
10384 ];
10385 unsafe {
10386 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10387 }
10388 return Ok(());
10389 }
10390 let f = self.func(kname);
10391 let cfg = LaunchConfig {
10392 grid_dim: (nrows as u32, 1, 1),
10393 block_dim: (rms_block(), 1, 1),
10394 shared_mem_bytes: 0,
10395 };
10396 let __s_b = self.gpu.stream();
10397 let mut b = __s_b.launch_builder(&f);
10398 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10399 unsafe {
10400 b.launch(cfg)?;
10401 }
10402 Ok(())
10403 }
10404
10405 pub fn rms_norm_decode(
10413 &self,
10414 x: &CudaSlice<f32>,
10415 w: &CudaSlice<f32>,
10416 dst: &mut CudaSlice<f32>,
10417 ncols: usize,
10418 nrows: usize,
10419 eps: f32,
10420 ) -> Result<(), Box<dyn std::error::Error>> {
10421 let f = self.func(if Self::norm_ilp_on() {
10422 "rms_norm_f32_v2"
10423 } else {
10424 "rms_norm_f32"
10425 });
10426 let cfg = LaunchConfig {
10427 grid_dim: (nrows as u32, 1, 1),
10428 block_dim: (1024, 1, 1),
10429 shared_mem_bytes: 0,
10430 };
10431 let (nc, e) = (ncols as i32, eps);
10432 let __s_b = self.gpu.stream();
10433 let mut b = __s_b.launch_builder(&f);
10434 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
10435 unsafe {
10436 b.launch(cfg)?;
10437 }
10438 Ok(())
10439 }
10440
10441 pub fn rms_norm_q8_1(
10445 &self,
10446 x: &CudaSlice<f32>,
10447 w: &CudaSlice<f32>,
10448 ncols: usize,
10449 nrows: usize,
10450 eps: f32,
10451 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10452 let nblk = ncols / 32;
10453 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10454 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10455 let (nc, e) = (ncols as i32, eps);
10456 if Self::pdl_on() {
10457 {
10458 use cudarc::driver::{DevicePtr, DevicePtrMut};
10459 let s = &self.gpu.stream();
10460 let (px, _g0) = x.device_ptr(s);
10461 let (pw, _g1) = w.device_ptr(s);
10462 let (pq, _g2) = q.device_ptr_mut(s);
10463 let (pd, _g3) = d.device_ptr_mut(s);
10464 let mut ps = [
10465 &px as *const _ as *mut std::ffi::c_void,
10466 &pw as *const _ as *mut _,
10467 &pq as *const _ as *mut _,
10468 &pd as *const _ as *mut _,
10469 &nc as *const _ as *mut _,
10470 &e as *const _ as *mut _,
10471 ];
10472 unsafe {
10473 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10474 }
10475 }
10476 return Ok((q, d));
10477 }
10478 let f = self.func("rms_norm_q8_1");
10479 let cfg = LaunchConfig {
10482 grid_dim: (nrows as u32, 1, 1),
10483 block_dim: (1024, 1, 1),
10484 shared_mem_bytes: 0,
10485 };
10486 let __s_b = self.gpu.stream();
10487 let mut b = __s_b.launch_builder(&f);
10488 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
10489 unsafe {
10490 b.launch(cfg)?;
10491 }
10492 Ok((q, d))
10493 }
10494
10495 pub fn rms_norm_q8_1_into(
10498 &self,
10499 x: &CudaSlice<f32>,
10500 w: &CudaSlice<f32>,
10501 ncols: usize,
10502 nrows: usize,
10503 eps: f32,
10504 q: &mut CudaSlice<i8>,
10505 d: &mut CudaSlice<f32>,
10506 ) -> Result<(), Box<dyn std::error::Error>> {
10507 let nblk = ncols / 32;
10508 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
10509 let (nc, e) = (ncols as i32, eps);
10510 if Self::pdl_on() {
10511 use cudarc::driver::{DevicePtr, DevicePtrMut};
10512 let s = &self.gpu.stream();
10513 let (px, _g0) = x.device_ptr(s);
10514 let (pw, _g1) = w.device_ptr(s);
10515 let (pq, _g2) = q.device_ptr_mut(s);
10516 let (pd, _g3) = d.device_ptr_mut(s);
10517 let mut ps = [
10518 &px as *const _ as *mut std::ffi::c_void,
10519 &pw as *const _ as *mut _,
10520 &pq as *const _ as *mut _,
10521 &pd as *const _ as *mut _,
10522 &nc as *const _ as *mut _,
10523 &e as *const _ as *mut _,
10524 ];
10525 unsafe {
10526 self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1), &mut ps)?;
10527 }
10528 return Ok(());
10529 }
10530 let f = self.func("rms_norm_q8_1");
10531 let cfg = LaunchConfig {
10532 grid_dim: (nrows as u32, 1, 1),
10533 block_dim: (1024, 1, 1),
10534 shared_mem_bytes: 0,
10535 };
10536 let __s_b = self.gpu.stream();
10537 let mut b = __s_b.launch_builder(&f);
10538 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
10539 unsafe {
10540 b.launch(cfg)?;
10541 }
10542 Ok(())
10543 }
10544
10545 pub fn quantize_q8_1_into(
10547 &self,
10548 x: &CudaSlice<f32>,
10549 m: usize,
10550 in_f: usize,
10551 q: &mut CudaSlice<i8>,
10552 d: &mut CudaSlice<f32>,
10553 ) -> Result<(), Box<dyn std::error::Error>> {
10554 let nblk = in_f / 32;
10555 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
10556 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
10557 let (inf, mi) = (in_f as i32, m as i32);
10558 if Self::pdl_on() && Self::pdl_wb_on() {
10559 use cudarc::driver::{DevicePtr, DevicePtrMut};
10560 let s = &self.gpu.stream();
10561 let (px, _g0) = x.device_ptr(s);
10562 let (pq, _g1) = q.device_ptr_mut(s);
10563 let (pd, _g2) = d.device_ptr_mut(s);
10564 let mut ps = [
10565 &px as *const _ as *mut std::ffi::c_void,
10566 &pq as *const _ as *mut _,
10567 &pd as *const _ as *mut _,
10568 &inf as *const _ as *mut _,
10569 &mi as *const _ as *mut _,
10570 ];
10571 unsafe {
10572 self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?;
10573 }
10574 return Ok(());
10575 }
10576 let f = self.func("quantize_q8_1");
10577 let __s_b = self.gpu.stream();
10578 let mut b = __s_b.launch_builder(&f);
10579 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
10580 unsafe {
10581 b.launch(cfg)?;
10582 }
10583 Ok(())
10584 }
10585
10586 pub fn add_rms_norm_q8_1(
10590 &self,
10591 a: &CudaSlice<f32>,
10592 b_in: &CudaSlice<f32>,
10593 w: &CudaSlice<f32>,
10594 res: &mut CudaSlice<f32>,
10595 ncols: usize,
10596 nrows: usize,
10597 eps: f32,
10598 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10599 let nblk = ncols / 32;
10600 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
10601 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
10602 let f = self.func("add_rms_norm_q8_1");
10603 let cfg = LaunchConfig {
10605 grid_dim: (nrows as u32, 1, 1),
10606 block_dim: (1024, 1, 1),
10607 shared_mem_bytes: 0,
10608 };
10609 let (nc, e) = (ncols as i32, eps);
10610 let __s_bld = self.gpu.stream();
10611 let mut bld = __s_bld.launch_builder(&f);
10612 bld.arg(a)
10613 .arg(b_in)
10614 .arg(w)
10615 .arg(res)
10616 .arg(&mut q)
10617 .arg(&mut d)
10618 .arg(&nc)
10619 .arg(&e);
10620 unsafe {
10621 bld.launch(cfg)?;
10622 }
10623 Ok((q, d))
10624 }
10625
10626 #[allow(clippy::too_many_arguments)]
10632 pub fn join_add_rms_norm_raw(
10633 &self,
10634 a0_raw: u64,
10635 a1_raw: u64,
10636 x: &CudaSlice<f32>,
10637 w: &CudaSlice<f32>,
10638 res: &mut CudaSlice<f32>,
10639 dst: &mut CudaSlice<f32>,
10640 ncols: usize,
10641 eps: f32,
10642 ) -> Result<(), Box<dyn std::error::Error>> {
10643 if a0_raw == 0 || a1_raw == 0 || x.len() < ncols || res.len() < ncols || dst.len() < ncols {
10644 return Err("join_add_rms_norm geometry".into());
10645 }
10646 let f = self.func("join_add_rms_norm_f32");
10647 let cfg = LaunchConfig {
10648 grid_dim: (1, 1, 1),
10649 block_dim: (rms_block(), 1, 1),
10650 shared_mem_bytes: 0,
10651 };
10652 let (nc, e) = (ncols as i32, eps);
10653 let __s_b = self.gpu.stream();
10654 let mut b = __s_b.launch_builder(&f);
10655 b.arg(&a0_raw)
10656 .arg(&a1_raw)
10657 .arg(x)
10658 .arg(w)
10659 .arg(&mut *res)
10660 .arg(&mut *dst)
10661 .arg(&nc)
10662 .arg(&e);
10663 unsafe {
10664 b.launch(cfg)?;
10665 }
10666 Ok(())
10667 }
10668
10669 pub fn add_rms_norm(
10670 &self,
10671 a: &CudaSlice<f32>,
10672 b: &CudaSlice<f32>,
10673 w: &CudaSlice<f32>,
10674 res: &mut CudaSlice<f32>,
10675 dst: &mut CudaSlice<f32>,
10676 ncols: usize,
10677 nrows: usize,
10678 eps: f32,
10679 ) -> Result<(), Box<dyn std::error::Error>> {
10680 let (nc, e) = (ncols as i32, eps);
10681 let kname = if Self::norm_ilp_on() {
10682 "add_rms_norm_f32_v2"
10683 } else {
10684 "add_rms_norm_f32"
10685 };
10686 if Self::pdl_on() && Self::pdl_wb_on() {
10687 use cudarc::driver::{DevicePtr, DevicePtrMut};
10688 let s = &self.gpu.stream();
10689 let (pa, _g0) = a.device_ptr(s);
10690 let (pb, _g1) = b.device_ptr(s);
10691 let (pw, _g2) = w.device_ptr(s);
10692 let (pr, _g3) = res.device_ptr_mut(s);
10693 let (pd, _g4) = dst.device_ptr_mut(s);
10694 let mut ps = [
10695 &pa as *const _ as *mut std::ffi::c_void,
10696 &pb as *const _ as *mut _,
10697 &pw as *const _ as *mut _,
10698 &pr as *const _ as *mut _,
10699 &pd as *const _ as *mut _,
10700 &nc as *const _ as *mut _,
10701 &e as *const _ as *mut _,
10702 ];
10703 unsafe {
10704 self.launch_pdl(kname, (nrows as u32, 1, 1), (rms_block(), 1, 1), &mut ps)?;
10705 }
10706 return Ok(());
10707 }
10708 let f = self.func(kname);
10709 let cfg = LaunchConfig {
10710 grid_dim: (nrows as u32, 1, 1),
10711 block_dim: (rms_block(), 1, 1),
10712 shared_mem_bytes: 0,
10713 };
10714 let __s_b2 = self.gpu.stream();
10715 let mut b2 = __s_b2.launch_builder(&f);
10716 b2.arg(a)
10717 .arg(b)
10718 .arg(w)
10719 .arg(&mut *res)
10720 .arg(&mut *dst)
10721 .arg(&nc)
10722 .arg(&e);
10723 unsafe {
10724 b2.launch(cfg)?;
10725 }
10726 Ok(())
10727 }
10728
10729 #[allow(clippy::too_many_arguments)]
10732 pub fn rms_pre_add_rms_norm(
10733 &self,
10734 a: &CudaSlice<f32>,
10735 wa: &CudaSlice<f32>,
10736 b: &CudaSlice<f32>,
10737 w: &CudaSlice<f32>,
10738 res: &mut CudaSlice<f32>,
10739 dst: &mut CudaSlice<f32>,
10740 ncols: usize,
10741 nrows: usize,
10742 eps: f32,
10743 ) -> Result<(), Box<dyn std::error::Error>> {
10744 let f = self.func("rms_pre_add_rms_norm_f32");
10745 let cfg = LaunchConfig {
10746 grid_dim: (nrows as u32, 1, 1),
10747 block_dim: (rms_block(), 1, 1),
10748 shared_mem_bytes: 0,
10749 };
10750 let (nc, e) = (ncols as i32, eps);
10751 let __s_b2 = self.gpu.stream();
10752 let mut b2 = __s_b2.launch_builder(&f);
10753 b2.arg(a)
10754 .arg(wa)
10755 .arg(b)
10756 .arg(w)
10757 .arg(&mut *res)
10758 .arg(&mut *dst)
10759 .arg(&nc)
10760 .arg(&e);
10761 unsafe {
10762 b2.launch(cfg)?;
10763 }
10764 Ok(())
10765 }
10766
10767 #[allow(clippy::too_many_arguments)]
10769 pub fn rms_pre_add_rms_norm_q8z(
10770 &self,
10771 a: &CudaSlice<f32>,
10772 wa: &CudaSlice<f32>,
10773 b: &CudaSlice<f32>,
10774 w: &CudaSlice<f32>,
10775 res: &mut CudaSlice<f32>,
10776 dst: &mut CudaSlice<f32>,
10777 ncols: usize,
10778 nrows: usize,
10779 eps: f32,
10780 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10781 debug_assert!(ncols % 128 == 0);
10782 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
10783 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
10784 let (nc, e) = (ncols as i32, eps);
10785 if Self::pdl_on() {
10786 {
10787 use cudarc::driver::{DevicePtr, DevicePtrMut};
10788 let s = &self.gpu.stream();
10789 let (pa, _g0) = a.device_ptr(s);
10790 let (pwa, _g1) = wa.device_ptr(s);
10791 let (pb, _g2) = b.device_ptr(s);
10792 let (pw, _g3) = w.device_ptr(s);
10793 let (pr, _g4) = res.device_ptr_mut(s);
10794 let (pdst, _g5) = dst.device_ptr_mut(s);
10795 let (pq, _g6) = out_q.device_ptr_mut(s);
10796 let (pd, _g7) = out_d.device_ptr_mut(s);
10797 let mut ps = [
10798 &pa as *const _ as *mut std::ffi::c_void,
10799 &pwa as *const _ as *mut _,
10800 &pb as *const _ as *mut _,
10801 &pw as *const _ as *mut _,
10802 &pr as *const _ as *mut _,
10803 &pdst as *const _ as *mut _,
10804 &pq as *const _ as *mut _,
10805 &pd as *const _ as *mut _,
10806 &nc as *const _ as *mut _,
10807 &e as *const _ as *mut _,
10808 ];
10809 unsafe {
10810 self.launch_pdl(
10811 "rms_pre_add_rms_norm_q8z_f32",
10812 (nrows as u32, 1, 1),
10813 (rms_block(), 1, 1),
10814 &mut ps,
10815 )?;
10816 }
10817 }
10818 return Ok((out_q, out_d));
10819 }
10820 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10821 let cfg = LaunchConfig {
10822 grid_dim: (nrows as u32, 1, 1),
10823 block_dim: (rms_block(), 1, 1),
10824 shared_mem_bytes: 0,
10825 };
10826 let __s_b2 = self.gpu.stream();
10827 let mut b2 = __s_b2.launch_builder(&f);
10828 b2.arg(a)
10829 .arg(wa)
10830 .arg(b)
10831 .arg(w)
10832 .arg(&mut *res)
10833 .arg(&mut *dst)
10834 .arg(&mut out_q)
10835 .arg(&mut out_d)
10836 .arg(&nc)
10837 .arg(&e);
10838 unsafe {
10839 b2.launch(cfg)?;
10840 }
10841 Ok((out_q, out_d))
10842 }
10843
10844 #[allow(clippy::too_many_arguments)]
10848 pub fn rms_pre_add_rms_norm_q8z_into(
10849 &self,
10850 a: &CudaSlice<f32>,
10851 wa: &CudaSlice<f32>,
10852 b: &CudaSlice<f32>,
10853 w: &CudaSlice<f32>,
10854 res: &mut CudaSlice<f32>,
10855 dst: &mut CudaSlice<f32>,
10856 ncols: usize,
10857 nrows: usize,
10858 eps: f32,
10859 out_q: &mut CudaSlice<i8>,
10860 out_d: &mut CudaSlice<f32>,
10861 ) -> Result<(), Box<dyn std::error::Error>> {
10862 debug_assert!(ncols % 128 == 0);
10863 let (nc, e) = (ncols as i32, eps);
10864 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
10865 let cfg = LaunchConfig {
10866 grid_dim: (nrows as u32, 1, 1),
10867 block_dim: (rms_block(), 1, 1),
10868 shared_mem_bytes: 0,
10869 };
10870 let __s_b = self.gpu.stream();
10871 let mut b2 = __s_b.launch_builder(&f);
10872 b2.arg(a)
10873 .arg(wa)
10874 .arg(b)
10875 .arg(w)
10876 .arg(&mut *res)
10877 .arg(&mut *dst)
10878 .arg(&mut *out_q)
10879 .arg(&mut *out_d)
10880 .arg(&nc)
10881 .arg(&e);
10882 unsafe {
10883 b2.launch(cfg)?;
10884 }
10885 Ok(())
10886 }
10887
10888 #[allow(clippy::too_many_arguments)]
10891 pub fn rms_pre_add_scale_rms_norm_q8_1_into(
10892 &self,
10893 a: &CudaSlice<f32>,
10894 wa: &CudaSlice<f32>,
10895 b_in: &CudaSlice<f32>,
10896 c: f32,
10897 w: &CudaSlice<f32>,
10898 res: &mut CudaSlice<f32>,
10899 ncols: usize,
10900 nrows: usize,
10901 eps: f32,
10902 out_q: &mut CudaSlice<i8>,
10903 out_d: &mut CudaSlice<f32>,
10904 ) -> Result<(), Box<dyn std::error::Error>> {
10905 debug_assert!(ncols % 128 == 0);
10906 let (nc, e2) = (ncols as i32, eps);
10907 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
10908 let cfg = LaunchConfig {
10909 grid_dim: (nrows as u32, 1, 1),
10910 block_dim: (rms_block(), 1, 1),
10911 shared_mem_bytes: 0,
10912 };
10913 let __s_b = self.gpu.stream();
10914 let mut b2 = __s_b.launch_builder(&f);
10915 b2.arg(a)
10916 .arg(wa)
10917 .arg(b_in)
10918 .arg(&c)
10919 .arg(w)
10920 .arg(&mut *res)
10921 .arg(&mut *out_q)
10922 .arg(&mut *out_d)
10923 .arg(&nc)
10924 .arg(&e2);
10925 unsafe {
10926 b2.launch(cfg)?;
10927 }
10928 Ok(())
10929 }
10930
10931 pub fn g4_pnfold_on() -> bool {
10939 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10940 *ON.get_or_init(|| {
10941 std::env::var("MEMRA_G4_PNFOLD")
10942 .map(|v| v != "0")
10943 .unwrap_or(true)
10944 })
10945 }
10946
10947 pub fn build_q4_out_concat3(
10951 &self,
10952 w0: &crate::model::GpuTensor,
10953 w1: &crate::model::GpuTensor,
10954 w2: &crate::model::GpuTensor,
10955 ) -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
10956 use crate::model::GpuTensor;
10957 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
10958 match w {
10959 GpuTensor::Quant {
10960 qtype,
10961 row_bytes,
10962 rp,
10963 ..
10964 } if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
10965 _ => None,
10966 }
10967 };
10968 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
10969 else {
10970 return Ok(None);
10971 };
10972 if rb0 != rb1
10973 || rb0 != rb2
10974 || w0.in_features() != w1.in_features()
10975 || w0.in_features() != w2.in_features()
10976 {
10977 return Ok(None);
10978 }
10979 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
10980 match w {
10981 crate::model::GpuTensor::Quant { bytes, .. } => bytes,
10982 _ => unreachable!(),
10983 }
10984 }
10985 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
10986 let total = rb0 * (o0 + o1 + o2);
10987 let mut cat = self.alloc_u8(total)?;
10988 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
10989 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
10990 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
10991 Ok(Some(GpuTensor::Quant {
10992 bytes: cat,
10993 qtype: QT_Q4_0,
10994 row_bytes: rb0,
10995 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64],
10996 scale: 1.0,
10997 rp: false,
10998 #[cfg(memra_cutlass)]
10999 cutlass: None,
11000 fp8: None,
11001 blk: None,
11002 rp4: None,
11003 f16: None,
11004 }))
11005 }
11006
11007 fn full_width_rope_only(
11025 kernel: &str,
11026 n_rot: usize,
11027 head_dim: usize,
11028 ) -> Result<(), Box<dyn std::error::Error>> {
11029 if n_rot == head_dim {
11030 return Ok(());
11031 }
11032 Err(format!(
11033 "{kernel}: PARTIAL ROTARY REFUSED — n_rot {n_rot} != head_dim {head_dim}. This fused \
11034 rms_norm+qkv+rope kernel carries no n_dims parameter and rotates the full head \
11035 width (half = ncols/2), so it would rotate dims {n_rot}..{head_dim} that must pass \
11036 through unrotated. Use the split path (rms_norm_qkv + rope_neox/rope_neox2 with \
11037 n_dims={n_rot}), or add an n_dims early-return to the kernel and widen this guard."
11038 )
11039 .into())
11040 }
11041
11042 #[allow(clippy::too_many_arguments)]
11046 pub fn rms_norm_qkv_rope_cat(
11047 &self,
11048 qkv: &CudaSlice<f32>,
11049 wq: &CudaSlice<f32>,
11050 wk: &CudaSlice<f32>,
11051 wv: &CudaSlice<f32>,
11052 q: &mut CudaSlice<f32>,
11053 k: &mut CudaSlice<f32>,
11054 v: &mut CudaSlice<f32>,
11055 head_dim: usize,
11056 n_rot: usize,
11057 rq: usize,
11058 rk: usize,
11059 pos: &CudaSlice<i32>,
11060 nh_q: usize,
11061 nh_k: usize,
11062 base: f32,
11063 freq_scale: f32,
11064 ff: Option<&CudaSlice<f32>>,
11065 eps: f32,
11066 ) -> Result<(), Box<dyn std::error::Error>> {
11067 Self::full_width_rope_only("rms_norm_qkv_rope_cat", n_rot, head_dim)?;
11068 let rows = rq + rk + rk;
11069 let theta_scale = base.powf(-2.0 / head_dim as f32);
11070 let (nc, rqi, rki, nhq, nhk) = (
11071 head_dim as i32,
11072 rq as i32,
11073 rk as i32,
11074 nh_q as i32,
11075 nh_k as i32,
11076 );
11077 if Self::pdl_on() {
11078 use cudarc::driver::{DevicePtr, DevicePtrMut};
11079 let s = &self.gpu.stream();
11080 let (pqkv, _g0) = qkv.device_ptr(s);
11081 let (pwq, _g1) = wq.device_ptr(s);
11082 let (pwk, _g2) = wk.device_ptr(s);
11083 let (pwv, _g3) = wv.device_ptr(s);
11084 let (pq, _g4) = q.device_ptr_mut(s);
11085 let (pk, _g5) = k.device_ptr_mut(s);
11086 let (pv, _g6) = v.device_ptr_mut(s);
11087 let (ppos, _g7) = pos.device_ptr(s);
11088 let (pff, _g8) = match ff {
11089 Some(t) => {
11090 let (p, g) = t.device_ptr(s);
11091 (p, Some(g))
11092 }
11093 None => (0, None),
11094 };
11095 let mut ps = [
11096 &pqkv as *const _ as *mut std::ffi::c_void,
11097 &pwq as *const _ as *mut _,
11098 &pwk as *const _ as *mut _,
11099 &pwv as *const _ as *mut _,
11100 &pq as *const _ as *mut _,
11101 &pk as *const _ as *mut _,
11102 &pv as *const _ as *mut _,
11103 &nc as *const _ as *mut _,
11104 &rqi as *const _ as *mut _,
11105 &rki as *const _ as *mut _,
11106 &ppos as *const _ as *mut _,
11107 &nhq as *const _ as *mut _,
11108 &nhk as *const _ as *mut _,
11109 &theta_scale as *const _ as *mut _,
11110 &freq_scale as *const _ as *mut _,
11111 &pff as *const _ as *mut _,
11112 &eps as *const _ as *mut _,
11113 ];
11114 unsafe {
11115 self.launch_pdl(
11116 "rms_norm_qkv_rope_cat_f32",
11117 (rows as u32, 1, 1),
11118 (rms_block(), 1, 1),
11119 &mut ps,
11120 )?;
11121 }
11122 return Ok(());
11123 }
11124 let f = self.func("rms_norm_qkv_rope_cat_f32");
11125 let cfg = LaunchConfig {
11126 grid_dim: (rows as u32, 1, 1),
11127 block_dim: (rms_block(), 1, 1),
11128 shared_mem_bytes: 0,
11129 };
11130 let __s_b = self.gpu.stream();
11131 let mut b = __s_b.launch_builder(&f);
11132 match ff {
11133 Some(t) => {
11134 b.arg(qkv)
11135 .arg(wq)
11136 .arg(wk)
11137 .arg(wv)
11138 .arg(&mut *q)
11139 .arg(&mut *k)
11140 .arg(&mut *v)
11141 .arg(&nc)
11142 .arg(&rqi)
11143 .arg(&rki)
11144 .arg(pos)
11145 .arg(&nhq)
11146 .arg(&nhk)
11147 .arg(&theta_scale)
11148 .arg(&freq_scale)
11149 .arg(t)
11150 .arg(&eps);
11151 unsafe {
11152 b.launch(cfg)?;
11153 }
11154 }
11155 None => {
11156 let null: u64 = 0;
11157 b.arg(qkv)
11158 .arg(wq)
11159 .arg(wk)
11160 .arg(wv)
11161 .arg(&mut *q)
11162 .arg(&mut *k)
11163 .arg(&mut *v)
11164 .arg(&nc)
11165 .arg(&rqi)
11166 .arg(&rki)
11167 .arg(pos)
11168 .arg(&nhq)
11169 .arg(&nhk)
11170 .arg(&theta_scale)
11171 .arg(&freq_scale)
11172 .arg(&null)
11173 .arg(&eps);
11174 unsafe {
11175 b.launch(cfg)?;
11176 }
11177 }
11178 }
11179 Ok(())
11180 }
11181
11182 #[allow(clippy::too_many_arguments)]
11186 pub fn rms_norm_qkv_rope(
11187 &self,
11188 q0: &CudaSlice<f32>,
11189 k0: &CudaSlice<f32>,
11190 v0: &CudaSlice<f32>,
11191 wq: &CudaSlice<f32>,
11192 wk: &CudaSlice<f32>,
11193 wv: &CudaSlice<f32>,
11194 q: &mut CudaSlice<f32>,
11195 k: &mut CudaSlice<f32>,
11196 v: &mut CudaSlice<f32>,
11197 head_dim: usize,
11198 n_rot: usize,
11199 rq: usize,
11200 rk: usize,
11201 pos: &CudaSlice<i32>,
11202 nh_q: usize,
11203 nh_k: usize,
11204 base: f32,
11205 freq_scale: f32,
11206 ff: Option<&CudaSlice<f32>>,
11207 eps: f32,
11208 ) -> Result<(), Box<dyn std::error::Error>> {
11209 Self::full_width_rope_only("rms_norm_qkv_rope", n_rot, head_dim)?;
11210 let f = self.func("rms_norm_qkv_rope_f32");
11211 let rows = rq + rk + rk; let cfg = LaunchConfig {
11213 grid_dim: (rows as u32, 1, 1),
11214 block_dim: (rms_block(), 1, 1),
11215 shared_mem_bytes: 0,
11216 };
11217 let theta_scale = base.powf(-2.0 / head_dim as f32);
11218 let (nc, rqi, rki, nhq, nhk) = (
11219 head_dim as i32,
11220 rq as i32,
11221 rk as i32,
11222 nh_q as i32,
11223 nh_k as i32,
11224 );
11225 let __s_b = self.gpu.stream();
11226 let mut b = __s_b.launch_builder(&f);
11227 match ff {
11228 Some(t) => {
11229 b.arg(q0)
11230 .arg(k0)
11231 .arg(v0)
11232 .arg(wq)
11233 .arg(wk)
11234 .arg(wv)
11235 .arg(&mut *q)
11236 .arg(&mut *k)
11237 .arg(&mut *v)
11238 .arg(&nc)
11239 .arg(&rqi)
11240 .arg(&rki)
11241 .arg(pos)
11242 .arg(&nhq)
11243 .arg(&nhk)
11244 .arg(&theta_scale)
11245 .arg(&freq_scale)
11246 .arg(t)
11247 .arg(&eps);
11248 unsafe {
11249 b.launch(cfg)?;
11250 }
11251 }
11252 None => {
11253 let null: u64 = 0;
11254 b.arg(q0)
11255 .arg(k0)
11256 .arg(v0)
11257 .arg(wq)
11258 .arg(wk)
11259 .arg(wv)
11260 .arg(&mut *q)
11261 .arg(&mut *k)
11262 .arg(&mut *v)
11263 .arg(&nc)
11264 .arg(&rqi)
11265 .arg(&rki)
11266 .arg(pos)
11267 .arg(&nhq)
11268 .arg(&nhk)
11269 .arg(&theta_scale)
11270 .arg(&freq_scale)
11271 .arg(&null)
11272 .arg(&eps);
11273 unsafe {
11274 b.launch(cfg)?;
11275 }
11276 }
11277 }
11278 Ok(())
11279 }
11280
11281 #[allow(clippy::too_many_arguments)]
11287 pub fn rms_norm_qkv_rope_append_dc(
11288 &self,
11289 q0: &CudaSlice<f32>,
11290 k0: &CudaSlice<f32>,
11291 v0: &CudaSlice<f32>,
11292 wq: &CudaSlice<f32>,
11293 wk: &CudaSlice<f32>,
11294 wv: &CudaSlice<f32>,
11295 q: &mut CudaSlice<f32>,
11296 k: &mut CudaSlice<f32>,
11297 v: &mut CudaSlice<f32>,
11298 head_dim: usize,
11299 n_rot: usize,
11300 rq: usize,
11301 rk: usize,
11302 pos: &CudaSlice<i32>,
11303 nh_q: usize,
11304 nh_k: usize,
11305 base: f32,
11306 freq_scale: f32,
11307 ff: Option<&CudaSlice<f32>>,
11308 eps: f32,
11309 kc: &mut CudaSlice<u8>,
11310 vc: &mut CudaSlice<u8>,
11311 t_dev: &CudaSlice<i32>,
11312 k_tok_bytes: usize,
11313 v_tok_bytes: usize,
11314 g: bool,
11315 ) -> Result<(), Box<dyn std::error::Error>> {
11316 Self::full_width_rope_only("rms_norm_qkv_rope_append_dc", n_rot, head_dim)?;
11317 let rows = rq + rk + rk;
11318 let theta_scale = base.powf(-2.0 / head_dim as f32);
11319 let (nc, rqi, rki, nhq, nhk) = (
11320 head_dim as i32,
11321 rq as i32,
11322 rk as i32,
11323 nh_q as i32,
11324 nh_k as i32,
11325 );
11326 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11327 if Self::pdl_on() && Self::pdl_wb_on() {
11328 use cudarc::driver::{DevicePtr, DevicePtrMut};
11329 let s = &self.gpu.stream();
11330 let (p0, _a0) = q0.device_ptr(s);
11331 let (p1, _a1) = k0.device_ptr(s);
11332 let (p2, _a2) = v0.device_ptr(s);
11333 let (pwq, _a3) = wq.device_ptr(s);
11334 let (pwk, _a4) = wk.device_ptr(s);
11335 let (pwv, _a5) = wv.device_ptr(s);
11336 let (pq, _a6) = q.device_ptr_mut(s);
11337 let (pk, _a7) = k.device_ptr_mut(s);
11338 let (pv, _a8) = v.device_ptr_mut(s);
11339 let (pp, _a9) = pos.device_ptr(s);
11340 let pff: u64 = match ff {
11341 Some(t) => {
11342 let (p, _gg) = t.device_ptr(s);
11343 p as u64
11344 }
11345 None => 0,
11346 };
11347 let (pkc, _a10) = kc.device_ptr_mut(s);
11348 let (pvc, _a11) = vc.device_ptr_mut(s);
11349 let (pt, _a12) = t_dev.device_ptr(s);
11350 let mut ps = [
11351 &p0 as *const _ as *mut std::ffi::c_void,
11352 &p1 as *const _ as *mut _,
11353 &p2 as *const _ as *mut _,
11354 &pwq as *const _ as *mut _,
11355 &pwk as *const _ as *mut _,
11356 &pwv as *const _ as *mut _,
11357 &pq as *const _ as *mut _,
11358 &pk as *const _ as *mut _,
11359 &pv as *const _ as *mut _,
11360 &nc as *const _ as *mut _,
11361 &rqi as *const _ as *mut _,
11362 &rki as *const _ as *mut _,
11363 &pp as *const _ as *mut _,
11364 &nhq as *const _ as *mut _,
11365 &nhk as *const _ as *mut _,
11366 &theta_scale as *const _ as *mut _,
11367 &freq_scale as *const _ as *mut _,
11368 &pff as *const _ as *mut _,
11369 &eps as *const _ as *mut _,
11370 &pkc as *const _ as *mut _,
11371 &pvc as *const _ as *mut _,
11372 &pt as *const _ as *mut _,
11373 &ktb as *const _ as *mut _,
11374 &vtb as *const _ as *mut _,
11375 ];
11376 unsafe {
11377 self.launch_pdl_flash(
11378 g,
11379 "rms_norm_qkv_rope_append_dc_f32",
11380 (rows as u32, 1, 1),
11381 (rms_block(), 1, 1),
11382 0,
11383 &mut ps,
11384 )?;
11385 }
11386 return Ok(());
11387 }
11388 let f = if g {
11389 self.func_g("rms_norm_qkv_rope_append_dc_f32")
11390 } else {
11391 self.func("rms_norm_qkv_rope_append_dc_f32")
11392 };
11393 let cfg = LaunchConfig {
11394 grid_dim: (rows as u32, 1, 1),
11395 block_dim: (rms_block(), 1, 1),
11396 shared_mem_bytes: 0,
11397 };
11398 let __s_b = self.gpu.stream();
11399 let mut b = __s_b.launch_builder(&f);
11400 match ff {
11401 Some(t) => {
11402 b.arg(q0)
11403 .arg(k0)
11404 .arg(v0)
11405 .arg(wq)
11406 .arg(wk)
11407 .arg(wv)
11408 .arg(&mut *q)
11409 .arg(&mut *k)
11410 .arg(&mut *v)
11411 .arg(&nc)
11412 .arg(&rqi)
11413 .arg(&rki)
11414 .arg(pos)
11415 .arg(&nhq)
11416 .arg(&nhk)
11417 .arg(&theta_scale)
11418 .arg(&freq_scale)
11419 .arg(t)
11420 .arg(&eps)
11421 .arg(&mut *kc)
11422 .arg(&mut *vc)
11423 .arg(t_dev)
11424 .arg(&ktb)
11425 .arg(&vtb);
11426 unsafe {
11427 b.launch(cfg)?;
11428 }
11429 }
11430 None => {
11431 let null: u64 = 0;
11432 b.arg(q0)
11433 .arg(k0)
11434 .arg(v0)
11435 .arg(wq)
11436 .arg(wk)
11437 .arg(wv)
11438 .arg(&mut *q)
11439 .arg(&mut *k)
11440 .arg(&mut *v)
11441 .arg(&nc)
11442 .arg(&rqi)
11443 .arg(&rki)
11444 .arg(pos)
11445 .arg(&nhq)
11446 .arg(&nhk)
11447 .arg(&theta_scale)
11448 .arg(&freq_scale)
11449 .arg(&null)
11450 .arg(&eps)
11451 .arg(&mut *kc)
11452 .arg(&mut *vc)
11453 .arg(t_dev)
11454 .arg(&ktb)
11455 .arg(&vtb);
11456 unsafe {
11457 b.launch(cfg)?;
11458 }
11459 }
11460 }
11461 Ok(())
11462 }
11463
11464 #[allow(clippy::too_many_arguments)]
11472 pub fn rms_norm_qkv_rope_append(
11473 &self,
11474 q0: &CudaSlice<f32>,
11475 k0: &CudaSlice<f32>,
11476 v0: &CudaSlice<f32>,
11477 wq: &CudaSlice<f32>,
11478 wk: &CudaSlice<f32>,
11479 wv: &CudaSlice<f32>,
11480 q: &mut CudaSlice<f32>,
11481 k: &mut CudaSlice<f32>,
11482 v: &mut CudaSlice<f32>,
11483 head_dim: usize,
11484 n_rot: usize,
11485 rq: usize,
11486 rk: usize,
11487 pos: &CudaSlice<i32>,
11488 nh_q: usize,
11489 nh_k: usize,
11490 base: f32,
11491 freq_scale: f32,
11492 ff: Option<&CudaSlice<f32>>,
11493 eps: f32,
11494 kc: &mut CudaSlice<u8>,
11495 vc: &mut CudaSlice<u8>,
11496 t: usize,
11497 k_tok_bytes: usize,
11498 v_tok_bytes: usize,
11499 g: bool,
11500 ) -> Result<(), Box<dyn std::error::Error>> {
11501 Self::full_width_rope_only("rms_norm_qkv_rope_append", n_rot, head_dim)?;
11502 let rows = rq + rk + rk;
11503 let theta_scale = base.powf(-2.0 / head_dim as f32);
11504 let (nc, rqi, rki, nhq, nhk) = (
11505 head_dim as i32,
11506 rq as i32,
11507 rk as i32,
11508 nh_q as i32,
11509 nh_k as i32,
11510 );
11511 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
11512 let ti = t as i32;
11513 if Self::pdl_on() && Self::pdl_wb_on() {
11514 use cudarc::driver::{DevicePtr, DevicePtrMut};
11515 let s = &self.gpu.stream();
11516 let (p0, _a0) = q0.device_ptr(s);
11517 let (p1, _a1) = k0.device_ptr(s);
11518 let (p2, _a2) = v0.device_ptr(s);
11519 let (pwq, _a3) = wq.device_ptr(s);
11520 let (pwk, _a4) = wk.device_ptr(s);
11521 let (pwv, _a5) = wv.device_ptr(s);
11522 let (pq, _a6) = q.device_ptr_mut(s);
11523 let (pk, _a7) = k.device_ptr_mut(s);
11524 let (pv, _a8) = v.device_ptr_mut(s);
11525 let (pp, _a9) = pos.device_ptr(s);
11526 let pff: u64 = match ff {
11527 Some(t) => {
11528 let (p, _gg) = t.device_ptr(s);
11529 p as u64
11530 }
11531 None => 0,
11532 };
11533 let (pkc, _a10) = kc.device_ptr_mut(s);
11534 let (pvc, _a11) = vc.device_ptr_mut(s);
11535 let mut ps = [
11536 &p0 as *const _ as *mut std::ffi::c_void,
11537 &p1 as *const _ as *mut _,
11538 &p2 as *const _ as *mut _,
11539 &pwq as *const _ as *mut _,
11540 &pwk as *const _ as *mut _,
11541 &pwv as *const _ as *mut _,
11542 &pq as *const _ as *mut _,
11543 &pk as *const _ as *mut _,
11544 &pv as *const _ as *mut _,
11545 &nc as *const _ as *mut _,
11546 &rqi as *const _ as *mut _,
11547 &rki as *const _ as *mut _,
11548 &pp as *const _ as *mut _,
11549 &nhq as *const _ as *mut _,
11550 &nhk as *const _ as *mut _,
11551 &theta_scale as *const _ as *mut _,
11552 &freq_scale as *const _ as *mut _,
11553 &pff as *const _ as *mut _,
11554 &eps as *const _ as *mut _,
11555 &pkc as *const _ as *mut _,
11556 &pvc as *const _ as *mut _,
11557 &ti as *const _ as *mut _,
11558 &ktb as *const _ as *mut _,
11559 &vtb as *const _ as *mut _,
11560 ];
11561 unsafe {
11562 self.launch_pdl_flash(
11563 g,
11564 "rms_norm_qkv_rope_append_f32",
11565 (rows as u32, 1, 1),
11566 (rms_block(), 1, 1),
11567 0,
11568 &mut ps,
11569 )?;
11570 }
11571 return Ok(());
11572 }
11573 let f = if g {
11574 self.func_g("rms_norm_qkv_rope_append_f32")
11575 } else {
11576 self.func("rms_norm_qkv_rope_append_f32")
11577 };
11578 let cfg = LaunchConfig {
11579 grid_dim: (rows as u32, 1, 1),
11580 block_dim: (rms_block(), 1, 1),
11581 shared_mem_bytes: 0,
11582 };
11583 let __s_b = self.gpu.stream();
11584 let mut b = __s_b.launch_builder(&f);
11585 let null: u64 = 0;
11586 b.arg(q0)
11587 .arg(k0)
11588 .arg(v0)
11589 .arg(wq)
11590 .arg(wk)
11591 .arg(wv)
11592 .arg(&mut *q)
11593 .arg(&mut *k)
11594 .arg(&mut *v)
11595 .arg(&nc)
11596 .arg(&rqi)
11597 .arg(&rki)
11598 .arg(pos)
11599 .arg(&nhq)
11600 .arg(&nhk)
11601 .arg(&theta_scale)
11602 .arg(&freq_scale);
11603 match ff {
11604 Some(t) => {
11605 b.arg(t);
11606 }
11607 None => {
11608 b.arg(&null);
11609 }
11610 }
11611 b.arg(&eps)
11612 .arg(&mut *kc)
11613 .arg(&mut *vc)
11614 .arg(&ti)
11615 .arg(&ktb)
11616 .arg(&vtb);
11617 unsafe {
11618 b.launch(cfg)?;
11619 }
11620 Ok(())
11621 }
11622
11623 pub fn add_q8_1(
11624 &self,
11625 a: &CudaSlice<f32>,
11626 b: &CudaSlice<f32>,
11627 res: &mut CudaSlice<f32>,
11628 ncols: usize,
11629 nrows: usize,
11630 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11631 debug_assert!(ncols % 128 == 0);
11632 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11633 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11634 let f = self.func("add_q8_1_f32");
11635 let cfg = LaunchConfig {
11636 grid_dim: (nrows as u32, 1, 1),
11637 block_dim: (rms_block(), 1, 1),
11638 shared_mem_bytes: 0,
11639 };
11640 let nc = ncols as i32;
11641 let __s_b2 = self.gpu.stream();
11642 let mut b2 = __s_b2.launch_builder(&f);
11643 b2.arg(a)
11644 .arg(b)
11645 .arg(&mut *res)
11646 .arg(&mut out_q)
11647 .arg(&mut out_d)
11648 .arg(&nc);
11649 unsafe {
11650 b2.launch(cfg)?;
11651 }
11652 Ok((out_q, out_d))
11653 }
11654
11655 pub fn rms_pre_add_q8_1(
11659 &self,
11660 a: &CudaSlice<f32>,
11661 wa: &CudaSlice<f32>,
11662 b: &CudaSlice<f32>,
11663 res: &mut CudaSlice<f32>,
11664 ncols: usize,
11665 nrows: usize,
11666 eps: f32,
11667 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11668 debug_assert!(ncols % 128 == 0);
11669 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11670 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11671 let f = self.func("rms_pre_add_q8_1_f32");
11672 let cfg = LaunchConfig {
11673 grid_dim: (nrows as u32, 1, 1),
11674 block_dim: (rms_block(), 1, 1),
11675 shared_mem_bytes: 0,
11676 };
11677 let (nc, ep) = (ncols as i32, eps);
11678 let __s_b2 = self.gpu.stream();
11679 let mut b2 = __s_b2.launch_builder(&f);
11680 b2.arg(a)
11681 .arg(wa)
11682 .arg(b)
11683 .arg(&mut *res)
11684 .arg(&mut out_q)
11685 .arg(&mut out_d)
11686 .arg(&nc)
11687 .arg(&ep);
11688 unsafe {
11689 b2.launch(cfg)?;
11690 }
11691 Ok((out_q, out_d))
11692 }
11693
11694 pub fn l2_v2_on(ncols: usize) -> bool {
11698 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
11699 }
11700
11701 pub fn l2_norm_pp(
11702 &self,
11703 x: &CudaSlice<f32>,
11704 dst: &mut CudaSlice<f32>,
11705 dst16: Option<&mut CudaSlice<u8>>,
11706 ncols: usize,
11707 nrows: usize,
11708 eps: f32,
11709 ) -> Result<(), Box<dyn std::error::Error>> {
11710 if Self::l2_v2_on(ncols) {
11711 let f = self.func("l2_norm_pp_v2_f32");
11712 let rows_per_block = 8u32; let cfg = LaunchConfig {
11714 grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1),
11715 block_dim: (256, 1, 1),
11716 shared_mem_bytes: 0,
11717 };
11718 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
11719 let d16: u64 = match dst16 {
11721 Some(d) => self.addr_u8(d),
11722 None => 0,
11723 };
11724 let __s_b = self.gpu.stream();
11725 let mut b = __s_b.launch_builder(&f);
11726 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
11727 unsafe {
11728 b.launch(cfg)?;
11729 }
11730 return Ok(());
11731 }
11732 self.l2_norm(x, dst, ncols, nrows, eps)
11733 }
11734
11735 pub fn l2_norm(
11736 &self,
11737 x: &CudaSlice<f32>,
11738 dst: &mut CudaSlice<f32>,
11739 ncols: usize,
11740 nrows: usize,
11741 eps: f32,
11742 ) -> Result<(), Box<dyn std::error::Error>> {
11743 let f = self.func("l2_norm_f32");
11744 let cfg = LaunchConfig {
11745 grid_dim: (nrows as u32, 1, 1),
11746 block_dim: (256, 1, 1),
11747 shared_mem_bytes: 0,
11748 };
11749 let (nc, e) = (ncols as i32, eps);
11750 let __s_b = self.gpu.stream();
11751 let mut b = __s_b.launch_builder(&f);
11752 b.arg(x).arg(dst).arg(&nc).arg(&e);
11753 unsafe {
11754 b.launch(cfg)?;
11755 }
11756 Ok(())
11757 }
11758
11759 pub fn l2_norm_decode(
11765 &self,
11766 x: &CudaSlice<f32>,
11767 dst: &mut CudaSlice<f32>,
11768 ncols: usize,
11769 nrows: usize,
11770 eps: f32,
11771 ) -> Result<(), Box<dyn std::error::Error>> {
11772 let f = self.func("l2_norm_f32");
11773 let cfg = LaunchConfig {
11774 grid_dim: (nrows as u32, 1, 1),
11775 block_dim: (32, 1, 1),
11776 shared_mem_bytes: 0,
11777 };
11778 let (nc, e) = (ncols as i32, eps);
11779 let __s_b = self.gpu.stream();
11780 let mut b = __s_b.launch_builder(&f);
11781 b.arg(x).arg(dst).arg(&nc).arg(&e);
11782 unsafe {
11783 b.launch(cfg)?;
11784 }
11785 Ok(())
11786 }
11787
11788 pub fn rope_neox(
11790 &self,
11791 x: &mut CudaSlice<f32>,
11792 pos: &CudaSlice<i32>,
11793 head_dim: usize,
11794 n_dims: usize,
11795 n_heads: usize,
11796 n_tokens: usize,
11797 freq_base: f32,
11798 freq_scale: f32,
11799 ) -> Result<(), Box<dyn std::error::Error>> {
11800 let f = self.func("rope_neox_f32");
11801 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11802 let grid = (n_heads * n_tokens) as u32;
11803 let cfg = LaunchConfig {
11804 grid_dim: (grid, 1, 1),
11805 block_dim: ((head_dim / 2) as u32, 1, 1),
11806 shared_mem_bytes: 0,
11807 };
11808 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11809 let __s_b = self.gpu.stream();
11810 let mut b = __s_b.launch_builder(&f);
11811 b.arg(x)
11812 .arg(pos)
11813 .arg(&hd)
11814 .arg(&nd)
11815 .arg(&nh)
11816 .arg(&theta_scale)
11817 .arg(&freq_scale);
11818 unsafe {
11819 b.launch(cfg)?;
11820 }
11821 Ok(())
11822 }
11823
11824 pub fn rope_neox_ff(
11826 &self,
11827 x: &mut CudaSlice<f32>,
11828 pos: &CudaSlice<i32>,
11829 head_dim: usize,
11830 n_dims: usize,
11831 n_heads: usize,
11832 n_tokens: usize,
11833 freq_base: f32,
11834 freq_scale: f32,
11835 ff: &CudaSlice<f32>,
11836 ) -> Result<(), Box<dyn std::error::Error>> {
11837 let f = self.func("rope_neox_ff_f32");
11838 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11839 let grid = (n_heads * n_tokens) as u32;
11840 let cfg = LaunchConfig {
11841 grid_dim: (grid, 1, 1),
11842 block_dim: ((head_dim / 2) as u32, 1, 1),
11843 shared_mem_bytes: 0,
11844 };
11845 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
11846 let __s_b = self.gpu.stream();
11847 let mut b = __s_b.launch_builder(&f);
11848 b.arg(x)
11849 .arg(pos)
11850 .arg(&hd)
11851 .arg(&nd)
11852 .arg(&nh)
11853 .arg(&theta_scale)
11854 .arg(&freq_scale)
11855 .arg(ff);
11856 unsafe {
11857 b.launch(cfg)?;
11858 }
11859 Ok(())
11860 }
11861
11862 #[allow(clippy::too_many_arguments)]
11864 pub fn rope_neox2(
11865 &self,
11866 q: &mut CudaSlice<f32>,
11867 k: &mut CudaSlice<f32>,
11868 pos: &CudaSlice<i32>,
11869 head_dim: usize,
11870 n_dims: usize,
11871 nh_q: usize,
11872 nh_k: usize,
11873 n_tokens: usize,
11874 freq_base: f32,
11875 freq_scale: f32,
11876 ff: Option<&CudaSlice<f32>>,
11877 ) -> Result<(), Box<dyn std::error::Error>> {
11878 let f = self.func("rope_neox2_f32");
11879 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
11880 let grid = ((nh_q + nh_k) * n_tokens) as u32;
11881 let cfg = LaunchConfig {
11882 grid_dim: (grid, 1, 1),
11883 block_dim: ((head_dim / 2) as u32, 1, 1),
11884 shared_mem_bytes: 0,
11885 };
11886 let (hd, nd, nq, nk, nt) = (
11887 head_dim as i32,
11888 n_dims as i32,
11889 nh_q as i32,
11890 nh_k as i32,
11891 n_tokens as i32,
11892 );
11893 let __s_b = self.gpu.stream();
11894 let mut b = __s_b.launch_builder(&f);
11895 b.arg(q)
11896 .arg(k)
11897 .arg(pos)
11898 .arg(&hd)
11899 .arg(&nd)
11900 .arg(&nq)
11901 .arg(&nk)
11902 .arg(&nt)
11903 .arg(&theta_scale)
11904 .arg(&freq_scale);
11905 match ff {
11906 Some(ffv) => {
11907 b.arg(ffv);
11908 unsafe {
11909 b.launch(cfg)?;
11910 }
11911 }
11912 None => {
11913 let null: u64 = 0;
11914 b.arg(&null);
11915 unsafe {
11916 b.launch(cfg)?;
11917 }
11918 }
11919 }
11920 Ok(())
11921 }
11922
11923 pub fn gelu_tanh_mul(
11925 &self,
11926 gate: &CudaSlice<f32>,
11927 up: &CudaSlice<f32>,
11928 dst: &mut CudaSlice<f32>,
11929 n: usize,
11930 ) -> Result<(), Box<dyn std::error::Error>> {
11931 let f = self.func("gelu_tanh_mul_f32");
11932 let cfg = LaunchConfig::for_num_elems(n as u32);
11933 let ni = n as i32;
11934 let __s_b = self.gpu.stream();
11935 let mut b = __s_b.launch_builder(&f);
11936 b.arg(gate).arg(up).arg(dst).arg(&ni);
11937 unsafe {
11938 b.launch(cfg)?;
11939 }
11940 Ok(())
11941 }
11942
11943 pub fn silu_mul(
11944 &self,
11945 gate: &CudaSlice<f32>,
11946 up: &CudaSlice<f32>,
11947 dst: &mut CudaSlice<f32>,
11948 n: usize,
11949 ) -> Result<(), Box<dyn std::error::Error>> {
11950 let f = self.func("silu_mul_f32");
11951 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
11953 let ni = n as i32;
11954 let __s_b = self.gpu.stream();
11955 let mut b = __s_b.launch_builder(&f);
11956 b.arg(gate).arg(up).arg(dst).arg(&ni);
11957 unsafe {
11958 b.launch(cfg)?;
11959 }
11960 Ok(())
11961 }
11962
11963 pub fn silu_mul_host_expf(
11965 &self,
11966 gate: &CudaSlice<f32>,
11967 up: &CudaSlice<f32>,
11968 dst: &mut CudaSlice<f32>,
11969 n: usize,
11970 ) -> Result<(), Box<dyn std::error::Error>> {
11971 let f = self.func("silu_mul_host_expf_f32");
11972 let cfg = LaunchConfig::for_num_elems(n as u32);
11973 let ni = n as i32;
11974 let __s_b = self.gpu.stream();
11975 let mut b = __s_b.launch_builder(&f);
11976 b.arg(gate).arg(up).arg(dst).arg(&ni);
11977 unsafe {
11978 b.launch(cfg)?;
11979 }
11980 Ok(())
11981 }
11982
11983 pub fn silu_clamped_mul_host_expf(
11985 &self,
11986 gate: &CudaSlice<f32>,
11987 up: &CudaSlice<f32>,
11988 limit: f32,
11989 dst: &mut CudaSlice<f32>,
11990 n: usize,
11991 ) -> Result<(), Box<dyn std::error::Error>> {
11992 if !limit.is_finite() || limit <= 0.0 {
11993 return Err(
11994 format!("Step routed-expert clamp limit must be positive, got {limit}").into(),
11995 );
11996 }
11997 let f = self.func("silu_clamped_mul_host_expf_f32");
11998 let cfg = LaunchConfig::for_num_elems(n as u32);
11999 let ni = n as i32;
12000 let __s_b = self.gpu.stream();
12001 let mut b = __s_b.launch_builder(&f);
12002 b.arg(gate).arg(up).arg(&limit).arg(dst).arg(&ni);
12003 unsafe {
12004 b.launch(cfg)?;
12005 }
12006 Ok(())
12007 }
12008
12009 pub fn silu_mul_f16out(
12012 &self,
12013 gate: &CudaSlice<f32>,
12014 up: &CudaSlice<f32>,
12015 dst: &mut CudaSlice<f32>,
12016 dst16: &mut CudaSlice<u8>,
12017 n: usize,
12018 ) -> Result<(), Box<dyn std::error::Error>> {
12019 let f = self.func("silu_mul_f16out_f32");
12020 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
12021 let ni = n as i32;
12022 let __s_b = self.gpu.stream();
12023 let mut b = __s_b.launch_builder(&f);
12024 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
12025 unsafe {
12026 b.launch(cfg)?;
12027 }
12028 Ok(())
12029 }
12030
12031 pub fn silu_mul_scaled(
12038 &self,
12039 gate: &CudaSlice<f32>,
12040 up: &CudaSlice<f32>,
12041 gs: f32,
12042 us: f32,
12043 dst: &mut CudaSlice<f32>,
12044 n: usize,
12045 ) -> Result<(), Box<dyn std::error::Error>> {
12046 let f = self.func("silu_mul_scaled_f32");
12047 let cfg = LaunchConfig::for_num_elems(n as u32);
12048 let ni = n as i32;
12049 let (gsf, usf) = (gs, us);
12050 let __s_b = self.gpu.stream();
12051 let mut b = __s_b.launch_builder(&f);
12052 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
12053 unsafe {
12054 b.launch(cfg)?;
12055 }
12056 Ok(())
12057 }
12058
12059 #[allow(clippy::too_many_arguments)]
12063 pub fn swigluoai_mul_scaled(
12064 &self,
12065 gate: &CudaSlice<f32>,
12066 up: &CudaSlice<f32>,
12067 gs: f32,
12068 us: f32,
12069 alpha: f32,
12070 limit: f32,
12071 dst: &mut CudaSlice<f32>,
12072 n: usize,
12073 ) -> Result<(), Box<dyn std::error::Error>> {
12074 let f = self.func("swigluoai_mul_scaled_f32");
12075 let cfg = LaunchConfig::for_num_elems(n as u32);
12076 let ni = n as i32;
12077 let __s_b = self.gpu.stream();
12078 let mut b = __s_b.launch_builder(&f);
12079 b.arg(gate)
12080 .arg(up)
12081 .arg(&gs)
12082 .arg(&us)
12083 .arg(&alpha)
12084 .arg(&limit)
12085 .arg(dst)
12086 .arg(&ni);
12087 unsafe {
12088 b.launch(cfg)?;
12089 }
12090 Ok(())
12091 }
12092
12093 pub fn silu_mul_scaled_q8_1(
12101 &self,
12102 gate: &CudaSlice<f32>,
12103 up: &CudaSlice<f32>,
12104 gs: f32,
12105 us: f32,
12106 n: usize,
12107 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
12108 let f = self.func("silu_mul_scaled_q8_1");
12109 let nblk = n / 32;
12110 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);
12114 let (gsf, usf, ni) = (gs, us, n as i32);
12115 let __s_b = self.gpu.stream();
12116 let mut b = __s_b.launch_builder(&f);
12117 b.arg(gate)
12118 .arg(up)
12119 .arg(&gsf)
12120 .arg(&usf)
12121 .arg(&mut aq)
12122 .arg(&mut ad)
12123 .arg(&ni);
12124 unsafe {
12125 b.launch(cfg)?;
12126 }
12127 Ok((aq, ad))
12128 }
12129
12130 pub fn add(
12131 &self,
12132 a: &CudaSlice<f32>,
12133 b_in: &CudaSlice<f32>,
12134 dst: &mut CudaSlice<f32>,
12135 n: usize,
12136 ) -> Result<(), Box<dyn std::error::Error>> {
12137 let f = self.func("add_f32");
12138 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
12140 let ni = n as i32;
12141 let __s_bld = self.gpu.stream();
12142 let mut bld = __s_bld.launch_builder(&f);
12143 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
12144 unsafe {
12145 bld.launch(cfg)?;
12146 }
12147 Ok(())
12148 }
12149
12150 pub fn mul(
12151 &self,
12152 a: &CudaSlice<f32>,
12153 b_in: &CudaSlice<f32>,
12154 dst: &mut CudaSlice<f32>,
12155 n: usize,
12156 ) -> Result<(), Box<dyn std::error::Error>> {
12157 let f = self.func("mul_f32");
12158 let cfg = LaunchConfig::for_num_elems(n as u32);
12159 let ni = n as i32;
12160 let __s_bld = self.gpu.stream();
12161 let mut bld = __s_bld.launch_builder(&f);
12162 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
12163 unsafe {
12164 bld.launch(cfg)?;
12165 }
12166 Ok(())
12167 }
12168
12169 pub fn matmul(
12172 &self,
12173 w: &crate::model::GpuTensor,
12174 x: &CudaSlice<f32>,
12175 m: usize,
12176 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12177 use crate::model::GpuTensor;
12178 let in_f = w.in_features();
12179 let out_f = w.out_features();
12180 #[allow(non_snake_case)]
12188 let GEMM_M_THRESHOLD = if self.verify_exact_on() {
12191 usize::MAX
12192 } else {
12193 16usize
12194 };
12195
12196 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
12221 if let Some(y) = self.try_fp8_gemm(w, x, m)? {
12222 return Ok(y);
12223 }
12224 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
12231 return Ok(y);
12232 }
12233 if let Some(y) = self.try_f16_gemm(w, x, m)? {
12236 return Ok(y);
12237 }
12238 }
12239 if let GpuTensor::Quant { qtype, .. } = w {
12254 if *qtype == QT_F8_E4M3_BLK {
12255 if m >= GEMM_M_THRESHOLD {
12256 if let Some(y) = self.try_e4m3_blk_prefill(w, x, m)? {
12257 return Ok(y);
12258 }
12259 }
12260 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12261 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12262 return Ok(y);
12263 }
12264 }
12265 }
12266 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
12267 return self.qmatvec_mmq(w, x, m);
12268 }
12269 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
12270 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12271 return self.qmatvec_gemm(w, &aq, &ad, m);
12272 }
12273 if m >= GEMM_M_THRESHOLD {
12276 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? {
12277 return Ok(y);
12278 }
12279 }
12280 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
12284 if m == 1 && fast {
12289 if let GpuTensor::Quant {
12290 bytes,
12291 qtype,
12292 row_bytes,
12293 rp,
12294 rp4,
12295 scale,
12296 ..
12297 } = w
12298 {
12299 if self.mmvq_supports(*qtype) {
12300 let (bytes, rp) = match rp4 {
12304 Some(m4) => (m4, true),
12305 None => (bytes, *rp),
12306 };
12307 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12308 return self.qmatvec_mmvq(
12309 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp,
12310 );
12311 }
12312 }
12313 }
12314 if (2..=16).contains(&m)
12330 && fast
12331 && std::env::var("MEMRA_NO_BATCHED").is_err()
12332 && (m <= 4 || Self::b8_enabled())
12333 {
12334 let m_ok = m <= 8
12344 || matches!(w, GpuTensor::Quant { qtype, .. }
12345 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || *qtype == QT_F8_E4M3
12346 || *qtype == QT_NVFP4 || *qtype == QT_Q4_K || *qtype == QT_Q5_K || *qtype == QT_Q8_0);
12347 if m_ok {
12348 if let GpuTensor::Quant {
12349 bytes,
12350 qtype,
12351 row_bytes,
12352 rp,
12353 rp4,
12354 ..
12355 } = w
12356 {
12357 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
12358 let (bytes, rp) = match rp4 {
12359 Some(m4) => (m4, true),
12360 None => (bytes, *rp),
12361 };
12362 let mcols = Self::batched_mcols(m);
12363 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12364 let mut y = self.qmatvec_mmvq_batched(
12365 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp,
12366 )?;
12367 if let GpuTensor::Quant { scale, .. } = w {
12368 if *scale != 1.0 {
12369 self.scale_inplace(&mut y, *scale, m * out_f)?;
12370 }
12371 }
12372 return Ok(y);
12373 }
12374 }
12375 }
12376 }
12377 if fast {
12383 if let GpuTensor::Quant {
12384 bytes,
12385 qtype,
12386 row_bytes,
12387 scale,
12388 ..
12389 } = w
12390 {
12391 if *qtype == QT_F8_E4M3 {
12392 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12393 return self.qmatvec_mmvq(
12394 bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, false,
12395 );
12396 }
12397 }
12398 }
12399 let mut y = match w {
12400 GpuTensor::Quant {
12401 bytes,
12402 qtype,
12403 row_bytes,
12404 ..
12405 } if fast && *qtype == QT_Q8_0 => {
12406 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12407 }
12408 GpuTensor::Quant {
12409 bytes,
12410 qtype,
12411 row_bytes,
12412 ..
12413 } if fast && *qtype == QT_Q4_K => {
12414 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12415 }
12416 GpuTensor::Quant {
12417 bytes,
12418 qtype,
12419 row_bytes,
12420 ..
12421 } if fast && *qtype == QT_Q6_K => {
12422 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12423 }
12424 GpuTensor::Quant {
12425 bytes,
12426 qtype,
12427 row_bytes,
12428 ..
12429 } if fast && *qtype == QT_Q5_K => {
12430 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12431 }
12432 GpuTensor::Quant {
12433 bytes,
12434 qtype,
12435 row_bytes,
12436 ..
12437 } if fast && *qtype == QT_Q3_K => {
12438 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12439 }
12440 GpuTensor::Quant {
12441 bytes,
12442 qtype,
12443 row_bytes,
12444 rp,
12445 ..
12446 } if fast && *qtype == QT_NVFP4 => self.qmatvec_dp4a_named(
12447 if *rp {
12448 "qmatvec_nvfp4_dp4a_rp"
12449 } else {
12450 "qmatvec_nvfp4_dp4a"
12451 },
12452 &bytes.slice(0..bytes.len()),
12453 x,
12454 m,
12455 in_f,
12456 out_f,
12457 *row_bytes,
12458 )?,
12459 GpuTensor::Quant {
12463 bytes,
12464 qtype,
12465 row_bytes,
12466 ..
12467 } if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() => {
12468 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?
12469 }
12470 GpuTensor::Quant {
12475 bytes,
12476 qtype,
12477 row_bytes,
12478 rp,
12479 ..
12480 } =>
12481 {
12484 self.qmatvec(
12485 bytes,
12486 x,
12487 m,
12488 in_f,
12489 out_f,
12490 if *rp && *qtype == QT_NVFP4 {
12491 QT_NVFP4_RP
12492 } else {
12493 *qtype
12494 },
12495 *row_bytes,
12496 )?
12497 }
12498 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
12499 GpuTensor::FloatBf16 { data, .. } => {
12502 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12509 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12510 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12511 y
12512 } else {
12513 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)?
12514 }
12515 }
12516 };
12517 if let GpuTensor::Quant { scale, .. } = w {
12519 if *scale != 1.0 {
12520 self.scale_inplace(&mut y, *scale, m * out_f)?;
12521 }
12522 }
12523 Ok(y)
12524 }
12525
12526 pub fn stage_a_raw_needed() -> bool {
12536 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12537 *ON.get_or_init(|| std::env::var("MEMRA_FAST").as_deref() == Ok("0"))
12538 }
12539
12540 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
12543 use crate::model::GpuTensor;
12544 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") {
12545 return false;
12546 }
12547 match w {
12548 GpuTensor::Quant { qtype, .. } => {
12555 matches!(
12556 *qtype,
12557 QT_Q8_0
12558 | QT_Q4_K
12559 | QT_Q6_K
12560 | QT_Q5_K
12561 | QT_Q3_K
12562 | QT_NVFP4
12563 | QT_F8_E4M3
12564 | QT_F8_E4M3_BLK
12565 | QT_Q4_0
12566 ) || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled())
12567 }
12568 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
12569 }
12570 }
12571
12572 pub fn matmul_pre(
12577 &self,
12578 w: &crate::model::GpuTensor,
12579 aq: &CudaSlice<i8>,
12580 ad: &CudaSlice<f32>,
12581 x_fallback: &CudaSlice<f32>,
12582 m: usize,
12583 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12584 use crate::model::GpuTensor;
12585 let x_raw_ok = x_fallback.len() >= m * w.in_features();
12591 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12594 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? {
12595 return Ok(y);
12596 }
12597 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? {
12600 return Ok(y);
12601 }
12602 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? {
12604 return Ok(y);
12605 }
12606 }
12607 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12613 if let Some(y) = self.try_e4m3_blk_prefill(w, x_fallback, m)? {
12614 return Ok(y);
12615 }
12616 }
12617 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12618 return Ok(y);
12619 }
12620 if m >= 16
12625 && w.out_features() >= 128
12626 && self.mmq_supports(w)
12627 && !self.verify_exact_on()
12628 && x_raw_ok
12629 {
12630 return self.qmatvec_mmq(w, x_fallback, m);
12631 }
12632 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
12635 if let Some(y) =
12636 self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())?
12637 {
12638 return Ok(y);
12639 }
12640 }
12641 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
12644 return self.qmatvec_gemm(w, aq, ad, m);
12645 }
12646 if !self.uses_q8_1_fast(w) {
12665 if !x_raw_ok {
12666 return Err(format!(
12667 "matmul_pre: q8_1-fast is off for this weight but x_fallback holds {} f32 \
12668 (need m*in_f = {}*{} = {}). This call site pre-quantized its activation and \
12669 dropped the f32, so there is nothing to fall back to — pass the real f32 \
12670 activation (see Engine::rms_norm_decode, which is bit-identical to \
12671 rms_norm_q8_1's reduction) or keep the weight on the q8_1 path.",
12672 x_fallback.len(),
12673 m,
12674 w.in_features(),
12675 m * w.in_features()
12676 )
12677 .into());
12678 }
12679 return self.matmul(w, x_fallback, m);
12680 }
12681 let in_f = w.in_features();
12682 let out_f = w.out_features();
12683 let (bytes, qtype, row_bytes, scale, rp) = match w {
12684 GpuTensor::Quant {
12685 bytes,
12686 qtype,
12687 row_bytes,
12688 scale,
12689 rp,
12690 ..
12691 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12692 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
12693 };
12694 let (mbytes, mrp) = match w {
12697 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12698 _ => (bytes, rp),
12699 };
12700 if m == 1 && self.mmvq_supports(qtype) {
12704 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
12705 }
12706 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12719 && std::env::var("MEMRA_NO_BATCHED").is_err()
12720 && (m <= 4 || Self::b8_enabled())
12721 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_NVFP4
12725 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_F8_E4M3 || qtype == QT_Q8_0)
12726 {
12727 let mcols = Self::batched_mcols(m);
12728 return self.qmatvec_mmvq_batched(
12729 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp,
12730 );
12731 }
12732 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
12738 let (b2, r2) = if qtype == QT_Q4_0 {
12739 (mbytes, mrp)
12740 } else {
12741 (bytes, rp)
12742 };
12743 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
12744 }
12745 let name = match qtype {
12746 QT_Q8_0 => "qmatvec_q8_0_dp4a",
12747 QT_Q4_K => "qmatvec_q4_K_dp4a",
12748 QT_Q6_K => "qmatvec_q6_K_dp4a",
12749 QT_Q5_K => "qmatvec_q5_K_dp4a",
12750 QT_Q3_K => "qmatvec_q3_K_dp4a",
12751 QT_NVFP4 => {
12752 if rp {
12753 "qmatvec_nvfp4_dp4a_rp"
12754 } else {
12755 "qmatvec_nvfp4_dp4a"
12756 }
12757 }
12758 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
12759 _ => unreachable!(),
12760 };
12761 let f = self.func(name);
12762 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
12764 grid_dim: (out_f as u32, m as u32, 1),
12765 block_dim: (128, 1, 1),
12766 shared_mem_bytes: 0,
12767 };
12768 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
12769 let __s_b = self.gpu.stream();
12770 let mut b = __s_b.launch_builder(&f);
12771 b.arg(bytes)
12772 .arg(aq)
12773 .arg(ad)
12774 .arg(&mut y)
12775 .arg(&inf)
12776 .arg(&outf)
12777 .arg(&mi)
12778 .arg(&rb);
12779 unsafe {
12780 b.launch(cfg)?;
12781 }
12782 if scale != 1.0 {
12783 self.scale_inplace(&mut y, scale, m * out_f)?;
12784 }
12785 Ok(y)
12786 }
12787
12788 pub fn matmul_decode_exact(
12796 &self,
12797 w: &crate::model::GpuTensor,
12798 x: &CudaSlice<f32>,
12799 m: usize,
12800 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12801 use crate::model::GpuTensor;
12802 if let GpuTensor::Float { data, .. } = w {
12810 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
12811 }
12812 if let GpuTensor::FloatBf16 { data, .. } = w {
12815 let (in_f, out_f) = (w.in_features(), w.out_features());
12816 if (1..=32).contains(&m) && Self::bf16_mmv_on() && in_f % 8 == 0 {
12819 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
12820 self.matvec_bf16_rows_into(data, x, &mut y, in_f, out_f, m)?;
12821 return Ok(y);
12822 }
12823 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true, None);
12824 }
12825 if !self.uses_q8_1_fast(w) {
12826 return self.matmul(w, x, m);
12827 }
12828 let in_f = w.in_features();
12829 let out_f = w.out_features();
12830 let (bytes, qtype, row_bytes, scale, rp) = match w {
12831 GpuTensor::Quant {
12832 bytes,
12833 qtype,
12834 row_bytes,
12835 scale,
12836 rp,
12837 ..
12838 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12839 _ => return self.matmul(w, x, m),
12840 };
12841 let (bytes, rp) = match w {
12844 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12845 _ => (bytes, rp),
12846 };
12847 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
12848 if let Some(y) = self.try_e4m3_blk_pre(w, &aq, &ad, m)? {
12852 return Ok(y);
12853 }
12854 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
12863 && std::env::var("MEMRA_NO_BATCHED").is_err()
12864 && (m <= 4 || Self::b8_enabled())
12865 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || qtype == QT_F8_E4M3
12868 || qtype == QT_NVFP4 || qtype == QT_Q4_K || qtype == QT_Q5_K || qtype == QT_Q8_0)
12869 {
12870 let mcols = Self::batched_mcols(m);
12871 return self.qmatvec_mmvq_batched(
12872 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12873 );
12874 }
12875 if self.mmvq_supports(qtype) {
12876 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12879 }
12880 self.matmul_pre(w, &aq, &ad, x, m)
12883 }
12884
12885 pub fn matmul_decode_exact_pre(
12895 &self,
12896 w: &crate::model::GpuTensor,
12897 aq: &CudaSlice<i8>,
12898 ad: &CudaSlice<f32>,
12899 m: usize,
12900 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
12901 use crate::model::GpuTensor;
12902 debug_assert!(
12903 self.uses_q8_1_fast(w),
12904 "matmul_decode_exact_pre: caller must guarantee q8_1-fast"
12905 );
12906 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
12908 return Ok(y);
12909 }
12910 let in_f = w.in_features();
12911 let out_f = w.out_features();
12912 let (bytes, qtype, row_bytes, scale, rp) = match w {
12913 GpuTensor::Quant {
12914 bytes,
12915 qtype,
12916 row_bytes,
12917 scale,
12918 rp,
12919 ..
12920 } => (bytes, *qtype, *row_bytes, *scale, *rp),
12921 _ => {
12922 return Err(
12923 "matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into(),
12924 );
12925 }
12926 };
12927 let (bytes, rp) = match w {
12929 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
12930 _ => (bytes, rp),
12931 };
12932 if (2..=16).contains(&m)
12934 && self.batched_supports(qtype)
12935 && self.mmvq_supports(qtype)
12936 && std::env::var("MEMRA_NO_BATCHED").is_err()
12937 && (m <= 4 || Self::b8_enabled())
12938 && (m <= 8
12939 || qtype == QT_Q4_0
12940 || qtype == QT_Q6_K
12941 || qtype == QT_F8_E4M3
12942 || qtype == QT_NVFP4
12943 || qtype == QT_Q4_K
12944 || qtype == QT_Q5_K
12945 || qtype == QT_Q8_0)
12946 {
12947 let mcols = Self::batched_mcols(m);
12948 return self.qmatvec_mmvq_batched(
12949 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp,
12950 );
12951 }
12952 if self.mmvq_supports(qtype) {
12953 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
12954 }
12955 let x0 = self.zeros(0)?;
12958 self.matmul_pre(w, aq, ad, &x0, m)
12959 }
12960
12961 pub fn matmul_decode_exact_dual_pre(
12970 &self,
12971 w0: &crate::model::GpuTensor,
12972 w1: &crate::model::GpuTensor,
12973 aq: &CudaSlice<i8>,
12974 ad: &CudaSlice<f32>,
12975 m: usize,
12976 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
12977 {
12978 use crate::model::GpuTensor;
12979 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12980 let on = *ON.get_or_init(|| {
12981 std::env::var("MEMRA_SPEC_DUAL_T")
12982 .map(|v| v != "0")
12983 .unwrap_or(true)
12984 });
12985 if !on
12986 || !(2..=7).contains(&m)
12987 || std::env::var("MEMRA_NO_BATCHED").is_ok()
12988 || !self.uses_q8_1_fast(w0)
12989 || !self.uses_q8_1_fast(w1)
12990 {
12991 return Ok(None);
12992 }
12993 if !self.mmvq_supports(QT_NVFP4) {
12998 return Ok(None);
12999 }
13000 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13001 if w1.in_features() != in_f || w1.out_features() != out_f {
13002 return Ok(None);
13003 }
13004 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
13005 (
13006 GpuTensor::Quant {
13007 bytes: b0,
13008 qtype: q0,
13009 row_bytes: rb0,
13010 scale: s0,
13011 rp: rp0,
13012 rp4: None,
13013 ..
13014 },
13015 GpuTensor::Quant {
13016 bytes: b1,
13017 qtype: q1,
13018 row_bytes: rb1,
13019 scale: s1,
13020 rp: rp1,
13021 rp4: None,
13022 ..
13023 },
13024 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13025 (b0, b1, *rb0, *s0, *s1, *rp0)
13026 }
13027 _ => return Ok(None),
13028 };
13029 if m > 4 && !(rp && Self::b8_enabled() && std::env::var("MEMRA_B567").as_deref() != Ok("0"))
13032 {
13033 return Ok(None);
13034 }
13035 let (y0, y1) =
13036 self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
13037 Ok(Some(((y0, s0), (y1, s1))))
13038 }
13039
13040 pub fn matmul_decode_exact_group4_pre(
13052 &self,
13053 ws: [&crate::model::GpuTensor; 4],
13054 aq: &CudaSlice<i8>,
13055 ad: &CudaSlice<f32>,
13056 m: usize,
13057 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
13058 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13059 let on = *ON.get_or_init(|| {
13060 std::env::var("MEMRA_TK_GDN_GROUP")
13061 .map(|v| v != "0")
13062 .unwrap_or(true)
13063 });
13064 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "GDN group4")
13065 }
13066
13067 pub fn matmul_decode_exact_group3_pre(
13072 &self,
13073 ws: [&crate::model::GpuTensor; 3],
13074 aq: &CudaSlice<i8>,
13075 ad: &CudaSlice<f32>,
13076 m: usize,
13077 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
13078 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13079 let on = *ON.get_or_init(|| {
13080 std::env::var("MEMRA_TK_FA_GROUP")
13081 .map(|v| v != "0")
13082 .unwrap_or(true)
13083 });
13084 self.matmul_decode_exact_group_pre(&ws, aq, ad, m, on, "FA group3")
13085 }
13086
13087 fn matmul_decode_exact_group_pre(
13091 &self,
13092 ws: &[&crate::model::GpuTensor],
13093 aq: &CudaSlice<i8>,
13094 ad: &CudaSlice<f32>,
13095 m: usize,
13096 on: bool,
13097 tag: &'static str,
13098 ) -> Result<Option<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
13099 use crate::model::GpuTensor;
13100 if !on
13101 || !(2..=16).contains(&m)
13102 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13103 || (m > 4 && !Self::b8_enabled())
13104 || !self.mmvq_supports(QT_NVFP4)
13105 || !self.batched_supports(QT_NVFP4)
13106 {
13107 return Ok(None);
13108 }
13109 let in_f = ws[0].in_features();
13110 let mut parts: Vec<(&CudaSlice<u8>, usize, f32)> = Vec::with_capacity(4);
13111 for w in ws {
13112 if !self.uses_q8_1_fast(w) || w.in_features() != in_f {
13113 return Ok(None);
13114 }
13115 match w {
13116 GpuTensor::Quant {
13117 bytes,
13118 qtype,
13119 scale,
13120 rp: true,
13121 rp4: None,
13122 ..
13123 } if *qtype == QT_NVFP4 && w.out_features() % 8 == 0 => {
13124 parts.push((bytes, w.out_features(), *scale));
13125 }
13126 _ => return Ok(None),
13127 }
13128 }
13129 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13131 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
13132 let mcols = if (5..=7).contains(&m) && b567 {
13133 m
13134 } else {
13135 Self::batched_mcols(m)
13136 };
13137 let kname: &'static str = match mcols {
13138 2 => "qmatvec_nvfp4_mmvq_group4_b2_rp",
13139 4 => "qmatvec_nvfp4_mmvq_group4_b4_rp",
13140 5 => "qmatvec_nvfp4_mmvq_group4_b5_rp",
13141 6 => "qmatvec_nvfp4_mmvq_group4_b6_rp",
13142 7 => "qmatvec_nvfp4_mmvq_group4_b7_rp",
13143 8 => "qmatvec_nvfp4_mmvq_group4_b8_rp",
13144 16 => "qmatvec_nvfp4_mmvq_group4_b16_rp",
13145 _ => return Ok(None),
13146 };
13147 if std::env::var("MEMRA_DEBUG").is_ok() {
13150 use std::sync::Mutex;
13151 static SEEN: Mutex<Vec<&'static str>> = Mutex::new(Vec::new());
13152 let mut seen = SEEN.lock().unwrap();
13153 if !seen.contains(&tag) {
13154 seen.push(tag);
13155 eprintln!("[memra] {tag} batched ENGAGED (m={m})");
13156 }
13157 }
13158 const ROWS_PER_BLOCK: u32 = 4; let rows_per_block = ROWS_PER_BLOCK * 2; let total: usize = parts.iter().map(|p| p.1).sum();
13161 let three = parts.len() == 3;
13162 let mut y0 = self.alloc_uninit::<f32>(m * parts[0].1)?;
13163 let mut y1 = self.alloc_uninit::<f32>(m * parts[1].1)?;
13164 let mut y2 = self.alloc_uninit::<f32>(m * parts[2].1)?;
13165 let mut y3 = self.alloc_uninit::<f32>(if three { 1 } else { m * parts[3].1 })?;
13168 let cfg = LaunchConfig {
13169 grid_dim: ((total as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
13170 block_dim: (32, ROWS_PER_BLOCK, 1),
13171 shared_mem_bytes: 0,
13172 };
13173 let (inf, mi) = (in_f as i32, m as i32);
13174 let (n0, n1, n2) = (parts[0].1 as i32, parts[1].1 as i32, parts[2].1 as i32);
13175 let n3 = if three { 0i32 } else { parts[3].1 as i32 };
13176 let (s0, s1, s2) = (parts[0].2, parts[1].2, parts[2].2);
13177 let s3 = if three { 1.0f32 } else { parts[3].2 };
13178 let w3 = if three { parts[0].0 } else { parts[3].0 };
13179 let f = self.func(kname);
13180 let __s_b = self.gpu.stream();
13181 let mut b = __s_b.launch_builder(&f);
13182 b.arg(parts[0].0)
13183 .arg(parts[1].0)
13184 .arg(parts[2].0)
13185 .arg(w3)
13186 .arg(aq)
13187 .arg(ad)
13188 .arg(&mut y0)
13189 .arg(&mut y1)
13190 .arg(&mut y2)
13191 .arg(&mut y3)
13192 .arg(&inf)
13193 .arg(&n0)
13194 .arg(&n1)
13195 .arg(&n2)
13196 .arg(&n3)
13197 .arg(&mi)
13198 .arg(&s0)
13199 .arg(&s1)
13200 .arg(&s2)
13201 .arg(&s3);
13202 unsafe {
13203 b.launch(cfg)?;
13204 }
13205 Ok(Some(if three {
13206 vec![y0, y1, y2]
13207 } else {
13208 vec![y0, y1, y2, y3]
13209 }))
13210 }
13211
13212 pub fn matmul_decode_exact_dual(
13228 &self,
13229 w0: &crate::model::GpuTensor,
13230 w1: &crate::model::GpuTensor,
13231 x: &CudaSlice<f32>,
13232 m: usize,
13233 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13234 use crate::model::GpuTensor;
13235 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13236 let on = *ON.get_or_init(|| {
13237 std::env::var("MEMRA_SPEC_DUAL_T")
13238 .map(|v| v != "0")
13239 .unwrap_or(true)
13240 });
13241 if !on
13242 || !(2..=4).contains(&m)
13243 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13244 || !self.uses_q8_1_fast(w0)
13245 || !self.uses_q8_1_fast(w1)
13246 {
13247 return Ok(None);
13248 }
13249 if !self.mmvq_supports(QT_NVFP4) {
13254 return Ok(None);
13255 }
13256 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13257 if w1.in_features() != in_f || w1.out_features() != out_f {
13258 return Ok(None);
13259 }
13260 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
13261 (
13262 GpuTensor::Quant {
13263 bytes: b0,
13264 qtype: q0,
13265 row_bytes: rb0,
13266 scale: s0,
13267 rp: rp0,
13268 rp4: None,
13269 ..
13270 },
13271 GpuTensor::Quant {
13272 bytes: b1,
13273 qtype: q1,
13274 row_bytes: rb1,
13275 scale: s1,
13276 rp: rp1,
13277 rp4: None,
13278 ..
13279 },
13280 ) if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 => {
13281 (b0, b1, *rb0, *s0, *s1, *rp0)
13282 }
13283 _ => return Ok(None),
13284 };
13285 if std::env::var("MEMRA_DEBUG").is_ok() {
13288 static ONCE: std::sync::Once = std::sync::Once::new();
13289 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
13290 }
13291 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
13292 let (y0, y1) =
13293 self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
13294 let mut y0 = y0;
13295 let mut y1 = y1;
13296 if s0 != 1.0 {
13297 self.scale_inplace(&mut y0, s0, m * out_f)?;
13298 }
13299 if s1 != 1.0 {
13300 self.scale_inplace(&mut y1, s1, m * out_f)?;
13301 }
13302 Ok(Some((y0, y1)))
13303 }
13304
13305 #[allow(clippy::too_many_arguments)]
13310 pub fn qmatvec_batched_dual_raw(
13311 &self,
13312 b0: &CudaSlice<u8>,
13313 b1: &CudaSlice<u8>,
13314 aq: &CudaSlice<i8>,
13315 ad: &CudaSlice<f32>,
13316 m: usize,
13317 in_f: usize,
13318 out_f: usize,
13319 row_bytes: usize,
13320 rp: bool,
13321 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
13322 const ROWS_PER_BLOCK: u32 = 4;
13323 let mcols = Self::batched_mcols(m);
13324 let tiny_rp1 = rp
13327 && mcols == 4
13328 && out_f <= 128
13329 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0");
13330 let (name, rows_per_block) = if tiny_rp1 {
13331 ("qmatvec_nvfp4_mmvq_dual_b4_rp", ROWS_PER_BLOCK)
13332 } else {
13333 match (mcols, rp, m) {
13334 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
13335 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
13336 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
13337 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
13338 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
13339 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
13340 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
13341 _ => {
13342 return Err(
13343 format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into(),
13344 );
13345 }
13346 }
13347 };
13348 let f = self.func(name);
13349 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
13350 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
13351 let cfg = LaunchConfig {
13352 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13353 block_dim: (32, ROWS_PER_BLOCK, 1),
13354 shared_mem_bytes: 0,
13355 };
13356 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
13357 let __s_b = self.gpu.stream();
13358 let mut b = __s_b.launch_builder(&f);
13359 b.arg(b0)
13360 .arg(b1)
13361 .arg(aq)
13362 .arg(ad)
13363 .arg(&mut y0)
13364 .arg(&mut y1)
13365 .arg(&inf)
13366 .arg(&outf)
13367 .arg(&mi)
13368 .arg(&rb);
13369 unsafe {
13370 b.launch(cfg)?;
13371 }
13372 Ok((y0, y1))
13373 }
13374
13375 pub fn matmul_pre_dual_noscale(
13387 &self,
13388 w0: &crate::model::GpuTensor,
13389 w1: &crate::model::GpuTensor,
13390 aq: &CudaSlice<i8>,
13391 ad: &CudaSlice<f32>,
13392 m: usize,
13393 ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>>
13394 {
13395 use crate::model::GpuTensor;
13396 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
13397 return Ok(None);
13398 }
13399 if !self.mmvq_supports(QT_NVFP4) {
13409 return Ok(None);
13410 }
13411 let (in_f, out_f) = (w0.in_features(), w0.out_features());
13412 if w1.in_features() != in_f || w1.out_features() != out_f {
13413 return Ok(None);
13414 }
13415 let no_mirror =
13428 |w: &crate::model::GpuTensor| !matches!(w, GpuTensor::Quant { rp4: Some(_), .. });
13429 if self.q8_ffn_fuse2_on()
13430 && no_mirror(w0)
13431 && no_mirror(w1)
13432 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
13433 {
13434 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
13435 return Ok(Some(((y0, 1.0), (y1, 1.0))));
13436 }
13437 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
13447 let (y0, y1) =
13448 self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2, 1.0, 1.0)?;
13449 return Ok(Some(((y0, p0.3), (y1, p1.3))));
13450 }
13451 let (b0, q0, rb0, s0, rp0) = match w0 {
13452 GpuTensor::Quant {
13453 bytes,
13454 qtype,
13455 row_bytes,
13456 scale,
13457 rp,
13458 ..
13459 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13460 _ => return Ok(None),
13461 };
13462 let (b1, q1, rb1, s1, rp1) = match w1 {
13463 GpuTensor::Quant {
13464 bytes,
13465 qtype,
13466 row_bytes,
13467 scale,
13468 rp,
13469 ..
13470 } => (bytes, *qtype, *row_bytes, *scale, *rp),
13471 _ => return Ok(None),
13472 };
13473 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 {
13474 return Ok(None);
13475 }
13476 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13478 let rows_per_block = ROWS_PER_BLOCK * RPW;
13479 let f = self.func(if rp0 {
13480 "qmatvec_nvfp4_mmvq_dual_mr2_rp"
13481 } else {
13482 "qmatvec_nvfp4_mmvq_dual_mr2"
13483 });
13484 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
13485 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
13486 let cfg = LaunchConfig {
13487 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
13488 block_dim: (32, ROWS_PER_BLOCK, 1),
13489 shared_mem_bytes: 0,
13490 };
13491 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
13492 let one = 1.0f32;
13495 let __s_b = self.gpu.stream();
13496 let mut b = __s_b.launch_builder(&f);
13497 b.arg(b0)
13498 .arg(b1)
13499 .arg(aq)
13500 .arg(ad)
13501 .arg(&mut y0)
13502 .arg(&mut y1)
13503 .arg(&inf)
13504 .arg(&outf)
13505 .arg(&mi)
13506 .arg(&rb)
13507 .arg(&one)
13508 .arg(&one);
13509 unsafe {
13510 b.launch(cfg)?;
13511 }
13512 Ok(Some(((y0, s0), (y1, s1))))
13513 }
13514
13515 #[allow(clippy::too_many_arguments)]
13523 pub fn matmul_nvfp4_fused3(
13524 &self,
13525 w0: &crate::model::GpuTensor,
13526 w1: &crate::model::GpuTensor,
13527 w2: &crate::model::GpuTensor,
13528 aq: &CudaSlice<i8>,
13529 ad: &CudaSlice<f32>,
13530 m: usize,
13531 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
13532 {
13533 use crate::model::GpuTensor;
13534 if !self.mmvq_supports(QT_NVFP4)
13541 || !self.uses_q8_1_fast(w0)
13542 || !self.uses_q8_1_fast(w1)
13543 || !self.uses_q8_1_fast(w2)
13544 {
13545 return Ok(None);
13546 }
13547 if (9..=16).contains(&m) {
13550 return Ok(
13551 match self.matmul_decode_exact_group3_pre([w0, w1, w2], aq, ad, m)? {
13552 Some(mut ys) => {
13553 let y2 = ys.pop().unwrap();
13554 let y1 = ys.pop().unwrap();
13555 let y0 = ys.pop().unwrap();
13556 Some((y0, y1, y2))
13557 }
13558 None => None,
13559 },
13560 );
13561 }
13562 if !(1..=8).contains(&m) {
13563 return Ok(None);
13564 }
13565 if m > 1 {
13566 let in_f = w0.in_features();
13567 if std::env::var("MEMRA_NVFP4_FUSED3B").as_deref() == Ok("0")
13568 || !self.batched_supports(QT_NVFP4)
13569 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13570 || (m > 4 && !Self::b8_enabled())
13571 || in_f % 512 != 0
13572 || in_f / 64 > 272
13573 {
13574 return Ok(None);
13575 }
13576 }
13577 let unpack = |w: &crate::model::GpuTensor| match w {
13578 GpuTensor::Quant {
13579 bytes,
13580 qtype,
13581 scale,
13582 rp,
13583 ..
13584 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13585 _ => None,
13586 };
13587 let (Some(p0), Some(p1), Some(p2)) = (unpack(w0), unpack(w1), unpack(w2)) else {
13588 return Ok(None);
13589 };
13590 let in_f = w0.in_features();
13591 if w1.in_features() != in_f || w2.in_features() != in_f {
13592 return Ok(None);
13593 }
13594 let (o0, o1, o2) = (w0.out_features(), w1.out_features(), w2.out_features());
13595 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13597 let rows_pb = ROWS_PER_BLOCK * RPW;
13598 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13599 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13600 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13601 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13602 let (inf, oi0, oi1, oi2, mi) = (in_f as i32, o0 as i32, o1 as i32, o2 as i32, m as i32);
13603 let (b0, b1, b2) = unsafe { (&*p0.0, &*p1.0, &*p2.0) };
13606 if m > 1 {
13607 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 {
13609 return Ok(None);
13610 }
13611 let f = self.func("qmatvec_nvfp4_mmvq_fused3_b8_rpsc");
13612 let cfg = LaunchConfig {
13613 grid_dim: (nb(o0) + nb(o1) + nb(o2), 1, 1),
13614 block_dim: (32, ROWS_PER_BLOCK, 1),
13615 shared_mem_bytes: 0,
13616 };
13617 let __s_b = self.gpu.stream();
13618 let mut b = __s_b.launch_builder(&f);
13619 b.arg(b0)
13620 .arg(b1)
13621 .arg(b2)
13622 .arg(aq)
13623 .arg(ad)
13624 .arg(&mut y0)
13625 .arg(&mut y1)
13626 .arg(&mut y2)
13627 .arg(&inf)
13628 .arg(&oi0)
13629 .arg(&oi1)
13630 .arg(&oi2)
13631 .arg(&mi);
13632 unsafe {
13633 b.launch(cfg)?;
13634 }
13635 return Ok(Some((y0, y1, y2)));
13636 }
13637 let f = self.func("qmatvec_nvfp4_mmvq_fused3_rp");
13638 let cfg = LaunchConfig {
13639 grid_dim: (nb(o0) + nb(o1) + nb(o2), m as u32, 1),
13640 block_dim: (32, ROWS_PER_BLOCK, 1),
13641 shared_mem_bytes: 0,
13642 };
13643 let __s_b = self.gpu.stream();
13644 let mut b = __s_b.launch_builder(&f);
13645 b.arg(b0)
13646 .arg(b1)
13647 .arg(b2)
13648 .arg(aq)
13649 .arg(ad)
13650 .arg(&mut y0)
13651 .arg(&mut y1)
13652 .arg(&mut y2)
13653 .arg(&inf)
13654 .arg(&oi0)
13655 .arg(&oi1)
13656 .arg(&oi2)
13657 .arg(&mi)
13658 .arg(&p0.1)
13659 .arg(&p1.1)
13660 .arg(&p2.1);
13661 unsafe {
13662 b.launch(cfg)?;
13663 }
13664 Ok(Some((y0, y1, y2)))
13665 }
13666
13667 pub fn matmul_nvfp4_fused2(
13676 &self,
13677 w0: &crate::model::GpuTensor,
13678 w1: &crate::model::GpuTensor,
13679 aq: &CudaSlice<i8>,
13680 ad: &CudaSlice<f32>,
13681 m: usize,
13682 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
13683 use crate::model::GpuTensor;
13684 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13685 let off =
13686 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13687 if off
13690 || m != 1
13691 || !self.mmvq_supports(QT_NVFP4)
13692 || !self.uses_q8_1_fast(w0)
13693 || !self.uses_q8_1_fast(w1)
13694 {
13695 return Ok(None);
13696 }
13697 let unpack = |w: &crate::model::GpuTensor| match w {
13698 GpuTensor::Quant {
13699 bytes,
13700 qtype,
13701 scale,
13702 rp,
13703 ..
13704 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13705 _ => None,
13706 };
13707 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13708 return Ok(None);
13709 };
13710 let in_f = w0.in_features();
13711 if w1.in_features() != in_f {
13712 return Ok(None);
13713 }
13714 let (o0, o1) = (w0.out_features(), w1.out_features());
13715 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13717 let rows_pb = ROWS_PER_BLOCK * RPW;
13718 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13719 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13720 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13721 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13722 let cfg = LaunchConfig {
13723 grid_dim: (nb(o0) + nb(o1), m as u32, 1),
13724 block_dim: (32, ROWS_PER_BLOCK, 1),
13725 shared_mem_bytes: 0,
13726 };
13727 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, m as i32);
13728 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13731 if Self::pdl_on() && Self::pdl_mmvq_on() && Self::pdl_nvfp4q8_on() {
13734 {
13735 use cudarc::driver::{DevicePtr, DevicePtrMut};
13736 let s = &self.gpu.stream();
13737 let (pw0, _g0) = b0.device_ptr(s);
13738 let (pw1, _g1) = b1.device_ptr(s);
13739 let (paq, _g2) = aq.device_ptr(s);
13740 let (pad, _g3) = ad.device_ptr(s);
13741 let (py0, _g4) = y0.device_ptr_mut(s);
13742 let (py1, _g5) = y1.device_ptr_mut(s);
13743 let (s0, s1) = (p0.1, p1.1);
13744 let mut ps = [
13745 &pw0 as *const _ as *mut std::ffi::c_void,
13746 &pw1 as *const _ as *mut _,
13747 &paq as *const _ as *mut _,
13748 &pad as *const _ as *mut _,
13749 &py0 as *const _ as *mut _,
13750 &py1 as *const _ as *mut _,
13751 &inf as *const _ as *mut _,
13752 &oi0 as *const _ as *mut _,
13753 &oi1 as *const _ as *mut _,
13754 &mi as *const _ as *mut _,
13755 &s0 as *const _ as *mut _,
13756 &s1 as *const _ as *mut _,
13757 ];
13758 unsafe {
13759 self.launch_pdl(
13760 "qmatvec_nvfp4_mmvq_fused2_rp",
13761 cfg.grid_dim,
13762 cfg.block_dim,
13763 &mut ps,
13764 )?;
13765 }
13766 }
13767 return Ok(Some((y0, y1)));
13768 }
13769 let __s_b = self.gpu.stream();
13770 let mut b = __s_b.launch_builder(&f);
13771 b.arg(b0)
13772 .arg(b1)
13773 .arg(aq)
13774 .arg(ad)
13775 .arg(&mut y0)
13776 .arg(&mut y1)
13777 .arg(&inf)
13778 .arg(&oi0)
13779 .arg(&oi1)
13780 .arg(&mi)
13781 .arg(&p0.1)
13782 .arg(&p1.1);
13783 unsafe {
13784 b.launch(cfg)?;
13785 }
13786 Ok(Some((y0, y1)))
13787 }
13788
13789 pub fn matmul_nvfp4_fused2_into(
13794 &self,
13795 w0: &crate::model::GpuTensor,
13796 w1: &crate::model::GpuTensor,
13797 aq: &CudaSlice<i8>,
13798 ad: &CudaSlice<f32>,
13799 y0: &mut CudaSlice<f32>,
13800 y1: &mut CudaSlice<f32>,
13801 ) -> Result<bool, Box<dyn std::error::Error>> {
13802 use crate::model::GpuTensor;
13803 static FUSED2_OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13804 let off =
13805 *FUSED2_OFF.get_or_init(|| std::env::var("MEMRA_NVFP4_FUSED2").as_deref() == Ok("0"));
13806 if off
13807 || !self.mmvq_supports(QT_NVFP4)
13808 || !self.uses_q8_1_fast(w0)
13809 || !self.uses_q8_1_fast(w1)
13810 {
13811 return Ok(false);
13812 }
13813 let unpack = |w: &crate::model::GpuTensor| match w {
13814 GpuTensor::Quant {
13815 bytes,
13816 qtype,
13817 scale,
13818 rp,
13819 ..
13820 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13821 _ => None,
13822 };
13823 let (Some(p0), Some(p1)) = (unpack(w0), unpack(w1)) else {
13824 return Ok(false);
13825 };
13826 let in_f = w0.in_features();
13827 if w1.in_features() != in_f {
13828 return Ok(false);
13829 }
13830 let (o0, o1) = (w0.out_features(), w1.out_features());
13831 if y0.len() < o0 || y1.len() < o1 {
13832 return Ok(false);
13833 }
13834 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13836 let rows_pb = ROWS_PER_BLOCK * RPW;
13837 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13838 let f = self.func("qmatvec_nvfp4_mmvq_fused2_rp");
13839 let cfg = LaunchConfig {
13840 grid_dim: (nb(o0) + nb(o1), 1, 1),
13841 block_dim: (32, ROWS_PER_BLOCK, 1),
13842 shared_mem_bytes: 0,
13843 };
13844 let (inf, oi0, oi1, mi) = (in_f as i32, o0 as i32, o1 as i32, 1i32);
13845 let (b0, b1) = unsafe { (&*p0.0, &*p1.0) };
13848 let __s_b = self.gpu.stream();
13849 let mut b = __s_b.launch_builder(&f);
13850 b.arg(b0)
13851 .arg(b1)
13852 .arg(aq)
13853 .arg(ad)
13854 .arg(&mut *y0)
13855 .arg(&mut *y1)
13856 .arg(&inf)
13857 .arg(&oi0)
13858 .arg(&oi1)
13859 .arg(&mi)
13860 .arg(&p0.1)
13861 .arg(&p1.1);
13862 unsafe {
13863 b.launch(cfg)?;
13864 }
13865 Ok(true)
13866 }
13867
13868 #[allow(clippy::type_complexity)]
13873 pub fn matmul_nvfp4_fused4(
13874 &self,
13875 w0: &crate::model::GpuTensor,
13876 w1: &crate::model::GpuTensor,
13877 w2: &crate::model::GpuTensor,
13878 w3: &crate::model::GpuTensor,
13879 aq: &CudaSlice<i8>,
13880 ad: &CudaSlice<f32>,
13881 m: usize,
13882 ) -> Result<
13883 Option<(
13884 CudaSlice<f32>,
13885 CudaSlice<f32>,
13886 CudaSlice<f32>,
13887 CudaSlice<f32>,
13888 )>,
13889 Box<dyn std::error::Error>,
13890 > {
13891 use crate::model::GpuTensor;
13892 if std::env::var("MEMRA_NVFP4_FUSED4").as_deref() == Ok("0")
13899 || !self.mmvq_supports(QT_NVFP4)
13900 || !self.uses_q8_1_fast(w0)
13901 || !self.uses_q8_1_fast(w1)
13902 || !self.uses_q8_1_fast(w2)
13903 || !self.uses_q8_1_fast(w3)
13904 {
13905 return Ok(None);
13906 }
13907 if (9..=16).contains(&m) {
13912 return Ok(
13913 match self.matmul_decode_exact_group4_pre([w0, w1, w2, w3], aq, ad, m)? {
13914 Some(mut ys) => {
13915 let y3 = ys.pop().unwrap();
13916 let y2 = ys.pop().unwrap();
13917 let y1 = ys.pop().unwrap();
13918 let y0 = ys.pop().unwrap();
13919 Some((y0, y1, y2, y3))
13920 }
13921 None => None,
13922 },
13923 );
13924 }
13925 if !(1..=8).contains(&m) {
13926 return Ok(None);
13927 }
13928 if m > 1 {
13929 let in_f = w0.in_features();
13932 if !self.batched_supports(QT_NVFP4)
13933 || std::env::var("MEMRA_NO_BATCHED").is_ok()
13934 || (m > 4 && !Self::b8_enabled())
13935 || in_f % 512 != 0
13936 || in_f / 64 > 272
13937 {
13938 return Ok(None);
13939 }
13940 }
13941 let unpack = |w: &crate::model::GpuTensor| match w {
13942 GpuTensor::Quant {
13943 bytes,
13944 qtype,
13945 scale,
13946 rp,
13947 ..
13948 } if *qtype == QT_NVFP4 && *rp => Some((bytes as *const CudaSlice<u8>, *scale)),
13949 _ => None,
13950 };
13951 let (Some(p0), Some(p1), Some(p2), Some(p3)) =
13952 (unpack(w0), unpack(w1), unpack(w2), unpack(w3))
13953 else {
13954 return Ok(None);
13955 };
13956 let in_f = w0.in_features();
13957 if w1.in_features() != in_f || w2.in_features() != in_f || w3.in_features() != in_f {
13958 return Ok(None);
13959 }
13960 let (o0, o1, o2, o3) = (
13961 w0.out_features(),
13962 w1.out_features(),
13963 w2.out_features(),
13964 w3.out_features(),
13965 );
13966 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
13968 let rows_pb = ROWS_PER_BLOCK * RPW;
13969 let nb = |o: usize| (o as u32).div_ceil(rows_pb);
13970 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
13971 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
13972 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
13973 let mut y3 = self.alloc_uninit::<f32>(m * o3)?;
13974 let (inf, oi0, oi1, oi2, oi3, mi) = (
13975 in_f as i32,
13976 o0 as i32,
13977 o1 as i32,
13978 o2 as i32,
13979 o3 as i32,
13980 m as i32,
13981 );
13982 let (b0, b1, b2, b3) = unsafe { (&*p0.0, &*p1.0, &*p2.0, &*p3.0) };
13985 if m > 1 {
13986 if p0.1 != 1.0 || p1.1 != 1.0 || p2.1 != 1.0 || p3.1 != 1.0 {
13989 return Ok(None);
13990 }
13991 let f = self.func("qmatvec_nvfp4_mmvq_fused4_b8_rpsc");
13992 let cfg = LaunchConfig {
13993 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), 1, 1),
13994 block_dim: (32, ROWS_PER_BLOCK, 1),
13995 shared_mem_bytes: 0,
13996 };
13997 let __s_b = self.gpu.stream();
13998 let mut b = __s_b.launch_builder(&f);
13999 b.arg(b0)
14000 .arg(b1)
14001 .arg(b2)
14002 .arg(b3)
14003 .arg(aq)
14004 .arg(ad)
14005 .arg(&mut y0)
14006 .arg(&mut y1)
14007 .arg(&mut y2)
14008 .arg(&mut y3)
14009 .arg(&inf)
14010 .arg(&oi0)
14011 .arg(&oi1)
14012 .arg(&oi2)
14013 .arg(&oi3)
14014 .arg(&mi);
14015 unsafe {
14016 b.launch(cfg)?;
14017 }
14018 return Ok(Some((y0, y1, y2, y3)));
14019 }
14020 let f = self.func("qmatvec_nvfp4_mmvq_fused4_rp");
14021 let cfg = LaunchConfig {
14022 grid_dim: (nb(o0) + nb(o1) + nb(o2) + nb(o3), m as u32, 1),
14023 block_dim: (32, ROWS_PER_BLOCK, 1),
14024 shared_mem_bytes: 0,
14025 };
14026 let __s_b = self.gpu.stream();
14027 let mut b = __s_b.launch_builder(&f);
14028 b.arg(b0)
14029 .arg(b1)
14030 .arg(b2)
14031 .arg(b3)
14032 .arg(aq)
14033 .arg(ad)
14034 .arg(&mut y0)
14035 .arg(&mut y1)
14036 .arg(&mut y2)
14037 .arg(&mut y3)
14038 .arg(&inf)
14039 .arg(&oi0)
14040 .arg(&oi1)
14041 .arg(&oi2)
14042 .arg(&oi3)
14043 .arg(&mi)
14044 .arg(&p0.1)
14045 .arg(&p1.1)
14046 .arg(&p2.1)
14047 .arg(&p3.1);
14048 unsafe {
14049 b.launch(cfg)?;
14050 }
14051 Ok(Some((y0, y1, y2, y3)))
14052 }
14053
14054 pub fn matmul_q8_fused2(
14062 &self,
14063 w0: &crate::model::GpuTensor,
14064 w1: &crate::model::GpuTensor,
14065 aq: &CudaSlice<i8>,
14066 ad: &CudaSlice<f32>,
14067 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14068 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14074 return Ok(Some(self.e4m3_fused2_core(
14075 p0.0,
14076 p1.0,
14077 aq,
14078 ad,
14079 w0.in_features(),
14080 p0.1,
14081 p1.1,
14082 p0.2,
14083 p0.3,
14084 p1.3,
14085 )?));
14086 }
14087 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14088 return Ok(None);
14089 };
14090 Ok(Some(self.q8_fused2_core(
14091 p0.0,
14092 p1.0,
14093 aq,
14094 ad,
14095 w0.in_features(),
14096 p0.1,
14097 p1.1,
14098 p0.2,
14099 )?))
14100 }
14101
14102 #[allow(clippy::too_many_arguments)]
14103 fn q8_fused2_core(
14104 &self,
14105 b0: &CudaSlice<u8>,
14106 b1: &CudaSlice<u8>,
14107 aq: &CudaSlice<i8>,
14108 ad: &CudaSlice<f32>,
14109 in_f: usize,
14110 out0: usize,
14111 out1: usize,
14112 row_bytes: usize,
14113 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14114 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14116 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14117 let f = self.func("qmatvec_q8_0_mmvq_fused2");
14118 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14119 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14120 let cfg = LaunchConfig {
14121 grid_dim: (nb0 + nb1, 1, 1),
14122 block_dim: (32, ROWS_PER_BLOCK, 1),
14123 shared_mem_bytes: 0,
14124 };
14125 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
14126 let __s_b = self.gpu.stream();
14127 let mut b = __s_b.launch_builder(&f);
14128 b.arg(b0)
14129 .arg(b1)
14130 .arg(aq)
14131 .arg(ad)
14132 .arg(&mut y0)
14133 .arg(&mut y1)
14134 .arg(&inf)
14135 .arg(&o0)
14136 .arg(&o1)
14137 .arg(&rbl);
14138 unsafe {
14139 b.launch(cfg)?;
14140 }
14141 Ok((y0, y1))
14142 }
14143
14144 pub fn matmul_q8_fused2_x(
14150 &self,
14151 w0: &crate::model::GpuTensor,
14152 w1: &crate::model::GpuTensor,
14153 x: &CudaSlice<f32>,
14154 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14155 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
14156 return Ok(None);
14157 }
14158 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
14159 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
14160 return Ok(Some(self.e4m3_fused2_core(
14161 p0.0,
14162 p1.0,
14163 &aq,
14164 &ad,
14165 w0.in_features(),
14166 p0.1,
14167 p1.1,
14168 p0.2,
14169 p0.3,
14170 p1.3,
14171 )?));
14172 }
14173 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
14174 return Ok(None);
14175 };
14176 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
14177 Ok(Some(self.q8_fused2_core(
14178 p0.0,
14179 p1.0,
14180 &aq,
14181 &ad,
14182 w0.in_features(),
14183 p0.1,
14184 p1.1,
14185 p0.2,
14186 )?))
14187 }
14188
14189 #[allow(clippy::too_many_arguments)]
14192 pub fn qmatvec_q8_fused2_raw(
14193 &self,
14194 b0: &CudaSlice<u8>,
14195 b1: &CudaSlice<u8>,
14196 x: &CudaSlice<f32>,
14197 in_f: usize,
14198 out0: usize,
14199 out1: usize,
14200 row_bytes: usize,
14201 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14202 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14203 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
14204 }
14205
14206 pub fn matmul_q4_fused3(
14212 &self,
14213 w0: &crate::model::GpuTensor,
14214 w1: &crate::model::GpuTensor,
14215 w2: &crate::model::GpuTensor,
14216 aq: &CudaSlice<i8>,
14217 ad: &CudaSlice<f32>,
14218 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14219 {
14220 use crate::model::GpuTensor;
14221 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14222 match w {
14223 GpuTensor::Quant {
14224 qtype, row_bytes, ..
14225 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14226 _ => None,
14227 }
14228 };
14229 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14230 return Ok(None);
14231 };
14232 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14233 return Ok(None);
14234 }
14235 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14239 match w {
14240 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14241 Some(m) => (m, true),
14242 None => (bytes, *rp),
14243 },
14244 _ => unreachable!(),
14245 }
14246 }
14247 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14248 if rp0 != rp1 || rp1 != rp2 {
14249 return Ok(None);
14250 }
14251 let rp = rp0;
14252 let rpb: u32 = 4;
14253 let mr1 = rp && Self::q40_mr1_on();
14257 let nb = |o: usize| {
14258 if mr1 {
14259 (o as u32).div_ceil(rpb)
14260 } else {
14261 (o as u32).div_ceil(2).div_ceil(rpb)
14262 }
14263 };
14264 let grid = nb(o0) + nb(o1) + nb(o2);
14265 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14266 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14267 let mut y2 = self.alloc_uninit::<f32>(o2)?;
14268 let f = self.func(if mr1 {
14269 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14270 } else if rp {
14271 "qmatvec_q4_0_mmvq_fused3_rp"
14272 } else {
14273 "qmatvec_q4_0_mmvq_fused3"
14274 });
14275 let cfg = LaunchConfig {
14276 grid_dim: (grid, 1, 1),
14277 block_dim: (32, rpb, 1),
14278 shared_mem_bytes: 0,
14279 };
14280 let inf = w0.in_features() as i32;
14281 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14282 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14283 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14286 {
14287 use cudarc::driver::{DevicePtr, DevicePtrMut};
14288 let s = &self.gpu.stream();
14289 let (p0, _g0) = b0.device_ptr(s);
14290 let (p1, _g1) = b1.device_ptr(s);
14291 let (p2, _g2) = b2.device_ptr(s);
14292 let (paq, _g3) = aq.device_ptr(s);
14293 let (pad, _g4) = ad.device_ptr(s);
14294 let (py0, _g5) = y0.device_ptr_mut(s);
14295 let (py1, _g6) = y1.device_ptr_mut(s);
14296 let (py2, _g7) = y2.device_ptr_mut(s);
14297 let mut ps = [
14298 &p0 as *const _ as *mut std::ffi::c_void,
14299 &p1 as *const _ as *mut _,
14300 &p2 as *const _ as *mut _,
14301 &paq as *const _ as *mut _,
14302 &pad as *const _ as *mut _,
14303 &py0 as *const _ as *mut _,
14304 &py1 as *const _ as *mut _,
14305 &py2 as *const _ as *mut _,
14306 &inf as *const _ as *mut _,
14307 &oo0 as *const _ as *mut _,
14308 &oo1 as *const _ as *mut _,
14309 &oo2 as *const _ as *mut _,
14310 &r0 as *const _ as *mut _,
14311 &r1 as *const _ as *mut _,
14312 &r2 as *const _ as *mut _,
14313 ];
14314 unsafe {
14315 self.launch_pdl(
14316 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14317 (grid, 1, 1),
14318 (32, rpb, 1),
14319 &mut ps,
14320 )?;
14321 }
14322 }
14323 return Ok(Some((y0, y1, y2)));
14324 }
14325 let __s_b = self.gpu.stream();
14326 let mut b = __s_b.launch_builder(&f);
14327 b.arg(b0)
14328 .arg(b1)
14329 .arg(b2)
14330 .arg(aq)
14331 .arg(ad)
14332 .arg(&mut y0)
14333 .arg(&mut y1)
14334 .arg(&mut y2)
14335 .arg(&inf)
14336 .arg(&oo0)
14337 .arg(&oo1)
14338 .arg(&oo2)
14339 .arg(&r0)
14340 .arg(&r1)
14341 .arg(&r2);
14342 unsafe {
14343 b.launch(cfg)?;
14344 }
14345 Ok(Some((y0, y1, y2)))
14346 }
14347
14348 #[allow(clippy::too_many_arguments)]
14351 pub fn matmul_q4_fused3_into(
14352 &self,
14353 w0: &crate::model::GpuTensor,
14354 w1: &crate::model::GpuTensor,
14355 w2: &crate::model::GpuTensor,
14356 aq: &CudaSlice<i8>,
14357 ad: &CudaSlice<f32>,
14358 y0: &mut CudaSlice<f32>,
14359 y1: &mut CudaSlice<f32>,
14360 y2: &mut CudaSlice<f32>,
14361 ) -> Result<bool, Box<dyn std::error::Error>> {
14362 use crate::model::GpuTensor;
14363 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14364 match w {
14365 GpuTensor::Quant {
14366 qtype, row_bytes, ..
14367 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14368 _ => None,
14369 }
14370 };
14371 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2)) else {
14372 return Ok(false);
14373 };
14374 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14375 return Ok(false);
14376 }
14377 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14378 match w {
14379 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14380 Some(m) => (m, true),
14381 None => (bytes, *rp),
14382 },
14383 _ => unreachable!(),
14384 }
14385 }
14386 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14387 if rp0 != rp1 || rp1 != rp2 {
14388 return Ok(false);
14389 }
14390 let rp = rp0;
14391 let rpb: u32 = 4;
14392 let mr1 = rp && Self::q40_mr1_on();
14393 let nb = |o: usize| {
14394 if mr1 {
14395 (o as u32).div_ceil(rpb)
14396 } else {
14397 (o as u32).div_ceil(2).div_ceil(rpb)
14398 }
14399 };
14400 let grid = nb(o0) + nb(o1) + nb(o2);
14401 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
14402 let f = self.func(if mr1 {
14403 "qmatvec_q4_0_mmvq_fused3_mr1_rp"
14404 } else if rp {
14405 "qmatvec_q4_0_mmvq_fused3_rp"
14406 } else {
14407 "qmatvec_q4_0_mmvq_fused3"
14408 });
14409 let cfg = LaunchConfig {
14410 grid_dim: (grid, 1, 1),
14411 block_dim: (32, rpb, 1),
14412 shared_mem_bytes: 0,
14413 };
14414 let inf = w0.in_features() as i32;
14415 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
14416 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
14417 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14419 use cudarc::driver::{DevicePtr, DevicePtrMut};
14420 let s = &self.gpu.stream();
14421 let (p0, _g0) = b0.device_ptr(s);
14422 let (p1, _g1) = b1.device_ptr(s);
14423 let (p2, _g2) = b2.device_ptr(s);
14424 let (paq, _g3) = aq.device_ptr(s);
14425 let (pad, _g4) = ad.device_ptr(s);
14426 let (py0, _g5) = y0.device_ptr_mut(s);
14427 let (py1, _g6) = y1.device_ptr_mut(s);
14428 let (py2, _g7) = y2.device_ptr_mut(s);
14429 let mut ps = [
14430 &p0 as *const _ as *mut std::ffi::c_void,
14431 &p1 as *const _ as *mut _,
14432 &p2 as *const _ as *mut _,
14433 &paq as *const _ as *mut _,
14434 &pad as *const _ as *mut _,
14435 &py0 as *const _ as *mut _,
14436 &py1 as *const _ as *mut _,
14437 &py2 as *const _ as *mut _,
14438 &inf as *const _ as *mut _,
14439 &oo0 as *const _ as *mut _,
14440 &oo1 as *const _ as *mut _,
14441 &oo2 as *const _ as *mut _,
14442 &r0 as *const _ as *mut _,
14443 &r1 as *const _ as *mut _,
14444 &r2 as *const _ as *mut _,
14445 ];
14446 unsafe {
14447 self.launch_pdl(
14448 "qmatvec_q4_0_mmvq_fused3_mr1_rp",
14449 (grid, 1, 1),
14450 (32, rpb, 1),
14451 &mut ps,
14452 )?;
14453 }
14454 return Ok(true);
14455 }
14456 let __s_b = self.gpu.stream();
14457 let mut b = __s_b.launch_builder(&f);
14458 b.arg(b0)
14459 .arg(b1)
14460 .arg(b2)
14461 .arg(aq)
14462 .arg(ad)
14463 .arg(&mut *y0)
14464 .arg(&mut *y1)
14465 .arg(&mut *y2)
14466 .arg(&inf)
14467 .arg(&oo0)
14468 .arg(&oo1)
14469 .arg(&oo2)
14470 .arg(&r0)
14471 .arg(&r1)
14472 .arg(&r2);
14473 unsafe {
14474 b.launch(cfg)?;
14475 }
14476 Ok(true)
14477 }
14478
14479 pub fn matmul_q4_fused2(
14481 &self,
14482 w0: &crate::model::GpuTensor,
14483 w1: &crate::model::GpuTensor,
14484 aq: &CudaSlice<i8>,
14485 ad: &CudaSlice<f32>,
14486 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14487 use crate::model::GpuTensor;
14488 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14489 match w {
14490 GpuTensor::Quant {
14491 qtype, row_bytes, ..
14492 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14493 _ => None,
14494 }
14495 };
14496 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14497 return Ok(None);
14498 };
14499 if w0.in_features() != w1.in_features() {
14500 return Ok(None);
14501 }
14502 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14504 match w {
14505 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14506 Some(m) => (m, true),
14507 None => (bytes, *rp),
14508 },
14509 _ => unreachable!(),
14510 }
14511 }
14512 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14513 if rp0 != rp1 {
14514 return Ok(None);
14515 }
14516 let rp = rp0;
14517 let rpb: u32 = 4;
14518 let mr1 = rp && Self::q40_mr1_on();
14520 let nb = |o: usize| {
14521 if mr1 {
14522 (o as u32).div_ceil(rpb)
14523 } else {
14524 (o as u32).div_ceil(2).div_ceil(rpb)
14525 }
14526 };
14527 let grid = nb(o0) + nb(o1);
14528 let mut y0 = self.alloc_uninit::<f32>(o0)?;
14529 let mut y1 = self.alloc_uninit::<f32>(o1)?;
14530 let f = self.func(if mr1 {
14531 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14532 } else if rp {
14533 "qmatvec_q4_0_mmvq_fused2_rp"
14534 } else {
14535 "qmatvec_q4_0_mmvq_fused2"
14536 });
14537 let cfg = LaunchConfig {
14538 grid_dim: (grid, 1, 1),
14539 block_dim: (32, rpb, 1),
14540 shared_mem_bytes: 0,
14541 };
14542 let inf = w0.in_features() as i32;
14543 let (oo0, oo1) = (o0 as i32, o1 as i32);
14544 let (r0, r1) = (rb0 as i64, rb1 as i64);
14545 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14547 {
14548 use cudarc::driver::{DevicePtr, DevicePtrMut};
14549 let s = &self.gpu.stream();
14550 let (p0, _g0) = b0.device_ptr(s);
14551 let (p1, _g1) = b1.device_ptr(s);
14552 let (paq, _g2) = aq.device_ptr(s);
14553 let (pad, _g3) = ad.device_ptr(s);
14554 let (py0, _g4) = y0.device_ptr_mut(s);
14555 let (py1, _g5) = y1.device_ptr_mut(s);
14556 let mut ps = [
14557 &p0 as *const _ as *mut std::ffi::c_void,
14558 &p1 as *const _ as *mut _,
14559 &paq as *const _ as *mut _,
14560 &pad as *const _ as *mut _,
14561 &py0 as *const _ as *mut _,
14562 &py1 as *const _ as *mut _,
14563 &inf as *const _ as *mut _,
14564 &oo0 as *const _ as *mut _,
14565 &oo1 as *const _ as *mut _,
14566 &r0 as *const _ as *mut _,
14567 &r1 as *const _ as *mut _,
14568 ];
14569 unsafe {
14570 self.launch_pdl(
14571 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14572 (grid, 1, 1),
14573 (32, rpb, 1),
14574 &mut ps,
14575 )?;
14576 }
14577 }
14578 return Ok(Some((y0, y1)));
14579 }
14580 let __s_b = self.gpu.stream();
14581 let mut b = __s_b.launch_builder(&f);
14582 b.arg(b0)
14583 .arg(b1)
14584 .arg(aq)
14585 .arg(ad)
14586 .arg(&mut y0)
14587 .arg(&mut y1)
14588 .arg(&inf)
14589 .arg(&oo0)
14590 .arg(&oo1)
14591 .arg(&r0)
14592 .arg(&r1);
14593 unsafe {
14594 b.launch(cfg)?;
14595 }
14596 Ok(Some((y0, y1)))
14597 }
14598
14599 pub fn matmul_q4_fused2_into(
14601 &self,
14602 w0: &crate::model::GpuTensor,
14603 w1: &crate::model::GpuTensor,
14604 aq: &CudaSlice<i8>,
14605 ad: &CudaSlice<f32>,
14606 y0: &mut CudaSlice<f32>,
14607 y1: &mut CudaSlice<f32>,
14608 ) -> Result<bool, Box<dyn std::error::Error>> {
14609 use crate::model::GpuTensor;
14610 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14611 match w {
14612 GpuTensor::Quant {
14613 qtype, row_bytes, ..
14614 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14615 _ => None,
14616 }
14617 };
14618 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else {
14619 return Ok(false);
14620 };
14621 if w0.in_features() != w1.in_features() {
14622 return Ok(false);
14623 }
14624 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14625 match w {
14626 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14627 Some(m) => (m, true),
14628 None => (bytes, *rp),
14629 },
14630 _ => unreachable!(),
14631 }
14632 }
14633 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14634 if rp0 != rp1 {
14635 return Ok(false);
14636 }
14637 let rp = rp0;
14638 let rpb: u32 = 4;
14639 let mr1 = rp && Self::q40_mr1_on();
14640 let nb = |o: usize| {
14641 if mr1 {
14642 (o as u32).div_ceil(rpb)
14643 } else {
14644 (o as u32).div_ceil(2).div_ceil(rpb)
14645 }
14646 };
14647 let grid = nb(o0) + nb(o1);
14648 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
14649 let f = self.func(if mr1 {
14650 "qmatvec_q4_0_mmvq_fused2_mr1_rp"
14651 } else if rp {
14652 "qmatvec_q4_0_mmvq_fused2_rp"
14653 } else {
14654 "qmatvec_q4_0_mmvq_fused2"
14655 });
14656 let cfg = LaunchConfig {
14657 grid_dim: (grid, 1, 1),
14658 block_dim: (32, rpb, 1),
14659 shared_mem_bytes: 0,
14660 };
14661 let inf = w0.in_features() as i32;
14662 let (oo0, oo1) = (o0 as i32, o1 as i32);
14663 let (r0, r1) = (rb0 as i64, rb1 as i64);
14664 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
14666 use cudarc::driver::{DevicePtr, DevicePtrMut};
14667 let s = &self.gpu.stream();
14668 let (p0, _g0) = b0.device_ptr(s);
14669 let (p1, _g1) = b1.device_ptr(s);
14670 let (paq, _g2) = aq.device_ptr(s);
14671 let (pad, _g3) = ad.device_ptr(s);
14672 let (py0, _g4) = y0.device_ptr_mut(s);
14673 let (py1, _g5) = y1.device_ptr_mut(s);
14674 let mut ps = [
14675 &p0 as *const _ as *mut std::ffi::c_void,
14676 &p1 as *const _ as *mut _,
14677 &paq as *const _ as *mut _,
14678 &pad as *const _ as *mut _,
14679 &py0 as *const _ as *mut _,
14680 &py1 as *const _ as *mut _,
14681 &inf as *const _ as *mut _,
14682 &oo0 as *const _ as *mut _,
14683 &oo1 as *const _ as *mut _,
14684 &r0 as *const _ as *mut _,
14685 &r1 as *const _ as *mut _,
14686 ];
14687 unsafe {
14688 self.launch_pdl(
14689 "qmatvec_q4_0_mmvq_fused2_mr1_rp",
14690 (grid, 1, 1),
14691 (32, rpb, 1),
14692 &mut ps,
14693 )?;
14694 }
14695 return Ok(true);
14696 }
14697 let __s_b = self.gpu.stream();
14698 let mut b = __s_b.launch_builder(&f);
14699 b.arg(b0)
14700 .arg(b1)
14701 .arg(aq)
14702 .arg(ad)
14703 .arg(&mut *y0)
14704 .arg(&mut *y1)
14705 .arg(&inf)
14706 .arg(&oo0)
14707 .arg(&oo1)
14708 .arg(&r0)
14709 .arg(&r1);
14710 unsafe {
14711 b.launch(cfg)?;
14712 }
14713 Ok(true)
14714 }
14715
14716 pub fn matmul_q4_fused2_batched(
14721 &self,
14722 w0: &crate::model::GpuTensor,
14723 w1: &crate::model::GpuTensor,
14724 aq: &CudaSlice<i8>,
14725 ad: &CudaSlice<f32>,
14726 m: usize,
14727 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
14728 use crate::model::GpuTensor;
14729 if m < 2 || m > 8 {
14730 return Ok(None);
14731 }
14732 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
14733 match w {
14734 GpuTensor::Quant {
14735 qtype, row_bytes, ..
14736 } if *qtype == QT_Q4_0 => Some((*row_bytes, w.out_features())),
14737 _ => None,
14738 }
14739 };
14740 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else {
14741 return Ok(None);
14742 };
14743 if w0.in_features() != w1.in_features() {
14744 return Ok(None);
14745 }
14746 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14747 match w {
14748 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14749 Some(mr) => (mr, true),
14750 None => (bytes, *rp),
14751 },
14752 _ => unreachable!(),
14753 }
14754 }
14755 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
14756 if !rp0 || !rp1 {
14757 return Ok(None);
14758 }
14759 let mcols = Self::batched_mcols(m);
14760 let rpb: u32 = 4;
14761 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14762 let grid = nb(o0) + nb(o1);
14763 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14764 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14765 let f = self.func(match mcols {
14766 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
14767 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
14768 _ => "qmatvec_q4_0_mmvq_b8_f2_rp",
14769 });
14770 let cfg = LaunchConfig {
14771 grid_dim: (grid, 1, 1),
14772 block_dim: (32, rpb, 1),
14773 shared_mem_bytes: 0,
14774 };
14775 let inf = w0.in_features() as i32;
14776 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
14777 let rb = rb0 as i64;
14778 let __s_b = self.gpu.stream();
14779 let mut b = __s_b.launch_builder(&f);
14780 b.arg(b0)
14781 .arg(b1)
14782 .arg(aq)
14783 .arg(ad)
14784 .arg(&mut y0)
14785 .arg(&mut y1)
14786 .arg(&inf)
14787 .arg(&oo0)
14788 .arg(&oo1)
14789 .arg(&mi)
14790 .arg(&rb);
14791 unsafe {
14792 b.launch(cfg)?;
14793 }
14794 Ok(Some((y0, y1)))
14795 }
14796
14797 #[allow(clippy::too_many_arguments)]
14800 pub fn matmul_q4_fused3_batched(
14801 &self,
14802 w0: &crate::model::GpuTensor,
14803 w1: &crate::model::GpuTensor,
14804 w2: &crate::model::GpuTensor,
14805 aq: &CudaSlice<i8>,
14806 ad: &CudaSlice<f32>,
14807 m: usize,
14808 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14809 {
14810 use crate::model::GpuTensor;
14811 if m < 2 || m > 8 {
14812 return Ok(None);
14813 }
14814 let q4 = |w: &GpuTensor| -> Option<usize> {
14815 match w {
14816 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
14817 _ => None,
14818 }
14819 };
14820 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else {
14821 return Ok(None);
14822 };
14823 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
14824 return Ok(None);
14825 }
14826 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
14827 match w {
14828 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
14829 Some(mr) => (mr, true),
14830 None => (bytes, *rp),
14831 },
14832 _ => unreachable!(),
14833 }
14834 }
14835 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
14836 if !rp0 || !rp1 || !rp2 {
14837 return Ok(None);
14838 }
14839 let mcols = Self::batched_mcols(m);
14840 let rpb: u32 = 4;
14841 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
14842 let grid = nb(o0) + nb(o1) + nb(o2);
14843 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
14844 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
14845 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
14846 let f = self.func(match mcols {
14847 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
14848 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
14849 _ => "qmatvec_q4_0_mmvq_b8_f3_rp",
14850 });
14851 let cfg = LaunchConfig {
14852 grid_dim: (grid, 1, 1),
14853 block_dim: (32, rpb, 1),
14854 shared_mem_bytes: 0,
14855 };
14856 let inf = w0.in_features() as i32;
14857 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
14858 let rb = 0i64;
14859 let __s_b = self.gpu.stream();
14860 let mut b = __s_b.launch_builder(&f);
14861 b.arg(b0)
14862 .arg(b1)
14863 .arg(b2)
14864 .arg(aq)
14865 .arg(ad)
14866 .arg(&mut y0)
14867 .arg(&mut y1)
14868 .arg(&mut y2)
14869 .arg(&inf)
14870 .arg(&oo0)
14871 .arg(&oo1)
14872 .arg(&oo2)
14873 .arg(&mi)
14874 .arg(&rb);
14875 unsafe {
14876 b.launch(cfg)?;
14877 }
14878 Ok(Some((y0, y1, y2)))
14879 }
14880
14881 pub fn matmul_q8_fused3(
14882 &self,
14883 w0: &crate::model::GpuTensor,
14884 w1: &crate::model::GpuTensor,
14885 w2: &crate::model::GpuTensor,
14886 aq: &CudaSlice<i8>,
14887 ad: &CudaSlice<f32>,
14888 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
14889 {
14890 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
14893 return Ok(Some(self.e4m3_fused3_core(
14894 p0.0,
14895 p1.0,
14896 p2.0,
14897 aq,
14898 ad,
14899 w0.in_features(),
14900 p0.1,
14901 p1.1,
14902 p2.1,
14903 p0.2,
14904 p0.3,
14905 p1.3,
14906 p2.3,
14907 )?));
14908 }
14909 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
14910 return Ok(None);
14911 };
14912 Ok(Some(self.q8_fused3_core(
14913 p0.0,
14914 p1.0,
14915 p2.0,
14916 aq,
14917 ad,
14918 w0.in_features(),
14919 p0.1,
14920 p1.1,
14921 p2.1,
14922 p0.2,
14923 )?))
14924 }
14925
14926 #[allow(clippy::too_many_arguments)]
14927 fn q8_fused3_core(
14928 &self,
14929 b0: &CudaSlice<u8>,
14930 b1: &CudaSlice<u8>,
14931 b2: &CudaSlice<u8>,
14932 aq: &CudaSlice<i8>,
14933 ad: &CudaSlice<f32>,
14934 in_f: usize,
14935 out0: usize,
14936 out1: usize,
14937 out2: usize,
14938 row_bytes: usize,
14939 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14940 const ROWS_PER_BLOCK: u32 = 4;
14941 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
14942 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
14943 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
14944 let f = self.func("qmatvec_q8_0_mmvq_fused3");
14945 let mut y0 = self.alloc_uninit::<f32>(out0)?;
14946 let mut y1 = self.alloc_uninit::<f32>(out1)?;
14947 let mut y2 = self.alloc_uninit::<f32>(out2)?;
14948 let cfg = LaunchConfig {
14949 grid_dim: (nb0 + nb1 + nb2, 1, 1),
14950 block_dim: (32, ROWS_PER_BLOCK, 1),
14951 shared_mem_bytes: 0,
14952 };
14953 let (inf, o0, o1, o2, rbl) = (
14954 in_f as i32,
14955 out0 as i32,
14956 out1 as i32,
14957 out2 as i32,
14958 row_bytes as i64,
14959 );
14960 let __s_b = self.gpu.stream();
14961 let mut b = __s_b.launch_builder(&f);
14962 b.arg(b0)
14963 .arg(b1)
14964 .arg(b2)
14965 .arg(aq)
14966 .arg(ad)
14967 .arg(&mut y0)
14968 .arg(&mut y1)
14969 .arg(&mut y2)
14970 .arg(&inf)
14971 .arg(&o0)
14972 .arg(&o1)
14973 .arg(&o2)
14974 .arg(&rbl);
14975 unsafe {
14976 b.launch(cfg)?;
14977 }
14978 Ok((y0, y1, y2))
14979 }
14980
14981 #[allow(clippy::too_many_arguments)]
14983 pub fn qmatvec_q8_fused3_raw(
14984 &self,
14985 b0: &CudaSlice<u8>,
14986 b1: &CudaSlice<u8>,
14987 b2: &CudaSlice<u8>,
14988 x: &CudaSlice<f32>,
14989 in_f: usize,
14990 out0: usize,
14991 out1: usize,
14992 out2: usize,
14993 row_bytes: usize,
14994 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
14995 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
14996 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
14997 }
14998
14999 pub fn matmul_q8_fused2_t(
15010 &self,
15011 w0: &crate::model::GpuTensor,
15012 w1: &crate::model::GpuTensor,
15013 aq: &CudaSlice<i8>,
15014 ad: &CudaSlice<f32>,
15015 m: usize,
15016 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
15017 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
15021 return Ok(None);
15022 }
15023 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
15026 if m > 4 && !Self::b8_enabled() {
15027 return Ok(None);
15028 }
15029 return Ok(Some(self.e4m3_fused2_t_core(
15030 p0.0,
15031 p1.0,
15032 aq,
15033 ad,
15034 m,
15035 w0.in_features(),
15036 p0.1,
15037 p1.1,
15038 p0.2,
15039 p0.3,
15040 p1.3,
15041 )?));
15042 }
15043 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else {
15044 return Ok(None);
15045 };
15046 Ok(Some(self.q8_fused2_t_core(
15047 p0.0,
15048 p1.0,
15049 aq,
15050 ad,
15051 m,
15052 w0.in_features(),
15053 p0.1,
15054 p1.1,
15055 p0.2,
15056 )?))
15057 }
15058
15059 #[allow(clippy::too_many_arguments)]
15060 fn q8_fused2_t_core(
15061 &self,
15062 b0: &CudaSlice<u8>,
15063 b1: &CudaSlice<u8>,
15064 aq: &CudaSlice<i8>,
15065 ad: &CudaSlice<f32>,
15066 m: usize,
15067 in_f: usize,
15068 out0: usize,
15069 out1: usize,
15070 row_bytes: usize,
15071 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15072 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15074 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15075 let f = self.func(match Self::batched_mcols(m) {
15076 2 => "qmatvec_q8_0_mmvq_fused2_b2",
15077 4 => "qmatvec_q8_0_mmvq_fused2_b4",
15078 _ => "qmatvec_q8_0_mmvq_fused2_b8",
15080 });
15081 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15082 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15083 let cfg = LaunchConfig {
15084 grid_dim: (nb0 + nb1, 1, 1),
15085 block_dim: (32, ROWS_PER_BLOCK, 1),
15086 shared_mem_bytes: 0,
15087 };
15088 let (inf, o0, o1, mi, rbl) = (
15089 in_f as i32,
15090 out0 as i32,
15091 out1 as i32,
15092 m as i32,
15093 row_bytes as i64,
15094 );
15095 let __s_b = self.gpu.stream();
15096 let mut b = __s_b.launch_builder(&f);
15097 b.arg(b0)
15098 .arg(b1)
15099 .arg(aq)
15100 .arg(ad)
15101 .arg(&mut y0)
15102 .arg(&mut y1)
15103 .arg(&inf)
15104 .arg(&o0)
15105 .arg(&o1)
15106 .arg(&mi)
15107 .arg(&rbl);
15108 unsafe {
15109 b.launch(cfg)?;
15110 }
15111 Ok((y0, y1))
15112 }
15113
15114 #[allow(clippy::too_many_arguments)]
15117 pub fn qmatvec_q8_fused2_t_raw(
15118 &self,
15119 b0: &CudaSlice<u8>,
15120 b1: &CudaSlice<u8>,
15121 x: &CudaSlice<f32>,
15122 m: usize,
15123 in_f: usize,
15124 out0: usize,
15125 out1: usize,
15126 row_bytes: usize,
15127 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15128 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15129 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
15130 }
15131
15132 #[allow(clippy::too_many_arguments)]
15135 pub fn matmul_q8_fused3_t(
15136 &self,
15137 w0: &crate::model::GpuTensor,
15138 w1: &crate::model::GpuTensor,
15139 w2: &crate::model::GpuTensor,
15140 aq: &CudaSlice<i8>,
15141 ad: &CudaSlice<f32>,
15142 m: usize,
15143 ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
15144 {
15145 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() {
15146 return Ok(None);
15147 }
15148 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
15149 return Ok(Some(self.e4m3_fused3_t_core(
15150 p0.0,
15151 p1.0,
15152 p2.0,
15153 aq,
15154 ad,
15155 m,
15156 w0.in_features(),
15157 p0.1,
15158 p1.1,
15159 p2.1,
15160 p0.2,
15161 p0.3,
15162 p1.3,
15163 p2.3,
15164 )?));
15165 }
15166 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else {
15167 return Ok(None);
15168 };
15169 Ok(Some(self.q8_fused3_t_core(
15170 p0.0,
15171 p1.0,
15172 p2.0,
15173 aq,
15174 ad,
15175 m,
15176 w0.in_features(),
15177 p0.1,
15178 p1.1,
15179 p2.1,
15180 p0.2,
15181 )?))
15182 }
15183
15184 #[allow(clippy::too_many_arguments)]
15185 fn q8_fused3_t_core(
15186 &self,
15187 b0: &CudaSlice<u8>,
15188 b1: &CudaSlice<u8>,
15189 b2: &CudaSlice<u8>,
15190 aq: &CudaSlice<i8>,
15191 ad: &CudaSlice<f32>,
15192 m: usize,
15193 in_f: usize,
15194 out0: usize,
15195 out1: usize,
15196 out2: usize,
15197 row_bytes: usize,
15198 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15199 const ROWS_PER_BLOCK: u32 = 4;
15200 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15201 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15202 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15203 let f = self.func(if Self::batched_mcols(m) == 2 {
15204 "qmatvec_q8_0_mmvq_fused3_b2"
15205 } else {
15206 "qmatvec_q8_0_mmvq_fused3_b4"
15207 });
15208 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15209 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15210 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15211 let cfg = LaunchConfig {
15212 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15213 block_dim: (32, ROWS_PER_BLOCK, 1),
15214 shared_mem_bytes: 0,
15215 };
15216 let (inf, o0, o1, o2, mi, rbl) = (
15217 in_f as i32,
15218 out0 as i32,
15219 out1 as i32,
15220 out2 as i32,
15221 m as i32,
15222 row_bytes as i64,
15223 );
15224 let __s_b = self.gpu.stream();
15225 let mut b = __s_b.launch_builder(&f);
15226 b.arg(b0)
15227 .arg(b1)
15228 .arg(b2)
15229 .arg(aq)
15230 .arg(ad)
15231 .arg(&mut y0)
15232 .arg(&mut y1)
15233 .arg(&mut y2)
15234 .arg(&inf)
15235 .arg(&o0)
15236 .arg(&o1)
15237 .arg(&o2)
15238 .arg(&mi)
15239 .arg(&rbl);
15240 unsafe {
15241 b.launch(cfg)?;
15242 }
15243 Ok((y0, y1, y2))
15244 }
15245
15246 #[allow(clippy::too_many_arguments)]
15248 pub fn qmatvec_q8_fused3_t_raw(
15249 &self,
15250 b0: &CudaSlice<u8>,
15251 b1: &CudaSlice<u8>,
15252 b2: &CudaSlice<u8>,
15253 x: &CudaSlice<f32>,
15254 m: usize,
15255 in_f: usize,
15256 out0: usize,
15257 out1: usize,
15258 out2: usize,
15259 row_bytes: usize,
15260 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15261 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15262 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
15263 }
15264
15265 pub fn q8_ffn_fuse2_on(&self) -> bool {
15269 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15270 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
15271 }
15272
15273 #[allow(clippy::type_complexity)]
15279 fn q8_fused_params<'w, const N: usize>(
15280 &self,
15281 ws: &[&'w crate::model::GpuTensor; N],
15282 ) -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
15283 use crate::model::GpuTensor;
15284 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
15285 return None;
15286 }
15287 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") {
15288 return None;
15289 }
15290 let in_f = ws[0].in_features();
15291 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
15292 for (i, w) in ws.iter().enumerate() {
15293 match w {
15294 GpuTensor::Quant {
15295 bytes,
15296 qtype,
15297 row_bytes,
15298 scale,
15299 ..
15300 } if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f => {
15301 out[i] = Some((bytes, w.out_features(), *row_bytes))
15302 }
15303 _ => return None,
15304 }
15305 }
15306 Some(out.map(|o| o.unwrap()))
15307 }
15308
15309 pub fn e4m3_dual_on(&self) -> bool {
15312 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
15313 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
15314 }
15315
15316 #[allow(clippy::type_complexity)]
15328 fn e4m3_fused_params<'w, const N: usize>(
15329 &self,
15330 ws: &[&'w crate::model::GpuTensor; N],
15331 ) -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
15332 use crate::model::GpuTensor;
15333 if !self.e4m3_dual_on() {
15334 return None;
15335 }
15336 let in_f = ws[0].in_features();
15337 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
15338 for (i, w) in ws.iter().enumerate() {
15339 match w {
15340 GpuTensor::Quant {
15341 bytes,
15342 qtype,
15343 row_bytes,
15344 scale,
15345 rp,
15346 rp4,
15347 ..
15348 } if *qtype == QT_F8_E4M3
15349 && w.in_features() == in_f
15350 && *row_bytes == in_f
15351 && !*rp
15352 && rp4.is_none() =>
15353 {
15354 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale))
15355 }
15356 _ => return None,
15357 }
15358 }
15359 Some(out.map(|o| o.unwrap()))
15360 }
15361
15362 #[allow(clippy::too_many_arguments)]
15366 fn e4m3_fused2_core(
15367 &self,
15368 b0: &CudaSlice<u8>,
15369 b1: &CudaSlice<u8>,
15370 aq: &CudaSlice<i8>,
15371 ad: &CudaSlice<f32>,
15372 in_f: usize,
15373 out0: usize,
15374 out1: usize,
15375 row_bytes: usize,
15376 ws0: f32,
15377 ws1: f32,
15378 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15379 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15381 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15382 let f = self.func("qmatvec_e4m3_mmvq_fused2");
15383 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15384 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15385 let cfg = LaunchConfig {
15386 grid_dim: (nb0 + nb1, 1, 1),
15387 block_dim: (32, ROWS_PER_BLOCK, 1),
15388 shared_mem_bytes: 0,
15389 };
15390 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
15391 let __s_b = self.gpu.stream();
15392 let mut b = __s_b.launch_builder(&f);
15393 b.arg(b0)
15394 .arg(b1)
15395 .arg(aq)
15396 .arg(ad)
15397 .arg(&mut y0)
15398 .arg(&mut y1)
15399 .arg(&inf)
15400 .arg(&o0)
15401 .arg(&o1)
15402 .arg(&rbl)
15403 .arg(&ws0)
15404 .arg(&ws1);
15405 unsafe {
15406 b.launch(cfg)?;
15407 }
15408 Ok((y0, y1))
15409 }
15410
15411 #[allow(clippy::too_many_arguments)]
15413 fn e4m3_fused3_core(
15414 &self,
15415 b0: &CudaSlice<u8>,
15416 b1: &CudaSlice<u8>,
15417 b2: &CudaSlice<u8>,
15418 aq: &CudaSlice<i8>,
15419 ad: &CudaSlice<f32>,
15420 in_f: usize,
15421 out0: usize,
15422 out1: usize,
15423 out2: usize,
15424 row_bytes: usize,
15425 ws0: f32,
15426 ws1: f32,
15427 ws2: f32,
15428 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15429 const ROWS_PER_BLOCK: u32 = 4;
15430 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15431 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15432 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15433 let f = self.func("qmatvec_e4m3_mmvq_fused3");
15434 let mut y0 = self.alloc_uninit::<f32>(out0)?;
15435 let mut y1 = self.alloc_uninit::<f32>(out1)?;
15436 let mut y2 = self.alloc_uninit::<f32>(out2)?;
15437 let cfg = LaunchConfig {
15438 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15439 block_dim: (32, ROWS_PER_BLOCK, 1),
15440 shared_mem_bytes: 0,
15441 };
15442 let (inf, o0, o1, o2, rbl) = (
15443 in_f as i32,
15444 out0 as i32,
15445 out1 as i32,
15446 out2 as i32,
15447 row_bytes as i64,
15448 );
15449 let __s_b = self.gpu.stream();
15450 let mut b = __s_b.launch_builder(&f);
15451 b.arg(b0)
15452 .arg(b1)
15453 .arg(b2)
15454 .arg(aq)
15455 .arg(ad)
15456 .arg(&mut y0)
15457 .arg(&mut y1)
15458 .arg(&mut y2)
15459 .arg(&inf)
15460 .arg(&o0)
15461 .arg(&o1)
15462 .arg(&o2)
15463 .arg(&rbl)
15464 .arg(&ws0)
15465 .arg(&ws1)
15466 .arg(&ws2);
15467 unsafe {
15468 b.launch(cfg)?;
15469 }
15470 Ok((y0, y1, y2))
15471 }
15472
15473 #[allow(clippy::too_many_arguments)]
15477 fn e4m3_fused2_t_core(
15478 &self,
15479 b0: &CudaSlice<u8>,
15480 b1: &CudaSlice<u8>,
15481 aq: &CudaSlice<i8>,
15482 ad: &CudaSlice<f32>,
15483 m: usize,
15484 in_f: usize,
15485 out0: usize,
15486 out1: usize,
15487 row_bytes: usize,
15488 ws0: f32,
15489 ws1: f32,
15490 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15491 const ROWS_PER_BLOCK: u32 = 4;
15492 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15493 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15494 let f = self.func(match Self::batched_mcols(m) {
15495 2 => "qmatvec_e4m3_mmvq_fused2_b2",
15496 4 => "qmatvec_e4m3_mmvq_fused2_b4",
15497 _ => "qmatvec_e4m3_mmvq_fused2_b8",
15498 });
15499 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15500 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15501 let cfg = LaunchConfig {
15502 grid_dim: (nb0 + nb1, 1, 1),
15503 block_dim: (32, ROWS_PER_BLOCK, 1),
15504 shared_mem_bytes: 0,
15505 };
15506 let (inf, o0, o1, mi, rbl) = (
15507 in_f as i32,
15508 out0 as i32,
15509 out1 as i32,
15510 m as i32,
15511 row_bytes as i64,
15512 );
15513 let __s_b = self.gpu.stream();
15514 let mut b = __s_b.launch_builder(&f);
15515 b.arg(b0)
15516 .arg(b1)
15517 .arg(aq)
15518 .arg(ad)
15519 .arg(&mut y0)
15520 .arg(&mut y1)
15521 .arg(&inf)
15522 .arg(&o0)
15523 .arg(&o1)
15524 .arg(&mi)
15525 .arg(&rbl);
15526 unsafe {
15527 b.launch(cfg)?;
15528 }
15529 if ws0 != 1.0 {
15530 self.scale_inplace(&mut y0, ws0, m * out0)?;
15531 }
15532 if ws1 != 1.0 {
15533 self.scale_inplace(&mut y1, ws1, m * out1)?;
15534 }
15535 Ok((y0, y1))
15536 }
15537
15538 #[allow(clippy::too_many_arguments)]
15540 fn e4m3_fused3_t_core(
15541 &self,
15542 b0: &CudaSlice<u8>,
15543 b1: &CudaSlice<u8>,
15544 b2: &CudaSlice<u8>,
15545 aq: &CudaSlice<i8>,
15546 ad: &CudaSlice<f32>,
15547 m: usize,
15548 in_f: usize,
15549 out0: usize,
15550 out1: usize,
15551 out2: usize,
15552 row_bytes: usize,
15553 ws0: f32,
15554 ws1: f32,
15555 ws2: f32,
15556 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15557 const ROWS_PER_BLOCK: u32 = 4;
15558 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
15559 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
15560 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
15561 let f = self.func(if Self::batched_mcols(m) == 2 {
15562 "qmatvec_e4m3_mmvq_fused3_b2"
15563 } else {
15564 "qmatvec_e4m3_mmvq_fused3_b4"
15565 });
15566 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
15567 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
15568 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
15569 let cfg = LaunchConfig {
15570 grid_dim: (nb0 + nb1 + nb2, 1, 1),
15571 block_dim: (32, ROWS_PER_BLOCK, 1),
15572 shared_mem_bytes: 0,
15573 };
15574 let (inf, o0, o1, o2, mi, rbl) = (
15575 in_f as i32,
15576 out0 as i32,
15577 out1 as i32,
15578 out2 as i32,
15579 m as i32,
15580 row_bytes as i64,
15581 );
15582 let __s_b = self.gpu.stream();
15583 let mut b = __s_b.launch_builder(&f);
15584 b.arg(b0)
15585 .arg(b1)
15586 .arg(b2)
15587 .arg(aq)
15588 .arg(ad)
15589 .arg(&mut y0)
15590 .arg(&mut y1)
15591 .arg(&mut y2)
15592 .arg(&inf)
15593 .arg(&o0)
15594 .arg(&o1)
15595 .arg(&o2)
15596 .arg(&mi)
15597 .arg(&rbl);
15598 unsafe {
15599 b.launch(cfg)?;
15600 }
15601 if ws0 != 1.0 {
15602 self.scale_inplace(&mut y0, ws0, m * out0)?;
15603 }
15604 if ws1 != 1.0 {
15605 self.scale_inplace(&mut y1, ws1, m * out1)?;
15606 }
15607 if ws2 != 1.0 {
15608 self.scale_inplace(&mut y2, ws2, m * out2)?;
15609 }
15610 Ok((y0, y1, y2))
15611 }
15612
15613 pub fn qmatvec_e4m3_blk_mmvq(
15623 &self,
15624 bytes: &CudaSlice<u8>,
15625 aq: &CudaSlice<i8>,
15626 ad: &CudaSlice<f32>,
15627 scales: &CudaSlice<f32>,
15628 m: usize,
15629 in_f: usize,
15630 out_f: usize,
15631 row_bytes: usize,
15632 scale_cols: usize,
15633 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15634 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_e4m3_blk_mmvq_into(
15636 bytes, aq, ad, scales, m, in_f, out_f, row_bytes, scale_cols, &mut y,
15637 )?;
15638 Ok(y)
15639 }
15640
15641 #[allow(clippy::too_many_arguments)]
15643 pub fn qmatvec_e4m3_blk_mmvq_into(
15644 &self,
15645 bytes: &CudaSlice<u8>,
15646 aq: &CudaSlice<i8>,
15647 ad: &CudaSlice<f32>,
15648 scales: &CudaSlice<f32>,
15649 m: usize,
15650 in_f: usize,
15651 out_f: usize,
15652 row_bytes: usize,
15653 scale_cols: usize,
15654 y: &mut CudaSlice<f32>,
15655 ) -> Result<(), Box<dyn std::error::Error>> {
15656 const ROWS_PER_BLOCK: u32 = 4; let f = self.func("qmatvec_e4m3_blk_mmvq");
15658 let cfg = LaunchConfig {
15659 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), m as u32, 1),
15660 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
15663 let (inf, outf, mi, rb, sc) = (
15664 in_f as i32,
15665 out_f as i32,
15666 m as i32,
15667 row_bytes as i64,
15668 scale_cols as i32,
15669 );
15670 let __s_b = self.gpu.stream();
15671 let mut b = __s_b.launch_builder(&f);
15672 b.arg(bytes)
15673 .arg(aq)
15674 .arg(ad)
15675 .arg(scales)
15676 .arg(&mut *y)
15677 .arg(&inf)
15678 .arg(&outf)
15679 .arg(&mi)
15680 .arg(&rb)
15681 .arg(&sc);
15682 unsafe {
15683 b.launch(cfg)?;
15684 }
15685 Ok(())
15686 }
15687
15688 #[allow(clippy::too_many_arguments)]
15694 pub fn qmatvec_e4m3_blk_mmvq_batched(
15695 &self,
15696 bytes: &CudaSlice<u8>,
15697 aq: &CudaSlice<i8>,
15698 ad: &CudaSlice<f32>,
15699 scales: &CudaSlice<f32>,
15700 m: usize,
15701 in_f: usize,
15702 out_f: usize,
15703 row_bytes: usize,
15704 scale_cols: usize,
15705 mcols: usize,
15706 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15707 const ROWS_PER_BLOCK: u32 = 4; debug_assert!(mcols >= m, "blk batched: mcols {mcols} < m {m}");
15709 let name = match mcols {
15710 2 => "qmatvec_e4m3_blk_mmvq_b2",
15711 4 => "qmatvec_e4m3_blk_mmvq_b4",
15712 8 => "qmatvec_e4m3_blk_mmvq_b8",
15713 16 => "qmatvec_e4m3_blk_mmvq_b16",
15714 _ => {
15715 return Err(
15716 format!("qmatvec_e4m3_blk_mmvq_batched: no kernel for mcols {mcols}").into(),
15717 );
15718 }
15719 };
15720 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
15721 let f = self.func(name);
15722 let cfg = LaunchConfig {
15723 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
15724 block_dim: (32, ROWS_PER_BLOCK, 1),
15725 shared_mem_bytes: 0,
15726 };
15727 let (inf, outf, mi, rb, sc) = (
15728 in_f as i32,
15729 out_f as i32,
15730 m as i32,
15731 row_bytes as i64,
15732 scale_cols as i32,
15733 );
15734 let __s_b = self.gpu.stream();
15735 let mut b = __s_b.launch_builder(&f);
15736 b.arg(bytes)
15737 .arg(aq)
15738 .arg(ad)
15739 .arg(scales)
15740 .arg(&mut y)
15741 .arg(&inf)
15742 .arg(&outf)
15743 .arg(&mi)
15744 .arg(&rb)
15745 .arg(&sc);
15746 unsafe {
15747 b.launch(cfg)?;
15748 }
15749 Ok(y)
15750 }
15751
15752 #[allow(clippy::too_many_arguments)]
15755 pub fn qmatvec_e4m3_blk_batched_raw(
15756 &self,
15757 bytes: &CudaSlice<u8>,
15758 x: &CudaSlice<f32>,
15759 scales: &CudaSlice<f32>,
15760 m: usize,
15761 in_f: usize,
15762 out_f: usize,
15763 row_bytes: usize,
15764 scale_cols: usize,
15765 mcols: usize,
15766 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15767 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15768 self.qmatvec_e4m3_blk_mmvq_batched(
15769 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols, mcols,
15770 )
15771 }
15772
15773 #[allow(clippy::too_many_arguments)]
15776 pub fn qmatvec_e4m3_blk_mmvq_raw(
15777 &self,
15778 bytes: &CudaSlice<u8>,
15779 x: &CudaSlice<f32>,
15780 scales: &CudaSlice<f32>,
15781 m: usize,
15782 in_f: usize,
15783 out_f: usize,
15784 row_bytes: usize,
15785 scale_cols: usize,
15786 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
15787 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15788 self.qmatvec_e4m3_blk_mmvq(
15789 bytes, &aq, &ad, scales, m, in_f, out_f, row_bytes, scale_cols,
15790 )
15791 }
15792
15793 #[allow(clippy::too_many_arguments)]
15796 pub fn qmatvec_e4m3_fused2_raw(
15797 &self,
15798 b0: &CudaSlice<u8>,
15799 b1: &CudaSlice<u8>,
15800 x: &CudaSlice<f32>,
15801 in_f: usize,
15802 out0: usize,
15803 out1: usize,
15804 row_bytes: usize,
15805 ws0: f32,
15806 ws1: f32,
15807 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15808 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15809 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
15810 }
15811
15812 #[allow(clippy::too_many_arguments)]
15813 pub fn qmatvec_e4m3_fused3_raw(
15814 &self,
15815 b0: &CudaSlice<u8>,
15816 b1: &CudaSlice<u8>,
15817 b2: &CudaSlice<u8>,
15818 x: &CudaSlice<f32>,
15819 in_f: usize,
15820 out0: usize,
15821 out1: usize,
15822 out2: usize,
15823 row_bytes: usize,
15824 ws0: f32,
15825 ws1: f32,
15826 ws2: f32,
15827 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15828 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
15829 self.e4m3_fused3_core(
15830 b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15831 )
15832 }
15833
15834 #[allow(clippy::too_many_arguments)]
15835 pub fn qmatvec_e4m3_fused2_t_raw(
15836 &self,
15837 b0: &CudaSlice<u8>,
15838 b1: &CudaSlice<u8>,
15839 x: &CudaSlice<f32>,
15840 m: usize,
15841 in_f: usize,
15842 out0: usize,
15843 out1: usize,
15844 row_bytes: usize,
15845 ws0: f32,
15846 ws1: f32,
15847 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15848 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15849 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
15850 }
15851
15852 #[allow(clippy::too_many_arguments)]
15853 pub fn qmatvec_e4m3_fused3_t_raw(
15854 &self,
15855 b0: &CudaSlice<u8>,
15856 b1: &CudaSlice<u8>,
15857 b2: &CudaSlice<u8>,
15858 x: &CudaSlice<f32>,
15859 m: usize,
15860 in_f: usize,
15861 out0: usize,
15862 out1: usize,
15863 out2: usize,
15864 row_bytes: usize,
15865 ws0: f32,
15866 ws1: f32,
15867 ws2: f32,
15868 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
15869 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
15870 self.e4m3_fused3_t_core(
15871 b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes, ws0, ws1, ws2,
15872 )
15873 }
15874
15875 fn try_e4m3_blk_pre(
15886 &self,
15887 w: &crate::model::GpuTensor,
15888 aq: &CudaSlice<i8>,
15889 ad: &CudaSlice<f32>,
15890 m: usize,
15891 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15892 use crate::model::GpuTensor;
15893 if let GpuTensor::Quant {
15894 bytes,
15895 qtype,
15896 row_bytes,
15897 blk: Some(g),
15898 ..
15899 } = w
15900 {
15901 if *qtype == QT_F8_E4M3_BLK {
15902 if (2..=16).contains(&m)
15908 && std::env::var("MEMRA_NO_BATCHED").is_err()
15909 && (m <= 4 || Self::b8_enabled())
15910 {
15911 let mcols = Self::batched_mcols(m);
15912 return Ok(Some(self.qmatvec_e4m3_blk_mmvq_batched(
15913 bytes,
15914 aq,
15915 ad,
15916 &g.scales,
15917 m,
15918 w.in_features(),
15919 w.out_features(),
15920 *row_bytes,
15921 g.cols,
15922 mcols,
15923 )?));
15924 }
15925 return Ok(Some(self.qmatvec_e4m3_blk_mmvq(
15926 bytes,
15927 aq,
15928 ad,
15929 &g.scales,
15930 m,
15931 w.in_features(),
15932 w.out_features(),
15933 *row_bytes,
15934 g.cols,
15935 )?));
15936 }
15937 }
15938 Ok(None)
15939 }
15940
15941 fn try_e4m3_blk_prefill(
15988 &self,
15989 w: &crate::model::GpuTensor,
15990 x: &CudaSlice<f32>,
15991 m: usize,
15992 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
15993 use crate::model::GpuTensor;
15994 let GpuTensor::Quant {
15995 bytes,
15996 qtype,
15997 blk: Some(g),
15998 ..
15999 } = w
16000 else {
16001 return Ok(None);
16002 };
16003 if *qtype != QT_F8_E4M3_BLK {
16004 return Ok(None);
16005 }
16006 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? {
16011 return Ok(Some(y));
16012 }
16013 let (in_f, out_f) = (w.in_features(), w.out_features());
16014 let slab = self.fp8_blk_dequant_q8_0_dev(bytes, &g.scales, out_f, in_f)?;
16015 let tmp = GpuTensor::Quant {
16016 bytes: slab,
16017 qtype: QT_Q8_0,
16018 row_bytes: in_f / 32 * 34,
16019 ne: vec![in_f as u64, out_f as u64],
16020 scale: 1.0,
16021 rp: false,
16022 #[cfg(memra_cutlass)]
16023 cutlass: None,
16024 fp8: None,
16025 blk: None,
16026 f16: None,
16027 rp4: None,
16028 };
16029 Ok(Some(self.matmul(&tmp, x, m)?))
16031 }
16032
16033 pub fn matmul_pre_noscale(
16034 &self,
16035 w: &crate::model::GpuTensor,
16036 aq: &CudaSlice<i8>,
16037 ad: &CudaSlice<f32>,
16038 m: usize,
16039 ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
16040 use crate::model::GpuTensor;
16041 if m == 1 {
16045 if let Some(y) = self.try_e4m3_blk_pre(w, aq, ad, m)? {
16046 return Ok(Some((y, 1.0)));
16047 }
16048 }
16049 if m != 1 || !self.uses_q8_1_fast(w) {
16051 return Ok(None);
16052 }
16053 let in_f = w.in_features();
16054 let out_f = w.out_features();
16055 let (bytes, qtype, row_bytes, scale, rp) = match w {
16056 GpuTensor::Quant {
16057 bytes,
16058 qtype,
16059 row_bytes,
16060 scale,
16061 rp,
16062 ..
16063 } => (bytes, *qtype, *row_bytes, *scale, *rp),
16064 _ => return Ok(None),
16065 };
16066 if self.mmvq_supports(qtype) {
16068 let (mbytes, mrp) = match w {
16070 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
16071 _ => (bytes, rp),
16072 };
16073 let y = self.qmatvec_mmvq(
16074 mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp,
16075 )?;
16076 return Ok(Some((y, scale)));
16077 }
16078 let name = match qtype {
16080 QT_Q8_0 => "qmatvec_q8_0_dp4a",
16081 QT_Q4_K => "qmatvec_q4_K_dp4a",
16082 QT_Q6_K => "qmatvec_q6_K_dp4a",
16083 QT_Q5_K => "qmatvec_q5_K_dp4a",
16084 QT_Q3_K => "qmatvec_q3_K_dp4a",
16085 QT_NVFP4 => {
16086 if rp {
16087 "qmatvec_nvfp4_dp4a_rp"
16088 } else {
16089 "qmatvec_nvfp4_dp4a"
16090 }
16091 }
16092 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
16093 _ => return Ok(None),
16094 };
16095 let f = self.func(name);
16096 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16097 let cfg = LaunchConfig {
16098 grid_dim: (out_f as u32, m as u32, 1),
16099 block_dim: (128, 1, 1),
16100 shared_mem_bytes: 0,
16101 };
16102 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16103 let __s_b = self.gpu.stream();
16104 let mut b = __s_b.launch_builder(&f);
16105 b.arg(bytes)
16106 .arg(aq)
16107 .arg(ad)
16108 .arg(&mut y)
16109 .arg(&inf)
16110 .arg(&outf)
16111 .arg(&mi)
16112 .arg(&rb);
16113 unsafe {
16114 b.launch(cfg)?;
16115 }
16116 Ok(Some((y, scale)))
16117 }
16118
16119 pub fn mmvq_supports(&self, qtype: i32) -> bool {
16122 if qtype == QT_F8_E4M3 {
16127 return true;
16128 }
16129 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") {
16130 return false;
16131 }
16132 matches!(
16133 qtype,
16134 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0
16135 )
16136 }
16137
16138 pub fn qmatvec_mmvq(
16143 &self,
16144 bytes: &CudaSlice<u8>,
16145 aq: &CudaSlice<i8>,
16146 ad: &CudaSlice<f32>,
16147 m: usize,
16148 in_f: usize,
16149 out_f: usize,
16150 qtype: i32,
16151 row_bytes: usize,
16152 scale: f32,
16153 rp: bool,
16154 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16155 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(
16157 bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y,
16158 )?;
16159 Ok(y)
16160 }
16161
16162 #[allow(clippy::too_many_arguments)]
16164 pub fn qmatvec_mmvq_into(
16165 &self,
16166 bytes: &CudaSlice<u8>,
16167 aq: &CudaSlice<i8>,
16168 ad: &CudaSlice<f32>,
16169 m: usize,
16170 in_f: usize,
16171 out_f: usize,
16172 qtype: i32,
16173 row_bytes: usize,
16174 scale: f32,
16175 rp: bool,
16176 y: &mut CudaSlice<f32>,
16177 ) -> Result<(), Box<dyn std::error::Error>> {
16178 debug_assert!(y.len() >= m * out_f);
16179 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0
16185 && rp
16186 && m == 1
16187 && out_f >= 64
16188 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
16189 && {
16190 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16191 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
16192 }
16193 {
16194 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
16195 let cfg = LaunchConfig {
16196 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
16197 block_dim: (32, 2, 1),
16198 shared_mem_bytes: 0,
16199 };
16200 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
16201 let __s_b = self.gpu.stream();
16202 let mut b = __s_b.launch_builder(&f);
16203 b.arg(bytes)
16204 .arg(aq)
16205 .arg(ad)
16206 .arg(&mut *y)
16207 .arg(&inf)
16208 .arg(&outf)
16209 .arg(&mi)
16210 .arg(&rb);
16211 unsafe {
16212 b.launch(cfg)?;
16213 }
16214 if scale != 1.0 {
16215 self.scale_inplace(y, scale, out_f)?;
16216 }
16217 return Ok(());
16218 }
16219 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) {
16228 2
16229 } else {
16230 1
16231 };
16232 if m == 1 && qtype == QT_Q4_0 {
16237 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16238 mr = *Q40MR.get_or_init(|| {
16241 std::env::var("MEMRA_Q40_MR")
16242 .ok()
16243 .and_then(|v| v.parse().ok())
16244 .unwrap_or(1)
16245 });
16246 }
16247 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
16258 let q5_force = q5_mode.as_deref() == Some("2");
16259 let q5_il = qtype == QT_Q5_K
16262 && m == 1
16263 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
16264 if q5_il && !q5_force && out_f > 65536 {
16265 mr = 1;
16266 }
16267 if qtype == QT_Q4_0 && rp && mr != 1 {
16270 mr = 2;
16271 }
16272 if qtype == QT_Q8_0 && rp {
16276 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
16277 mr = *Q80MR.get_or_init(|| {
16278 std::env::var("MEMRA_Q80_MR")
16279 .ok()
16280 .and_then(|v| v.parse().ok())
16281 .unwrap_or(1)
16282 });
16283 }
16284 let name = match (qtype, mr, rp) {
16285 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
16286 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
16287 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
16288 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
16289 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
16290 (QT_Q5_K, 2, _) => {
16291 if q5_il {
16292 "qmatvec_q5_K_mmvq_mr2_il"
16293 } else {
16294 "qmatvec_q5_K_mmvq_mr2"
16295 }
16296 }
16297 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
16298 (QT_Q8_0, _, true)
16303 if in_f % 1024 == 0 && {
16304 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16305 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
16306 } =>
16307 {
16308 "qmatvec_q8_0_mmvq_rpca"
16309 }
16310 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
16311 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
16312 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
16316 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
16317 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
16318 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
16319 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
16320 (QT_Q5_K, _, _) => {
16321 if q5_il {
16322 "qmatvec_q5_K_mmvq_il"
16323 } else {
16324 "qmatvec_q5_K_mmvq"
16325 }
16326 }
16327 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
16328 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
16329 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
16330 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
16331 };
16332 let f = self.func(name);
16333 let rows_per_block = ROWS_PER_BLOCK * mr;
16335 let cfg = LaunchConfig {
16336 grid_dim: (
16337 (out_f as u32 + rows_per_block - 1) / rows_per_block,
16338 m as u32,
16339 1,
16340 ),
16341 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
16344 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16345 let __s_b = self.gpu.stream();
16346 let mut b = __s_b.launch_builder(&f);
16347 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
16352 if Self::pdl_on()
16355 && Self::pdl_mmvq_on()
16356 && Self::pdl_nvfp4q8_on()
16357 && name == "qmatvec_nvfp4_mmvq_mr2_rp"
16358 {
16359 use cudarc::driver::{DevicePtr, DevicePtrMut};
16360 let s = &self.gpu.stream();
16361 let (pw, _g0) = bytes.device_ptr(s);
16362 let (paq, _g1) = aq.device_ptr(s);
16363 let (pad, _g2) = ad.device_ptr(s);
16364 let (py, _g3) = y.device_ptr_mut(s);
16365 let mut ps = [
16366 &pw as *const _ as *mut std::ffi::c_void,
16367 &paq as *const _ as *mut _,
16368 &pad as *const _ as *mut _,
16369 &py as *const _ as *mut _,
16370 &inf as *const _ as *mut _,
16371 &outf as *const _ as *mut _,
16372 &mi as *const _ as *mut _,
16373 &rb as *const _ as *mut _,
16374 &scale as *const _ as *mut _,
16375 ];
16376 unsafe {
16377 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16378 }
16379 return Ok(());
16380 }
16381 b.arg(bytes)
16382 .arg(aq)
16383 .arg(ad)
16384 .arg(&mut *y)
16385 .arg(&inf)
16386 .arg(&outf)
16387 .arg(&mi)
16388 .arg(&rb)
16389 .arg(&scale);
16390 unsafe {
16391 b.launch(cfg)?;
16392 }
16393 } else if Self::pdl_on()
16394 && Self::pdl_mmvq_on()
16395 && (matches!(
16396 name,
16397 "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq" | "qmatvec_q6_K_mmvq_rp"
16398 ) || (Self::pdl_nvfp4q8_on()
16399 && matches!(name, "qmatvec_q8_0_mmvq_rp" | "qmatvec_q8_0_mmvq_mr2_rp")))
16400 {
16401 {
16405 use cudarc::driver::{DevicePtr, DevicePtrMut};
16406 let s = &self.gpu.stream();
16407 let (pw, _g0) = bytes.device_ptr(s);
16408 let (paq, _g1) = aq.device_ptr(s);
16409 let (pad, _g2) = ad.device_ptr(s);
16410 let (py, _g3) = y.device_ptr_mut(s);
16411 let mut ps = [
16412 &pw as *const _ as *mut std::ffi::c_void,
16413 &paq as *const _ as *mut _,
16414 &pad as *const _ as *mut _,
16415 &py as *const _ as *mut _,
16416 &inf as *const _ as *mut _,
16417 &outf as *const _ as *mut _,
16418 &mi as *const _ as *mut _,
16419 &rb as *const _ as *mut _,
16420 ];
16421 unsafe {
16422 self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?;
16423 }
16424 }
16425 if scale != 1.0 {
16426 self.scale_inplace(y, scale, m * out_f)?;
16427 }
16428 } else {
16429 b.arg(bytes)
16430 .arg(aq)
16431 .arg(ad)
16432 .arg(&mut *y)
16433 .arg(&inf)
16434 .arg(&outf)
16435 .arg(&mi)
16436 .arg(&rb);
16437 unsafe {
16438 b.launch(cfg)?;
16439 }
16440 if scale != 1.0 {
16441 self.scale_inplace(y, scale, m * out_f)?;
16442 }
16443 }
16444 Ok(())
16445 }
16446
16447 pub fn qmatvec_mmvq_raw(
16451 &self,
16452 bytes: &CudaSlice<u8>,
16453 x: &CudaSlice<f32>,
16454 m: usize,
16455 in_f: usize,
16456 out_f: usize,
16457 qtype: i32,
16458 row_bytes: usize,
16459 rp: bool,
16460 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16461 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
16462 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
16463 }
16464
16465 pub fn batched_supports(&self, qtype: i32) -> bool {
16469 matches!(
16470 qtype,
16471 QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0
16472 )
16473 }
16474
16475 pub fn iq_fast_enabled() -> bool {
16483 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16484 *ON.get_or_init(|| {
16485 std::env::var("MEMRA_IQ_FAST")
16486 .map(|v| v != "0")
16487 .unwrap_or(true)
16488 })
16489 }
16490
16491 pub fn b8_enabled() -> bool {
16494 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16495 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
16496 }
16497
16498 pub fn batched_mcols(m: usize) -> usize {
16500 if m == 2 {
16501 2
16502 } else if m <= 4 {
16503 4
16504 } else if m <= 8 {
16505 8
16506 } else {
16507 16
16508 }
16509 }
16510
16511 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
16516 Some(match (qtype, mcols) {
16517 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2",
16518 (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
16519 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
16520 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
16526 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2",
16527 (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
16528 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
16529 (QT_Q4_K, 16) => "qmatvec_q4_K_mmvq_b16",
16532 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2",
16533 (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
16534 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
16535 (QT_Q5_K, 16) => "qmatvec_q5_K_mmvq_b16",
16538 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2",
16539 (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
16540 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8",
16541 (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
16542 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2",
16543 (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
16544 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
16545 (QT_NVFP4, 16) => "qmatvec_nvfp4_mmvq_b16",
16549 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2",
16550 (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
16551 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
16552 (QT_F8_E4M3, 16) => "qmatvec_e4m3_mmvq_b16",
16556 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2",
16557 (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
16558 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8",
16559 (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
16560 _ => return None,
16561 })
16562 }
16563
16564 pub fn sm_count(&self) -> i32 {
16599 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16600 *SMS.get_or_init(|| {
16601 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16602 self.gpu
16603 .ctx
16604 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16605 .unwrap_or(82)
16606 })
16607 }
16608
16609 pub fn batched_variant(
16610 &self,
16611 _m: usize,
16612 in_f: usize,
16613 out_f: usize,
16614 qtype: i32,
16615 row_bytes: usize,
16616 mcols: usize,
16617 rp: bool,
16618 ) -> &'static str {
16619 if qtype == QT_Q8_0 {
16624 return if rp { "rp" } else { "base" };
16625 }
16626 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16627 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
16628 Ok("base") => "base",
16629 Ok("pf") => "pf",
16630 Ok("r2") => "r2",
16631 Ok("r2w8") => "r2w8",
16632 Ok("pfr2") => "pfr2",
16633 Ok("ca") => "ca",
16634 Ok("car2") => "car2",
16635 Ok("rp") => "rp",
16638 Ok("rpr2") => "rpr2",
16639 Ok("rpr2w8") => "rpr2w8",
16640 Ok("rpca") => "rpca",
16643 Ok("rpcar2") => "rpcar2",
16644 Ok("rpsc") => "rpsc",
16651 Ok("rpms") => "rpms",
16652 Ok("rpmsc") => "rpmsc",
16653 Ok("rpks") => "rpks",
16654 Ok("rpksc") => "rpksc",
16655 _ => "auto",
16656 });
16657 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
16661 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16666 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
16667 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
16668 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
16669 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
16670 let sms = *SMS.get_or_init(|| {
16671 use cudarc::driver::sys::CUdevice_attribute_enum as A;
16672 self.gpu
16673 .ctx
16674 .attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)
16675 .unwrap_or(82)
16676 });
16677 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
16697 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16700 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
16701 Ok("base") => "base",
16702 Ok("r2") => "r2",
16703 Ok("r2w8") => "r2w8",
16704 _ => "auto",
16705 });
16706 let variant: &'static str = if qtype == QT_Q4_0 {
16707 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
16711 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
16712 Ok("base") => "base",
16718 Ok("r2") => "r2",
16719 Ok("ms") => "ms",
16720 Ok("sm") => "sm",
16721 Ok("la") => "la",
16722 _ => "auto",
16723 });
16724 let v = if q40 != "auto" {
16725 q40
16726 } else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 {
16727 "r2"
16728 } else {
16729 "base"
16730 };
16731 if rp {
16736 match v {
16737 "ms" => "r2ms_rp",
16738 "sm" => "r2sm_rp",
16739 "la" => "r2la_rp",
16740 "r2" => "r2_rp",
16741 _ => "rp",
16742 }
16743 } else if matches!(v, "ms" | "sm" | "la") {
16744 "r2"
16745 } else {
16746 v
16747 }
16748 } else if qtype != QT_NVFP4 && !kq_r2 {
16749 "base"
16750 } else if kq_r2 && rp {
16751 "rp"
16755 } else if kq_r2 {
16756 if kq_bv != "auto" {
16759 if kq_bv == "r2w8" && mcols != 4 {
16760 "r2"
16761 } else {
16762 kq_bv
16763 }
16764 } else if bv != "auto" {
16765 match bv {
16766 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
16767 "r2w8" | "rpr2w8" => {
16768 if mcols != 4 {
16769 "r2"
16770 } else {
16771 "r2w8"
16772 }
16773 }
16774 _ => "base", }
16776 } else {
16777 let blocks = (out_f + 7) / 8;
16778 let waves = blocks as f64 / (7 * sms as usize) as f64;
16779 let filled = blocks >= 4 * sms as usize;
16780 let use_r2 = if qtype == QT_Q4_K {
16781 filled
16782 } else {
16783 waves >= 2.0
16784 };
16785 if use_r2 { "r2" } else { "base" }
16786 }
16787 } else if bv != "auto" {
16788 let v = if bv == "r2w8" && mcols == 2 {
16793 "r2"
16794 } else if bv == "ca" && (!ca_ok || mcols == 8) {
16795 "pf"
16796 } else if bv == "car2" && (!ca_ok || mcols == 8) {
16797 "r2"
16798 } else if bv == "pfr2" && mcols == 8 {
16799 "r2"
16800 } else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 {
16801 "rpr2"
16802 }
16803 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
16805 if mcols == 8 { "rpr2w8" } else { "rpr2" }
16806 } else if bv == "rpcar2" && mcols == 2 {
16807 "rpca"
16808 }
16809 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok {
16812 "rpr2"
16813 } else if (bv == "rpks" || bv == "rpksc") && !ks_ok {
16814 "rpr2"
16815 } else {
16816 bv
16817 };
16818 if rp {
16819 match v {
16820 "base" | "pf" | "ca" | "rp" => "rp",
16821 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
16822 "r2w8" | "rpr2w8" => {
16823 if mcols == 2 {
16824 "rpr2"
16825 } else {
16826 "rpr2w8"
16827 }
16828 }
16829 other => other, }
16831 } else {
16832 v
16833 }
16834 } else if mcols == 8 {
16835 if rp {
16846 if sc_ok { "rpsc" } else { "rpr2w8" }
16847 } else {
16848 "r2w8"
16849 }
16850 } else if mcols >= 4 {
16851 let blocks = (out_f + 7) / 8;
16855 let r7 = 7 * sms as usize;
16856 let r8 = 8 * sms as usize;
16857 let waves = blocks as f64 / r7 as f64;
16858 let filled = blocks >= 4 * sms as usize;
16859 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
16863 if rp { "rpr2w8" } else { "r2w8" }
16867 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
16868 if rp { "rpr2" } else { "r2" }
16871 } else {
16872 if rp { "rp" } else { "pf" }
16876 }
16877 } else if in_f >= 6144 {
16878 if rp { "rpr2" } else { "r2" }
16882 } else if rp {
16883 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
16888 if sc_ok && waves >= 0.9 && waves <= 1.1 {
16889 "rpsc"
16890 } else {
16891 "rp"
16892 }
16893 } else {
16894 "base"
16895 };
16896 variant
16897 }
16898
16899 pub fn qmatvec_mmvq_batched(
16900 &self,
16901 bytes: &CudaSlice<u8>,
16902 aq: &CudaSlice<i8>,
16903 ad: &CudaSlice<f32>,
16904 m: usize,
16905 in_f: usize,
16906 out_f: usize,
16907 qtype: i32,
16908 row_bytes: usize,
16909 mcols: usize,
16910 scale: f32,
16911 rp: bool,
16912 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
16913 const ROWS_PER_BLOCK: u32 = 4;
16914 let forced: Option<&'static str> = {
16919 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
16920 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
16921 .as_deref()
16922 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
16923 };
16924 let variant = match forced {
16925 Some(v) if !rp || v.contains("rp") => v,
16926 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
16927 };
16928 let base_name = Self::batched_kernel_name(qtype, mcols).ok_or_else(|| {
16929 format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}")
16930 })?;
16931 let variant = if mcols == 16 {
16935 if rp { "rp" } else { "base" }
16936 } else {
16937 variant
16938 };
16939 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
16946 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
16947 if b567
16948 && qtype == QT_NVFP4
16949 && rp
16950 && mcols == 8
16951 && (5..=7).contains(&m)
16952 && matches!(variant, "rpsc" | "rpr2w8")
16953 {
16954 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
16955 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
16957 let cfg = LaunchConfig {
16958 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
16959 block_dim: (32, ROWS_PER_BLOCK, 1),
16960 shared_mem_bytes: 0,
16961 };
16962 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
16963 let __s_b = self.gpu.stream();
16964 let mut b = __s_b.launch_builder(&f);
16965 b.arg(bytes)
16966 .arg(aq)
16967 .arg(ad)
16968 .arg(&mut y)
16969 .arg(&inf)
16970 .arg(&outf)
16971 .arg(&mi)
16972 .arg(&rb);
16973 unsafe {
16974 b.launch(cfg)?;
16975 }
16976 if scale != 1.0 {
16977 self.scale_inplace(&mut y, scale, m * out_f)?;
16978 }
16979 return Ok(y);
16980 }
16981 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
16982 "base" => (base_name.into(), ROWS_PER_BLOCK),
16983 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
16984 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
16985 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
16986 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
16990 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
16991 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
16992 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
16993 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
16994 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
16995 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
16996 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
16998 debug_assert!(
16999 !rp || name.contains("_rp"),
17000 "rp weight dispatched to a GGUF-layout kernel"
17001 );
17002 let f = self.func(&name);
17003 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
17004 let smem = if name.contains("_r2sm_rp") {
17006 (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32
17007 } else {
17008 0
17009 };
17010 let cfg = LaunchConfig {
17011 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
17012 block_dim: (32, ROWS_PER_BLOCK, 1),
17013 shared_mem_bytes: smem,
17014 };
17015 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17016 let __s_b = self.gpu.stream();
17017 let mut b = __s_b.launch_builder(&f);
17018 b.arg(bytes)
17019 .arg(aq)
17020 .arg(ad)
17021 .arg(&mut y)
17022 .arg(&inf)
17023 .arg(&outf)
17024 .arg(&mi)
17025 .arg(&rb);
17026 unsafe {
17027 b.launch(cfg)?;
17028 }
17029 if scale != 1.0 {
17030 self.scale_inplace(&mut y, scale, m * out_f)?;
17031 }
17032 Ok(y)
17033 }
17034
17035 pub fn qmatvec_batched_raw(
17039 &self,
17040 bytes: &CudaSlice<u8>,
17041 x: &CudaSlice<f32>,
17042 m: usize,
17043 in_f: usize,
17044 out_f: usize,
17045 qtype: i32,
17046 row_bytes: usize,
17047 mcols: usize,
17048 rp: bool,
17049 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17050 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17051 self.qmatvec_mmvq_batched(
17052 bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp,
17053 )
17054 }
17055
17056 pub fn qmatvec_nvfp4_batched_raw(
17058 &self,
17059 bytes: &CudaSlice<u8>,
17060 x: &CudaSlice<f32>,
17061 m: usize,
17062 in_f: usize,
17063 out_f: usize,
17064 row_bytes: usize,
17065 mcols: usize,
17066 rp: bool,
17067 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17068 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
17069 }
17070
17071 fn try_fp4_gemm(
17075 &self,
17076 w: &crate::model::GpuTensor,
17077 x: &CudaSlice<f32>,
17078 m: usize,
17079 in_f: usize,
17080 out_f: usize,
17081 ) -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17082 use crate::model::GpuTensor;
17083 if cfg!(memra_portable_cuda) {
17084 return Ok(None);
17085 }
17086 if std::env::var("MEMRA_FP4").is_ok() {
17089 refuse_portable_force("MEMRA_FP4", "the sm_120a mxf4 block-scale MMA");
17090 }
17091 if std::env::var("MEMRA_FP4").is_err() {
17092 return Ok(None);
17093 }
17094 #[cfg(memra_cutlass)]
17103 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
17104 if let GpuTensor::Quant {
17105 bytes,
17106 qtype,
17107 scale,
17108 row_bytes,
17109 cutlass,
17110 ..
17111 } = w
17112 {
17113 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
17114 if let Some(cw) = cutlass {
17115 let y = self.cutlass_fp4_gemm(
17117 &cw.b_packed,
17118 &cw.sfb_swizzled,
17119 x,
17120 *scale,
17121 m,
17122 out_f,
17123 in_f,
17124 )?;
17125 return Ok(Some(y));
17126 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
17127 let (b_packed, sfb_sw) =
17132 self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
17133 let y =
17134 self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
17135 return Ok(Some(y));
17136 }
17137 }
17138 }
17139 }
17140 if let GpuTensor::Quant {
17141 bytes,
17142 qtype,
17143 row_bytes,
17144 scale,
17145 rp,
17146 ..
17147 } = w
17148 {
17149 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
17152 let y =
17153 self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
17154 return Ok(Some(y));
17155 }
17156 }
17157 Ok(None)
17158 }
17159
17160 pub fn rms_norm_f16out(
17164 &self,
17165 x: &CudaSlice<f32>,
17166 w: &CudaSlice<f32>,
17167 dst: &mut CudaSlice<f32>,
17168 dst16: &mut CudaSlice<u8>,
17169 ncols: usize,
17170 nrows: usize,
17171 eps: f32,
17172 ) -> Result<(), Box<dyn std::error::Error>> {
17173 let f = self.func("rms_norm_f16out_f32");
17174 let cfg = LaunchConfig {
17175 grid_dim: (nrows as u32, 1, 1),
17176 block_dim: (rms_block(), 1, 1),
17177 shared_mem_bytes: 0,
17178 };
17179 let (nc, e) = (ncols as i32, eps);
17180 let __s_b = self.gpu.stream();
17181 let mut b = __s_b.launch_builder(&f);
17182 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
17183 unsafe {
17184 b.launch(cfg)?;
17185 }
17186 Ok(())
17187 }
17188
17189 #[allow(clippy::too_many_arguments)]
17192 pub fn add_rms_norm_f16out(
17193 &self,
17194 a: &CudaSlice<f32>,
17195 b: &CudaSlice<f32>,
17196 w: &CudaSlice<f32>,
17197 res: &mut CudaSlice<f32>,
17198 dst: &mut CudaSlice<f32>,
17199 dst16: &mut CudaSlice<u8>,
17200 ncols: usize,
17201 nrows: usize,
17202 eps: f32,
17203 ) -> Result<(), Box<dyn std::error::Error>> {
17204 let f = self.func("add_rms_norm_f16out_f32");
17205 let cfg = LaunchConfig {
17206 grid_dim: (nrows as u32, 1, 1),
17207 block_dim: (rms_block(), 1, 1),
17208 shared_mem_bytes: 0,
17209 };
17210 let (nc, e) = (ncols as i32, eps);
17211 let __s_lb = self.gpu.stream();
17212 let mut lb = __s_lb.launch_builder(&f);
17213 lb.arg(a)
17214 .arg(b)
17215 .arg(w)
17216 .arg(res)
17217 .arg(dst)
17218 .arg(dst16)
17219 .arg(&nc)
17220 .arg(&e);
17221 unsafe {
17222 lb.launch(cfg)?;
17223 }
17224 Ok(())
17225 }
17226
17227 pub fn matmul_group_xh(
17230 &self,
17231 ws: &[&crate::model::GpuTensor],
17232 x: &CudaSlice<f32>,
17233 xh: &CudaSlice<u8>,
17234 m: usize,
17235 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17236 let mut out = Vec::with_capacity(ws.len());
17237 let in_f = ws[0].in_features();
17238 for w in ws {
17239 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
17240 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
17241 out.push(y);
17242 continue;
17243 }
17244 }
17245 out.push(self.matmul(w, x, m)?);
17246 }
17247 Ok(out)
17248 }
17249
17250 pub fn gdn_pad_mask(
17253 &self,
17254 beta: &mut CudaSlice<f32>,
17255 g_log: &mut CudaSlice<f32>,
17256 len_d: &CudaSlice<i32>,
17257 h: usize,
17258 t: usize,
17259 ) -> Result<(), Box<dyn std::error::Error>> {
17260 let f = self.func("gdn_pad_mask_f32");
17261 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
17262 let (hi, ti) = (h as i32, t as i32);
17263 let __s_b = self.gpu.stream();
17264 let mut b = __s_b.launch_builder(&f);
17265 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
17266 unsafe {
17267 b.launch(cfg)?;
17268 }
17269 Ok(())
17270 }
17271
17272 pub fn row_gather_dev(
17275 &self,
17276 src: &CudaSlice<f32>,
17277 dst: &mut CudaSlice<f32>,
17278 len_d: &CudaSlice<i32>,
17279 ncols: usize,
17280 ) -> Result<(), Box<dyn std::error::Error>> {
17281 let f = self.func("row_gather_dev_f32");
17282 let cfg = LaunchConfig::for_num_elems(ncols as u32);
17283 let nc = ncols as i32;
17284 let __s_b = self.gpu.stream();
17285 let mut b = __s_b.launch_builder(&f);
17286 b.arg(src).arg(dst).arg(len_d).arg(&nc);
17287 unsafe {
17288 b.launch(cfg)?;
17289 }
17290 Ok(())
17291 }
17292
17293 pub fn matmul_group(
17300 &self,
17301 ws: &[&crate::model::GpuTensor],
17302 x: &CudaSlice<f32>,
17303 m: usize,
17304 ) -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
17305 use crate::model::GpuTensor;
17306 let mut out = Vec::with_capacity(ws.len());
17307 let any_mirror = ws
17308 .iter()
17309 .any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
17310 if m >= 16 && any_mirror && !self.verify_exact_on() {
17311 let in_f = ws[0].in_features();
17312 let xh = self.f16_act(x, m * in_f, in_f)?;
17313 for w in ws {
17314 if w.in_features() == in_f {
17315 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
17316 out.push(y);
17317 continue;
17318 }
17319 }
17320 out.push(self.matmul(w, x, m)?);
17321 }
17322 return Ok(out);
17323 }
17324 for w in ws {
17325 out.push(self.matmul(w, x, m)?);
17326 }
17327 Ok(out)
17328 }
17329
17330 pub fn matmul_group_multi(
17337 &self,
17338 ws: &[&crate::model::GpuTensor],
17339 xs: &[&CudaSlice<f32>],
17340 ms: &[usize],
17341 ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
17342 assert_eq!(xs.len(), ms.len());
17343 let in_f = ws[0].in_features();
17344 let total: usize = ms.iter().sum();
17345 let mut xcat = self.uninit(total * in_f)?;
17346 let mut off = 0usize;
17347 for (x, &m) in xs.iter().zip(ms) {
17348 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
17349 off += m;
17350 }
17351 let ys = self.matmul_group(ws, &xcat, total)?;
17352 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
17353 for (w, y) in ws.iter().zip(ys) {
17354 let out_f = w.out_features();
17355 let mut off = 0usize;
17356 for (s, &m) in ms.iter().enumerate() {
17357 let mut ys_s = self.uninit(m * out_f)?;
17358 let src = y.slice(off * out_f..(off + m) * out_f);
17359 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
17360 out[s].push(ys_s);
17361 off += m;
17362 }
17363 }
17364 Ok(out)
17365 }
17366
17367 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
17377 use crate::model::GpuTensor;
17378 if !legacy_quant_gemm_allowed(
17379 cfg!(memra_portable_cuda),
17380 cfg!(memra_hopper_mma),
17381 std::env::var_os("MEMRA_NO_GEMM").is_some(),
17382 ) {
17383 return false;
17384 }
17385 match w {
17386 GpuTensor::Quant { qtype, .. } => {
17387 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
17388 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0)
17389 }
17390 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
17391 }
17392 }
17393
17394 pub fn qmatvec_gemm(
17401 &self,
17402 w: &crate::model::GpuTensor,
17403 aq: &CudaSlice<i8>,
17404 ad: &CudaSlice<f32>,
17405 m: usize,
17406 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17407 use crate::model::GpuTensor;
17408 let in_f = w.in_features();
17409 let out_f = w.out_features();
17410 let (bytes, qtype, row_bytes, scale, rp) = match w {
17411 GpuTensor::Quant {
17412 bytes,
17413 qtype,
17414 row_bytes,
17415 scale,
17416 rp,
17417 ..
17418 } => (bytes, *qtype, *row_bytes, *scale, *rp),
17419 _ => unreachable!("gemm_supports guaranteed Quant"),
17420 };
17421 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
17427 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
17428 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
17429 if scale != 1.0 {
17430 self.scale_inplace(&mut y, scale, m * out_f)?;
17431 }
17432 return Ok(y);
17433 }
17434 }
17435 let name = match qtype {
17436 QT_Q8_0 => "qmatvec_gemm_q8_0",
17437 QT_Q4_K => "qmatvec_gemm_q4_K",
17438 QT_Q4_0 => {
17439 if rp {
17440 "qmatvec_gemm_q4_0_rp"
17441 } else {
17442 "qmatvec_gemm_q4_0"
17443 }
17444 }
17445 QT_Q5_K => "qmatvec_gemm_q5_K",
17446 QT_Q6_K => "qmatvec_gemm_q6_K",
17447 QT_NVFP4 => {
17448 if rp {
17449 "qmatvec_gemm_nvfp4_rp"
17450 } else {
17451 "qmatvec_gemm_nvfp4"
17452 }
17453 }
17454 _ => unreachable!(),
17455 };
17456 let f = self.func(name);
17457 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);
17462 let k1_tile = if is_k1 {
17464 k1_launch_override().unwrap_or((128, 128, 8))
17465 } else {
17466 (128, 128, 8)
17467 };
17468 let (bm, bn): (u32, u32) = if is_k1 {
17469 (k1_tile.0, k1_tile.1)
17470 } else {
17471 (64, 256)
17472 };
17473 let warps: u32 = if is_k1 {
17474 k1_tile.2
17475 } else {
17476 match qtype {
17477 QT_NVFP4 => 8,
17478 _ => 4,
17479 }
17480 };
17481 let cfg = LaunchConfig {
17482 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17483 block_dim: (32, warps, 1),
17484 shared_mem_bytes: 0,
17485 };
17486 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17487 let __s_b = self.gpu.stream();
17488 let mut b = __s_b.launch_builder(&f);
17489 b.arg(bytes)
17490 .arg(aq)
17491 .arg(ad)
17492 .arg(&mut y)
17493 .arg(&inf)
17494 .arg(&outf)
17495 .arg(&mi)
17496 .arg(&rb);
17497 unsafe {
17498 b.launch(cfg)?;
17499 }
17500 if scale != 1.0 {
17501 self.scale_inplace(&mut y, scale, m * out_f)?;
17502 }
17503 Ok(y)
17504 }
17505
17506 pub fn qmatvec_gemm_raw(
17511 &self,
17512 bytes: &CudaSlice<u8>,
17513 x: &CudaSlice<f32>,
17514 m: usize,
17515 in_f: usize,
17516 out_f: usize,
17517 qtype: i32,
17518 row_bytes: usize,
17519 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17520 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
17521 let name = match qtype {
17522 QT_Q8_0 => "qmatvec_gemm_q8_0",
17523 QT_Q4_K => "qmatvec_gemm_q4_K",
17524 QT_Q4_0 => "qmatvec_gemm_q4_0",
17525 QT_Q5_K => "qmatvec_gemm_q5_K",
17526 QT_Q6_K => "qmatvec_gemm_q6_K",
17527 QT_NVFP4 => "qmatvec_gemm_nvfp4",
17528 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
17529 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
17530 };
17531 let f = self.func(name);
17532 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);
17536 let k1_tile = if is_k1 {
17538 k1_launch_override().unwrap_or((128, 128, 8))
17539 } else {
17540 (128, 128, 8)
17541 };
17542 let (bm, bn): (u32, u32) = if is_k1 {
17543 (k1_tile.0, k1_tile.1)
17544 } else {
17545 (64, 256)
17546 };
17547 let warps: u32 = if is_k1 {
17548 k1_tile.2
17549 } else {
17550 match qtype {
17551 QT_NVFP4 | QT_NVFP4_RP => 8,
17552 _ => 4,
17553 }
17554 };
17555 let cfg = LaunchConfig {
17556 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
17557 block_dim: (32, warps, 1),
17558 shared_mem_bytes: 0,
17559 };
17560 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
17561 let __s_b = self.gpu.stream();
17562 let mut b = __s_b.launch_builder(&f);
17563 b.arg(bytes)
17564 .arg(&aq)
17565 .arg(&ad)
17566 .arg(&mut y)
17567 .arg(&inf)
17568 .arg(&outf)
17569 .arg(&mi)
17570 .arg(&rb);
17571 unsafe {
17572 b.launch(cfg)?;
17573 }
17574 Ok(y)
17575 }
17576
17577 pub fn qmatvec_gemm_q8_0_wgmma_raw(
17584 &self,
17585 rp4: &CudaSlice<u8>,
17586 aq: &CudaSlice<i8>,
17587 ad: &CudaSlice<f32>,
17588 m: usize,
17589 in_f: usize,
17590 out_f: usize,
17591 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17592 assert!(
17593 out_f % 64 == 0 && in_f % 32 == 0,
17594 "wgmma GEMM needs out_f%64==0, in_f%32==0"
17595 );
17596 let f = self.func("qmatvec_gemm_q8_0_wgmma");
17597 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
17599 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
17600 block_dim: (128, 1, 1),
17601 shared_mem_bytes: 0,
17602 };
17603 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
17604 let __s_b = self.gpu.stream();
17605 let mut b = __s_b.launch_builder(&f);
17606 b.arg(rp4)
17607 .arg(aq)
17608 .arg(ad)
17609 .arg(&mut y)
17610 .arg(&inf)
17611 .arg(&outf)
17612 .arg(&mi);
17613 unsafe {
17614 b.launch(cfg)?;
17615 }
17616 Ok(y)
17617 }
17618
17619 pub fn scale_inplace(
17621 &self,
17622 y: &mut CudaSlice<f32>,
17623 s: f32,
17624 n: usize,
17625 ) -> Result<(), Box<dyn std::error::Error>> {
17626 let f = self.func("scale_f32");
17627 let cfg = LaunchConfig::for_num_elems(n as u32);
17628 let (sf, ni) = (s, n as i32);
17629 let __s_b = self.gpu.stream();
17630 let mut b = __s_b.launch_builder(&f);
17631 b.arg(y).arg(&sf).arg(&ni);
17632 unsafe {
17633 b.launch(cfg)?;
17634 }
17635 Ok(())
17636 }
17637
17638 pub fn bf16_to_f32(
17643 &self,
17644 data: &cudarc::driver::CudaView<'_, u8>,
17645 n: usize,
17646 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17647 let mut out = self.alloc_uninit::<f32>(n)?;
17648 let f = self.func("bf16_to_f32");
17649 let cfg = LaunchConfig::for_num_elems(n as u32);
17650 let ni = n as i32;
17651 let __s_b = self.gpu.stream();
17652 let mut b = __s_b.launch_builder(&f);
17653 b.arg(data).arg(&mut out).arg(&ni);
17654 unsafe {
17655 b.launch(cfg)?;
17656 }
17657 Ok(out)
17658 }
17659
17660 fn linear_bf16_chunked(
17667 &self,
17668 x: &CudaSlice<f32>,
17669 data: &CudaSlice<u8>,
17670 m: usize,
17671 in_f: usize,
17672 out_f: usize,
17673 exact: bool,
17674 canonical_chunk_rows: Option<usize>,
17675 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17676 static EXP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17680 static EXP_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17681 static EXP_WBYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
17682 let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
17683 let started = timing.then(std::time::Instant::now);
17684 let result =
17685 self.linear_bf16_chunked_inner(x, data, m, in_f, out_f, exact, canonical_chunk_rows);
17686 if let Some(started) = started {
17687 use std::sync::atomic::Ordering;
17688 self.stream().synchronize()?;
17689 let ns = EXP_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed)
17690 + started.elapsed().as_nanos() as u64;
17691 let wb = EXP_WBYTES.fetch_add((in_f * out_f * 2) as u64, Ordering::Relaxed)
17692 + (in_f * out_f * 2) as u64;
17693 let calls = EXP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
17694 if calls % 1024 == 0 {
17695 eprintln!(
17696 "[bf16-expand-timing] calls={calls} total_ms={:.1} avg_us={:.1} \
17697 weight_gb={:.2}",
17698 ns as f64 / 1.0e6,
17699 ns as f64 / calls as f64 / 1.0e3,
17700 wb as f64 / 1.0e9,
17701 );
17702 }
17703 }
17704 result
17705 }
17706
17707 pub(crate) fn bf16_mmv_on() -> bool {
17712 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
17713 *ON.get_or_init(|| std::env::var("MEMRA_BF16_MMV").as_deref() == Ok("1"))
17714 }
17715
17716 fn matvec_bf16(
17719 &self,
17720 data: &CudaSlice<u8>,
17721 x: &CudaSlice<f32>,
17722 in_f: usize,
17723 out_f: usize,
17724 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
17725 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 {
17726 return Err(format!(
17727 "matvec_bf16 geometry bytes={} x={} in={in_f} out={out_f}",
17728 data.len(),
17729 x.len()
17730 )
17731 .into());
17732 }
17733 let mut y = self.alloc_uninit::<f32>(out_f)?;
17734 let f = self.func("matvec_bf16_f32acc");
17735 let cfg = LaunchConfig {
17736 grid_dim: (out_f as u32, 1, 1),
17737 block_dim: (mmv_block(), 1, 1),
17738 shared_mem_bytes: 0,
17739 };
17740 let ini = in_f as i32;
17741 let __s_bld = self.gpu.stream();
17742 let mut bld = __s_bld.launch_builder(&f);
17743 bld.arg(data).arg(x).arg(&mut y).arg(&ini);
17744 unsafe {
17745 bld.launch(cfg)?;
17746 }
17747 Ok(y)
17748 }
17749
17750 #[allow(clippy::too_many_arguments)]
17754 #[allow(clippy::too_many_arguments)]
17759 #[allow(clippy::too_many_arguments)]
17764 pub fn qk_norm_rope_append_inc_dcw_rows(
17765 &self,
17766 q_raw_t: &CudaSlice<f32>,
17767 k_raw_t: &CudaSlice<f32>,
17768 v_raw_t: &CudaSlice<f32>,
17769 qw: &CudaSlice<f32>,
17770 kw: &CudaSlice<f32>,
17771 q_out_t: &mut CudaSlice<f32>,
17772 k_out_t: &mut CudaSlice<f32>,
17773 tab: &CudaSlice<u64>,
17774 pos_t: &CudaSlice<i32>,
17775 same_session: bool,
17776 t: usize,
17777 kv_dim_k: usize,
17778 kv_dim_v: usize,
17779 k_tok_bytes: usize,
17780 v_tok_bytes: usize,
17781 head_dim: usize,
17782 n_dims: usize,
17783 nh_q: usize,
17784 nh_k: usize,
17785 eps: f32,
17786 freq_base: f32,
17787 freq_scale: f32,
17788 ff: Option<&CudaSlice<f32>>,
17789 ) -> Result<(), Box<dyn std::error::Error>> {
17790 if head_dim != 128
17791 || kv_dim_v != kv_dim_k
17792 || kv_dim_k != nh_k * head_dim
17793 || t == 0
17794 || t > 32
17795 || tab.len() < t * 6
17796 || pos_t.len() < t
17797 || q_raw_t.len() < t * nh_q * head_dim
17798 || k_raw_t.len() < t * nh_k * head_dim
17799 || v_raw_t.len() < t * kv_dim_v
17800 || q_out_t.len() < t * nh_q * head_dim
17801 || k_out_t.len() < t * nh_k * head_dim
17802 {
17803 return Err(format!(
17804 "qk_norm_rope_append_inc_rows geometry head_dim={head_dim} t={t} \
17805 nh_q={nh_q} nh_k={nh_k}"
17806 )
17807 .into());
17808 }
17809 let f = self.func("qk_norm_rope_append_inc_dcw_rows");
17810 let same_t: i32 = if same_session { t as i32 } else { 0 };
17811 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17812 let cfg = LaunchConfig {
17813 grid_dim: ((nh_q + nh_k) as u32, 1, t as u32),
17814 block_dim: (128, 1, 1),
17815 shared_mem_bytes: 0,
17816 };
17817 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17818 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17819 let (hd, nd, nq, nk) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32);
17820 let null: u64 = 0;
17821 let __s_b = self.gpu.stream();
17822 let mut b = __s_b.launch_builder(&f);
17823 b.arg(q_raw_t)
17824 .arg(k_raw_t)
17825 .arg(v_raw_t)
17826 .arg(qw)
17827 .arg(kw)
17828 .arg(q_out_t)
17829 .arg(k_out_t)
17830 .arg(tab)
17831 .arg(pos_t)
17832 .arg(&same_t)
17833 .arg(&kvk)
17834 .arg(&kvv)
17835 .arg(&ktb)
17836 .arg(&vtb)
17837 .arg(&hd)
17838 .arg(&nd)
17839 .arg(&nq)
17840 .arg(&nk)
17841 .arg(&eps)
17842 .arg(&theta_scale)
17843 .arg(&freq_scale);
17844 match ff {
17845 Some(freqs) => {
17846 b.arg(freqs);
17847 }
17848 None => {
17849 b.arg(&null);
17850 }
17851 }
17852 unsafe {
17853 b.launch(cfg)?;
17854 }
17855 Ok(())
17856 }
17857
17858 pub fn qk_norm_rope_append_inc_dcw(
17859 &self,
17860 q_raw: &CudaSlice<f32>,
17861 k_raw: &CudaSlice<f32>,
17862 v_raw: &CudaSlice<f32>,
17863 qw: &CudaSlice<f32>,
17864 kw: &CudaSlice<f32>,
17865 q_out: &mut CudaSlice<f32>,
17866 k_out: &mut CudaSlice<f32>,
17867 pos: &CudaSlice<i32>,
17868 k_plane: &mut CudaSlice<u8>,
17869 v_plane: &mut CudaSlice<u8>,
17870 len_dev: &CudaSlice<i32>,
17873 base_dev: Option<&CudaSlice<i32>>,
17874 done_ctr: &mut CudaSlice<u32>,
17875 kv_dim_k: usize,
17876 kv_dim_v: usize,
17877 k_tok_bytes: usize,
17878 v_tok_bytes: usize,
17879 head_dim: usize,
17880 n_dims: usize,
17881 nh_q: usize,
17882 nh_k: usize,
17883 eps: f32,
17884 freq_base: f32,
17885 freq_scale: f32,
17886 ff: Option<&CudaSlice<f32>>,
17887 ) -> Result<(), Box<dyn std::error::Error>> {
17888 if head_dim != 128
17889 || kv_dim_v != kv_dim_k
17890 || kv_dim_k != nh_k * head_dim
17891 || q_raw.len() < nh_q * head_dim
17892 || k_raw.len() < nh_k * head_dim
17893 || v_raw.len() < kv_dim_v
17894 || q_out.len() < nh_q * head_dim
17895 || k_out.len() < nh_k * head_dim
17896 || pos.is_empty()
17897 || done_ctr.is_empty()
17898 {
17899 return Err(format!(
17900 "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}"
17901 )
17902 .into());
17903 }
17904 let f = self.func("qk_norm_rope_append_inc_dcw");
17905 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17906 let cfg = LaunchConfig {
17907 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17908 block_dim: (128, 1, 1),
17909 shared_mem_bytes: 0,
17910 };
17911 let (kvk, kvv) = (kv_dim_k as i32, kv_dim_v as i32);
17912 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
17913 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
17914 let null: u64 = 0;
17915 let __s_b = self.gpu.stream();
17916 let mut b = __s_b.launch_builder(&f);
17917 b.arg(q_raw)
17918 .arg(k_raw)
17919 .arg(v_raw)
17920 .arg(qw)
17921 .arg(kw)
17922 .arg(q_out)
17923 .arg(k_out)
17924 .arg(pos)
17925 .arg(&mut *k_plane)
17926 .arg(&mut *v_plane)
17927 .arg(len_dev);
17928 match base_dev {
17929 Some(base) => {
17930 b.arg(base);
17931 }
17932 None => {
17933 b.arg(&null);
17934 }
17935 }
17936 b.arg(&mut *done_ctr)
17937 .arg(&kvk)
17938 .arg(&kvv)
17939 .arg(&ktb)
17940 .arg(&vtb)
17941 .arg(&hd)
17942 .arg(&nd)
17943 .arg(&nq)
17944 .arg(&eps)
17945 .arg(&theta_scale)
17946 .arg(&freq_scale);
17947 match ff {
17948 Some(freqs) => {
17949 b.arg(freqs);
17950 }
17951 None => {
17952 b.arg(&null);
17953 }
17954 }
17955 unsafe {
17956 b.launch(cfg)?;
17957 }
17958 Ok(())
17959 }
17960
17961 pub fn qk_norm_rope_into(
17962 &self,
17963 q_raw: &CudaSlice<f32>,
17964 k_raw: &CudaSlice<f32>,
17965 qw: &CudaSlice<f32>,
17966 kw: &CudaSlice<f32>,
17967 q_out: &mut CudaSlice<f32>,
17968 k_out: &mut CudaSlice<f32>,
17969 pos: &CudaSlice<i32>,
17970 head_dim: usize,
17971 n_dims: usize,
17972 nh_q: usize,
17973 nh_k: usize,
17974 eps: f32,
17975 freq_base: f32,
17976 freq_scale: f32,
17977 ff: Option<&CudaSlice<f32>>,
17978 ) -> Result<(), Box<dyn std::error::Error>> {
17979 if head_dim > 512
17980 || q_raw.len() < nh_q * head_dim
17981 || k_raw.len() < nh_k * head_dim
17982 || q_out.len() < nh_q * head_dim
17983 || k_out.len() < nh_k * head_dim
17984 || qw.len() < head_dim
17985 || kw.len() < head_dim
17986 || pos.is_empty()
17987 {
17988 return Err(format!(
17989 "qk_norm_rope geometry head_dim={head_dim} nh_q={nh_q} nh_k={nh_k}"
17990 )
17991 .into());
17992 }
17993 let f = self.func("qk_norm_rope_f32");
17994 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
17995 let cfg = LaunchConfig {
17996 grid_dim: ((nh_q + nh_k) as u32, 1, 1),
17997 block_dim: (128, 1, 1),
17998 shared_mem_bytes: 0,
17999 };
18000 let (hd, nd, nq) = (head_dim as i32, n_dims as i32, nh_q as i32);
18001 let __s_b = self.gpu.stream();
18002 let mut b = __s_b.launch_builder(&f);
18003 b.arg(q_raw)
18004 .arg(k_raw)
18005 .arg(qw)
18006 .arg(kw)
18007 .arg(q_out)
18008 .arg(k_out)
18009 .arg(pos)
18010 .arg(&hd)
18011 .arg(&nd)
18012 .arg(&nq)
18013 .arg(&eps)
18014 .arg(&theta_scale)
18015 .arg(&freq_scale);
18016 match ff {
18017 Some(ffv) => {
18018 b.arg(ffv);
18019 unsafe {
18020 b.launch(cfg)?;
18021 }
18022 }
18023 None => {
18024 let null: u64 = 0;
18025 b.arg(&null);
18026 unsafe {
18027 b.launch(cfg)?;
18028 }
18029 }
18030 }
18031 Ok(())
18032 }
18033
18034 #[allow(clippy::too_many_arguments)]
18037 pub fn matvec_f32_b4_into(
18038 &self,
18039 w: [&CudaSlice<f32>; 4],
18040 x: &CudaSlice<f32>,
18041 y: &mut CudaSlice<f32>,
18042 block_cols: usize,
18043 out_f: usize,
18044 ) -> Result<(), Box<dyn std::error::Error>> {
18045 if block_cols % 4 != 0
18046 || x.len() < 4 * block_cols
18047 || y.len() < out_f
18048 || w.iter().any(|w| w.len() != out_f * block_cols)
18049 {
18050 return Err(format!(
18051 "matvec_f32_b4 geometry block_cols={block_cols} out={out_f} x={}",
18052 x.len()
18053 )
18054 .into());
18055 }
18056 let f = self.func("matvec_f32_b4");
18057 let cfg = LaunchConfig {
18058 grid_dim: (out_f as u32, 1, 1),
18059 block_dim: (128, 1, 1),
18060 shared_mem_bytes: 0,
18061 };
18062 let (bc, of) = (block_cols as i32, out_f as i32);
18063 let __s_b = self.gpu.stream();
18064 let mut b = __s_b.launch_builder(&f);
18065 b.arg(w[0])
18066 .arg(w[1])
18067 .arg(w[2])
18068 .arg(w[3])
18069 .arg(x)
18070 .arg(y)
18071 .arg(&bc)
18072 .arg(&of);
18073 unsafe {
18074 b.launch(cfg)?;
18075 }
18076 Ok(())
18077 }
18078
18079 pub fn axpy_rows_seq_into(
18082 &self,
18083 x: &CudaSlice<f32>,
18084 w: &CudaSlice<f32>,
18085 y: &mut CudaSlice<f32>,
18086 width: usize,
18087 n_rows: usize,
18088 ) -> Result<(), Box<dyn std::error::Error>> {
18089 if x.len() < n_rows * width || w.len() < n_rows || y.len() < width {
18090 return Err(format!(
18091 "axpy_rows_seq geometry x={} w={} y={} width={width} rows={n_rows}",
18092 x.len(),
18093 w.len(),
18094 y.len()
18095 )
18096 .into());
18097 }
18098 let f = self.func("axpy_rows_seq_f32");
18099 let cfg = LaunchConfig::for_num_elems(width as u32);
18100 let (wi, nr) = (width as i32, n_rows as i32);
18101 let __s_b = self.gpu.stream();
18102 let mut b = __s_b.launch_builder(&f);
18103 b.arg(x).arg(w).arg(y).arg(&wi).arg(&nr);
18104 unsafe {
18105 b.launch(cfg)?;
18106 }
18107 Ok(())
18108 }
18109
18110 #[allow(clippy::too_many_arguments)]
18114 pub fn axpy_rows_seq_md_off_into(
18115 &self,
18116 x: &CudaSlice<f32>,
18117 w_route: &CudaSlice<f32>,
18118 md: &CudaSlice<f32>,
18119 sel: &CudaSlice<i32>,
18120 y: &mut CudaSlice<f32>,
18121 width: usize,
18122 n_rows: usize,
18123 row0: usize,
18124 ) -> Result<(), Box<dyn std::error::Error>> {
18125 if x.len() < (row0 + n_rows) * width
18126 || w_route.len() < row0 + n_rows
18127 || sel.len() < row0 + n_rows
18128 || y.len() < width
18129 {
18130 return Err(format!(
18131 "axpy_rows_seq_md_off geometry x={} w={} sel={} y={} width={width} \
18132 rows={n_rows} row0={row0}",
18133 x.len(),
18134 w_route.len(),
18135 sel.len(),
18136 y.len()
18137 )
18138 .into());
18139 }
18140 let f = self.func("axpy_rows_seq_md_off_f32");
18141 let cfg = LaunchConfig::for_num_elems(width as u32);
18142 let (wi, nr, r0) = (width as i32, n_rows as i32, row0 as i32);
18143 let __s_b = self.gpu.stream();
18144 let mut b = __s_b.launch_builder(&f);
18145 b.arg(x)
18146 .arg(w_route)
18147 .arg(md)
18148 .arg(sel)
18149 .arg(y)
18150 .arg(&wi)
18151 .arg(&nr)
18152 .arg(&r0);
18153 unsafe {
18154 b.launch(cfg)?;
18155 }
18156 Ok(())
18157 }
18158
18159 #[allow(clippy::too_many_arguments)]
18164 pub fn qmatvec_nvfp4_sel_gu_tcol_into(
18165 &self,
18166 gate_bank: &CudaSlice<u8>,
18167 up_bank: &CudaSlice<u8>,
18168 sel: &CudaSlice<i32>,
18169 aq: &CudaSlice<i8>,
18170 ad: &CudaSlice<f32>,
18171 yg: &mut CudaSlice<f32>,
18172 yu: &mut CudaSlice<f32>,
18173 n_sel: usize,
18174 n_sel_col: usize,
18175 in_f: usize,
18176 out_f: usize,
18177 row_bytes: usize,
18178 expert_stride: usize,
18179 act_row_stride: usize,
18180 ad_row_stride: usize,
18181 ) -> Result<(), Box<dyn std::error::Error>> {
18182 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0");
18183 if yg.len() < n_sel * out_f
18184 || yu.len() < n_sel * out_f
18185 || sel.len() < n_sel
18186 || n_sel_col == 0
18187 || n_sel % n_sel_col != 0
18188 {
18189 return Err("NVFP4 gu tcol geometry".into());
18190 }
18191 let f = self.func("qmatvec_nvfp4_dp4a_sel_v2_gu_tcol");
18192 let cfg = LaunchConfig {
18193 grid_dim: ((2 * out_f) as u32, n_sel as u32, 1),
18194 block_dim: (128, 1, 1),
18195 shared_mem_bytes: 0,
18196 };
18197 let (inf, outf, ns, nsc) = (in_f as i32, out_f as i32, n_sel as i32, n_sel_col as i32);
18198 let (rb, es) = (row_bytes as i64, expert_stride as i64);
18199 let (ars, adrs) = (act_row_stride as i64, ad_row_stride as i64);
18200 let __s_b = self.gpu.stream();
18201 let mut b = __s_b.launch_builder(&f);
18202 b.arg(gate_bank)
18203 .arg(up_bank)
18204 .arg(sel)
18205 .arg(aq)
18206 .arg(ad)
18207 .arg(yg)
18208 .arg(yu)
18209 .arg(&inf)
18210 .arg(&outf)
18211 .arg(&ns)
18212 .arg(&rb)
18213 .arg(&es)
18214 .arg(&ars)
18215 .arg(&adrs)
18216 .arg(&nsc);
18217 unsafe {
18218 b.launch(cfg)?;
18219 }
18220 Ok(())
18221 }
18222
18223 #[allow(clippy::too_many_arguments)]
18226 pub fn axpy_rows_seq_md_into(
18227 &self,
18228 x: &CudaSlice<f32>,
18229 w_route: &CudaSlice<f32>,
18230 md: &CudaSlice<f32>,
18231 sel: &CudaSlice<i32>,
18232 y: &mut CudaSlice<f32>,
18233 width: usize,
18234 n_rows: usize,
18235 ) -> Result<(), Box<dyn std::error::Error>> {
18236 if x.len() < n_rows * width
18237 || w_route.len() < n_rows
18238 || sel.len() < n_rows
18239 || y.len() < width
18240 {
18241 return Err(format!(
18242 "axpy_rows_seq_md geometry x={} w={} sel={} y={} width={width} rows={n_rows}",
18243 x.len(),
18244 w_route.len(),
18245 sel.len(),
18246 y.len()
18247 )
18248 .into());
18249 }
18250 let f = self.func("axpy_rows_seq_md_f32");
18251 let cfg = LaunchConfig::for_num_elems(width as u32);
18252 let (wi, nr) = (width as i32, n_rows as i32);
18253 let __s_b = self.gpu.stream();
18254 let mut b = __s_b.launch_builder(&f);
18255 b.arg(x)
18256 .arg(w_route)
18257 .arg(md)
18258 .arg(sel)
18259 .arg(y)
18260 .arg(&wi)
18261 .arg(&nr);
18262 unsafe {
18263 b.launch(cfg)?;
18264 }
18265 Ok(())
18266 }
18267
18268 #[allow(clippy::too_many_arguments)]
18270 #[allow(clippy::too_many_arguments)]
18274 pub fn matvec_bf16_qkvg_tcol_into(
18275 &self,
18276 wq: &CudaSlice<u8>,
18277 wk: &CudaSlice<u8>,
18278 wv: &CudaSlice<u8>,
18279 wg: &CudaSlice<u8>,
18280 x_t: &CudaSlice<f32>,
18281 yq: &mut CudaSlice<f32>,
18282 yk: &mut CudaSlice<f32>,
18283 yv: &mut CudaSlice<f32>,
18284 yg: &mut CudaSlice<f32>,
18285 in_f: usize,
18286 out_q: usize,
18287 out_kv: usize,
18288 out_g: usize,
18289 t: usize,
18290 ) -> Result<(), Box<dyn std::error::Error>> {
18291 if t == 0
18292 || t > 8
18293 || in_f % 8 != 0
18294 || x_t.len() < t * in_f
18295 || yq.len() < t * out_q
18296 || yk.len() < t * out_kv
18297 || yv.len() < t * out_kv
18298 || (out_g > 0 && yg.len() < t * out_g)
18299 {
18300 return Err("matvec_bf16_qkvg_tcol geometry".into());
18301 }
18302 let grid = out_q + 2 * out_kv + out_g;
18303 let cfg = LaunchConfig {
18304 grid_dim: (grid as u32, 1, 1),
18305 block_dim: (mmv_block(), 1, 1),
18306 shared_mem_bytes: 0,
18307 };
18308 let (ini, oq, okv, og, ti) = (
18309 in_f as i32,
18310 out_q as i32,
18311 out_kv as i32,
18312 out_g as i32,
18313 t as i32,
18314 );
18315 let __s_b = self.gpu.stream();
18316 let f = self.func("matvec_bf16_qkvg_tcol");
18322 let mut b = __s_b.launch_builder(&f);
18323 b.arg(wq)
18324 .arg(wk)
18325 .arg(wv)
18326 .arg(wg)
18327 .arg(x_t)
18328 .arg(yq)
18329 .arg(yk)
18330 .arg(yv)
18331 .arg(yg)
18332 .arg(&ini)
18333 .arg(&oq)
18334 .arg(&okv)
18335 .arg(&og)
18336 .arg(&ti);
18337 unsafe {
18338 b.launch(cfg)?;
18339 }
18340 Ok(())
18341 }
18342
18343 pub fn matvec_bf16_qkvg_into(
18344 &self,
18345 wq: &CudaSlice<u8>,
18346 wk: &CudaSlice<u8>,
18347 wv: &CudaSlice<u8>,
18348 wg: &CudaSlice<u8>,
18349 x: &CudaSlice<f32>,
18350 yq: &mut CudaSlice<f32>,
18351 yk: &mut CudaSlice<f32>,
18352 yv: &mut CudaSlice<f32>,
18353 yg: &mut CudaSlice<f32>,
18354 in_f: usize,
18355 out_q: usize,
18356 out_kv: usize,
18357 out_g: usize,
18358 ) -> Result<(), Box<dyn std::error::Error>> {
18359 if in_f % 8 != 0
18360 || wq.len() != out_q * in_f * 2
18361 || wk.len() != out_kv * in_f * 2
18362 || wv.len() != out_kv * in_f * 2
18363 || wg.len() < out_g * in_f * 2
18364 || x.len() < in_f
18365 || yq.len() < out_q
18366 || yk.len() < out_kv
18367 || yv.len() < out_kv
18368 || (out_g > 0 && yg.len() < out_g)
18369 {
18370 return Err(format!(
18371 "fused bf16 QKV geometry in={in_f} out_q={out_q} out_kv={out_kv} out_g={out_g}"
18372 )
18373 .into());
18374 }
18375 let f = self.func("matvec_bf16_qkvg");
18376 let cfg = LaunchConfig {
18377 grid_dim: ((out_q + 2 * out_kv + out_g) as u32, 1, 1),
18378 block_dim: (mmv_block(), 1, 1),
18379 shared_mem_bytes: 0,
18380 };
18381 let (inf, oq, okv, og) = (in_f as i32, out_q as i32, out_kv as i32, out_g as i32);
18382 let __s_b = self.gpu.stream();
18383 let mut b = __s_b.launch_builder(&f);
18384 b.arg(wq)
18385 .arg(wk)
18386 .arg(wv)
18387 .arg(wg)
18388 .arg(x)
18389 .arg(yq)
18390 .arg(yk)
18391 .arg(yv)
18392 .arg(yg)
18393 .arg(&inf)
18394 .arg(&oq)
18395 .arg(&okv)
18396 .arg(&og);
18397 unsafe {
18398 b.launch(cfg)?;
18399 }
18400 Ok(())
18401 }
18402
18403 pub fn matvec_bf16_b4_into(
18405 &self,
18406 w: [&CudaSlice<u8>; 4],
18407 x: &CudaSlice<f32>,
18408 y: &mut CudaSlice<f32>,
18409 block_cols: usize,
18410 out_f: usize,
18411 ) -> Result<(), Box<dyn std::error::Error>> {
18412 if block_cols % 8 != 0
18413 || x.len() < 4 * block_cols
18414 || y.len() < out_f
18415 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18416 {
18417 return Err(format!(
18418 "bf16 b4 geometry block_cols={block_cols} out={out_f} x={}",
18419 x.len()
18420 )
18421 .into());
18422 }
18423 static B4_X2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
18426 let x2 = *B4_X2.get_or_init(|| std::env::var("MEMRA_B4_X2").as_deref() == Ok("1"));
18427 let f = self.func(if x2 {
18428 "matvec_bf16_b4_x2"
18429 } else {
18430 "matvec_bf16_b4"
18431 });
18432 let grid = if x2 { out_f.div_ceil(2) } else { out_f };
18433 let cfg = LaunchConfig {
18434 grid_dim: (grid as u32, 1, 1),
18435 block_dim: (mmv_block(), 1, 1),
18436 shared_mem_bytes: 0,
18437 };
18438 let (bc, of) = (block_cols as i32, out_f as i32);
18439 let __s_b = self.gpu.stream();
18440 let mut b = __s_b.launch_builder(&f);
18441 b.arg(w[0])
18442 .arg(w[1])
18443 .arg(w[2])
18444 .arg(w[3])
18445 .arg(x)
18446 .arg(y)
18447 .arg(&bc)
18448 .arg(&of);
18449 unsafe {
18450 b.launch(cfg)?;
18451 }
18452 Ok(())
18453 }
18454
18455 pub fn matvec_bf16_b4_tcol_into(
18461 &self,
18462 w: [&CudaSlice<u8>; 4],
18463 x_t: &CudaSlice<f32>,
18464 y_t: &mut CudaSlice<f32>,
18465 block_cols: usize,
18466 out_f: usize,
18467 t: usize,
18468 ) -> Result<(), Box<dyn std::error::Error>> {
18469 if block_cols % 8 != 0
18470 || t == 0
18471 || t > 8
18472 || x_t.len() < t * 4 * block_cols
18473 || y_t.len() < t * out_f
18474 || w.iter().any(|w| w.len() != out_f * block_cols * 2)
18475 {
18476 return Err(format!(
18477 "bf16 b4 tcol geometry block_cols={block_cols} out={out_f} t={t} x={}",
18478 x_t.len()
18479 )
18480 .into());
18481 }
18482 if std::env::var("MEMRA_B4_X2").as_deref() == Ok("1") {
18483 return Err(
18484 "b4 tcol verify is qualified against the plain b4 kernel only \
18485 (MEMRA_B4_X2=1 is a different t=1 program)"
18486 .into(),
18487 );
18488 }
18489 let cfg = LaunchConfig {
18493 grid_dim: (out_f as u32, 1, 1),
18494 block_dim: (mmv_block(), 1, 1),
18495 shared_mem_bytes: 0,
18496 };
18497 let (bc, of, ti) = (block_cols as i32, out_f as i32, t as i32);
18498 let __s_b = self.gpu.stream();
18499 let f = self.func("matvec_bf16_b4_tcol");
18500 let mut b = __s_b.launch_builder(&f);
18501 b.arg(w[0])
18502 .arg(w[1])
18503 .arg(w[2])
18504 .arg(w[3])
18505 .arg(x_t)
18506 .arg(y_t)
18507 .arg(&bc)
18508 .arg(&of)
18509 .arg(&ti);
18510 unsafe {
18511 b.launch(cfg)?;
18512 }
18513 Ok(())
18514 }
18515
18516 pub fn q8_0_row_bytes(in_f: usize) -> usize {
18519 in_f / 32 * 34
18520 }
18521
18522 pub fn encode_q8_0_from_bf16(
18526 &self,
18527 w_bf16: &CudaSlice<u8>,
18528 out: &mut CudaSlice<u8>,
18529 in_f: usize,
18530 out_f: usize,
18531 ) -> Result<(), Box<dyn std::error::Error>> {
18532 if in_f % 32 != 0
18533 || w_bf16.len() < in_f * out_f * 2
18534 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18535 {
18536 return Err(format!(
18537 "encode_q8_0_from_bf16 geometry in={in_f} out={out_f} src={} dst={}",
18538 w_bf16.len(),
18539 out.len()
18540 )
18541 .into());
18542 }
18543 let f = self.func("encode_q8_0_rows_from_bf16");
18544 const PAIRS_PER_BLOCK: u32 = 4;
18547 let pairs = (out_f * (in_f / 32)) as u64;
18548 let cfg = LaunchConfig {
18549 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18550 block_dim: (32, PAIRS_PER_BLOCK, 1),
18551 shared_mem_bytes: 0,
18552 };
18553 let (ini, outi) = (in_f as i32, out_f as i32);
18554 let __s_b = self.gpu.stream();
18555 let mut b = __s_b.launch_builder(&f);
18556 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18557 unsafe {
18558 b.launch(cfg)?;
18559 }
18560 Ok(())
18561 }
18562
18563 pub fn encode_q8_0_from_bf16_view(
18567 &self,
18568 w_bf16: &cudarc::driver::CudaView<'_, u8>,
18569 out: &mut CudaSlice<u8>,
18570 in_f: usize,
18571 out_f: usize,
18572 ) -> Result<(), Box<dyn std::error::Error>> {
18573 if in_f % 32 != 0
18574 || w_bf16.len() < in_f * out_f * 2
18575 || out.len() < out_f * Self::q8_0_row_bytes(in_f)
18576 {
18577 return Err(format!(
18578 "encode_q8_0_from_bf16_view geometry in={in_f} out={out_f} src={} dst={}",
18579 w_bf16.len(),
18580 out.len()
18581 )
18582 .into());
18583 }
18584 let f = self.func("encode_q8_0_rows_from_bf16");
18585 const PAIRS_PER_BLOCK: u32 = 4;
18586 let pairs = (out_f * (in_f / 32)) as u64;
18587 let cfg = LaunchConfig {
18588 grid_dim: ((pairs.div_ceil(PAIRS_PER_BLOCK as u64)) as u32, 1, 1),
18589 block_dim: (32, PAIRS_PER_BLOCK, 1),
18590 shared_mem_bytes: 0,
18591 };
18592 let (ini, outi) = (in_f as i32, out_f as i32);
18593 let __s_b = self.gpu.stream();
18594 let mut b = __s_b.launch_builder(&f);
18595 b.arg(w_bf16).arg(out).arg(&ini).arg(&outi);
18596 unsafe {
18597 b.launch(cfg)?;
18598 }
18599 Ok(())
18600 }
18601
18602 #[allow(clippy::too_many_arguments)]
18607 pub fn qmatvec_q8_0_qkv_rp_into(
18608 &self,
18609 wq: &CudaSlice<u8>,
18610 wk: &CudaSlice<u8>,
18611 wv: &CudaSlice<u8>,
18612 aq: &CudaSlice<i8>,
18613 ad: &CudaSlice<f32>,
18614 yq: &mut CudaSlice<f32>,
18615 yk: &mut CudaSlice<f32>,
18616 yv: &mut CudaSlice<f32>,
18617 in_f: usize,
18618 out_q: usize,
18619 out_kv: usize,
18620 ) -> Result<(), Box<dyn std::error::Error>> {
18621 const ROWS_PER_BLOCK: u32 = 4; let rows = out_q + 2 * out_kv;
18623 let nblk = in_f / 32;
18624 if in_f % 32 != 0
18625 || aq.len() < in_f
18626 || ad.len() < nblk
18627 || yq.len() < out_q
18628 || yk.len() < out_kv
18629 || yv.len() < out_kv
18630 || wq.len() < out_q * nblk * 34
18631 || wk.len() < out_kv * nblk * 34
18632 || wv.len() < out_kv * nblk * 34
18633 {
18634 return Err(
18635 format!("q8_0 qkv rp geometry in={in_f} out_q={out_q} out_kv={out_kv}").into(),
18636 );
18637 }
18638 let f = self.func("qmatvec_q8_0_qkv_rp");
18639 let cfg = LaunchConfig {
18640 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18641 block_dim: (32, ROWS_PER_BLOCK, 1),
18642 shared_mem_bytes: 0,
18643 };
18644 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18645 let __s_b = self.gpu.stream();
18646 let mut b = __s_b.launch_builder(&f);
18647 b.arg(wq)
18648 .arg(wk)
18649 .arg(wv)
18650 .arg(aq)
18651 .arg(ad)
18652 .arg(yq)
18653 .arg(yk)
18654 .arg(yv)
18655 .arg(&ini)
18656 .arg(&oq)
18657 .arg(&okv);
18658 unsafe {
18659 b.launch(cfg)?;
18660 }
18661 Ok(())
18662 }
18663
18664 #[allow(clippy::too_many_arguments)]
18668 pub fn qmatvec_q8_0_b4_rp_into(
18669 &self,
18670 w: [&CudaSlice<u8>; 4],
18671 aq: &CudaSlice<i8>,
18672 ad: &CudaSlice<f32>,
18673 y: &mut CudaSlice<f32>,
18674 block_cols: usize,
18675 out_f: usize,
18676 ) -> Result<(), Box<dyn std::error::Error>> {
18677 const ROWS_PER_BLOCK: u32 = 4; let nblk = block_cols / 32;
18679 if block_cols % 32 != 0
18680 || aq.len() < 4 * block_cols
18681 || ad.len() < 4 * nblk
18682 || y.len() < out_f
18683 || w.iter().any(|p| p.len() < out_f * nblk * 34)
18684 {
18685 return Err(format!("q8_0 b4 rp geometry block_cols={block_cols} out={out_f}").into());
18686 }
18687 let f = self.func("qmatvec_q8_0_b4_rp");
18688 let cfg = LaunchConfig {
18689 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18690 block_dim: (32, ROWS_PER_BLOCK, 1),
18691 shared_mem_bytes: 0,
18692 };
18693 let (bc, of) = (block_cols as i32, out_f as i32);
18694 let __s_b = self.gpu.stream();
18695 let mut b = __s_b.launch_builder(&f);
18696 b.arg(w[0])
18697 .arg(w[1])
18698 .arg(w[2])
18699 .arg(w[3])
18700 .arg(aq)
18701 .arg(ad)
18702 .arg(y)
18703 .arg(&bc)
18704 .arg(&of);
18705 unsafe {
18706 b.launch(cfg)?;
18707 }
18708 Ok(())
18709 }
18710
18711 fn matvec_bf16_via_q8_mirror_t(
18714 &self,
18715 data: &CudaSlice<u8>,
18716 x: &CudaSlice<f32>,
18717 y: &mut CudaSlice<f32>,
18718 in_f: usize,
18719 out_f: usize,
18720 t: usize,
18721 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18722 use cudarc::driver::DevicePtr;
18723 let key = {
18724 let s = self.gpu.stream();
18725 let (p, _g) = data.device_ptr(&s);
18726 (p as u64, in_f as u32, out_f as u32)
18727 };
18728 {
18729 let mut mirrors = self
18730 .w8_mirrors
18731 .lock()
18732 .map_err(|_| "w8 mirror map is poisoned")?;
18733 if !mirrors.contains_key(&key) {
18734 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18735 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18736 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18737 mirrors.insert(key, planar);
18738 }
18739 }
18740 let nblk = in_f / 32;
18741 let akey = in_f * 64 + t.min(32);
18743 {
18744 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18745 if !act.contains_key(&akey) {
18746 let aq = self.alloc_i8_uninit(32 * in_f)?;
18747 let ad = self.alloc_uninit::<f32>(32 * nblk)?;
18748 act.insert(akey, (aq, ad));
18749 }
18750 let (aq, ad) = act.get_mut(&akey).expect("just inserted");
18751 self.quantize_q8_1_into(x, t, in_f, aq, ad)?;
18752 }
18753 let mirrors = self
18754 .w8_mirrors
18755 .lock()
18756 .map_err(|_| "w8 mirror map is poisoned")?;
18757 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18758 let mirror = mirrors.get(&key).expect("built above");
18759 let (aq, ad) = act.get(&akey).expect("built above");
18760 const ROWS_PER_BLOCK: u32 = 4;
18761 let (ini, of) = (in_f as i32, out_f as i32);
18762 if q8t_wonce_on() && t <= 32 {
18766 let f = self.func(if t <= 8 {
18767 "qmatvec_q8_0_rows_tw"
18768 } else {
18769 "qmatvec_q8_0_rows_tw32"
18770 });
18771 let cfg = LaunchConfig {
18772 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18773 block_dim: (32, ROWS_PER_BLOCK, 1),
18774 shared_mem_bytes: 0,
18775 };
18776 let ti = t as i32;
18777 let __s_b = self.gpu.stream();
18778 let mut b = __s_b.launch_builder(&f);
18779 b.arg(mirror)
18780 .arg(aq)
18781 .arg(ad)
18782 .arg(&mut *y)
18783 .arg(&ini)
18784 .arg(&of)
18785 .arg(&ti);
18786 unsafe {
18787 b.launch(cfg)?;
18788 }
18789 return Ok(Some(()));
18790 }
18791 let f = self.func("qmatvec_q8_0_rows_t");
18792 let cfg = LaunchConfig {
18793 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
18794 block_dim: (32, ROWS_PER_BLOCK, 1),
18795 shared_mem_bytes: 0,
18796 };
18797 let __s_b = self.gpu.stream();
18798 let mut b = __s_b.launch_builder(&f);
18799 b.arg(mirror)
18800 .arg(aq)
18801 .arg(ad)
18802 .arg(&mut *y)
18803 .arg(&ini)
18804 .arg(&of);
18805 unsafe {
18806 b.launch(cfg)?;
18807 }
18808 Ok(Some(()))
18809 }
18810
18811 fn matvec_bf16_via_q8_mirror(
18814 &self,
18815 data: &CudaSlice<u8>,
18816 x: &CudaSlice<f32>,
18817 y: &mut CudaSlice<f32>,
18818 in_f: usize,
18819 out_f: usize,
18820 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
18821 use cudarc::driver::DevicePtr;
18822 let key = {
18823 let s = self.gpu.stream();
18824 let (p, _g) = data.device_ptr(&s);
18825 (p as u64, in_f as u32, out_f as u32)
18826 };
18827 {
18828 let mut mirrors = self
18829 .w8_mirrors
18830 .lock()
18831 .map_err(|_| "w8 mirror map is poisoned")?;
18832 if !mirrors.contains_key(&key) {
18833 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
18834 self.encode_q8_0_from_bf16(data, &mut interleaved, in_f, out_f)?;
18835 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
18836 mirrors.insert(key, planar);
18837 if std::env::var("MEMRA_W8_TRACE").as_deref() == Ok("1") {
18843 eprintln!(
18844 "[w8-mirror] built in_f={in_f} out_f={out_f} mirrors={}",
18845 mirrors.len()
18846 );
18847 }
18848 }
18849 }
18850 let nblk = in_f / 32;
18851 {
18852 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18853 if !act.contains_key(&in_f) {
18854 let aq = self.alloc_uninit::<i8>(in_f)?;
18855 let ad = self.alloc_uninit::<f32>(nblk)?;
18856 act.insert(in_f, (aq, ad));
18857 }
18858 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
18859 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
18860 }
18861 let mirrors = self
18862 .w8_mirrors
18863 .lock()
18864 .map_err(|_| "w8 mirror map is poisoned")?;
18865 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
18866 let mirror = mirrors.get(&key).expect("built above");
18867 let (aq, ad) = act.get(&in_f).expect("built above");
18868 self.qmatvec_mmvq_into(
18869 mirror,
18870 aq,
18871 ad,
18872 1,
18873 in_f,
18874 out_f,
18875 QT_Q8_0,
18876 Self::q8_0_row_bytes(in_f),
18877 1.0,
18878 true,
18879 y,
18880 )?;
18881 Ok(Some(()))
18882 }
18883
18884 #[allow(clippy::too_many_arguments)]
18889 pub fn qmatvec_q8_0_qkv_rp_t_into(
18890 &self,
18891 wq: &CudaSlice<u8>,
18892 wk: &CudaSlice<u8>,
18893 wv: &CudaSlice<u8>,
18894 aq: &CudaSlice<i8>,
18895 ad: &CudaSlice<f32>,
18896 yq: &mut CudaSlice<f32>,
18897 yk: &mut CudaSlice<f32>,
18898 yv: &mut CudaSlice<f32>,
18899 in_f: usize,
18900 out_q: usize,
18901 out_kv: usize,
18902 t: usize,
18903 ) -> Result<(), Box<dyn std::error::Error>> {
18904 const ROWS_PER_BLOCK: u32 = 4;
18905 let rows = out_q + 2 * out_kv;
18906 let nblk = in_f / 32;
18907 if in_f % 32 != 0
18908 || t == 0
18909 || aq.len() < t * in_f
18910 || ad.len() < t * nblk
18911 || yq.len() < t * out_q
18912 || yk.len() < t * out_kv
18913 || yv.len() < t * out_kv
18914 {
18915 return Err(format!("q8_0 qkv rp_t geometry in={in_f} t={t}").into());
18916 }
18917 let (ini, oq, okv) = (in_f as i32, out_q as i32, out_kv as i32);
18918 if q8t_wonce_on() && t <= 32 {
18922 let f = self.func(if t <= 8 {
18923 "qmatvec_q8_0_qkv_rp_tw"
18924 } else {
18925 "qmatvec_q8_0_qkv_rp_tw32"
18926 });
18927 let cfg = LaunchConfig {
18928 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
18929 block_dim: (32, ROWS_PER_BLOCK, 1),
18930 shared_mem_bytes: 0,
18931 };
18932 let ti = t as i32;
18933 let __s_b = self.gpu.stream();
18934 let mut b = __s_b.launch_builder(&f);
18935 b.arg(wq)
18936 .arg(wk)
18937 .arg(wv)
18938 .arg(aq)
18939 .arg(ad)
18940 .arg(yq)
18941 .arg(yk)
18942 .arg(yv)
18943 .arg(&ini)
18944 .arg(&oq)
18945 .arg(&okv)
18946 .arg(&ti);
18947 unsafe {
18948 b.launch(cfg)?;
18949 }
18950 return Ok(());
18951 }
18952 let f = self.func("qmatvec_q8_0_qkv_rp_t");
18953 let cfg = LaunchConfig {
18954 grid_dim: ((rows as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
18955 block_dim: (32, ROWS_PER_BLOCK, 1),
18956 shared_mem_bytes: 0,
18957 };
18958 let __s_b = self.gpu.stream();
18959 let mut b = __s_b.launch_builder(&f);
18960 b.arg(wq)
18961 .arg(wk)
18962 .arg(wv)
18963 .arg(aq)
18964 .arg(ad)
18965 .arg(yq)
18966 .arg(yk)
18967 .arg(yv)
18968 .arg(&ini)
18969 .arg(&oq)
18970 .arg(&okv);
18971 unsafe {
18972 b.launch(cfg)?;
18973 }
18974 Ok(())
18975 }
18976
18977 #[allow(clippy::too_many_arguments)]
18980 pub fn qmatvec_q8_0_b4_rp_t_into(
18981 &self,
18982 w: [&CudaSlice<u8>; 4],
18983 aq: &CudaSlice<i8>,
18984 ad: &CudaSlice<f32>,
18985 y: &mut CudaSlice<f32>,
18986 block_cols: usize,
18987 out_f: usize,
18988 t: usize,
18989 ) -> Result<(), Box<dyn std::error::Error>> {
18990 const ROWS_PER_BLOCK: u32 = 4;
18991 let nblk = block_cols / 32;
18992 if block_cols % 32 != 0
18993 || t == 0
18994 || aq.len() < t * 4 * block_cols
18995 || ad.len() < t * 4 * nblk
18996 || y.len() < t * out_f
18997 {
18998 return Err(format!("q8_0 b4 rp_t geometry cols={block_cols} t={t}").into());
18999 }
19000 let (bc, of) = (block_cols as i32, out_f as i32);
19001 if q8t_wonce_on() && t <= 32 {
19004 let f = self.func(if t <= 8 {
19005 "qmatvec_q8_0_b4_rp_tw"
19006 } else {
19007 "qmatvec_q8_0_b4_rp_tw32"
19008 });
19009 let cfg = LaunchConfig {
19010 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), 1, 1),
19011 block_dim: (32, ROWS_PER_BLOCK, 1),
19012 shared_mem_bytes: 0,
19013 };
19014 let ti = t as i32;
19015 let __s_b = self.gpu.stream();
19016 let mut b = __s_b.launch_builder(&f);
19017 b.arg(w[0])
19018 .arg(w[1])
19019 .arg(w[2])
19020 .arg(w[3])
19021 .arg(aq)
19022 .arg(ad)
19023 .arg(y)
19024 .arg(&bc)
19025 .arg(&of)
19026 .arg(&ti);
19027 unsafe {
19028 b.launch(cfg)?;
19029 }
19030 return Ok(());
19031 }
19032 let f = self.func("qmatvec_q8_0_b4_rp_t");
19033 let cfg = LaunchConfig {
19034 grid_dim: ((out_f as u32).div_ceil(ROWS_PER_BLOCK), t as u32, 1),
19035 block_dim: (32, ROWS_PER_BLOCK, 1),
19036 shared_mem_bytes: 0,
19037 };
19038 let __s_b = self.gpu.stream();
19039 let mut b = __s_b.launch_builder(&f);
19040 b.arg(w[0])
19041 .arg(w[1])
19042 .arg(w[2])
19043 .arg(w[3])
19044 .arg(aq)
19045 .arg(ad)
19046 .arg(y)
19047 .arg(&bc)
19048 .arg(&of);
19049 unsafe {
19050 b.launch(cfg)?;
19051 }
19052 Ok(())
19053 }
19054
19055 fn matvec_bf16_view_via_q8_mirror(
19065 &self,
19066 data: &cudarc::driver::CudaView<'_, u8>,
19067 x: &CudaSlice<f32>,
19068 y: &mut CudaSlice<f32>,
19069 in_f: usize,
19070 out_f: usize,
19071 ) -> Result<Option<()>, Box<dyn std::error::Error>> {
19072 use cudarc::driver::DevicePtr;
19073 let key = {
19074 let s = self.gpu.stream();
19075 let (p, _g) = data.device_ptr(&s);
19076 (p as u64, in_f as u32, out_f as u32)
19077 };
19078 {
19079 let mut mirrors = self
19080 .w8_mirrors
19081 .lock()
19082 .map_err(|_| "w8 mirror map is poisoned")?;
19083 if !mirrors.contains_key(&key) {
19084 let mut interleaved = self.alloc_u8_uninit(out_f * Self::q8_0_row_bytes(in_f))?;
19085 self.encode_q8_0_from_bf16_view(data, &mut interleaved, in_f, out_f)?;
19086 let planar = self.build_q8_rp4_raw(&interleaved, in_f, out_f)?;
19087 mirrors.insert(key, planar);
19088 eprintln!("[w8-view] mirror built in_f={in_f} out_f={out_f}");
19092 }
19093 }
19094 let nblk = in_f / 32;
19095 {
19096 let mut act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
19097 if !act.contains_key(&in_f) {
19098 let aq = self.alloc_uninit::<i8>(in_f)?;
19099 let ad = self.alloc_uninit::<f32>(nblk)?;
19100 act.insert(in_f, (aq, ad));
19101 }
19102 let (aq, ad) = act.get_mut(&in_f).expect("just inserted");
19103 self.quantize_q8_1_into(x, 1, in_f, aq, ad)?;
19104 }
19105 let mirrors = self
19106 .w8_mirrors
19107 .lock()
19108 .map_err(|_| "w8 mirror map is poisoned")?;
19109 let act = self.w8_act.lock().map_err(|_| "w8 act map is poisoned")?;
19110 let mirror = mirrors.get(&key).expect("built above");
19111 let (aq, ad) = act.get(&in_f).expect("built above");
19112 self.qmatvec_mmvq_into(
19113 mirror,
19114 aq,
19115 ad,
19116 1,
19117 in_f,
19118 out_f,
19119 QT_Q8_0,
19120 Self::q8_0_row_bytes(in_f),
19121 1.0,
19122 true,
19123 y,
19124 )?;
19125 Ok(Some(()))
19126 }
19127
19128 pub fn matvec_bf16_into(
19129 &self,
19130 data: &CudaSlice<u8>,
19131 x: &CudaSlice<f32>,
19132 y: &mut CudaSlice<f32>,
19133 in_f: usize,
19134 out_f: usize,
19135 ) -> Result<(), Box<dyn std::error::Error>> {
19136 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
19137 return Err(format!(
19138 "matvec_bf16_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
19139 data.len(),
19140 x.len(),
19141 y.len()
19142 )
19143 .into());
19144 }
19145 if step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19152 if let Some(()) = self.matvec_bf16_via_q8_mirror(data, x, y, in_f, out_f)? {
19153 return Ok(());
19154 }
19155 }
19156 static X4: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
19160 let x4 = *X4.get_or_init(|| std::env::var("MEMRA_DOWN_X4").as_deref() == Ok("1"))
19161 && in_f <= 2048;
19162 if x4 {
19163 let f = self.func("matvec_bf16_f32acc_x4");
19164 let cfg = LaunchConfig {
19165 grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
19166 block_dim: (mmv_block(), 1, 1),
19167 shared_mem_bytes: 0,
19168 };
19169 let (ini, outi) = (in_f as i32, out_f as i32);
19170 let __s_b = self.gpu.stream();
19171 let mut b = __s_b.launch_builder(&f);
19172 b.arg(data).arg(x).arg(y).arg(&ini).arg(&outi);
19173 unsafe {
19174 b.launch(cfg)?;
19175 }
19176 return Ok(());
19177 }
19178 let f = self.func("matvec_bf16_f32acc");
19179 let cfg = LaunchConfig {
19180 grid_dim: (out_f as u32, 1, 1),
19181 block_dim: (mmv_block(), 1, 1),
19182 shared_mem_bytes: 0,
19183 };
19184 let ini = in_f as i32;
19185 let __s_b = self.gpu.stream();
19186 let mut b = __s_b.launch_builder(&f);
19187 b.arg(data).arg(x).arg(y).arg(&ini);
19188 unsafe {
19189 b.launch(cfg)?;
19190 }
19191 Ok(())
19192 }
19193
19194 pub fn matvec_bf16_view_into(
19197 &self,
19198 data: &cudarc::driver::CudaView<'_, u8>,
19199 x: &CudaSlice<f32>,
19200 y: &mut CudaSlice<f32>,
19201 in_f: usize,
19202 out_f: usize,
19203 ) -> Result<(), Box<dyn std::error::Error>> {
19204 if data.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y.len() < out_f {
19205 return Err(format!(
19206 "matvec_bf16_view_into geometry bytes={} x={} y={} in={in_f} out={out_f}",
19207 data.len(),
19208 x.len(),
19209 y.len()
19210 )
19211 .into());
19212 }
19213 if w8_view_on() && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19214 if let Some(()) = self.matvec_bf16_view_via_q8_mirror(data, x, y, in_f, out_f)? {
19215 return Ok(());
19216 }
19217 }
19218 let f = self.func("matvec_bf16_f32acc");
19219 let cfg = LaunchConfig {
19220 grid_dim: (out_f as u32, 1, 1),
19221 block_dim: (mmv_block(), 1, 1),
19222 shared_mem_bytes: 0,
19223 };
19224 let ini = in_f as i32;
19225 let __s_b = self.gpu.stream();
19226 let mut b = __s_b.launch_builder(&f);
19227 b.arg(data).arg(x).arg(y).arg(&ini);
19228 unsafe {
19229 b.launch(cfg)?;
19230 }
19231 Ok(())
19232 }
19233
19234 pub fn matvec_bf16_raw_out(
19237 &self,
19238 w: &CudaSlice<u8>,
19239 x: &CudaSlice<f32>,
19240 y_raw: u64,
19241 in_f: usize,
19242 out_f: usize,
19243 ) -> Result<(), Box<dyn std::error::Error>> {
19244 if w.len() != in_f * out_f * 2 || x.len() < in_f || in_f % 8 != 0 || y_raw == 0 {
19245 return Err("matvec_bf16_raw_out geometry".into());
19246 }
19247 let f = self.func("matvec_bf16_f32acc");
19248 let cfg = LaunchConfig {
19249 grid_dim: (out_f as u32, 1, 1),
19250 block_dim: (mmv_block(), 1, 1),
19251 shared_mem_bytes: 0,
19252 };
19253 let ini = in_f as i32;
19254 let __s_b = self.gpu.stream();
19255 let mut b = __s_b.launch_builder(&f);
19256 b.arg(w).arg(x).arg(&y_raw).arg(&ini);
19257 unsafe {
19258 b.launch(cfg)?;
19259 }
19260 Ok(())
19261 }
19262
19263 pub fn add3_raw(
19267 &self,
19268 a: &CudaSlice<f32>,
19269 b: &CudaSlice<f32>,
19270 sh_raw: u64,
19271 scale_raw: u64,
19272 dst: &mut CudaSlice<f32>,
19273 n: usize,
19274 ) -> Result<(), Box<dyn std::error::Error>> {
19275 if a.len() < n || b.len() < n || dst.len() < n || sh_raw == 0 || scale_raw == 0 {
19276 return Err("add3_raw geometry".into());
19277 }
19278 let f = self.func("add3_f32");
19279 let cfg = LaunchConfig {
19280 grid_dim: ((n as u32).div_ceil(256), 1, 1),
19281 block_dim: (256, 1, 1),
19282 shared_mem_bytes: 0,
19283 };
19284 let ni = n as i32;
19285 let __s_b = self.gpu.stream();
19286 let mut bld = __s_b.launch_builder(&f);
19287 bld.arg(a)
19288 .arg(b)
19289 .arg(&sh_raw)
19290 .arg(&scale_raw)
19291 .arg(dst)
19292 .arg(&ni);
19293 unsafe {
19294 bld.launch(cfg)?;
19295 }
19296 Ok(())
19297 }
19298
19299 pub fn matvec_bf16_down_addscale_into(
19302 &self,
19303 w: &CudaSlice<u8>,
19304 x: &CudaSlice<f32>,
19305 scale: &CudaSlice<f32>,
19306 dst: &mut CudaSlice<f32>,
19307 in_f: usize,
19308 out_f: usize,
19309 ) -> Result<(), Box<dyn std::error::Error>> {
19310 if w.len() != in_f * out_f * 2
19311 || x.len() < in_f
19312 || in_f % 8 != 0
19313 || dst.len() < out_f
19314 || scale.is_empty()
19315 {
19316 return Err("matvec_bf16_down_addscale geometry".into());
19317 }
19318 let f = self.func("matvec_bf16_down_addscale");
19319 let cfg = LaunchConfig {
19320 grid_dim: (out_f as u32, 1, 1),
19321 block_dim: (mmv_block(), 1, 1),
19322 shared_mem_bytes: 0,
19323 };
19324 let ini = in_f as i32;
19325 let __s_b = self.gpu.stream();
19326 let mut b = __s_b.launch_builder(&f);
19327 b.arg(w).arg(x).arg(scale).arg(dst).arg(&ini);
19328 unsafe {
19329 b.launch(cfg)?;
19330 }
19331 Ok(())
19332 }
19333
19334 #[allow(clippy::too_many_arguments)]
19338 pub fn matvec_bf16_dual_silu_rows_into(
19339 &self,
19340 wg: &CudaSlice<u8>,
19341 wu: &CudaSlice<u8>,
19342 x: &CudaSlice<f32>,
19343 act: &mut CudaSlice<f32>,
19344 in_f: usize,
19345 out_f: usize,
19346 limit: Option<f32>,
19347 t: usize,
19348 ) -> Result<(), Box<dyn std::error::Error>> {
19349 if x.len() < t * in_f || act.len() < t * out_f || t == 0 || t > 32 {
19350 return Err("matvec_bf16_dual_silu_rows geometry".into());
19351 }
19352 let f = self.func("matvec_bf16_dual_silu_rows");
19353 let cfg = LaunchConfig {
19354 grid_dim: (out_f as u32, t as u32, 1),
19355 block_dim: (mmv_block(), 1, 1),
19356 shared_mem_bytes: 0,
19357 };
19358 let (ini, outi) = (in_f as i32, out_f as i32);
19359 let lim = limit.unwrap_or(0.0);
19360 let __s_b = self.gpu.stream();
19361 let mut b = __s_b.launch_builder(&f);
19362 b.arg(wg)
19363 .arg(wu)
19364 .arg(x)
19365 .arg(&mut *act)
19366 .arg(&ini)
19367 .arg(&outi)
19368 .arg(&lim);
19369 unsafe {
19370 b.launch(cfg)?;
19371 }
19372 Ok(())
19373 }
19374
19375 pub fn matvec_bf16_rows_into(
19377 &self,
19378 w: &CudaSlice<u8>,
19379 x: &CudaSlice<f32>,
19380 y: &mut CudaSlice<f32>,
19381 in_f: usize,
19382 out_f: usize,
19383 t: usize,
19384 ) -> Result<(), Box<dyn std::error::Error>> {
19385 if x.len() < t * in_f || y.len() < t * out_f || t == 0 || t > 32 || in_f % 8 != 0 {
19386 return Err("matvec_bf16_rows geometry".into());
19387 }
19388 if t >= 2 && t <= 32 && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19393 if let Some(()) = self.matvec_bf16_via_q8_mirror_t(w, x, y, in_f, out_f, t)? {
19394 return Ok(());
19395 }
19396 }
19397 if t == 1 && step_tp_w8_on() && w8_hybrid_on() && in_f % 32 == 0 && out_f >= 64 {
19403 if let Some(()) = self.matvec_bf16_via_q8_mirror(w, x, y, in_f, out_f)? {
19404 return Ok(());
19405 }
19406 }
19407 let f = self.func("matvec_bf16_f32acc_x4_rows");
19408 let cfg = LaunchConfig {
19409 grid_dim: (out_f.div_ceil(4) as u32, t as u32, 1),
19410 block_dim: (mmv_block(), 1, 1),
19411 shared_mem_bytes: 0,
19412 };
19413 let (ini, outi) = (in_f as i32, out_f as i32);
19414 let __s_b = self.gpu.stream();
19415 let mut b = __s_b.launch_builder(&f);
19416 b.arg(w).arg(x).arg(&mut *y).arg(&ini).arg(&outi);
19417 unsafe {
19418 b.launch(cfg)?;
19419 }
19420 Ok(())
19421 }
19422
19423 pub fn matvec_bf16_dual_silu_into(
19424 &self,
19425 wg: &CudaSlice<u8>,
19426 wu: &CudaSlice<u8>,
19427 x: &CudaSlice<f32>,
19428 act: &mut CudaSlice<f32>,
19429 in_f: usize,
19430 out_f: usize,
19431 limit: Option<f32>,
19432 ) -> Result<(), Box<dyn std::error::Error>> {
19433 if wg.len() != in_f * out_f * 2
19434 || wu.len() != in_f * out_f * 2
19435 || x.len() < in_f
19436 || in_f % 8 != 0
19437 || act.len() < out_f
19438 {
19439 return Err("matvec_bf16_dual_silu geometry".into());
19440 }
19441 let f = self.func("matvec_bf16_dual_silu");
19442 let cfg = LaunchConfig {
19443 grid_dim: (out_f as u32, 1, 1),
19444 block_dim: (mmv_block(), 1, 1),
19445 shared_mem_bytes: 0,
19446 };
19447 let (ini, outi) = (in_f as i32, out_f as i32);
19448 let lim = limit.unwrap_or(0.0);
19449 let __s_b = self.gpu.stream();
19450 let mut b = __s_b.launch_builder(&f);
19451 b.arg(wg)
19452 .arg(wu)
19453 .arg(x)
19454 .arg(act)
19455 .arg(&ini)
19456 .arg(&outi)
19457 .arg(&lim);
19458 unsafe {
19459 b.launch(cfg)?;
19460 }
19461 Ok(())
19462 }
19463
19464 #[allow(clippy::too_many_arguments)]
19467 pub fn matvec_bf16_dual_view_into(
19468 &self,
19469 wg: &cudarc::driver::CudaView<'_, u8>,
19470 wu: &cudarc::driver::CudaView<'_, u8>,
19471 x: &CudaSlice<f32>,
19472 yg: &mut CudaSlice<f32>,
19473 yu: &mut CudaSlice<f32>,
19474 in_f: usize,
19475 out_f: usize,
19476 ) -> Result<(), Box<dyn std::error::Error>> {
19477 if wg.len() != in_f * out_f * 2
19478 || wu.len() != in_f * out_f * 2
19479 || x.len() < in_f
19480 || in_f % 8 != 0
19481 || yg.len() < out_f
19482 || yu.len() < out_f
19483 {
19484 return Err(format!(
19485 "matvec_bf16_dual_view_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
19486 wg.len(),
19487 wu.len(),
19488 x.len()
19489 )
19490 .into());
19491 }
19492 let f = self.func("matvec_bf16_dual");
19493 let cfg = LaunchConfig {
19494 grid_dim: ((2 * out_f) as u32, 1, 1),
19495 block_dim: (mmv_block(), 1, 1),
19496 shared_mem_bytes: 0,
19497 };
19498 let (ini, outi) = (in_f as i32, out_f as i32);
19499 let __s_b = self.gpu.stream();
19500 let mut b = __s_b.launch_builder(&f);
19501 b.arg(wg)
19502 .arg(wu)
19503 .arg(x)
19504 .arg(yg)
19505 .arg(yu)
19506 .arg(&ini)
19507 .arg(&outi);
19508 unsafe {
19509 b.launch(cfg)?;
19510 }
19511 Ok(())
19512 }
19513
19514 #[allow(clippy::too_many_arguments)]
19516 pub fn matvec_bf16_dual_into(
19517 &self,
19518 wg: &CudaSlice<u8>,
19519 wu: &CudaSlice<u8>,
19520 x: &CudaSlice<f32>,
19521 yg: &mut CudaSlice<f32>,
19522 yu: &mut CudaSlice<f32>,
19523 in_f: usize,
19524 out_f: usize,
19525 ) -> Result<(), Box<dyn std::error::Error>> {
19526 if wg.len() != in_f * out_f * 2
19527 || wu.len() != in_f * out_f * 2
19528 || x.len() < in_f
19529 || in_f % 8 != 0
19530 || yg.len() < out_f
19531 || yu.len() < out_f
19532 {
19533 return Err(format!(
19534 "matvec_bf16_dual_into geometry wg={} wu={} x={} in={in_f} out={out_f}",
19535 wg.len(),
19536 wu.len(),
19537 x.len()
19538 )
19539 .into());
19540 }
19541 let f = self.func("matvec_bf16_dual");
19542 let cfg = LaunchConfig {
19543 grid_dim: ((2 * out_f) as u32, 1, 1),
19544 block_dim: (mmv_block(), 1, 1),
19545 shared_mem_bytes: 0,
19546 };
19547 let (ini, outi) = (in_f as i32, out_f as i32);
19548 let __s_b = self.gpu.stream();
19549 let mut b = __s_b.launch_builder(&f);
19550 b.arg(wg)
19551 .arg(wu)
19552 .arg(x)
19553 .arg(yg)
19554 .arg(yu)
19555 .arg(&ini)
19556 .arg(&outi);
19557 unsafe {
19558 b.launch(cfg)?;
19559 }
19560 Ok(())
19561 }
19562
19563 pub(crate) fn matvec_bf16_dual(
19566 &self,
19567 wg: &CudaSlice<u8>,
19568 wu: &CudaSlice<u8>,
19569 x: &CudaSlice<f32>,
19570 in_f: usize,
19571 out_f: usize,
19572 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
19573 if wg.len() != in_f * out_f * 2
19574 || wu.len() != in_f * out_f * 2
19575 || x.len() < in_f
19576 || in_f % 8 != 0
19577 {
19578 return Err(format!(
19579 "matvec_bf16_dual geometry wg={} wu={} x={} in={in_f} out={out_f}",
19580 wg.len(),
19581 wu.len(),
19582 x.len()
19583 )
19584 .into());
19585 }
19586 let mut yg = self.alloc_uninit::<f32>(out_f)?;
19587 let mut yu = self.alloc_uninit::<f32>(out_f)?;
19588 let f = self.func("matvec_bf16_dual");
19589 let cfg = LaunchConfig {
19590 grid_dim: ((2 * out_f) as u32, 1, 1),
19591 block_dim: (mmv_block(), 1, 1),
19592 shared_mem_bytes: 0,
19593 };
19594 let (ini, outi) = (in_f as i32, out_f as i32);
19595 let __s_b = self.gpu.stream();
19596 let mut b = __s_b.launch_builder(&f);
19597 b.arg(wg)
19598 .arg(wu)
19599 .arg(x)
19600 .arg(&mut yg)
19601 .arg(&mut yu)
19602 .arg(&ini)
19603 .arg(&outi);
19604 unsafe {
19605 b.launch(cfg)?;
19606 }
19607 Ok((yg, yu))
19608 }
19609
19610 #[allow(clippy::too_many_arguments)]
19611 fn linear_bf16_chunked_inner(
19612 &self,
19613 x: &CudaSlice<f32>,
19614 data: &CudaSlice<u8>,
19615 m: usize,
19616 in_f: usize,
19617 out_f: usize,
19618 exact: bool,
19619 canonical_chunk_rows: Option<usize>,
19620 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19621 const CHUNK_BYTES: usize = 256 << 20;
19622 if m == 1
19625 && !exact
19626 && canonical_chunk_rows.is_none()
19627 && in_f % 8 == 0
19628 && Self::bf16_mmv_on()
19629 {
19630 return self.matvec_bf16(data, x, in_f, out_f);
19631 }
19632 if m >= 16
19637 && !exact
19638 && canonical_chunk_rows.is_none()
19639 && data.len() == in_f * out_f * 2
19640 && crate::f16_ffi::pp_bf16_enabled()
19641 {
19642 if let Some(y) = self.bf16_tc_gemm(data, x, m, in_f, out_f)? {
19645 return Ok(y);
19646 }
19647 }
19648 let row_bytes = in_f
19649 .checked_mul(std::mem::size_of::<f32>())
19650 .ok_or("BF16 chunk row byte count overflow")?;
19651 if row_bytes == 0 || out_f == 0 {
19652 return Err("BF16 chunk dimensions must be nonzero".into());
19653 }
19654 let max_chunk_rows = (CHUNK_BYTES / row_bytes).max(1).min(out_f);
19655 let chunk_rows = match canonical_chunk_rows {
19656 Some(rows) if rows == 0 => {
19657 return Err("canonical BF16 chunk rows must be nonzero".into());
19658 }
19659 Some(rows) if rows > max_chunk_rows => {
19660 return Err(format!(
19661 "canonical BF16 chunk rows {rows} exceed the {max_chunk_rows}-row scratch limit"
19662 )
19663 .into());
19664 }
19665 Some(rows) if out_f % rows != 0 => {
19666 return Err(format!(
19667 "BF16 output width {out_f} is not divisible by canonical {rows}-row chunks"
19668 )
19669 .into());
19670 }
19671 Some(rows) => rows,
19672 None => max_chunk_rows,
19673 };
19674 if chunk_rows >= out_f {
19675 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
19676 return if exact {
19677 self.linear_decode_exact(x, &wf32, m, in_f, out_f)
19678 } else {
19679 self.linear(x, &wf32, m, in_f, out_f)
19680 };
19681 }
19682 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19683 let mut r0 = 0usize;
19684 while r0 < out_f {
19685 let rows = chunk_rows.min(out_f - r0);
19686 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
19687 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
19688 let yc = if exact {
19689 self.linear_decode_exact(x, &wf32, m, in_f, rows)?
19690 } else {
19691 self.linear(x, &wf32, m, in_f, rows)?
19692 };
19693 for mi in 0..m {
19695 let src = yc.slice(mi * rows..(mi + 1) * rows);
19696 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
19697 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
19698 }
19699 r0 += rows;
19700 }
19701 Ok(y)
19702 }
19703
19704 pub fn linear_bf16_resident(
19708 &self,
19709 x: &CudaSlice<f32>,
19710 data: &CudaSlice<u8>,
19711 m: usize,
19712 in_f: usize,
19713 out_f: usize,
19714 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19715 if data.len() != in_f * out_f * 2 {
19716 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19717 }
19718 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, None)
19719 }
19720
19721 pub fn linear_bf16_resident_canonical_rows(
19727 &self,
19728 x: &CudaSlice<f32>,
19729 data: &CudaSlice<u8>,
19730 m: usize,
19731 in_f: usize,
19732 out_f: usize,
19733 canonical_chunk_rows: usize,
19734 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19735 if data.len() != in_f * out_f * 2 {
19736 return Err(format!("resident BF16 bytes {} != {out_f}x{in_f}x2", data.len()).into());
19737 }
19738 self.linear_bf16_chunked(x, data, m, in_f, out_f, false, Some(canonical_chunk_rows))
19739 }
19740
19741 pub fn linear_f32_resident_canonical_rows(
19746 &self,
19747 x: &CudaSlice<f32>,
19748 data: &CudaSlice<f32>,
19749 m: usize,
19750 in_f: usize,
19751 out_f: usize,
19752 canonical_chunk_rows: usize,
19753 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19754 self.linear_f32_resident_canonical_rows_inner(
19755 x,
19756 data,
19757 m,
19758 in_f,
19759 out_f,
19760 canonical_chunk_rows,
19761 false,
19762 )
19763 }
19764
19765 pub fn linear_f32_resident_canonical_rows_strided(
19771 &self,
19772 x: &CudaSlice<f32>,
19773 data: &CudaSlice<f32>,
19774 m: usize,
19775 in_f: usize,
19776 out_f: usize,
19777 canonical_chunk_rows: usize,
19778 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19779 self.linear_f32_resident_canonical_rows_inner(
19780 x,
19781 data,
19782 m,
19783 in_f,
19784 out_f,
19785 canonical_chunk_rows,
19786 true,
19787 )
19788 }
19789
19790 fn linear_f32_resident_canonical_rows_inner(
19791 &self,
19792 x: &CudaSlice<f32>,
19793 data: &CudaSlice<f32>,
19794 m: usize,
19795 in_f: usize,
19796 out_f: usize,
19797 canonical_chunk_rows: usize,
19798 strided_output: bool,
19799 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19800 if data.len() != in_f * out_f {
19801 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19802 }
19803 if canonical_chunk_rows == 0
19804 || canonical_chunk_rows > out_f
19805 || out_f % canonical_chunk_rows != 0
19806 {
19807 return Err(format!(
19808 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19809 )
19810 .into());
19811 }
19812 if canonical_chunk_rows == out_f {
19813 return self.linear(x, data, m, in_f, out_f);
19814 }
19815
19816 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
19817 let input = x.slice(0..x.len());
19818 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19819 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19820 if m == 1 {
19821 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19822 self.linear_device_into(
19823 &input,
19824 &weights,
19825 &mut destination,
19826 1,
19827 in_f,
19828 canonical_chunk_rows,
19829 )?;
19830 continue;
19831 }
19832 let chunk = self.linear_device(&input, &weights, m, in_f, canonical_chunk_rows)?;
19833 if strided_output {
19834 self.place_rows_strided(&chunk, &mut y, canonical_chunk_rows, m, out_f, r0)?;
19835 } else {
19836 for token in 0..m {
19837 let source = chunk
19838 .slice(token * canonical_chunk_rows..(token + 1) * canonical_chunk_rows);
19839 let mut destination =
19840 y.slice_mut(token * out_f + r0..token * out_f + r0 + canonical_chunk_rows);
19841 self.gpu.stream().memcpy_dtod(&source, &mut destination)?;
19842 }
19843 }
19844 }
19845 Ok(y)
19846 }
19847
19848 pub fn linear_f32_resident_canonical_rows_t1_into(
19853 &self,
19854 x: &CudaSlice<f32>,
19855 data: &CudaSlice<f32>,
19856 y: &mut CudaSlice<f32>,
19857 in_f: usize,
19858 out_f: usize,
19859 canonical_chunk_rows: usize,
19860 ) -> Result<(), Box<dyn std::error::Error>> {
19861 if data.len() != in_f * out_f {
19862 return Err(format!("resident F32 values {} != {out_f}x{in_f}", data.len()).into());
19863 }
19864 if y.len() != out_f || x.len() != in_f {
19865 return Err(format!(
19866 "resident F32 t1 shapes x={} y={} != in {in_f} out {out_f}",
19867 x.len(),
19868 y.len()
19869 )
19870 .into());
19871 }
19872 if canonical_chunk_rows == 0
19873 || canonical_chunk_rows > out_f
19874 || out_f % canonical_chunk_rows != 0
19875 {
19876 return Err(format!(
19877 "invalid canonical F32 chunk rows {canonical_chunk_rows} for output width {out_f}"
19878 )
19879 .into());
19880 }
19881 let input = x.slice(0..x.len());
19882 for r0 in (0..out_f).step_by(canonical_chunk_rows) {
19883 let weights = data.slice(r0 * in_f..(r0 + canonical_chunk_rows) * in_f);
19884 let mut destination = y.slice_mut(r0..r0 + canonical_chunk_rows);
19885 self.linear_device_into(
19886 &input,
19887 &weights,
19888 &mut destination,
19889 1,
19890 in_f,
19891 canonical_chunk_rows,
19892 )?;
19893 }
19894 Ok(())
19895 }
19896
19897 pub fn linear_t1_into(
19900 &self,
19901 x: &cudarc::driver::CudaView<'_, f32>,
19902 w: &cudarc::driver::CudaView<'_, f32>,
19903 y: &mut cudarc::driver::CudaViewMut<'_, f32>,
19904 in_f: usize,
19905 out_f: usize,
19906 ) -> Result<(), Box<dyn std::error::Error>> {
19907 self.linear_device_into(x, w, y, 1, in_f, out_f)
19908 }
19909
19910 pub fn linear_decode_exact(
19917 &self,
19918 x: &CudaSlice<f32>,
19919 w: &CudaSlice<f32>,
19920 m_tokens: usize,
19921 in_f: usize,
19922 out_f: usize,
19923 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19924 if m_tokens == 1 {
19925 return self.linear(x, w, 1, in_f, out_f);
19926 }
19927 let xv = self.view(x, m_tokens * in_f);
19928 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
19929 for t in 0..m_tokens {
19930 let row = xv.slice(t * in_f..(t + 1) * in_f);
19931 let mut xr = self.alloc_uninit::<f32>(in_f)?;
19932 self.copy_view_into(&mut xr, 0, &row, in_f)?;
19933 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
19934 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
19935 }
19936 Ok(y)
19937 }
19938
19939 pub fn linear(
19940 &self,
19941 x: &CudaSlice<f32>,
19942 w: &CudaSlice<f32>,
19943 m_tokens: usize,
19944 in_f: usize,
19945 out_f: usize,
19946 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
19947 self.linear_device(x, w, m_tokens, in_f, out_f)
19948 }
19949
19950 fn linear_device<I>(
19951 &self,
19952 x: &I,
19953 w: &I,
19954 m_tokens: usize,
19955 in_f: usize,
19956 out_f: usize,
19957 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>>
19958 where
19959 I: cudarc::driver::DevicePtr<f32>,
19960 {
19961 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)?;
19963 Ok(c)
19964 }
19965
19966 fn linear_device_into<I, O>(
19967 &self,
19968 x: &I,
19969 w: &I,
19970 c: &mut O,
19971 m_tokens: usize,
19972 in_f: usize,
19973 out_f: usize,
19974 ) -> Result<(), Box<dyn std::error::Error>>
19975 where
19976 I: cudarc::driver::DevicePtr<f32>,
19977 O: cudarc::driver::DevicePtrMut<f32>,
19978 {
19979 use cudarc::cublaslt::{Matmul, MatmulConfig};
19980 let cfg = MatmulConfig {
19981 transa: true,
19982 transb: false,
19983 transc: false,
19984 m: out_f as u64,
19985 n: m_tokens as u64,
19986 k: in_f as u64,
19987 alpha: 1.0,
19988 lda: in_f as i64,
19989 ldb: in_f as i64,
19990 beta: 0.0,
19991 ldc: out_f as i64,
19992 stride_a: None,
19993 stride_b: None,
19994 stride_c: None,
19995 stride_bias: None,
19996 batch_size: None,
19997 };
19998 let blas = self.gpu.blas();
19999 unsafe {
20000 blas.matmul(cfg, w, x, c, None, None)?;
20001 }
20002 Ok(())
20003 }
20004
20005 pub fn sdpa_naive(
20014 &self,
20015 q: &CudaSlice<f32>,
20016 k: &CudaSlice<f32>,
20017 v: &CudaSlice<f32>,
20018 o: &mut CudaSlice<f32>,
20019 head_dim: usize,
20020 n_head: usize,
20021 n_head_kv: usize,
20022 t: usize,
20023 t_kv: usize,
20024 scale: f32,
20025 causal: bool,
20026 ) -> Result<(), Box<dyn std::error::Error>> {
20027 if t_kv * 4 > SDPA_NAIVE_SMEM_MAX {
20028 return self.sdpa_naive_gmem(
20029 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20030 );
20031 }
20032 let f = self.func("sdpa_naive_f32");
20033 let cfg = LaunchConfig {
20034 grid_dim: (n_head as u32, t as u32, 1),
20035 block_dim: (128, 1, 1),
20036 shared_mem_bytes: (t_kv * 4) as u32,
20037 };
20038 let (hd, nh, nhkv, ti, tkvi, cz) = (
20039 head_dim as i32,
20040 n_head as i32,
20041 n_head_kv as i32,
20042 t as i32,
20043 t_kv as i32,
20044 causal as i32,
20045 );
20046 let __s_b = self.gpu.stream();
20047 let mut b = __s_b.launch_builder(&f);
20048 b.arg(q)
20049 .arg(k)
20050 .arg(v)
20051 .arg(o)
20052 .arg(&hd)
20053 .arg(&nh)
20054 .arg(&nhkv)
20055 .arg(&ti)
20056 .arg(&tkvi)
20057 .arg(&scale)
20058 .arg(&cz);
20059 unsafe {
20060 b.launch(cfg)?;
20061 }
20062 Ok(())
20063 }
20064
20065 #[allow(clippy::too_many_arguments)]
20074 pub fn sdpa_naive_gmem(
20075 &self,
20076 q: &CudaSlice<f32>,
20077 k: &CudaSlice<f32>,
20078 v: &CudaSlice<f32>,
20079 o: &mut CudaSlice<f32>,
20080 head_dim: usize,
20081 n_head: usize,
20082 n_head_kv: usize,
20083 t: usize,
20084 t_kv: usize,
20085 scale: f32,
20086 causal: bool,
20087 ) -> Result<(), Box<dyn std::error::Error>> {
20088 let ws_len = n_head
20089 .checked_mul(t)
20090 .and_then(|x| x.checked_mul(t_kv))
20091 .ok_or("sdpa_naive_gmem: scores workspace size overflow")?;
20092 let ws_bytes = ws_len
20093 .checked_mul(std::mem::size_of::<f32>())
20094 .ok_or("sdpa_naive_gmem: scores workspace byte count overflow")?;
20095 if ws_bytes > SDPA_NAIVE_GMEM_WS_MAX {
20096 return Err(format!(
20097 "sdpa_naive_gmem: scores workspace {ws_bytes} bytes (heads {n_head} x T {t} x \
20098 T_kv {t_kv}) exceeds the {SDPA_NAIVE_GMEM_WS_MAX}-byte guard — this shape \
20099 needs a tiled/flash kernel, not the naive oracle"
20100 )
20101 .into());
20102 }
20103 let mut scores = self.uninit(ws_len)?;
20104 let f = self.func("sdpa_naive_gmem_f32");
20105 let cfg = LaunchConfig {
20106 grid_dim: (n_head as u32, t as u32, 1),
20107 block_dim: (128, 1, 1),
20108 shared_mem_bytes: 0,
20109 };
20110 let (hd, nh, nhkv, ti, tkvi, cz) = (
20111 head_dim as i32,
20112 n_head as i32,
20113 n_head_kv as i32,
20114 t as i32,
20115 t_kv as i32,
20116 causal as i32,
20117 );
20118 let __s_b = self.gpu.stream();
20119 let mut b = __s_b.launch_builder(&f);
20120 b.arg(q)
20121 .arg(k)
20122 .arg(v)
20123 .arg(o)
20124 .arg(&mut scores)
20125 .arg(&hd)
20126 .arg(&nh)
20127 .arg(&nhkv)
20128 .arg(&ti)
20129 .arg(&tkvi)
20130 .arg(&scale)
20131 .arg(&cz);
20132 unsafe {
20133 b.launch(cfg)?;
20134 }
20135 Ok(())
20136 }
20137
20138 #[allow(clippy::too_many_arguments)]
20143 pub fn sdpa_naive_island(
20144 &self,
20145 q: &CudaSlice<f32>,
20146 k: &CudaSlice<f32>,
20147 v: &CudaSlice<f32>,
20148 o: &mut CudaSlice<f32>,
20149 span_id: &CudaSlice<i32>,
20150 head_dim: usize,
20151 n_head: usize,
20152 n_head_kv: usize,
20153 t: usize,
20154 t_kv: usize,
20155 scale: f32,
20156 window: usize,
20157 ) -> Result<(), Box<dyn std::error::Error>> {
20158 let f = self.func("sdpa_naive_island_f32");
20159 let cfg = LaunchConfig {
20160 grid_dim: (n_head as u32, t as u32, 1),
20161 block_dim: (128, 1, 1),
20162 shared_mem_bytes: (t_kv * 4) as u32,
20163 };
20164 let (hd, nh, nhkv, ti, tkvi, wi) = (
20165 head_dim as i32,
20166 n_head as i32,
20167 n_head_kv as i32,
20168 t as i32,
20169 t_kv as i32,
20170 window as i32,
20171 );
20172 let __s_b = self.gpu.stream();
20173 let mut b = __s_b.launch_builder(&f);
20174 b.arg(q)
20175 .arg(k)
20176 .arg(v)
20177 .arg(o)
20178 .arg(span_id)
20179 .arg(&hd)
20180 .arg(&nh)
20181 .arg(&nhkv)
20182 .arg(&ti)
20183 .arg(&tkvi)
20184 .arg(&scale)
20185 .arg(&wi);
20186 unsafe {
20187 b.launch(cfg)?;
20188 }
20189 Ok(())
20190 }
20191
20192 #[allow(clippy::too_many_arguments)]
20194 pub fn sdpa_naive_w(
20195 &self,
20196 q: &CudaSlice<f32>,
20197 k: &CudaSlice<f32>,
20198 v: &CudaSlice<f32>,
20199 o: &mut CudaSlice<f32>,
20200 head_dim: usize,
20201 n_head: usize,
20202 n_head_kv: usize,
20203 t: usize,
20204 t_kv: usize,
20205 scale: f32,
20206 causal: bool,
20207 window: usize,
20208 ) -> Result<(), Box<dyn std::error::Error>> {
20209 let f = self.func("sdpa_naive_w_f32");
20210 let cfg = LaunchConfig {
20211 grid_dim: (n_head as u32, t as u32, 1),
20212 block_dim: (128, 1, 1),
20213 shared_mem_bytes: (t_kv * 4) as u32,
20214 };
20215 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20216 head_dim as i32,
20217 n_head as i32,
20218 n_head_kv as i32,
20219 t as i32,
20220 t_kv as i32,
20221 causal as i32,
20222 window as i32,
20223 );
20224 let __s_b = self.gpu.stream();
20225 let mut b = __s_b.launch_builder(&f);
20226 b.arg(q)
20227 .arg(k)
20228 .arg(v)
20229 .arg(o)
20230 .arg(&hd)
20231 .arg(&nh)
20232 .arg(&nhkv)
20233 .arg(&ti)
20234 .arg(&tkvi)
20235 .arg(&scale)
20236 .arg(&cz)
20237 .arg(&wi);
20238 unsafe {
20239 b.launch(cfg)?;
20240 }
20241 Ok(())
20242 }
20243
20244 #[allow(clippy::too_many_arguments)]
20254 pub fn sdpa_naive_w_lo(
20255 &self,
20256 q: &CudaSlice<f32>,
20257 k: &CudaSlice<f32>,
20258 v: &CudaSlice<f32>,
20259 o: &mut CudaSlice<f32>,
20260 head_dim: usize,
20261 n_head: usize,
20262 n_head_kv: usize,
20263 t: usize,
20264 t_kv: usize,
20265 scale: f32,
20266 causal: bool,
20267 window: usize,
20268 ) -> Result<(), Box<dyn std::error::Error>> {
20269 let kv_lo = if window > 0 {
20270 (t_kv - t + 1).saturating_sub(window)
20271 } else {
20272 0
20273 };
20274 let smem = (t_kv - kv_lo) * 4;
20275 if smem > 48 * 1024 {
20276 return Err(format!(
20277 "sdpa_naive_w_lo: window {window} + T {t} rows need {smem} bytes of dynamic \
20278 shared memory (> 48KB launch bound) — this kernel clips the OLD side only; \
20279 a window this wide needs the multi-pass long-ctx kernel"
20280 )
20281 .into());
20282 }
20283 let f = self.func("sdpa_naive_w_lo_f32");
20284 let cfg = LaunchConfig {
20285 grid_dim: (n_head as u32, t as u32, 1),
20286 block_dim: (128, 1, 1),
20287 shared_mem_bytes: smem as u32,
20288 };
20289 let (hd, nh, nhkv, ti, tkvi, cz, wi, lo) = (
20290 head_dim as i32,
20291 n_head as i32,
20292 n_head_kv as i32,
20293 t as i32,
20294 t_kv as i32,
20295 causal as i32,
20296 window as i32,
20297 kv_lo as i32,
20298 );
20299 let __s_b = self.gpu.stream();
20300 let mut b = __s_b.launch_builder(&f);
20301 b.arg(q)
20302 .arg(k)
20303 .arg(v)
20304 .arg(o)
20305 .arg(&hd)
20306 .arg(&nh)
20307 .arg(&nhkv)
20308 .arg(&ti)
20309 .arg(&tkvi)
20310 .arg(&scale)
20311 .arg(&cz)
20312 .arg(&wi)
20313 .arg(&lo);
20314 unsafe {
20315 b.launch(cfg)?;
20316 }
20317 Ok(())
20318 }
20319
20320 pub fn sdpa_naive_view(
20322 &self,
20323 q: &CudaSlice<f32>,
20324 k: &cudarc::driver::CudaView<f32>,
20325 v: &cudarc::driver::CudaView<f32>,
20326 o: &mut CudaSlice<f32>,
20327 head_dim: usize,
20328 n_head: usize,
20329 n_head_kv: usize,
20330 t: usize,
20331 t_kv: usize,
20332 scale: f32,
20333 causal: bool,
20334 ) -> Result<(), Box<dyn std::error::Error>> {
20335 let f = self.func("sdpa_naive_f32");
20336 let cfg = LaunchConfig {
20337 grid_dim: (n_head as u32, t as u32, 1),
20338 block_dim: (128, 1, 1),
20339 shared_mem_bytes: (t_kv * 4) as u32,
20340 };
20341 let (hd, nh, nhkv, ti, tkvi, cz) = (
20342 head_dim as i32,
20343 n_head as i32,
20344 n_head_kv as i32,
20345 t as i32,
20346 t_kv as i32,
20347 causal as i32,
20348 );
20349 let __s_b = self.gpu.stream();
20350 let mut b = __s_b.launch_builder(&f);
20351 b.arg(q)
20352 .arg(k)
20353 .arg(v)
20354 .arg(o)
20355 .arg(&hd)
20356 .arg(&nh)
20357 .arg(&nhkv)
20358 .arg(&ti)
20359 .arg(&tkvi)
20360 .arg(&scale)
20361 .arg(&cz);
20362 unsafe {
20363 b.launch(cfg)?;
20364 }
20365 Ok(())
20366 }
20367
20368 #[allow(clippy::too_many_arguments)]
20376 pub fn fa_dequant_kv_view_f32(
20377 &self,
20378 k: &cudarc::driver::CudaView<u8>,
20379 v: &cudarc::driver::CudaView<u8>,
20380 kf: &mut CudaSlice<f32>,
20381 vf: &mut CudaSlice<f32>,
20382 kv_dim_k: usize,
20383 kv_dim_v: usize,
20384 t_kv: usize,
20385 k_tok_bytes: usize,
20386 v_tok_bytes: usize,
20387 g: bool,
20388 ) -> Result<(), Box<dyn std::error::Error>> {
20389 let f = if g {
20390 self.func_g("fa_dequant_kv_ws_f32")
20391 } else {
20392 self.func("fa_dequant_kv_ws_f32")
20393 };
20394 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
20395 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20396 let cfg = LaunchConfig {
20397 grid_dim: (nblk.max(1), 1, 1),
20398 block_dim: (256, 1, 1),
20399 shared_mem_bytes: 0,
20400 };
20401 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
20402 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
20403 let __s_b = self.gpu.stream();
20404 let mut b = __s_b.launch_builder(&f);
20405 b.arg(k)
20406 .arg(v)
20407 .arg(&mut *kf)
20408 .arg(&mut *vf)
20409 .arg(&kdk)
20410 .arg(&kdv)
20411 .arg(&tkvi)
20412 .arg(&ktb)
20413 .arg(&vtb);
20414 unsafe {
20415 b.launch(cfg)?;
20416 }
20417 Ok(())
20418 }
20419
20420 #[allow(clippy::too_many_arguments)]
20421 pub fn sdpa_naive_quantized_view(
20422 &self,
20423 q: &CudaSlice<f32>,
20424 k: &cudarc::driver::CudaView<u8>,
20425 v: &cudarc::driver::CudaView<u8>,
20426 o: &mut CudaSlice<f32>,
20427 head_dim: usize,
20428 n_head: usize,
20429 n_head_kv: usize,
20430 t: usize,
20431 t_kv: usize,
20432 scale: f32,
20433 causal: bool,
20434 k_tok_bytes: usize,
20435 v_tok_bytes: usize,
20436 ) -> Result<(), Box<dyn std::error::Error>> {
20437 let kv_dim = n_head_kv * head_dim;
20438 let mut kf = self.uninit(t_kv * kv_dim)?;
20439 let mut vf = self.uninit(t_kv * kv_dim)?;
20440 let f = self.func("fa_dequant_kv_ws_f32");
20441 let total = (2 * t_kv * kv_dim) as u64;
20442 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20443 let cfg = LaunchConfig {
20444 grid_dim: (nblk.max(1), 1, 1),
20445 block_dim: (256, 1, 1),
20446 shared_mem_bytes: 0,
20447 };
20448 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
20449 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
20450 let __s_b = self.gpu.stream();
20451 let mut b = __s_b.launch_builder(&f);
20452 b.arg(k)
20453 .arg(v)
20454 .arg(&mut kf)
20455 .arg(&mut vf)
20456 .arg(&kv_dim_i)
20457 .arg(&kv_dim_i)
20458 .arg(&t_kv_i)
20459 .arg(&k_tok_bytes_i)
20460 .arg(&v_tok_bytes_i);
20461 unsafe { b.launch(cfg)? };
20462 self.sdpa_naive(
20463 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20464 )
20465 }
20466
20467 #[allow(clippy::too_many_arguments)]
20479 pub fn sdpa_naive_w_quantized_view(
20480 &self,
20481 q: &CudaSlice<f32>,
20482 k: &cudarc::driver::CudaView<u8>,
20483 v: &cudarc::driver::CudaView<u8>,
20484 o: &mut CudaSlice<f32>,
20485 head_dim: usize,
20486 n_head: usize,
20487 n_head_kv: usize,
20488 t: usize,
20489 t_kv: usize,
20490 scale: f32,
20491 causal: bool,
20492 window: usize,
20493 k_tok_bytes: usize,
20494 v_tok_bytes: usize,
20495 ) -> Result<(), Box<dyn std::error::Error>> {
20496 let kv_dim = n_head_kv * head_dim;
20497 let mut kf = self.uninit(t_kv * kv_dim)?;
20498 let mut vf = self.uninit(t_kv * kv_dim)?;
20499 let f = self.func("fa_dequant_kv_ws_f32");
20500 let total = (2 * t_kv * kv_dim) as u64;
20501 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
20502 let cfg = LaunchConfig {
20503 grid_dim: (nblk.max(1), 1, 1),
20504 block_dim: (256, 1, 1),
20505 shared_mem_bytes: 0,
20506 };
20507 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
20508 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
20509 let __s_b = self.gpu.stream();
20510 let mut b = __s_b.launch_builder(&f);
20511 b.arg(k)
20512 .arg(v)
20513 .arg(&mut kf)
20514 .arg(&mut vf)
20515 .arg(&kv_dim_i)
20516 .arg(&kv_dim_i)
20517 .arg(&t_kv_i)
20518 .arg(&k_tok_bytes_i)
20519 .arg(&v_tok_bytes_i);
20520 unsafe { b.launch(cfg)? };
20521 self.sdpa_naive_w(
20522 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20523 )
20524 }
20525
20526 pub fn fa_prefill(
20530 &self,
20531 q: &CudaSlice<f32>,
20532 k: &CudaSlice<f32>,
20533 v: &CudaSlice<f32>,
20534 o: &mut CudaSlice<f32>,
20535 head_dim: usize,
20536 n_head: usize,
20537 n_head_kv: usize,
20538 t: usize,
20539 t_kv: usize,
20540 scale: f32,
20541 causal: bool,
20542 ) -> Result<(), Box<dyn std::error::Error>> {
20543 if portable_mma_gated() {
20544 return self.sdpa_naive(
20545 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
20546 );
20547 }
20548 let fa3_on = head_dim == 256
20556 && causal
20557 && t == t_kv
20558 && match std::env::var("MEMRA_FA3").as_deref() {
20559 Ok("0") => false,
20560 Ok("1") => {
20564 refuse_portable_force("MEMRA_FA3=1", "the sm_90a fa3/bf16 kernels");
20565 true
20566 }
20567 _ => cfg!(memra_hopper_mma),
20568 };
20569 if fa3_on {
20570 let n = t * n_head * head_dim;
20571 let nkv = t * n_head_kv * head_dim;
20572 let mut q16 = self.alloc_u8_uninit(n * 2)?;
20573 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
20574 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
20575 self.f32_to_bf16_into(q, &mut q16, n)?;
20576 self.f32_to_bf16_into(k, &mut k16, nkv)?;
20577 self.f32_to_bf16_into(v, &mut v16, nkv)?;
20578 let rc = {
20579 use cudarc::driver::{DevicePtr, DevicePtrMut};
20580 let stream = self.gpu.stream();
20581 let (qp, _g1) = q16.device_ptr(&stream);
20582 let (kp, _g2) = k16.device_ptr(&stream);
20583 let (vp, _g3) = v16.device_ptr(&stream);
20584 let (op, _g4) = o.device_ptr_mut(&stream);
20585 unsafe {
20586 memra_fa3_prefill(
20587 qp as *const core::ffi::c_void,
20588 kp as *const core::ffi::c_void,
20589 vp as *const core::ffi::c_void,
20590 op as *mut f32,
20591 t as i32,
20592 n_head as i32,
20593 n_head_kv as i32,
20594 head_dim as i32,
20595 scale,
20596 stream.cu_stream() as *mut core::ffi::c_void,
20597 )
20598 }
20599 };
20600 if rc != 0 {
20601 return Err(format!("memra_fa3_prefill rc={rc}").into());
20602 }
20603 return Ok(());
20604 }
20605 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20610 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
20611 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
20612 const BLOCK_Q: usize = 64;
20613 const BKX: usize = 32;
20614 let f = self.func("fa_prefill_bf16_p1");
20615 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
20616 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
20617 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20618 f.set_attribute(
20619 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20620 shmem as i32,
20621 )?;
20622 let cfg = LaunchConfig {
20623 grid_dim: (
20624 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20625 n_head as u32,
20626 1,
20627 ),
20628 block_dim: (32, 4, 1),
20629 shared_mem_bytes: shmem,
20630 };
20631 let (hd, nh, nhkv, ti, tkvi, cz) = (
20632 head_dim as i32,
20633 n_head as i32,
20634 n_head_kv as i32,
20635 t as i32,
20636 t_kv as i32,
20637 causal as i32,
20638 );
20639 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
20640 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
20641 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
20642 let __s_b = self.gpu.stream();
20643 let mut b = __s_b.launch_builder(&f);
20644 b.arg(&qb)
20645 .arg(&kb)
20646 .arg(&vb)
20647 .arg(o)
20648 .arg(&hd)
20649 .arg(&nh)
20650 .arg(&nhkv)
20651 .arg(&ti)
20652 .arg(&tkvi)
20653 .arg(&scale)
20654 .arg(&cz);
20655 unsafe {
20656 b.launch(cfg)?;
20657 }
20658 return Ok(());
20659 }
20660 const BK: usize = 32;
20666 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
20669 let (block_q, warps, w2_sfx): (usize, u32, &str) =
20670 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
20671 let hd_sfx = fa_hd_suffix(head_dim)?;
20675 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20676 let bf16kv = !floor && !w2 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
20681 let (kb16, vb16) = if bf16kv {
20682 let n = t_kv * n_head_kv * head_dim;
20683 let mut kb = self.alloc_u8_uninit(n * 2)?;
20684 let mut vb = self.alloc_u8_uninit(n * 2)?;
20685 let fcv = self.func("f32_to_bf16_bulk");
20686 let ni = n as i64;
20687 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
20688 let __s_b = self.gpu.stream();
20689 let mut b = __s_b.launch_builder(&fcv);
20690 b.arg(k).arg(&mut kb).arg(&ni);
20691 unsafe {
20692 b.launch(cfgc)?;
20693 }
20694 let __s_b = self.gpu.stream();
20695 let mut b = __s_b.launch_builder(&fcv);
20696 b.arg(v).arg(&mut vb).arg(&ni);
20697 unsafe {
20698 b.launch(cfgc)?;
20699 }
20700 (Some(kb), Some(vb))
20701 } else {
20702 (None, None)
20703 };
20704 let f = self.func(&if bf16kv {
20705 format!("fa_prefill_bf16kv_pp{hd_sfx}")
20706 } else {
20707 format!(
20708 "fa_prefill_f32{}{}{hd_sfx}",
20709 if floor { "" } else { "_pp" },
20710 if floor { "" } else { w2_sfx }
20711 )
20712 });
20713 let kv_stages = if bf16kv { 2 } else { 1 };
20716 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
20717 + 4 * (block_q * BK + 2 * block_q)) as u32;
20718 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20719 f.set_attribute(
20720 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20721 shmem as i32,
20722 )?;
20723 let cfg = LaunchConfig {
20724 grid_dim: (
20725 (t as u32 + block_q as u32 - 1) / block_q as u32,
20726 n_head as u32,
20727 1,
20728 ),
20729 block_dim: (32, warps, 1),
20730 shared_mem_bytes: shmem,
20731 };
20732 let (hd, nh, nhkv, ti, tkvi, cz) = (
20733 head_dim as i32,
20734 n_head as i32,
20735 n_head_kv as i32,
20736 t as i32,
20737 t_kv as i32,
20738 causal as i32,
20739 );
20740 let __s_b = self.gpu.stream();
20741 let mut b = __s_b.launch_builder(&f);
20742 b.arg(q);
20743 match (&kb16, &vb16) {
20744 (Some(kb), Some(vb)) => {
20745 b.arg(kb).arg(vb);
20746 }
20747 _ => {
20748 b.arg(k).arg(v);
20749 }
20750 }
20751 b.arg(o)
20752 .arg(&hd)
20753 .arg(&nh)
20754 .arg(&nhkv)
20755 .arg(&ti)
20756 .arg(&tkvi)
20757 .arg(&scale)
20758 .arg(&cz);
20759 unsafe {
20760 b.launch(cfg)?;
20761 }
20762 Ok(())
20763 }
20764
20765 #[allow(clippy::too_many_arguments)]
20769 pub fn fa_prefill_w(
20770 &self,
20771 q: &CudaSlice<f32>,
20772 k: &CudaSlice<f32>,
20773 v: &CudaSlice<f32>,
20774 o: &mut CudaSlice<f32>,
20775 head_dim: usize,
20776 n_head: usize,
20777 n_head_kv: usize,
20778 t: usize,
20779 t_kv: usize,
20780 scale: f32,
20781 causal: bool,
20782 window: usize,
20783 ) -> Result<(), Box<dyn std::error::Error>> {
20784 if portable_mma_gated() {
20787 return self.sdpa_naive_w(
20788 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal, window,
20789 );
20790 }
20791 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20795 let faw_f32 =
20796 *FAW_F32.get_or_init(|| std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32"));
20797 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
20798 self.fa_prefill_w_arm(
20799 q,
20800 k,
20801 v,
20802 o,
20803 head_dim,
20804 n_head,
20805 n_head_kv,
20806 t,
20807 t_kv,
20808 scale,
20809 causal,
20810 window,
20811 floor || faw_f32,
20812 floor,
20813 )
20814 }
20815
20816 #[allow(clippy::too_many_arguments)]
20819 pub fn fa_prefill_w_pre(
20820 &self,
20821 qb: &CudaSlice<u8>,
20822 kb: &CudaSlice<u8>,
20823 vb: &CudaSlice<u8>,
20824 o: &mut CudaSlice<f32>,
20825 head_dim: usize,
20826 n_head: usize,
20827 n_head_kv: usize,
20828 t: usize,
20829 t_kv: usize,
20830 scale: f32,
20831 causal: bool,
20832 window: usize,
20833 v_f16: bool,
20834 ) -> Result<(), Box<dyn std::error::Error>> {
20835 const BLOCK_Q: usize = 64;
20836 const BK: usize = 32;
20837 debug_assert_eq!(head_dim, 256);
20838 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20839 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
20840 if hp {
20841 const BLOCK_QH: usize = 32;
20842 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
20845 let vh: &CudaSlice<u8> = if v_f16 {
20846 vb
20847 } else {
20848 let n = t_kv * n_head_kv * head_dim;
20849 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
20850 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
20851 }
20852 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
20853 vguard.as_ref().unwrap()
20854 };
20855 let f = self.func("fa_prefill_w_bf16_p1h2");
20856 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20857 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20858 f.set_attribute(
20859 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20860 shmem as i32,
20861 )?;
20862 let cfg = LaunchConfig {
20863 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20864 block_dim: (32, 4, 1),
20865 shared_mem_bytes: shmem,
20866 };
20867 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20868 head_dim as i32,
20869 n_head as i32,
20870 n_head_kv as i32,
20871 t as i32,
20872 t_kv as i32,
20873 causal as i32,
20874 window as i32,
20875 );
20876 let __s_b = self.gpu.stream();
20877 let mut b = __s_b.launch_builder(&f);
20878 b.arg(qb)
20879 .arg(kb)
20880 .arg(vh)
20881 .arg(o)
20882 .arg(&hd)
20883 .arg(&nh)
20884 .arg(&nhkv)
20885 .arg(&ti)
20886 .arg(&tkvi)
20887 .arg(&scale)
20888 .arg(&cz)
20889 .arg(&wi);
20890 unsafe {
20891 b.launch(cfg)?;
20892 }
20893 return Ok(());
20894 }
20895 let f = self.func("fa_prefill_w_bf16_p1");
20896 let shmem =
20897 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
20898 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20899 f.set_attribute(
20900 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20901 shmem as i32,
20902 )?;
20903 let cfg = LaunchConfig {
20904 grid_dim: (
20905 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
20906 n_head as u32,
20907 1,
20908 ),
20909 block_dim: (32, 4, 1),
20910 shared_mem_bytes: shmem,
20911 };
20912 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20913 head_dim as i32,
20914 n_head as i32,
20915 n_head_kv as i32,
20916 t as i32,
20917 t_kv as i32,
20918 causal as i32,
20919 window as i32,
20920 );
20921 let __s_b = self.gpu.stream();
20922 let mut b = __s_b.launch_builder(&f);
20923 b.arg(qb)
20924 .arg(kb)
20925 .arg(vb)
20926 .arg(o)
20927 .arg(&hd)
20928 .arg(&nh)
20929 .arg(&nhkv)
20930 .arg(&ti)
20931 .arg(&tkvi)
20932 .arg(&scale)
20933 .arg(&cz)
20934 .arg(&wi);
20935 unsafe {
20936 b.launch(cfg)?;
20937 }
20938 Ok(())
20939 }
20940
20941 #[allow(clippy::too_many_arguments)]
20943 pub fn fa_prefill_w_arm(
20944 &self,
20945 q: &CudaSlice<f32>,
20946 k: &CudaSlice<f32>,
20947 v: &CudaSlice<f32>,
20948 o: &mut CudaSlice<f32>,
20949 head_dim: usize,
20950 n_head: usize,
20951 n_head_kv: usize,
20952 t: usize,
20953 t_kv: usize,
20954 scale: f32,
20955 causal: bool,
20956 window: usize,
20957 f32_stage: bool,
20958 floor: bool,
20959 ) -> Result<(), Box<dyn std::error::Error>> {
20960 const BLOCK_Q: usize = 64;
20961 const BK: usize = 32;
20962 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
20963 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
20967 let p1 = !floor
20968 && !f32_stage
20969 && *P1_ON.get_or_init(|| {
20970 std::env::var("MEMRA_FAW_P1")
20971 .map(|v| v != "0")
20972 .unwrap_or(true)
20973 });
20974 let hp =
20975 p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
20976 if hp {
20977 const BLOCK_QH: usize = 32;
20978 let f = self.func("fa_prefill_w_bf16_p1h2");
20979 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK) + 4 * (2 * BLOCK_QH)) as u32;
20980 use cudarc::driver::sys::CUfunction_attribute_enum as A;
20981 f.set_attribute(
20982 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
20983 shmem as i32,
20984 )?;
20985 let cfg = LaunchConfig {
20986 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
20987 block_dim: (32, 4, 1),
20988 shared_mem_bytes: shmem,
20989 };
20990 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
20991 head_dim as i32,
20992 n_head as i32,
20993 n_head_kv as i32,
20994 t as i32,
20995 t_kv as i32,
20996 causal as i32,
20997 window as i32,
20998 );
20999 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21000 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21001 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
21002 let __s_b = self.gpu.stream();
21003 let mut b = __s_b.launch_builder(&f);
21004 b.arg(&qb)
21005 .arg(&kb)
21006 .arg(&vh)
21007 .arg(o)
21008 .arg(&hd)
21009 .arg(&nh)
21010 .arg(&nhkv)
21011 .arg(&ti)
21012 .arg(&tkvi)
21013 .arg(&scale)
21014 .arg(&cz)
21015 .arg(&wi);
21016 unsafe {
21017 b.launch(cfg)?;
21018 }
21019 return Ok(());
21020 }
21021 if p1 {
21022 let f = self.func("fa_prefill_w_bf16_p1");
21023 let shmem =
21024 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21025 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21026 f.set_attribute(
21027 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21028 shmem as i32,
21029 )?;
21030 let cfg = LaunchConfig {
21031 grid_dim: (
21032 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21033 n_head as u32,
21034 1,
21035 ),
21036 block_dim: (32, 4, 1),
21037 shared_mem_bytes: shmem,
21038 };
21039 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
21040 head_dim as i32,
21041 n_head as i32,
21042 n_head_kv as i32,
21043 t as i32,
21044 t_kv as i32,
21045 causal as i32,
21046 window as i32,
21047 );
21048 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21049 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21050 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21051 let __s_b = self.gpu.stream();
21052 let mut b = __s_b.launch_builder(&f);
21053 b.arg(&qb)
21054 .arg(&kb)
21055 .arg(&vb)
21056 .arg(o)
21057 .arg(&hd)
21058 .arg(&nh)
21059 .arg(&nhkv)
21060 .arg(&ti)
21061 .arg(&tkvi)
21062 .arg(&scale)
21063 .arg(&cz)
21064 .arg(&wi);
21065 unsafe {
21066 b.launch(cfg)?;
21067 }
21068 return Ok(());
21069 }
21070 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21073 let g4 = !floor
21074 && !f32_stage
21075 && n_head_kv == 1
21076 && n_head % 4 == 0
21077 && *G4_ON.get_or_init(|| {
21078 std::env::var("MEMRA_FAW_G4")
21079 .map(|v| v != "0")
21080 .unwrap_or(true)
21081 });
21082 if g4 {
21083 const SP_M: usize = 16;
21084 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21087 let o2 = *O2_ON.get_or_init(|| {
21088 std::env::var("MEMRA_FAW_O2")
21089 .map(|v| v != "0")
21090 .unwrap_or(true)
21091 });
21092 let f = self.func(if o2 {
21093 "fa_prefill_w_bf16_g4o2"
21094 } else {
21095 "fa_prefill_w_bf16_g4"
21096 });
21097 let shmem = if o2 {
21098 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
21099 } else {
21100 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M))
21101 as u32
21102 };
21103 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21104 f.set_attribute(
21105 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21106 shmem as i32,
21107 )?;
21108 let cfg = LaunchConfig {
21109 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
21110 block_dim: (32, 4, 1),
21111 shared_mem_bytes: shmem,
21112 };
21113 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
21114 head_dim as i32,
21115 n_head as i32,
21116 n_head_kv as i32,
21117 t as i32,
21118 t_kv as i32,
21119 causal as i32,
21120 window as i32,
21121 );
21122 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21123 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21124 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21125 let __s_b = self.gpu.stream();
21126 let mut b = __s_b.launch_builder(&f);
21127 b.arg(&qb)
21128 .arg(&kb)
21129 .arg(&vb)
21130 .arg(o)
21131 .arg(&hd)
21132 .arg(&nh)
21133 .arg(&nhkv)
21134 .arg(&ti)
21135 .arg(&tkvi)
21136 .arg(&scale)
21137 .arg(&cz)
21138 .arg(&wi);
21139 unsafe {
21140 b.launch(cfg)?;
21141 }
21142 return Ok(());
21143 }
21144 let f = self.func(if floor {
21145 "fa_prefill_w_f32"
21146 } else if f32_stage {
21147 "fa_prefill_w_f32_pp"
21148 } else {
21149 "fa_prefill_w_bf16_pp"
21150 });
21151 let shmem =
21152 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21153 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21154 f.set_attribute(
21155 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21156 shmem as i32,
21157 )?;
21158 let cfg = LaunchConfig {
21159 grid_dim: (
21160 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21161 n_head as u32,
21162 1,
21163 ),
21164 block_dim: (32, 4, 1),
21165 shared_mem_bytes: shmem,
21166 };
21167 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (
21168 head_dim as i32,
21169 n_head as i32,
21170 n_head_kv as i32,
21171 t as i32,
21172 t_kv as i32,
21173 causal as i32,
21174 window as i32,
21175 );
21176 if f32_stage {
21177 let __s_b = self.gpu.stream();
21178 let mut b = __s_b.launch_builder(&f);
21179 b.arg(q)
21180 .arg(k)
21181 .arg(v)
21182 .arg(o)
21183 .arg(&hd)
21184 .arg(&nh)
21185 .arg(&nhkv)
21186 .arg(&ti)
21187 .arg(&tkvi)
21188 .arg(&scale)
21189 .arg(&cz)
21190 .arg(&wi);
21191 unsafe {
21192 b.launch(cfg)?;
21193 }
21194 } else {
21195 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21196 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21197 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21198 let __s_b = self.gpu.stream();
21199 let mut b = __s_b.launch_builder(&f);
21200 b.arg(&qb)
21201 .arg(&kb)
21202 .arg(&vb)
21203 .arg(o)
21204 .arg(&hd)
21205 .arg(&nh)
21206 .arg(&nhkv)
21207 .arg(&ti)
21208 .arg(&tkvi)
21209 .arg(&scale)
21210 .arg(&cz)
21211 .arg(&wi);
21212 unsafe {
21213 b.launch(cfg)?;
21214 }
21215 }
21216 Ok(())
21217 }
21218
21219 #[allow(clippy::too_many_arguments)]
21223 pub fn fa_prefill_hd512(
21224 &self,
21225 q: &CudaSlice<f32>,
21226 k: &CudaSlice<f32>,
21227 v: &CudaSlice<f32>,
21228 o: &mut CudaSlice<f32>,
21229 head_dim: usize,
21230 n_head: usize,
21231 n_head_kv: usize,
21232 t: usize,
21233 t_kv: usize,
21234 scale: f32,
21235 causal: bool,
21236 ) -> Result<(), Box<dyn std::error::Error>> {
21237 if portable_mma_gated() {
21239 return self.sdpa_naive(
21240 q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
21241 );
21242 }
21243 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21249 let f32_stage =
21250 *F32_STAGE.get_or_init(|| std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32"));
21251 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
21255 let sp = !f32_stage
21256 && *SP_ON.get_or_init(|| {
21257 std::env::var("MEMRA_FA512_SP")
21258 .map(|v| v != "0")
21259 .unwrap_or(true)
21260 });
21261 self.fa_prefill_hd512_arm(
21262 q,
21263 k,
21264 v,
21265 o,
21266 head_dim,
21267 n_head,
21268 n_head_kv,
21269 t,
21270 t_kv,
21271 scale,
21272 causal,
21273 f32_stage,
21274 sp,
21275 sp && fa_f16pv_on(),
21276 )
21277 }
21278
21279 #[allow(clippy::too_many_arguments)]
21281 pub fn fa_prefill_hd512_pre(
21282 &self,
21283 qb: &CudaSlice<u8>,
21284 kb: &CudaSlice<u8>,
21285 vb: &CudaSlice<u8>,
21286 o: &mut CudaSlice<f32>,
21287 head_dim: usize,
21288 n_head: usize,
21289 n_head_kv: usize,
21290 t: usize,
21291 t_kv: usize,
21292 scale: f32,
21293 causal: bool,
21294 v_f16: bool,
21295 ) -> Result<(), Box<dyn std::error::Error>> {
21296 debug_assert_eq!(head_dim, 512);
21297 const SP_M: usize = 16;
21298 const BKS: usize = 32;
21299 let f16pv = fa_f16pv_on();
21303 let nw = if f16pv { fa512_wide_warps() } else { 2 };
21304 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
21305 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
21306 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
21307 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
21308 let n = t_kv * n_head_kv * head_dim;
21310 let need = n * 2;
21311 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
21312 *vguard = Some(self.alloc_uninit::<u8>(need)?);
21313 }
21314 let dst = vguard.as_mut().unwrap();
21315 self.bf16_to_f16_into(vb, n, dst)?;
21316 vguard.as_ref().unwrap()
21317 } else {
21318 vb
21319 };
21320 let f = self.func(if hp {
21321 "fa_prefill_bf16_hd512_sp16h2"
21322 } else {
21323 match (f16pv, nw) {
21324 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
21325 (true, _) => "fa_prefill_bf16_hd512_sp16",
21326 _ => "fa_prefill_bf16_hd512_sp",
21327 }
21328 });
21329 let (nwarp, npart) = if hp {
21330 (4usize, 4usize)
21331 } else if nw > 2 {
21332 (nw, nw)
21333 } else {
21334 (2, 1)
21335 };
21336 let shmem = if hp {
21338 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS) + 4 * (2 * npart * SP_M * BKS + 2 * SP_M))
21339 as u32
21340 } else {
21341 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
21342 + 4 * (npart * SP_M * BKS + SP_M)) as u32
21343 };
21344 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21345 f.set_attribute(
21346 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21347 shmem as i32,
21348 )?;
21349 let grid_y = if hp {
21350 (n_head / 2) as u32
21351 } else {
21352 n_head as u32
21353 };
21354 let cfg = LaunchConfig {
21355 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
21356 block_dim: (32, nwarp as u32, 1),
21357 shared_mem_bytes: shmem,
21358 };
21359 let (hd, nh, nhkv, ti, tkvi, cz) = (
21360 head_dim as i32,
21361 n_head as i32,
21362 n_head_kv as i32,
21363 t as i32,
21364 t_kv as i32,
21365 causal as i32,
21366 );
21367 let __s_b = self.gpu.stream();
21368 let mut b = __s_b.launch_builder(&f);
21369 b.arg(qb)
21370 .arg(kb)
21371 .arg(vref)
21372 .arg(o)
21373 .arg(&hd)
21374 .arg(&nh)
21375 .arg(&nhkv)
21376 .arg(&ti)
21377 .arg(&tkvi)
21378 .arg(&scale)
21379 .arg(&cz);
21380 unsafe {
21381 b.launch(cfg)?;
21382 }
21383 Ok(())
21384 }
21385
21386 #[allow(clippy::too_many_arguments)]
21389 pub fn fa_prefill_hd512_arm(
21390 &self,
21391 q: &CudaSlice<f32>,
21392 k: &CudaSlice<f32>,
21393 v: &CudaSlice<f32>,
21394 o: &mut CudaSlice<f32>,
21395 head_dim: usize,
21396 n_head: usize,
21397 n_head_kv: usize,
21398 t: usize,
21399 t_kv: usize,
21400 scale: f32,
21401 causal: bool,
21402 f32_stage: bool,
21403 sp: bool,
21404 f16pv: bool,
21405 ) -> Result<(), Box<dyn std::error::Error>> {
21406 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
21407 if sp && !f32_stage {
21408 const SP_M: usize = 16;
21412 const BKS: usize = 32;
21413 let nw = if f16pv { fa512_wide_warps() } else { 2 };
21414 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
21415 let f = self.func(if hp {
21416 "fa_prefill_bf16_hd512_sp16h2"
21417 } else {
21418 match (f16pv, nw) {
21419 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
21420 (true, _) => "fa_prefill_bf16_hd512_sp16",
21421 _ => "fa_prefill_bf16_hd512_sp",
21422 }
21423 });
21424 let (nwarp, npart) = if hp {
21425 (4usize, 4usize)
21426 } else if nw > 2 {
21427 (nw, nw)
21428 } else {
21429 (2, 1)
21430 };
21431 let shmem = if hp {
21432 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
21433 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
21434 } else {
21435 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
21436 + 4 * (npart * SP_M * BKS + SP_M)) as u32
21437 };
21438 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21439 f.set_attribute(
21440 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21441 shmem as i32,
21442 )?;
21443 let grid_y = if hp {
21444 (n_head / 2) as u32
21445 } else {
21446 n_head as u32
21447 };
21448 let cfg = LaunchConfig {
21449 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
21450 block_dim: (32, nwarp as u32, 1),
21451 shared_mem_bytes: shmem,
21452 };
21453 let (hd, nh, nhkv, ti, tkvi, cz) = (
21454 head_dim as i32,
21455 n_head as i32,
21456 n_head_kv as i32,
21457 t as i32,
21458 t_kv as i32,
21459 causal as i32,
21460 );
21461 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21462 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21463 let vb = if f16pv {
21464 self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?
21465 } else {
21466 self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?
21467 };
21468 let __s_b = self.gpu.stream();
21469 let mut b = __s_b.launch_builder(&f);
21470 b.arg(&qb)
21471 .arg(&kb)
21472 .arg(&vb)
21473 .arg(o)
21474 .arg(&hd)
21475 .arg(&nh)
21476 .arg(&nhkv)
21477 .arg(&ti)
21478 .arg(&tkvi)
21479 .arg(&scale)
21480 .arg(&cz);
21481 unsafe {
21482 b.launch(cfg)?;
21483 }
21484 return Ok(());
21485 }
21486 const BLOCK_Q: usize = 32;
21487 const BK: usize = 32;
21488 const HALF: usize = 256;
21489 let f = self.func(if f32_stage {
21490 "fa_prefill_f32_hd512"
21491 } else {
21492 "fa_prefill_bf16_hd512"
21493 });
21494 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
21496 + 4 * BLOCK_Q) as u32;
21497 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21498 f.set_attribute(
21499 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21500 shmem as i32,
21501 )?;
21502 let cfg = LaunchConfig {
21503 grid_dim: (
21504 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21505 n_head as u32,
21506 2,
21507 ),
21508 block_dim: (32, 2, 1),
21509 shared_mem_bytes: shmem,
21510 };
21511 let (hd, nh, nhkv, ti, tkvi, cz) = (
21512 head_dim as i32,
21513 n_head as i32,
21514 n_head_kv as i32,
21515 t as i32,
21516 t_kv as i32,
21517 causal as i32,
21518 );
21519 if f32_stage {
21520 let __s_b = self.gpu.stream();
21521 let mut b = __s_b.launch_builder(&f);
21522 b.arg(q)
21523 .arg(k)
21524 .arg(v)
21525 .arg(o)
21526 .arg(&hd)
21527 .arg(&nh)
21528 .arg(&nhkv)
21529 .arg(&ti)
21530 .arg(&tkvi)
21531 .arg(&scale)
21532 .arg(&cz);
21533 unsafe {
21534 b.launch(cfg)?;
21535 }
21536 } else {
21537 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
21538 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
21539 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
21540 let __s_b = self.gpu.stream();
21541 let mut b = __s_b.launch_builder(&f);
21542 b.arg(&qb)
21543 .arg(&kb)
21544 .arg(&vb)
21545 .arg(o)
21546 .arg(&hd)
21547 .arg(&nh)
21548 .arg(&nhkv)
21549 .arg(&ti)
21550 .arg(&tkvi)
21551 .arg(&scale)
21552 .arg(&cz);
21553 unsafe {
21554 b.launch(cfg)?;
21555 }
21556 }
21557 Ok(())
21558 }
21559
21560 #[allow(clippy::too_many_arguments)]
21564 pub fn rope_neox2_bf16e(
21565 &self,
21566 q: &mut CudaSlice<f32>,
21567 k: &mut CudaSlice<f32>,
21568 qb: &mut CudaSlice<u8>,
21569 kb: &mut CudaSlice<u8>,
21570 pos: &CudaSlice<i32>,
21571 head_dim: usize,
21572 n_dims: usize,
21573 nh_q: usize,
21574 nh_k: usize,
21575 n_tokens: usize,
21576 base: f32,
21577 freq_scale: f32,
21578 ff: Option<&CudaSlice<f32>>,
21579 ) -> Result<(), Box<dyn std::error::Error>> {
21580 let f = self.func("rope_neox2_bf16e_f32");
21581 let rows = ((nh_q + nh_k) * n_tokens) as u32;
21582 let cfg = LaunchConfig {
21583 grid_dim: (rows, 1, 1),
21584 block_dim: ((head_dim / 2) as u32, 1, 1),
21585 shared_mem_bytes: 0,
21586 };
21587 let theta_scale = base.powf(-2.0 / n_dims as f32);
21588 let (hd, nd, nhq, nhk, nt) = (
21589 head_dim as i32,
21590 n_dims as i32,
21591 nh_q as i32,
21592 nh_k as i32,
21593 n_tokens as i32,
21594 );
21595 let __s_b = self.gpu.stream();
21596 let mut b = __s_b.launch_builder(&f);
21597 match ff {
21598 Some(t) => {
21599 b.arg(&mut *q)
21600 .arg(&mut *k)
21601 .arg(&mut *qb)
21602 .arg(&mut *kb)
21603 .arg(pos)
21604 .arg(&hd)
21605 .arg(&nd)
21606 .arg(&nhq)
21607 .arg(&nhk)
21608 .arg(&nt)
21609 .arg(&theta_scale)
21610 .arg(&freq_scale)
21611 .arg(t);
21612 unsafe {
21613 b.launch(cfg)?;
21614 }
21615 }
21616 None => {
21617 let null: u64 = 0;
21618 b.arg(&mut *q)
21619 .arg(&mut *k)
21620 .arg(&mut *qb)
21621 .arg(&mut *kb)
21622 .arg(pos)
21623 .arg(&hd)
21624 .arg(&nd)
21625 .arg(&nhq)
21626 .arg(&nhk)
21627 .arg(&nt)
21628 .arg(&theta_scale)
21629 .arg(&freq_scale)
21630 .arg(&null);
21631 unsafe {
21632 b.launch(cfg)?;
21633 }
21634 }
21635 }
21636 Ok(())
21637 }
21638
21639 pub fn f32_to_bf16(
21642 &self,
21643 x: &CudaSlice<f32>,
21644 n: usize,
21645 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21646 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
21647 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21648 let f = self.func("f32_to_bf16_flat");
21649 let n_i = n as i64;
21650 let cfg = LaunchConfig {
21651 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21652 block_dim: (256, 1, 1),
21653 shared_mem_bytes: 0,
21654 };
21655 let __s_b = self.gpu.stream();
21656 let mut b = __s_b.launch_builder(&f);
21657 b.arg(x).arg(&mut y).arg(&n_i);
21658 unsafe {
21659 b.launch(cfg)?;
21660 }
21661 Ok(y)
21662 }
21663
21664 pub fn f32_to_f16(
21665 &self,
21666 x: &CudaSlice<f32>,
21667 n: usize,
21668 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21669 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
21670 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21671 let f = self.func("f32_to_f16_flat");
21672 let n_i = n as i64;
21673 let cfg = LaunchConfig {
21674 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
21675 block_dim: (256, 1, 1),
21676 shared_mem_bytes: 0,
21677 };
21678 let __s_b = self.gpu.stream();
21679 let mut b = __s_b.launch_builder(&f);
21680 b.arg(x).arg(&mut y).arg(&n_i);
21681 unsafe {
21682 b.launch(cfg)?;
21683 }
21684 Ok(y)
21685 }
21686
21687 pub fn bf16_to_f16(
21689 &self,
21690 xb: &CudaSlice<u8>,
21691 n: usize,
21692 ) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
21693 let mut y = self.alloc_uninit::<u8>(n * 2)?;
21694 self.bf16_to_f16_into(xb, n, &mut y)?;
21695 Ok(y)
21696 }
21697
21698 pub fn bf16_to_f16_into(
21700 &self,
21701 xb: &CudaSlice<u8>,
21702 n: usize,
21703 y: &mut CudaSlice<u8>,
21704 ) -> Result<(), Box<dyn std::error::Error>> {
21705 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
21706 assert!(y.len() >= n * 2);
21707 let f = self.func("bf16_to_f16_flat");
21708 let n2 = (n / 2) as i64;
21709 let cfg = LaunchConfig {
21710 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
21711 block_dim: (256, 1, 1),
21712 shared_mem_bytes: 0,
21713 };
21714 let __s_b = self.gpu.stream();
21715 let mut b = __s_b.launch_builder(&f);
21716 b.arg(xb).arg(y).arg(&n2);
21717 unsafe {
21718 b.launch(cfg)?;
21719 }
21720 Ok(())
21721 }
21722
21723 #[allow(clippy::too_many_arguments)]
21728 pub fn fa_prefill_vl8(
21729 &self,
21730 seqs: &[FaSeqVl],
21731 head_dim: usize,
21732 n_head: usize,
21733 n_head_kv: usize,
21734 scale: f32,
21735 ) -> Result<(), Box<dyn std::error::Error>> {
21736 const BK: usize = 32;
21737 let b = seqs.len();
21738 assert!(b >= 1 && b <= 8);
21739 let mut packed = [FaSeqVl::default(); 8];
21740 packed[..b].copy_from_slice(seqs);
21741 let v = FaVl8(packed);
21742 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21743 let ept = (n_head_kv * head_dim) as i32;
21744 {
21745 let f = self.func("fa_mirror_vl");
21746 let max_n = (max_t as i64) * ept as i64;
21747 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
21748 for which in 0..2i32 {
21749 let cfg = LaunchConfig {
21750 grid_dim: (blocks, 1, b as u32),
21751 block_dim: (256, 1, 1),
21752 shared_mem_bytes: 0,
21753 };
21754 let __s_lb = self.gpu.stream();
21755 let mut lb = __s_lb.launch_builder(&f);
21756 lb.arg(&v).arg(&ept).arg(&which);
21757 unsafe {
21758 lb.launch(cfg)?;
21759 }
21760 }
21761 }
21762 let hd_sfx = fa_hd_suffix(head_dim)?;
21763 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
21764 let block_q = 64usize;
21765 let kv_stages = 2usize;
21766 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
21767 + 4 * (block_q * BK + 2 * block_q)) as u32;
21768 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21769 f.set_attribute(
21770 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21771 shmem as i32,
21772 )?;
21773 let cfg = LaunchConfig {
21774 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
21775 block_dim: (32, 4, 1),
21776 shared_mem_bytes: shmem,
21777 };
21778 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21779 let __s_lb = self.gpu.stream();
21780 let mut lb = __s_lb.launch_builder(&f);
21781 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
21782 unsafe {
21783 lb.launch(cfg)?;
21784 }
21785 Ok(())
21786 }
21787
21788 #[allow(clippy::too_many_arguments)]
21792 pub fn attn_pre_vl8(
21793 &self,
21794 seqs: &[AttnPreVl],
21795 wq: &CudaSlice<f32>,
21796 wk: &CudaSlice<f32>,
21797 head_dim: usize,
21798 rope_dims: usize,
21799 n_head: usize,
21800 n_head_kv: usize,
21801 eps: f32,
21802 freq_base: f32,
21803 freq_scale: f32,
21804 kv_dim_k: usize,
21805 kv_dim_v: usize,
21806 k_tok_bytes: usize,
21807 v_tok_bytes: usize,
21808 ) -> Result<(), Box<dyn std::error::Error>> {
21809 let b = seqs.len();
21810 assert!(b >= 1 && b <= 8);
21811 let mut packed = [AttnPreVl::default(); 8];
21812 packed[..b].copy_from_slice(seqs);
21813 let v = AttnPreVl8(packed);
21814 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
21815 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
21816 {
21817 let f = self.func("q_gate_split_vl");
21818 let n = max_t * (n_head * head_dim) as u32;
21819 let cfg = LaunchConfig {
21820 grid_dim: (n.div_ceil(256), 1, b as u32),
21821 block_dim: (256, 1, 1),
21822 shared_mem_bytes: 0,
21823 };
21824 let __s_lb = self.gpu.stream();
21825 let mut lb = __s_lb.launch_builder(&f);
21826 lb.arg(&v).arg(&hd).arg(&nh);
21827 unsafe {
21828 lb.launch(cfg)?;
21829 }
21830 }
21831 {
21832 let f = self.func("attn_rms_vl");
21833 let cfg = LaunchConfig {
21834 grid_dim: (max_t * n_head as u32, 2, b as u32),
21835 block_dim: (rms_block(), 1, 1),
21836 shared_mem_bytes: 0,
21837 };
21838 let __s_lb = self.gpu.stream();
21839 let mut lb = __s_lb.launch_builder(&f);
21840 lb.arg(&v)
21841 .arg(wq)
21842 .arg(wk)
21843 .arg(&hd)
21844 .arg(&nh)
21845 .arg(&nhkv)
21846 .arg(&eps);
21847 unsafe {
21848 lb.launch(cfg)?;
21849 }
21850 }
21851 {
21852 let f = self.func("attn_rope_vl");
21853 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
21854 let nd = rope_dims as i32;
21855 let cfg = LaunchConfig {
21856 grid_dim: (max_t * n_head as u32, 2, b as u32),
21857 block_dim: ((head_dim / 2) as u32, 1, 1),
21858 shared_mem_bytes: 0,
21859 };
21860 let __s_lb = self.gpu.stream();
21861 let mut lb = __s_lb.launch_builder(&f);
21862 lb.arg(&v)
21863 .arg(&hd)
21864 .arg(&nd)
21865 .arg(&nh)
21866 .arg(&nhkv)
21867 .arg(&theta_scale)
21868 .arg(&freq_scale);
21869 unsafe {
21870 lb.launch(cfg)?;
21871 }
21872 }
21873 {
21874 let f = self.func("append_kv_vl");
21875 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
21876 let cfg = LaunchConfig {
21877 grid_dim: (nblk, max_t, b as u32),
21878 block_dim: (32, 1, 1),
21879 shared_mem_bytes: 0,
21880 };
21881 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
21882 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21883 let __s_lb = self.gpu.stream();
21884 let mut lb = __s_lb.launch_builder(&f);
21885 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
21886 unsafe {
21887 lb.launch(cfg)?;
21888 }
21889 }
21890 Ok(())
21891 }
21892
21893 pub fn fa_prefill_view(
21898 &self,
21899 q: &CudaSlice<f32>,
21900 k: &cudarc::driver::CudaView<u8>,
21901 v: &cudarc::driver::CudaView<u8>,
21902 o: &mut CudaSlice<f32>,
21903 head_dim: usize,
21904 n_head: usize,
21905 n_head_kv: usize,
21906 t: usize,
21907 t_kv: usize,
21908 scale: f32,
21909 causal: bool,
21910 k_tok_bytes: usize,
21911 v_tok_bytes: usize,
21912 g: bool,
21913 ) -> Result<(), Box<dyn std::error::Error>> {
21914 if portable_mma_gated() {
21915 return self.sdpa_naive_quantized_view(
21916 q,
21917 k,
21918 v,
21919 o,
21920 head_dim,
21921 n_head,
21922 n_head_kv,
21923 t,
21924 t_kv,
21925 scale,
21926 causal,
21927 k_tok_bytes,
21928 v_tok_bytes,
21929 );
21930 }
21931 const BLOCK_Q: usize = 64;
21932 const BK: usize = 32;
21933 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
21936 let f = if g {
21937 self.func_g(&name)
21938 } else {
21939 self.func(&name)
21940 };
21941 let shmem =
21942 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
21943 use cudarc::driver::sys::CUfunction_attribute_enum as A;
21944 f.set_attribute(
21945 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
21946 shmem as i32,
21947 )?;
21948 let cfg = LaunchConfig {
21949 grid_dim: (
21950 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
21951 n_head as u32,
21952 1,
21953 ),
21954 block_dim: (32, 4, 1),
21955 shared_mem_bytes: shmem,
21956 };
21957 let (hd, nh, nhkv, ti, tkvi, cz) = (
21958 head_dim as i32,
21959 n_head as i32,
21960 n_head_kv as i32,
21961 t as i32,
21962 t_kv as i32,
21963 causal as i32,
21964 );
21965 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
21966 let __s_b = self.gpu.stream();
21967 let mut b = __s_b.launch_builder(&f);
21968 b.arg(q)
21969 .arg(k)
21970 .arg(v)
21971 .arg(o)
21972 .arg(&hd)
21973 .arg(&nh)
21974 .arg(&nhkv)
21975 .arg(&ti)
21976 .arg(&tkvi)
21977 .arg(&scale)
21978 .arg(&cz)
21979 .arg(&ktb)
21980 .arg(&vtb);
21981 unsafe {
21982 b.launch(cfg)?;
21983 }
21984 Ok(())
21985 }
21986
21987 #[allow(clippy::too_many_arguments)]
21997 pub fn fa_prefill_view_ws(
21998 &self,
21999 q: &CudaSlice<f32>,
22000 k: &cudarc::driver::CudaView<u8>,
22001 v: &cudarc::driver::CudaView<u8>,
22002 o: &mut CudaSlice<f32>,
22003 head_dim: usize,
22004 n_head: usize,
22005 n_head_kv: usize,
22006 t: usize,
22007 t_kv: usize,
22008 scale: f32,
22009 causal: bool,
22010 k_tok_bytes: usize,
22011 v_tok_bytes: usize,
22012 g: bool,
22013 ) -> Result<(), Box<dyn std::error::Error>> {
22014 if portable_mma_gated() {
22015 return self.sdpa_naive_quantized_view(
22016 q,
22017 k,
22018 v,
22019 o,
22020 head_dim,
22021 n_head,
22022 n_head_kv,
22023 t,
22024 t_kv,
22025 scale,
22026 causal,
22027 k_tok_bytes,
22028 v_tok_bytes,
22029 );
22030 }
22031 const BLOCK_Q: usize = 64;
22032 const BK: usize = 32;
22033 let kv_dim_k = n_head_kv * head_dim;
22034 let kv_dim_v = n_head_kv * head_dim;
22035 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
22037 let mut guard = self.prime_deqw_ws.lock().unwrap();
22039 let need_grow = match guard.as_ref() {
22040 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
22041 None => true,
22042 };
22043 if need_grow {
22044 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
22045 let (ck, cv) = guard
22046 .as_ref()
22047 .map(|(a, b)| (a.len(), b.len()))
22048 .unwrap_or((0, 0));
22049 *guard = Some((
22050 self.alloc_u8(grow(ck, k_ws_bytes))?,
22051 self.alloc_u8(grow(cv, v_ws_bytes))?,
22052 ));
22053 }
22054 let (kw, vw) = guard.as_mut().unwrap();
22055 {
22057 let f = if g {
22059 self.func_g("fa_dequant_kv_ws_bf16")
22060 } else {
22061 self.func("fa_dequant_kv_ws_bf16")
22062 };
22063 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
22064 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
22065 let cfg = LaunchConfig {
22066 grid_dim: (nblk.max(1), 1, 1),
22067 block_dim: (256, 1, 1),
22068 shared_mem_bytes: 0,
22069 };
22070 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
22071 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22072 let __s_b = self.gpu.stream();
22073 let mut b = __s_b.launch_builder(&f);
22074 b.arg(k)
22075 .arg(v)
22076 .arg(&mut *kw)
22077 .arg(&mut *vw)
22078 .arg(&kdk)
22079 .arg(&kdv)
22080 .arg(&tkvi)
22081 .arg(&ktb)
22082 .arg(&vtb);
22083 unsafe {
22084 b.launch(cfg)?;
22085 }
22086 }
22087 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
22095 .map(|v| v != "0")
22096 .unwrap_or(true);
22097 {
22098 let hd_sfx = fa_hd_suffix(head_dim)?;
22099 let f = self.func(&format!(
22100 "fa_prefill_qw{}{hd_sfx}",
22101 if db { "_db" } else { "" }
22102 ));
22103 let shmem = if db {
22104 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
22106 } else {
22107 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
22108 };
22109 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22110 f.set_attribute(
22111 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22112 shmem as i32,
22113 )?;
22114 let cfg = LaunchConfig {
22115 grid_dim: (
22116 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
22117 n_head as u32,
22118 1,
22119 ),
22120 block_dim: (32, 4, 1),
22121 shared_mem_bytes: shmem,
22122 };
22123 let (hd, nh, nhkv, ti, tkvi, cz) = (
22124 head_dim as i32,
22125 n_head as i32,
22126 n_head_kv as i32,
22127 t as i32,
22128 t_kv as i32,
22129 causal as i32,
22130 );
22131 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
22132 let __s_b = self.gpu.stream();
22133 let mut b = __s_b.launch_builder(&f);
22134 b.arg(q)
22135 .arg(&*kw)
22136 .arg(&*vw)
22137 .arg(o)
22138 .arg(&hd)
22139 .arg(&nh)
22140 .arg(&nhkv)
22141 .arg(&ti)
22142 .arg(&tkvi)
22143 .arg(&scale)
22144 .arg(&cz)
22145 .arg(&kdk)
22146 .arg(&kdv);
22147 unsafe {
22148 b.launch(cfg)?;
22149 }
22150 }
22151 Ok(())
22152 }
22153
22154 #[allow(clippy::too_many_arguments)]
22170 pub fn fa_prefill_view_ws_w_hd128(
22171 &self,
22172 q: &CudaSlice<f32>,
22173 k: &cudarc::driver::CudaView<u8>,
22174 v: &cudarc::driver::CudaView<u8>,
22175 o: &mut CudaSlice<f32>,
22176 head_dim: usize,
22177 n_head: usize,
22178 n_head_kv: usize,
22179 t: usize,
22180 t_kv: usize,
22181 scale: f32,
22182 causal: bool,
22183 window: usize,
22184 k_tok_bytes: usize,
22185 v_tok_bytes: usize,
22186 ) -> Result<(), Box<dyn std::error::Error>> {
22187 assert_eq!(
22188 head_dim, 128,
22189 "fa_prefill_view_ws_w_hd128: only the hd128 twin is stamped"
22190 );
22191 if portable_mma_gated() {
22192 return self.sdpa_naive_w_quantized_view(
22193 q,
22194 k,
22195 v,
22196 o,
22197 head_dim,
22198 n_head,
22199 n_head_kv,
22200 t,
22201 t_kv,
22202 scale,
22203 causal,
22204 window,
22205 k_tok_bytes,
22206 v_tok_bytes,
22207 );
22208 }
22209 const BLOCK_Q: usize = 64;
22210 const BK: usize = 32;
22211 let kv_dim_k = n_head_kv * head_dim;
22212 let kv_dim_v = n_head_kv * head_dim;
22213 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
22215 let mut guard = self.prime_deqw_ws.lock().unwrap();
22216 let need_grow = match guard.as_ref() {
22217 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
22218 None => true,
22219 };
22220 if need_grow {
22221 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
22222 let (ck, cv) = guard
22223 .as_ref()
22224 .map(|(a, b)| (a.len(), b.len()))
22225 .unwrap_or((0, 0));
22226 *guard = Some((
22227 self.alloc_u8(grow(ck, k_ws_bytes))?,
22228 self.alloc_u8(grow(cv, v_ws_bytes))?,
22229 ));
22230 }
22231 let (kw, vw) = guard.as_mut().unwrap();
22232 {
22235 let f = self.func("fa_dequant_kv_ws_bf16");
22236 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
22237 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
22238 let cfg = LaunchConfig {
22239 grid_dim: (nblk.max(1), 1, 1),
22240 block_dim: (256, 1, 1),
22241 shared_mem_bytes: 0,
22242 };
22243 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
22244 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22245 let __s_b = self.gpu.stream();
22246 let mut b = __s_b.launch_builder(&f);
22247 b.arg(k)
22248 .arg(v)
22249 .arg(&mut *kw)
22250 .arg(&mut *vw)
22251 .arg(&kdk)
22252 .arg(&kdv)
22253 .arg(&tkvi)
22254 .arg(&ktb)
22255 .arg(&vtb);
22256 unsafe {
22257 b.launch(cfg)?;
22258 }
22259 }
22260 let db = std::env::var("MEMRA_PRIME_DEQW_DB")
22262 .map(|v| v != "0")
22263 .unwrap_or(true);
22264 {
22265 let f = self.func(if db {
22266 "fa_prefill_qw_db_w_hd128"
22267 } else {
22268 "fa_prefill_qw_w_hd128"
22269 });
22270 let shmem = if db {
22271 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
22272 } else {
22273 (2 * (2 * BK * head_dim + BLOCK_Q * BK) + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
22274 };
22275 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22276 f.set_attribute(
22277 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22278 shmem as i32,
22279 )?;
22280 let cfg = LaunchConfig {
22281 grid_dim: (
22282 (t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32,
22283 n_head as u32,
22284 1,
22285 ),
22286 block_dim: (32, 4, 1),
22287 shared_mem_bytes: shmem,
22288 };
22289 let (hd, nh, nhkv, ti, tkvi, cz) = (
22290 head_dim as i32,
22291 n_head as i32,
22292 n_head_kv as i32,
22293 t as i32,
22294 t_kv as i32,
22295 causal as i32,
22296 );
22297 let (kdk, kdv, wnd) = (kv_dim_k as i32, kv_dim_v as i32, window as i32);
22298 let __s_b = self.gpu.stream();
22299 let mut b = __s_b.launch_builder(&f);
22300 b.arg(q)
22301 .arg(&*kw)
22302 .arg(&*vw)
22303 .arg(o)
22304 .arg(&hd)
22305 .arg(&nh)
22306 .arg(&nhkv)
22307 .arg(&ti)
22308 .arg(&tkvi)
22309 .arg(&scale)
22310 .arg(&cz)
22311 .arg(&kdk)
22312 .arg(&kdv)
22313 .arg(&wnd);
22314 unsafe {
22315 b.launch(cfg)?;
22316 }
22317 }
22318 Ok(())
22319 }
22320
22321 pub fn fa_decode(
22325 &self,
22326 q: &CudaSlice<f32>,
22327 k: &cudarc::driver::CudaView<u8>,
22328 v: &cudarc::driver::CudaView<u8>,
22329 o: &mut CudaSlice<f32>,
22330 head_dim: usize,
22331 n_head: usize,
22332 n_head_kv: usize,
22333 t_kv: usize,
22334 scale: f32,
22335 k_tok_bytes: usize,
22336 v_tok_bytes: usize,
22337 ) -> Result<(), Box<dyn std::error::Error>> {
22338 self.fa_decode_kvmod(
22339 q,
22340 k,
22341 v,
22342 o,
22343 head_dim,
22344 n_head,
22345 n_head_kv,
22346 t_kv,
22347 scale,
22348 k_tok_bytes,
22349 v_tok_bytes,
22350 false,
22351 )
22352 }
22353
22354 #[allow(clippy::too_many_arguments)]
22358 #[allow(clippy::too_many_arguments)]
22362 #[allow(clippy::too_many_arguments)]
22363 fn fa_decode_scalar_unified(
22364 &self,
22365 q: &cudarc::driver::CudaView<f32>,
22366 k: &cudarc::driver::CudaView<u8>,
22367 v: &cudarc::driver::CudaView<u8>,
22368 o: &mut cudarc::driver::CudaViewMut<f32>,
22369 head_dim: usize,
22370 n_head: usize,
22371 n_head_kv: usize,
22372 t_kv_host: usize,
22373 t_kv_dev: Option<&CudaSlice<i32>>,
22374 scale: f32,
22375 n_splits: usize,
22376 split_keys: usize,
22377 k_tok_bytes: usize,
22378 v_tok_bytes: usize,
22379 g: bool,
22380 part_o: &mut CudaSlice<f32>,
22381 part_m: &mut CudaSlice<f32>,
22382 part_l: &mut CudaSlice<f32>,
22383 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
22384 ) -> Result<(), Box<dyn std::error::Error>> {
22385 let f = if g {
22386 self.func_g("fa_decode_f32")
22387 } else {
22388 self.fa_func("fa_decode_f32", head_dim)
22389 };
22390 let cfg = LaunchConfig {
22391 grid_dim: (n_head as u32, n_splits as u32, 1),
22392 block_dim: (head_dim as u32, 1, 1),
22393 shared_mem_bytes: (4 * (head_dim + 32)) as u32,
22394 };
22395 let (hd, nh, nhkv, nsp) = (
22396 head_dim as i32,
22397 n_head as i32,
22398 n_head_kv as i32,
22399 n_splits as i32,
22400 );
22401 let (ktb, vtb, tkvi, ski) = (
22402 k_tok_bytes as i64,
22403 v_tok_bytes as i64,
22404 t_kv_host as i32,
22405 split_keys as i32,
22406 );
22407 let __s_b = self.gpu.stream();
22408 let mut b = __s_b.launch_builder(&f);
22409 match t_kv_dev {
22410 Some(d) => {
22411 b.arg(q)
22412 .arg(k)
22413 .arg(v)
22414 .arg(&mut *part_o)
22415 .arg(&mut *part_m)
22416 .arg(&mut *part_l)
22417 .arg(&hd)
22418 .arg(&nh)
22419 .arg(&nhkv)
22420 .arg(&tkvi)
22421 .arg(d)
22422 .arg(&scale)
22423 .arg(&nsp)
22424 .arg(&ski)
22425 .arg(&ktb)
22426 .arg(&vtb);
22427 unsafe {
22428 b.launch(cfg)?;
22429 }
22430 }
22431 None => {
22432 let null: u64 = 0;
22433 b.arg(q)
22434 .arg(k)
22435 .arg(v)
22436 .arg(&mut *part_o)
22437 .arg(&mut *part_m)
22438 .arg(&mut *part_l)
22439 .arg(&hd)
22440 .arg(&nh)
22441 .arg(&nhkv)
22442 .arg(&tkvi)
22443 .arg(&null)
22444 .arg(&scale)
22445 .arg(&nsp)
22446 .arg(&ski)
22447 .arg(&ktb)
22448 .arg(&vtb);
22449 unsafe {
22450 b.launch(cfg)?;
22451 }
22452 }
22453 }
22454 let cfg2 = LaunchConfig {
22455 grid_dim: (n_head as u32, 1, 1),
22456 block_dim: (head_dim as u32, 1, 1),
22457 shared_mem_bytes: 0,
22458 };
22459 if let Some((oq, od)) = q8_out {
22460 let fc = if g {
22462 self.func_g("fa_decode_combine_q8_1")
22463 } else {
22464 self.fa_func("fa_decode_combine_q8_1", head_dim)
22465 };
22466 let __s_b2 = self.gpu.stream();
22467 let mut b2 = __s_b2.launch_builder(&fc);
22468 b2.arg(&*part_o)
22469 .arg(&*part_m)
22470 .arg(&*part_l)
22471 .arg(oq)
22472 .arg(od)
22473 .arg(&hd)
22474 .arg(&nh)
22475 .arg(&nsp);
22476 unsafe {
22477 b2.launch(cfg2)?;
22478 }
22479 return Ok(());
22480 }
22481 let fc = if g {
22482 self.func_g("fa_decode_combine_f32")
22483 } else {
22484 self.fa_func("fa_decode_combine_f32", head_dim)
22485 };
22486 let __s_b2 = self.gpu.stream();
22487 let mut b2 = __s_b2.launch_builder(&fc);
22488 b2.arg(&*part_o)
22489 .arg(&*part_m)
22490 .arg(&*part_l)
22491 .arg(o)
22492 .arg(&hd)
22493 .arg(&nh)
22494 .arg(&nsp);
22495 unsafe {
22496 b2.launch(cfg2)?;
22497 }
22498 Ok(())
22499 }
22500
22501 pub fn fa_decode_kvmod(
22502 &self,
22503 q: &CudaSlice<f32>,
22504 k: &cudarc::driver::CudaView<u8>,
22505 v: &cudarc::driver::CudaView<u8>,
22506 o: &mut CudaSlice<f32>,
22507 head_dim: usize,
22508 n_head: usize,
22509 n_head_kv: usize,
22510 t_kv: usize,
22511 scale: f32,
22512 k_tok_bytes: usize,
22513 v_tok_bytes: usize,
22514 g: bool,
22515 ) -> Result<(), Box<dyn std::error::Error>> {
22516 let q_view = q.as_view();
22517 let mut o_view = o.as_view_mut();
22518 self.fa_decode_kvmod_view(
22519 &q_view,
22520 k,
22521 v,
22522 &mut o_view,
22523 head_dim,
22524 n_head,
22525 n_head_kv,
22526 t_kv,
22527 scale,
22528 k_tok_bytes,
22529 v_tok_bytes,
22530 g,
22531 )
22532 }
22533
22534 #[allow(clippy::too_many_arguments)]
22539 pub fn fa_decode_kvmod_view(
22540 &self,
22541 q: &cudarc::driver::CudaView<f32>,
22542 k: &cudarc::driver::CudaView<u8>,
22543 v: &cudarc::driver::CudaView<u8>,
22544 o: &mut cudarc::driver::CudaViewMut<f32>,
22545 head_dim: usize,
22546 n_head: usize,
22547 n_head_kv: usize,
22548 t_kv: usize,
22549 scale: f32,
22550 k_tok_bytes: usize,
22551 v_tok_bytes: usize,
22552 g: bool,
22553 ) -> Result<(), Box<dyn std::error::Error>> {
22554 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
22575 if g && head_dim == 256 && !fa_v4_at(t_kv) {
22579 fa_vec = false;
22580 }
22581 let sp = fa_split_keys(t_kv, n_head_kv);
22582 let n_splits = if fa_vec {
22583 ((t_kv + sp - 1) / sp).max(1)
22584 } else {
22585 ((t_kv + 255) / 256).max(1)
22586 };
22587 let o_len = n_head * n_splits * head_dim;
22588 let ml_len = n_head * n_splits;
22589 let mut part_guard = self.fa_part_pool.lock().unwrap();
22590 if part_guard
22591 .as_ref()
22592 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22593 .unwrap_or(true)
22594 {
22595 let old = part_guard.take();
22606 let (co, cm) = old
22607 .as_ref()
22608 .map(|pp| (pp.0.len(), pp.1.len()))
22609 .unwrap_or((0, 0));
22610 if let Some(old) = old {
22611 self.fa_part_retired.lock().unwrap().push(old);
22612 }
22613 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22614 eprintln!(
22615 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22616 co, o_len, cm, ml_len
22617 );
22618 }
22619 *part_guard =
22620 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
22621 }
22622 let pg = part_guard.as_mut().unwrap();
22623 self.gpu
22624 .stream()
22625 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22626 self.gpu
22627 .stream()
22628 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22629 self.gpu
22630 .stream()
22631 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22632 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22633 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
22634 let (hd, nh, nhkv, tkvi, nsp) = (
22635 head_dim as i32,
22636 n_head as i32,
22637 n_head_kv as i32,
22638 t_kv as i32,
22639 n_splits as i32,
22640 );
22641 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22642 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
22646 let fa512_min = fa512_min_tkv();
22651 let deep = fa_vec
22654 && head_dim == 256
22655 && fa_v4_at(t_kv)
22656 && !g
22657 && fa_deep_at(t_kv)
22658 && !matches!(fa_v4_mode(), "noB3" | "stage");
22659 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
22660 let gqa = (n_head / n_head_kv).max(1) as u32;
22663 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
22664 (
22665 fv,
22666 LaunchConfig {
22667 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22668 block_dim: (32, gqa, 1),
22669 shared_mem_bytes: 0,
22670 },
22671 )
22672 } else if fa_vec && head_dim <= 256 {
22673 let gqa = (n_head / n_head_kv).max(1) as u32;
22674 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
22685 let smem_tkv = *SMEM_TKV.get_or_init(|| {
22686 std::env::var("MEMRA_FA_SMEM_TKV")
22687 .ok()
22688 .and_then(|v| v.parse().ok())
22689 .unwrap_or_else(|| {
22690 FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
22691 })
22692 });
22693 if fa_v4_at(t_kv) && head_dim == 256 {
22694 let v4name = match fa_v4_mode() {
22698 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
22701 _ => "fa_decode_vec_q_v4",
22702 };
22703 let fv = if g {
22704 self.func_g(v4name)
22705 } else {
22706 self.func(v4name)
22707 };
22708 let shmem = (if deep { 12160 } else { 11520 }
22711 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
22712 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22713 fv.set_attribute(
22714 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22715 shmem as i32,
22716 )?;
22717 (
22718 fv,
22719 LaunchConfig {
22720 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22721 block_dim: (32, gqa, 1),
22722 shared_mem_bytes: shmem,
22723 },
22724 )
22725 } else if fa_v3_active(head_dim) {
22726 let fv = if g {
22729 self.func_g("fa_decode_vec_q_v3")
22730 } else {
22731 self.func("fa_decode_vec_q_v3")
22732 };
22733 let shmem = (32 * head_dim * 2) as u32; (
22735 fv,
22736 LaunchConfig {
22737 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22738 block_dim: (32, gqa, 1),
22739 shared_mem_bytes: shmem,
22740 },
22741 )
22742 } else if fa_v2_on() {
22743 let fv = if g {
22747 self.func_g("fa_decode_vec_q_v2")
22748 } else {
22749 self.func("fa_decode_vec_q_v2")
22750 };
22751 let shmem = (2 * 32 * head_dim * 2) as u32; (
22753 fv,
22754 LaunchConfig {
22755 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22756 block_dim: (32, gqa, 1),
22757 shared_mem_bytes: shmem,
22758 },
22759 )
22760 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g && !(head_dim == 512 && Self::gkv_on())
22761 {
22762 let fv = if g {
22766 self.func_g("fa_decode_vec_q_smem")
22767 } else {
22768 self.func("fa_decode_vec_q_smem")
22769 };
22770 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
22772 fv.set_attribute(
22773 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22774 shmem as i32,
22775 )?;
22776 (
22777 fv,
22778 LaunchConfig {
22779 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22780 block_dim: (32, gqa, 1),
22781 shared_mem_bytes: shmem,
22782 },
22783 )
22784 } else {
22785 let fv = if g {
22788 self.func_g("fa_decode_vec_q")
22789 } else {
22790 self.func("fa_decode_vec_q")
22791 };
22792 (
22793 fv,
22794 LaunchConfig {
22795 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
22796 block_dim: (32, gqa, 1),
22797 shared_mem_bytes: 0,
22798 },
22799 )
22800 }
22801 } else {
22802 return self.fa_decode_scalar_unified(
22805 q,
22806 k,
22807 v,
22808 o,
22809 head_dim,
22810 n_head,
22811 n_head_kv,
22812 t_kv,
22813 None,
22814 scale,
22815 n_splits,
22816 if fa_vec { sp } else { 256 },
22817 k_tok_bytes,
22818 v_tok_bytes,
22819 g,
22820 part_o,
22821 part_m,
22822 part_l,
22823 None,
22824 );
22825 };
22826 let __s_b = self.gpu.stream();
22827 let mut b = __s_b.launch_builder(&f);
22828 b.arg(q)
22829 .arg(k)
22830 .arg(v)
22831 .arg(&mut *part_o)
22832 .arg(&mut *part_m)
22833 .arg(&mut *part_l)
22834 .arg(&hd)
22835 .arg(&nh)
22836 .arg(&nhkv)
22837 .arg(&tkvi)
22838 .arg(&scale)
22839 .arg(&nsp)
22840 .arg(&ktb)
22841 .arg(&vtb);
22842 unsafe {
22843 b.launch(cfg)?;
22844 }
22845 let (fc, cfg2) = (
22848 if g {
22849 self.func_g("fa_decode_combine_f32")
22850 } else {
22851 self.fa_func("fa_decode_combine_f32", head_dim)
22852 },
22853 LaunchConfig {
22854 grid_dim: (n_head as u32, 1, 1),
22855 block_dim: (head_dim as u32, 1, 1),
22856 shared_mem_bytes: 0,
22857 },
22858 );
22859 let __s_b2 = self.gpu.stream();
22860 let mut b2 = __s_b2.launch_builder(&fc);
22861 b2.arg(&*part_o)
22862 .arg(&*part_m)
22863 .arg(&*part_l)
22864 .arg(o)
22865 .arg(&hd)
22866 .arg(&nh)
22867 .arg(&nsp);
22868 unsafe {
22869 b2.launch(cfg2)?;
22870 }
22871 Ok(())
22872 }
22873
22874 #[allow(clippy::too_many_arguments)]
22885 pub fn fa_decode_batch_seqs_v4(
22886 &self,
22887 q: &CudaSlice<f32>,
22888 kv_ptrs: &cudarc::driver::CudaView<u64>,
22889 pos_seq: &CudaSlice<i32>,
22890 o: &mut CudaSlice<f32>,
22891 head_dim: usize,
22892 n_head: usize,
22893 n_head_kv: usize,
22894 b_n: usize,
22895 t_kv_max: usize,
22896 scale: f32,
22897 split_keys: usize,
22898 k_tok_bytes: usize,
22899 v_tok_bytes: usize,
22900 ) -> Result<(), Box<dyn std::error::Error>> {
22901 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
22902 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
22903 let o_len = b_n * n_head * n_splits_max * head_dim;
22904 let ml_len = b_n * n_head * n_splits_max;
22905 let mut part_guard = self.fa_part_pool.lock().unwrap();
22906 if part_guard
22907 .as_ref()
22908 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
22909 .unwrap_or(true)
22910 {
22911 let old = part_guard.take();
22922 let (co, cm) = old
22923 .as_ref()
22924 .map(|pp| (pp.0.len(), pp.1.len()))
22925 .unwrap_or((0, 0));
22926 if let Some(old) = old {
22927 self.fa_part_retired.lock().unwrap().push(old);
22928 }
22929 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
22930 eprintln!(
22931 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
22932 co, o_len, cm, ml_len
22933 );
22934 }
22935 *part_guard =
22936 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
22937 }
22938 let pg = part_guard.as_mut().unwrap();
22939 self.gpu
22940 .stream()
22941 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
22942 self.gpu
22943 .stream()
22944 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
22945 self.gpu
22946 .stream()
22947 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
22948 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
22949 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
22950 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
22951 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
22952 let gqa = (n_head / n_head_kv).max(1) as u32;
22953 let f = self.func("fa_decode_vec_q_seqs_v4");
22954 let shmem = (11520 + 32 * head_dim * 2) as u32;
22956 use cudarc::driver::sys::CUfunction_attribute_enum as A;
22957 f.set_attribute(
22958 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
22959 shmem as i32,
22960 )?;
22961 let cfg = LaunchConfig {
22962 grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
22963 block_dim: (32, gqa, 1),
22964 shared_mem_bytes: shmem,
22965 };
22966 {
22967 let __s_b = self.gpu.stream();
22968 let mut b = __s_b.launch_builder(&f);
22969 b.arg(q)
22970 .arg(kv_ptrs)
22971 .arg(pos_seq)
22972 .arg(&mut *part_o)
22973 .arg(&mut *part_m)
22974 .arg(&mut *part_l)
22975 .arg(&hd)
22976 .arg(&nh)
22977 .arg(&nhkv)
22978 .arg(&scale)
22979 .arg(&nspm)
22980 .arg(&spk)
22981 .arg(&ktb)
22982 .arg(&vtb);
22983 unsafe {
22984 b.launch(cfg)?;
22985 }
22986 }
22987 let fc = self.func("fa_decode_combine_seqs");
22988 let cfg2 = LaunchConfig {
22989 grid_dim: (n_head as u32, b_n as u32, 1),
22990 block_dim: (head_dim as u32, 1, 1),
22991 shared_mem_bytes: 0,
22992 };
22993 let __s_b2 = self.gpu.stream();
22994 let mut b2 = __s_b2.launch_builder(&fc);
22995 b2.arg(&*part_o)
22996 .arg(&*part_m)
22997 .arg(&*part_l)
22998 .arg(o)
22999 .arg(&hd)
23000 .arg(&nh)
23001 .arg(pos_seq)
23002 .arg(&nspm)
23003 .arg(&spk);
23004 unsafe {
23005 b2.launch(cfg2)?;
23006 }
23007 Ok(())
23008 }
23009
23010 #[allow(clippy::too_many_arguments)]
23017 pub fn append_kv_quantized_seqs(
23018 &self,
23019 k_rows: &CudaSlice<f32>,
23020 v_rows: &CudaSlice<f32>,
23021 kv_ptrs: &cudarc::driver::CudaView<u64>,
23022 pos_seq: &CudaSlice<i32>,
23023 b_n: usize,
23024 kv_dim_k: usize,
23025 kv_dim_v: usize,
23026 k_tok_bytes: usize,
23027 v_tok_bytes: usize,
23028 ) -> Result<(), Box<dyn std::error::Error>> {
23029 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
23030 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
23031 let cfg = LaunchConfig {
23032 grid_dim: (nblk, b_n as u32, 1),
23033 block_dim: (32, 1, 1),
23034 shared_mem_bytes: 0,
23035 };
23036 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
23037 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23038 let __s_b = self.gpu.stream();
23039 let mut b = __s_b.launch_builder(&f);
23040 b.arg(k_rows)
23041 .arg(v_rows)
23042 .arg(kv_ptrs)
23043 .arg(pos_seq)
23044 .arg(&kdk)
23045 .arg(&kdv)
23046 .arg(&ktb)
23047 .arg(&vtb);
23048 unsafe {
23049 b.launch(cfg)?;
23050 }
23051 Ok(())
23052 }
23053
23054 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
23060 std::env::var("MEMRA_NO_FA_VEC").is_err()
23061 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
23062 && base_len + 1 >= fa_vec_min_tkv()
23063 && head_dim <= 256
23064 && head_dim % 32 == 0
23065 }
23066
23067 #[allow(clippy::too_many_arguments)]
23076 pub fn fa_decode_rows(
23077 &self,
23078 q: &CudaSlice<f32>,
23079 k: &cudarc::driver::CudaView<u8>,
23080 v: &cudarc::driver::CudaView<u8>,
23081 o: &mut CudaSlice<f32>,
23082 head_dim: usize,
23083 n_head: usize,
23084 n_head_kv: usize,
23085 base_len: usize,
23086 t: usize,
23087 scale: f32,
23088 k_tok_bytes: usize,
23089 v_tok_bytes: usize,
23090 base_dev: Option<(&CudaSlice<i32>, i32)>,
23094 kv_shared: bool,
23097 g: bool,
23101 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23104 ) -> Result<(), Box<dyn std::error::Error>> {
23105 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
23106 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
23113 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23114 let v = *SP512.get_or_init(|| {
23117 std::env::var("MEMRA_FA_SP512")
23118 .ok()
23119 .and_then(|x| x.parse().ok())
23120 .unwrap_or(0)
23121 });
23122 sp = if v >= 8 {
23123 v
23124 } else {
23125 FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
23126 };
23127 }
23128 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23129 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23130 let gqa = (n_head / n_head_kv).max(1) as u32;
23131 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
23142 groups.push((0, t, sp));
23143 } else {
23144 let mut r0 = 0usize;
23145 while r0 < t {
23146 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
23147 let mut r1 = r0 + 1;
23148 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g {
23149 r1 += 1;
23150 }
23151 groups.push((r0, r1 - r0, sp_g));
23152 r0 = r1;
23153 }
23154 }
23155 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23159 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
23160 std::env::var("MEMRA_FA_SMEM_TKV")
23161 .ok()
23162 .and_then(|v| v.parse().ok())
23163 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23164 });
23165 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
23166 let v3 = fa_v3_active(head_dim);
23167 let smem_rows =
23168 head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
23169 let _ = kv_shared;
23174 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
23177 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23191 let tb512 = head_dim == 512
23193 && sp <= 32
23194 && n_head / n_head_kv.max(1) <= 16
23195 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
23196 let fname = if tb512 {
23197 "fa_decode_vec_q_rows_v4_512_tb"
23198 } else if i2 {
23199 "fa_decode_vec_q_rows_dpl16_i2"
23200 } else if head_dim == 512 {
23201 "fa_decode_vec_q_rows_dpl16"
23202 }
23203 else if v4 {
23205 "fa_decode_vec_q_rows_v4"
23206 } else if v3 {
23207 "fa_decode_vec_q_rows_v3"
23208 } else if fa_v2_on() {
23209 "fa_decode_vec_q_rows_v2"
23210 } else if smem_rows {
23211 "fa_decode_vec_q_rows_smem"
23212 } else {
23213 "fa_decode_vec_q_rows"
23214 };
23215 let f = if head_dim == 512 {
23216 self.fa_func(fname, head_dim)
23217 } else if g {
23218 self.func_g(if smem_rows {
23226 "fa_decode_vec_q_rows"
23227 } else {
23228 fname
23229 })
23230 } else {
23231 self.func(fname)
23232 };
23233 let shmem = if tb512 {
23234 let gk = Self::gkv_on();
23236 let sh =
23237 (8192 + 1024 + 32 * 512 + 32 * 64 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
23238 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23239 f.set_attribute(
23240 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23241 sh as i32,
23242 )?;
23243 sh
23244 } else if v4 || v3 || smem_rows || fa_v2_on() {
23245 let sh = (if v4 {
23247 11520 + 32 * head_dim * if g { 1 } else { 2 }
23248 } else if v3 {
23249 32 * head_dim * 2
23250 } else {
23251 2 * 32 * head_dim * 2
23252 }) as u32;
23253 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23254 f.set_attribute(
23255 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23256 sh as i32,
23257 )?;
23258 sh
23259 } else {
23260 0
23261 };
23262 for &(r0, t_g, sp_g) in &groups {
23266 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
23267 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
23268 let base_i = (base_len + r0) as i32;
23269 let o_len = t_g * n_head * n_splits_g * head_dim;
23270 let ml_len = t_g * n_head * n_splits_g;
23271 let mut part_guard = self.fa_part_pool.lock().unwrap();
23272 if part_guard
23273 .as_ref()
23274 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23275 .unwrap_or(true)
23276 {
23277 let old = part_guard.take();
23288 let (co, cm) = old
23289 .as_ref()
23290 .map(|pp| (pp.0.len(), pp.1.len()))
23291 .unwrap_or((0, 0));
23292 if let Some(old) = old {
23293 self.fa_part_retired.lock().unwrap().push(old);
23294 }
23295 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23296 eprintln!(
23297 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23298 co, o_len, cm, ml_len
23299 );
23300 }
23301 *part_guard =
23302 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23303 }
23304 let pg = part_guard.as_mut().unwrap();
23305 self.gpu
23306 .stream()
23307 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23308 self.gpu
23309 .stream()
23310 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23311 self.gpu
23312 .stream()
23313 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23314 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23315 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
23316 let qv = self.view(q, t * n_head * head_dim);
23317 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
23318 let cfg = LaunchConfig {
23319 grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
23320 block_dim: (32, gqa, 1),
23321 shared_mem_bytes: shmem,
23322 };
23323 {
23324 let __s_b = self.gpu.stream();
23325 let mut b = __s_b.launch_builder(&f);
23326 if tb512 {
23327 let (bd, plus) =
23329 base_dev.expect("hd512 rows twin requires a device base counter");
23330 let plus_g = plus + r0 as i32;
23331 let nr = t_g as i32;
23332 if Self::pdl_on() && Self::pdl_wb_on() {
23333 use cudarc::driver::{DevicePtr, DevicePtrMut};
23335 let s = &self.gpu.stream();
23336 let (pq, _b0) = q_g.device_ptr(s);
23337 let (pk, _b1) = k.device_ptr(s);
23338 let (pv, _b2) = v.device_ptr(s);
23339 let (po, _b3) = part_o.device_ptr_mut(s);
23340 let (pm, _b4) = part_m.device_ptr_mut(s);
23341 let (pl, _b5) = part_l.device_ptr_mut(s);
23342 let (pb, _b6) = bd.device_ptr(s);
23343 let mut ps = [
23344 &pq as *const _ as *mut std::ffi::c_void,
23345 &pk as *const _ as *mut _,
23346 &pv as *const _ as *mut _,
23347 &po as *const _ as *mut _,
23348 &pm as *const _ as *mut _,
23349 &pl as *const _ as *mut _,
23350 &hd as *const _ as *mut _,
23351 &nh as *const _ as *mut _,
23352 &nhkv as *const _ as *mut _,
23353 &pb as *const _ as *mut _,
23354 &plus_g as *const _ as *mut _,
23355 &scale as *const _ as *mut _,
23356 &nspm as *const _ as *mut _,
23357 &spk as *const _ as *mut _,
23358 &ktb as *const _ as *mut _,
23359 &vtb as *const _ as *mut _,
23360 &nr as *const _ as *mut _,
23361 ];
23362 unsafe {
23363 self.launch_pdl_flash(
23364 Self::gkv_on(),
23365 "fa_decode_vec_q_rows_v4_512_tb",
23366 (n_head_kv as u32, n_splits_g as u32, 1),
23367 (32, gqa, 1),
23368 shmem,
23369 &mut ps,
23370 )?;
23371 }
23372 } else {
23373 let cfg_tb = LaunchConfig {
23374 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
23375 block_dim: (32, gqa, 1),
23376 shared_mem_bytes: shmem,
23377 };
23378 b.arg(&q_g)
23379 .arg(k)
23380 .arg(v)
23381 .arg(&mut *part_o)
23382 .arg(&mut *part_m)
23383 .arg(&mut *part_l)
23384 .arg(&hd)
23385 .arg(&nh)
23386 .arg(&nhkv)
23387 .arg(bd)
23388 .arg(&plus_g)
23389 .arg(&scale)
23390 .arg(&nspm)
23391 .arg(&spk)
23392 .arg(&ktb)
23393 .arg(&vtb)
23394 .arg(&nr);
23395 unsafe {
23396 b.launch(cfg_tb)?;
23397 }
23398 }
23399 } else if head_dim == 512 {
23400 let (bd, plus) =
23401 base_dev.expect("hd512 rows twin requires a device base counter");
23402 let plus_g = plus + r0 as i32;
23403 b.arg(&q_g)
23404 .arg(k)
23405 .arg(v)
23406 .arg(&mut *part_o)
23407 .arg(&mut *part_m)
23408 .arg(&mut *part_l)
23409 .arg(&hd)
23410 .arg(&nh)
23411 .arg(&nhkv)
23412 .arg(bd)
23413 .arg(&plus_g)
23414 .arg(&scale)
23415 .arg(&nspm)
23416 .arg(&spk)
23417 .arg(&ktb)
23418 .arg(&vtb);
23419 unsafe {
23420 b.launch(cfg)?;
23421 }
23422 } else {
23423 b.arg(&q_g)
23424 .arg(k)
23425 .arg(v)
23426 .arg(&mut *part_o)
23427 .arg(&mut *part_m)
23428 .arg(&mut *part_l)
23429 .arg(&hd)
23430 .arg(&nh)
23431 .arg(&nhkv)
23432 .arg(&base_i)
23433 .arg(&scale)
23434 .arg(&nspm)
23435 .arg(&spk)
23436 .arg(&ktb)
23437 .arg(&vtb);
23438 unsafe {
23439 b.launch(cfg)?;
23440 }
23441 }
23442 }
23443 let cfg2 = LaunchConfig {
23444 grid_dim: (n_head as u32, t_g as u32, 1),
23445 block_dim: (head_dim as u32, 1, 1),
23446 shared_mem_bytes: 0,
23447 };
23448 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
23449 if head_dim == 512 {
23450 let (bd, plus) = base_dev.unwrap();
23453 let plus_g = plus + r0 as i32;
23454 if let Some((oq, od)) = q8_out.as_mut() {
23455 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
23457 if Self::pdl_on() && Self::pdl_wb_on() {
23458 use cudarc::driver::{DevicePtr, DevicePtrMut};
23460 let s = &self.gpu.stream();
23461 let (po, _g0) = part_o.device_ptr(s);
23462 let (pm, _g1) = part_m.device_ptr(s);
23463 let (pl, _g2) = part_l.device_ptr(s);
23464 let (pq, _g3) = oq.device_ptr_mut(s);
23465 let (pd, _g4) = od.device_ptr_mut(s);
23466 let (pb, _g5) = bd.device_ptr(s);
23467 let mut ps = [
23468 &po as *const _ as *mut std::ffi::c_void,
23469 &pm as *const _ as *mut _,
23470 &pl as *const _ as *mut _,
23471 &pq as *const _ as *mut _,
23472 &pd as *const _ as *mut _,
23473 &hd as *const _ as *mut _,
23474 &nh as *const _ as *mut _,
23475 &pb as *const _ as *mut _,
23476 &plus_g as *const _ as *mut _,
23477 &nspm as *const _ as *mut _,
23478 &spk as *const _ as *mut _,
23479 ];
23480 unsafe {
23481 self.launch_pdl_flash(
23482 Self::gkv_on(),
23483 "fa_decode_combine_rows_dc_q8_1",
23484 cfg2.grid_dim,
23485 cfg2.block_dim,
23486 0,
23487 &mut ps,
23488 )?;
23489 }
23490 continue;
23491 }
23492 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
23493 let __s_b2 = self.gpu.stream();
23494 let mut b2 = __s_b2.launch_builder(&fc);
23495 b2.arg(&*part_o)
23496 .arg(&*part_m)
23497 .arg(&*part_l)
23498 .arg(&mut **oq)
23499 .arg(&mut **od)
23500 .arg(&hd)
23501 .arg(&nh)
23502 .arg(bd)
23503 .arg(&plus_g)
23504 .arg(&nspm)
23505 .arg(&spk);
23506 unsafe {
23507 b2.launch(cfg2)?;
23508 }
23509 continue;
23510 }
23511 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
23512 let __s_b2 = self.gpu.stream();
23513 let mut b2 = __s_b2.launch_builder(&fc);
23514 b2.arg(&*part_o)
23515 .arg(&*part_m)
23516 .arg(&*part_l)
23517 .arg(&mut o_g)
23518 .arg(&hd)
23519 .arg(&nh)
23520 .arg(bd)
23521 .arg(&plus_g)
23522 .arg(&nspm)
23523 .arg(&spk);
23524 unsafe {
23525 b2.launch(cfg2)?;
23526 }
23527 } else {
23528 assert!(
23531 q8_out.is_none(),
23532 "rows q8 emit requires the hd512 dc combine"
23533 );
23534 let fc = self.func("fa_decode_combine_rows");
23535 let __s_b2 = self.gpu.stream();
23536 let mut b2 = __s_b2.launch_builder(&fc);
23537 b2.arg(&*part_o)
23538 .arg(&*part_m)
23539 .arg(&*part_l)
23540 .arg(&mut o_g)
23541 .arg(&hd)
23542 .arg(&nh)
23543 .arg(&base_i)
23544 .arg(&nspm)
23545 .arg(&spk);
23546 unsafe {
23547 b2.launch(cfg2)?;
23548 }
23549 }
23550 }
23551 Ok(())
23552 }
23553
23554 #[allow(clippy::too_many_arguments)]
23558 pub fn fa_decode_rows_w(
23559 &self,
23560 q: &CudaSlice<f32>,
23561 k: &cudarc::driver::CudaView<u8>,
23562 v: &cudarc::driver::CudaView<u8>,
23563 o: &mut CudaSlice<f32>,
23564 head_dim: usize,
23565 n_head: usize,
23566 n_head_kv: usize,
23567 base_dev: &CudaSlice<i32>,
23568 base_plus: i32,
23569 t: usize,
23570 scale: f32,
23571 window: usize,
23572 k_tok_bytes: usize,
23573 v_tok_bytes: usize,
23574 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
23575 ) -> Result<(), Box<dyn std::error::Error>> {
23576 debug_assert!(head_dim == 256);
23581 let sp = {
23589 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23590 let v = *SPW.get_or_init(|| {
23591 std::env::var("MEMRA_FA_SPW")
23592 .ok()
23593 .and_then(|x| x.parse().ok())
23594 .unwrap_or(0)
23595 });
23596 if v >= 8 {
23597 v
23598 } else {
23599 FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)
23600 }
23601 };
23602 let n_splits_max = (window + sp - 1) / sp;
23603 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23604 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
23605 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23606 let gqa = (n_head / n_head_kv).max(1) as u32;
23607 let o_len = t * n_head * n_splits_max * head_dim;
23608 let ml_len = t * n_head * n_splits_max;
23609 let mut part_guard = self.fa_part_pool.lock().unwrap();
23610 if part_guard
23611 .as_ref()
23612 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23613 .unwrap_or(true)
23614 {
23615 let old = part_guard.take();
23626 let (co, cm) = old
23627 .as_ref()
23628 .map(|pp| (pp.0.len(), pp.1.len()))
23629 .unwrap_or((0, 0));
23630 if let Some(old) = old {
23631 self.fa_part_retired.lock().unwrap().push(old);
23632 }
23633 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
23634 eprintln!(
23635 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
23636 co, o_len, cm, ml_len
23637 );
23638 }
23639 *part_guard =
23640 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
23641 }
23642 let pg = part_guard.as_mut().unwrap();
23643 self.gpu
23644 .stream()
23645 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
23646 self.gpu
23647 .stream()
23648 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
23649 self.gpu
23650 .stream()
23651 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
23652 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
23653 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
23659 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
23660 std::env::var("MEMRA_FA_SMEM_TKV")
23661 .ok()
23662 .and_then(|v| v.parse().ok())
23663 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
23664 });
23665 use cudarc::driver::sys::CUfunction_attribute_enum as A;
23671 let wg = Self::wkv_on();
23676 let sp2 =
23679 gqa <= 4 && fa_v4_at(window) && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
23680 if sp2 {
23681 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23682 if Self::pdl_on() && Self::pdl_wb_on() {
23683 use cudarc::driver::{DevicePtr, DevicePtrMut};
23685 let s = &self.gpu.stream();
23686 let (pq, _b0) = q.device_ptr(s);
23687 let (pk, _b1) = k.device_ptr(s);
23688 let (pv, _b2) = v.device_ptr(s);
23689 let (po, _b3) = part_o.device_ptr_mut(s);
23690 let (pm, _b4) = part_m.device_ptr_mut(s);
23691 let (pl, _b5) = part_l.device_ptr_mut(s);
23692 let (pb, _b6) = base_dev.device_ptr(s);
23693 let mut ps = [
23694 &pq as *const _ as *mut std::ffi::c_void,
23695 &pk as *const _ as *mut _,
23696 &pv as *const _ as *mut _,
23697 &po as *const _ as *mut _,
23698 &pm as *const _ as *mut _,
23699 &pl as *const _ as *mut _,
23700 &hd as *const _ as *mut _,
23701 &nh as *const _ as *mut _,
23702 &nhkv as *const _ as *mut _,
23703 &pb as *const _ as *mut _,
23704 &base_plus as *const _ as *mut _,
23705 &scale as *const _ as *mut _,
23706 &nspm as *const _ as *mut _,
23707 &spk as *const _ as *mut _,
23708 &ktb as *const _ as *mut _,
23709 &vtb as *const _ as *mut _,
23710 &wini as *const _ as *mut _,
23711 ];
23712 unsafe {
23713 self.launch_pdl_flash(
23714 wg,
23715 "fa_decode_vec_q_rows_v4_w_sp",
23716 (n_head_kv as u32, n_splits_max as u32, t as u32),
23717 (32, gqa + 1, 1),
23718 sh,
23719 &mut ps,
23720 )?;
23721 }
23722 } else {
23723 let f = if wg {
23724 self.func_g("fa_decode_vec_q_rows_v4_w_sp")
23725 } else {
23726 self.func("fa_decode_vec_q_rows_v4_w_sp")
23727 };
23728 f.set_attribute(
23729 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23730 sh as i32,
23731 )?;
23732 let cfg = LaunchConfig {
23733 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23734 block_dim: (32, gqa + 1, 1),
23735 shared_mem_bytes: sh,
23736 };
23737 let __s_b = self.gpu.stream();
23738 let mut b = __s_b.launch_builder(&f);
23739 b.arg(q)
23740 .arg(k)
23741 .arg(v)
23742 .arg(&mut *part_o)
23743 .arg(&mut *part_m)
23744 .arg(&mut *part_l)
23745 .arg(&hd)
23746 .arg(&nh)
23747 .arg(&nhkv)
23748 .arg(base_dev)
23749 .arg(&base_plus)
23750 .arg(&scale)
23751 .arg(&nspm)
23752 .arg(&spk)
23753 .arg(&ktb)
23754 .arg(&vtb)
23755 .arg(&wini);
23756 unsafe {
23757 b.launch(cfg)?;
23758 }
23759 }
23760 } else {
23761 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
23762 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
23764 use cudarc::driver::{DevicePtr, DevicePtrMut};
23765 let s = &self.gpu.stream();
23766 let (pq, _b0) = q.device_ptr(s);
23767 let (pk, _b1) = k.device_ptr(s);
23768 let (pv, _b2) = v.device_ptr(s);
23769 let (po, _b3) = part_o.device_ptr_mut(s);
23770 let (pm, _b4) = part_m.device_ptr_mut(s);
23771 let (pl, _b5) = part_l.device_ptr_mut(s);
23772 let (pb, _b6) = base_dev.device_ptr(s);
23773 let mut ps = [
23774 &pq as *const _ as *mut std::ffi::c_void,
23775 &pk as *const _ as *mut _,
23776 &pv as *const _ as *mut _,
23777 &po as *const _ as *mut _,
23778 &pm as *const _ as *mut _,
23779 &pl as *const _ as *mut _,
23780 &hd as *const _ as *mut _,
23781 &nh as *const _ as *mut _,
23782 &nhkv as *const _ as *mut _,
23783 &pb as *const _ as *mut _,
23784 &base_plus as *const _ as *mut _,
23785 &scale as *const _ as *mut _,
23786 &nspm as *const _ as *mut _,
23787 &spk as *const _ as *mut _,
23788 &ktb as *const _ as *mut _,
23789 &vtb as *const _ as *mut _,
23790 &wini as *const _ as *mut _,
23791 ];
23792 unsafe {
23793 self.launch_pdl_flash(
23794 wg,
23795 "fa_decode_vec_q_rows_v4_w",
23796 (n_head_kv as u32, n_splits_max as u32, t as u32),
23797 (32, gqa, 1),
23798 sh,
23799 &mut ps,
23800 )?;
23801 }
23802 } else {
23803 let pick = |name: &str| {
23804 if wg {
23805 self.func_g(name)
23806 } else {
23807 self.func(name)
23808 }
23809 };
23810 let (f, sh) = if fa_v4_at(window) {
23811 let f = pick("fa_decode_vec_q_rows_v4_w");
23812 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
23813 } else if smem_tkv > 0 && window >= smem_tkv {
23814 (
23817 pick("fa_decode_vec_q_rows_smem_w"),
23818 (2 * 32 * head_dim * 2) as u32,
23819 )
23820 } else {
23821 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
23822 };
23823 f.set_attribute(
23824 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
23825 sh as i32,
23826 )?;
23827 let cfg = LaunchConfig {
23828 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
23829 block_dim: (32, gqa, 1),
23830 shared_mem_bytes: sh,
23831 };
23832 let __s_b = self.gpu.stream();
23833 let mut b = __s_b.launch_builder(&f);
23834 b.arg(q)
23835 .arg(k)
23836 .arg(v)
23837 .arg(&mut *part_o)
23838 .arg(&mut *part_m)
23839 .arg(&mut *part_l)
23840 .arg(&hd)
23841 .arg(&nh)
23842 .arg(&nhkv)
23843 .arg(base_dev)
23844 .arg(&base_plus)
23845 .arg(&scale)
23846 .arg(&nspm)
23847 .arg(&spk)
23848 .arg(&ktb)
23849 .arg(&vtb)
23850 .arg(&wini);
23851 unsafe {
23852 b.launch(cfg)?;
23853 }
23854 }
23855 }
23856 let cfg2 = LaunchConfig {
23857 grid_dim: (n_head as u32, t as u32, 1),
23858 block_dim: (head_dim as u32, 1, 1),
23859 shared_mem_bytes: 0,
23860 };
23861 if let Some((oq, od)) = q8_out {
23862 if Self::pdl_on() && Self::pdl_wb_on() {
23865 use cudarc::driver::{DevicePtr, DevicePtrMut};
23867 let s = &self.gpu.stream();
23868 let (po, _g0) = part_o.device_ptr(s);
23869 let (pm, _g1) = part_m.device_ptr(s);
23870 let (pl, _g2) = part_l.device_ptr(s);
23871 let (pq, _g3) = oq.device_ptr_mut(s);
23872 let (pd, _g4) = od.device_ptr_mut(s);
23873 let mut ps = [
23874 &po as *const _ as *mut std::ffi::c_void,
23875 &pm as *const _ as *mut _,
23876 &pl as *const _ as *mut _,
23877 &pq as *const _ as *mut _,
23878 &pd as *const _ as *mut _,
23879 &hd as *const _ as *mut _,
23880 &nh as *const _ as *mut _,
23881 &nspm as *const _ as *mut _,
23882 &spk as *const _ as *mut _,
23883 &wini as *const _ as *mut _,
23884 ];
23885 unsafe {
23886 self.launch_pdl_flash(
23887 wg,
23888 "fa_decode_combine_rows_w_q8_1",
23889 cfg2.grid_dim,
23890 cfg2.block_dim,
23891 0,
23892 &mut ps,
23893 )?;
23894 }
23895 return Ok(());
23896 }
23897 let fc = if wg {
23898 self.func_g("fa_decode_combine_rows_w_q8_1")
23899 } else {
23900 self.func("fa_decode_combine_rows_w_q8_1")
23901 };
23902 let __s_b2 = self.gpu.stream();
23903 let mut b2 = __s_b2.launch_builder(&fc);
23904 b2.arg(&*part_o)
23905 .arg(&*part_m)
23906 .arg(&*part_l)
23907 .arg(oq)
23908 .arg(od)
23909 .arg(&hd)
23910 .arg(&nh)
23911 .arg(&nspm)
23912 .arg(&spk)
23913 .arg(&wini);
23914 unsafe {
23915 b2.launch(cfg2)?;
23916 }
23917 return Ok(());
23918 }
23919 let fc = if wg {
23920 self.func_g("fa_decode_combine_rows_w")
23921 } else {
23922 self.func("fa_decode_combine_rows_w")
23923 };
23924 let __s_b2 = self.gpu.stream();
23925 let mut b2 = __s_b2.launch_builder(&fc);
23926 b2.arg(&*part_o)
23927 .arg(&*part_m)
23928 .arg(&*part_l)
23929 .arg(o)
23930 .arg(&hd)
23931 .arg(&nh)
23932 .arg(&nspm)
23933 .arg(&spk)
23934 .arg(&wini);
23935 unsafe {
23936 b2.launch(cfg2)?;
23937 }
23938 Ok(())
23939 }
23940
23941 #[allow(clippy::too_many_arguments)]
23947 pub fn fa_decode_rows_dc(
23948 &self,
23949 q: &CudaSlice<f32>,
23950 k: &cudarc::driver::CudaView<u8>,
23951 v: &cudarc::driver::CudaView<u8>,
23952 o: &mut CudaSlice<f32>,
23953 head_dim: usize,
23954 n_head: usize,
23955 n_head_kv: usize,
23956 base_dev: &CudaSlice<i32>,
23957 t_kv_upper: usize,
23958 t: usize,
23959 scale: f32,
23960 k_tok_bytes: usize,
23961 v_tok_bytes: usize,
23962 base_plus: i32,
23963 g: bool,
23964 ) -> Result<(), Box<dyn std::error::Error>> {
23965 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
23966 assert!(
23967 v4 || fa_v3_active(head_dim),
23968 "stream fa rows requires the v3 or v4 lane"
23969 );
23970 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
23971 if v4 {
23972 let sp = fa_split_keys(t_kv_upper, n_head_kv);
23973 let n_splits_max = (t_kv_upper + sp - 1) / sp;
23974 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
23975 let (nspm, spk) = (n_splits_max as i32, sp as i32);
23976 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
23977 let gqa = (n_head / n_head_kv).max(1) as u32;
23978 let o_len = t * n_head * n_splits_max * head_dim;
23979 let ml_len = t * n_head * n_splits_max;
23980 let mut part_guard = self.fa_part_pool.lock().unwrap();
23981 if part_guard
23982 .as_ref()
23983 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
23984 .unwrap_or(true)
23985 {
23986 let old = part_guard.take();
23997 let (co, cm) = old
23998 .as_ref()
23999 .map(|pp| (pp.0.len(), pp.1.len()))
24000 .unwrap_or((0, 0));
24001 if let Some(old) = old {
24002 self.fa_part_retired.lock().unwrap().push(old);
24003 }
24004 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
24005 eprintln!(
24006 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
24007 co, o_len, cm, ml_len
24008 );
24009 }
24010 *part_guard =
24011 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24012 }
24013 let pg = part_guard.as_mut().unwrap();
24014 self.gpu
24015 .stream()
24016 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24017 self.gpu
24018 .stream()
24019 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24020 self.gpu
24021 .stream()
24022 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24023 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24024 let f = if g {
24025 self.func_g("fa_decode_vec_q_rows_v4_dc")
24026 } else {
24027 self.func("fa_decode_vec_q_rows_v4_dc")
24028 };
24029 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
24030 use cudarc::driver::sys::CUfunction_attribute_enum as A;
24031 f.set_attribute(
24032 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24033 sh as i32,
24034 )?;
24035 let cfg = LaunchConfig {
24036 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
24037 block_dim: (32, gqa, 1),
24038 shared_mem_bytes: sh,
24039 };
24040 let __s_b = self.gpu.stream();
24041 let mut b = __s_b.launch_builder(&f);
24042 b.arg(q)
24043 .arg(k)
24044 .arg(v)
24045 .arg(&mut *part_o)
24046 .arg(&mut *part_m)
24047 .arg(&mut *part_l)
24048 .arg(&hd)
24049 .arg(&nh)
24050 .arg(&nhkv)
24051 .arg(base_dev)
24052 .arg(&base_plus)
24053 .arg(&scale)
24054 .arg(&nspm)
24055 .arg(&spk)
24056 .arg(&ktb)
24057 .arg(&vtb);
24058 unsafe {
24059 b.launch(cfg)?;
24060 }
24061 let fc = self.func("fa_decode_combine_rows_dc");
24062 let cfg2 = LaunchConfig {
24063 grid_dim: (n_head as u32, t as u32, 1),
24064 block_dim: (head_dim as u32, 1, 1),
24065 shared_mem_bytes: 0,
24066 };
24067 let __s_b2 = self.gpu.stream();
24068 let mut b2 = __s_b2.launch_builder(&fc);
24069 b2.arg(&*part_o)
24070 .arg(&*part_m)
24071 .arg(&*part_l)
24072 .arg(o)
24073 .arg(&hd)
24074 .arg(&nh)
24075 .arg(base_dev)
24076 .arg(&base_plus)
24077 .arg(&nspm)
24078 .arg(&spk);
24079 unsafe {
24080 b2.launch(cfg2)?;
24081 }
24082 return Ok(());
24083 }
24084 let sp = fa_split_keys(t_kv_upper, n_head_kv);
24085 let n_splits_max = (t_kv_upper + sp - 1) / sp;
24086 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24087 let (nspm, spk) = (n_splits_max as i32, sp as i32);
24088 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24089 let gqa = (n_head / n_head_kv).max(1) as u32;
24090 let o_len = t * n_head * n_splits_max * head_dim;
24091 let ml_len = t * n_head * n_splits_max;
24092 let mut part_guard = self.fa_part_pool.lock().unwrap();
24093 if part_guard
24094 .as_ref()
24095 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24096 .unwrap_or(true)
24097 {
24098 let old = part_guard.take();
24109 let (co, cm) = old
24110 .as_ref()
24111 .map(|pp| (pp.0.len(), pp.1.len()))
24112 .unwrap_or((0, 0));
24113 if let Some(old) = old {
24114 self.fa_part_retired.lock().unwrap().push(old);
24115 }
24116 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
24117 eprintln!(
24118 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
24119 co, o_len, cm, ml_len
24120 );
24121 }
24122 *part_guard =
24123 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24124 }
24125 let pg = part_guard.as_mut().unwrap();
24126 self.gpu
24127 .stream()
24128 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24129 self.gpu
24130 .stream()
24131 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24132 self.gpu
24133 .stream()
24134 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24135 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24136 let f = self.func("fa_decode_vec_q_rows_v3_dc");
24137 let sh = (32 * head_dim * 2) as u32;
24138 use cudarc::driver::sys::CUfunction_attribute_enum as A;
24139 f.set_attribute(
24140 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24141 sh as i32,
24142 )?;
24143 let cfg = LaunchConfig {
24144 grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
24145 block_dim: (32, gqa, 1),
24146 shared_mem_bytes: sh,
24147 };
24148 let __s_b = self.gpu.stream();
24149 let mut b = __s_b.launch_builder(&f);
24150 b.arg(q)
24151 .arg(k)
24152 .arg(v)
24153 .arg(&mut *part_o)
24154 .arg(&mut *part_m)
24155 .arg(&mut *part_l)
24156 .arg(&hd)
24157 .arg(&nh)
24158 .arg(&nhkv)
24159 .arg(base_dev)
24160 .arg(&scale)
24161 .arg(&nspm)
24162 .arg(&spk)
24163 .arg(&ktb)
24164 .arg(&vtb);
24165 unsafe {
24166 b.launch(cfg)?;
24167 }
24168 let fc = self.func("fa_decode_combine_rows_dc");
24169 let cfg2 = LaunchConfig {
24170 grid_dim: (n_head as u32, t as u32, 1),
24171 block_dim: (head_dim as u32, 1, 1),
24172 shared_mem_bytes: 0,
24173 };
24174 let plus0 = 0i32;
24175 let __s_b2 = self.gpu.stream();
24176 let mut b2 = __s_b2.launch_builder(&fc);
24177 b2.arg(&*part_o)
24178 .arg(&*part_m)
24179 .arg(&*part_l)
24180 .arg(o)
24181 .arg(&hd)
24182 .arg(&nh)
24183 .arg(base_dev)
24184 .arg(&plus0)
24185 .arg(&nspm)
24186 .arg(&spk);
24187 unsafe {
24188 b2.launch(cfg2)?;
24189 }
24190 Ok(())
24191 }
24192
24193 pub fn fa_decode_dc(
24204 &self,
24205 q: &CudaSlice<f32>,
24206 k: &cudarc::driver::CudaView<u8>,
24207 v: &cudarc::driver::CudaView<u8>,
24208 o: &mut CudaSlice<f32>,
24209 head_dim: usize,
24210 n_head: usize,
24211 n_head_kv: usize,
24212 t_kv_dev: &CudaSlice<i32>,
24213 bucket_max: usize,
24214 scale: f32,
24215 k_tok_bytes: usize,
24216 v_tok_bytes: usize,
24217 g: bool,
24218 ) -> Result<(), Box<dyn std::error::Error>> {
24219 self.fa_decode_dc_q8(
24220 q,
24221 k,
24222 v,
24223 o,
24224 head_dim,
24225 n_head,
24226 n_head_kv,
24227 t_kv_dev,
24228 bucket_max,
24229 scale,
24230 k_tok_bytes,
24231 v_tok_bytes,
24232 g,
24233 None,
24234 )
24235 }
24236
24237 #[allow(clippy::too_many_arguments)]
24240 pub fn fa_decode_dc_q8(
24241 &self,
24242 q: &CudaSlice<f32>,
24243 k: &cudarc::driver::CudaView<u8>,
24244 v: &cudarc::driver::CudaView<u8>,
24245 o: &mut CudaSlice<f32>,
24246 head_dim: usize,
24247 n_head: usize,
24248 n_head_kv: usize,
24249 t_kv_dev: &CudaSlice<i32>,
24250 bucket_max: usize,
24251 scale: f32,
24252 k_tok_bytes: usize,
24253 v_tok_bytes: usize,
24254 g: bool,
24255 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>,
24256 ) -> Result<(), Box<dyn std::error::Error>> {
24257 let mut fa_vec =
24265 std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24266 if g && head_dim == 256 && !fa_v4_at(bucket_max) {
24267 fa_vec = false;
24268 } let sp = fa_split_keys(bucket_max, n_head_kv);
24270 let n_splits = if fa_vec {
24271 ((bucket_max + sp - 1) / sp).max(1)
24272 } else {
24273 ((bucket_max + 255) / 256).max(1)
24274 };
24275 let o_len = n_head * n_splits * head_dim;
24276 let ml_len = n_head * n_splits;
24277 let mut part_guard = self.fa_part_pool.lock().unwrap();
24278 if part_guard
24279 .as_ref()
24280 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24281 .unwrap_or(true)
24282 {
24283 let old = part_guard.take();
24294 let (co, cm) = old
24295 .as_ref()
24296 .map(|pp| (pp.0.len(), pp.1.len()))
24297 .unwrap_or((0, 0));
24298 if let Some(old) = old {
24299 self.fa_part_retired.lock().unwrap().push(old);
24300 }
24301 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
24302 eprintln!(
24303 "[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)",
24304 co, o_len, cm, ml_len
24305 );
24306 }
24307 *part_guard =
24308 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24309 }
24310 let pg = part_guard.as_mut().unwrap();
24311 self.gpu
24312 .stream()
24313 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
24314 self.gpu
24315 .stream()
24316 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
24317 self.gpu
24318 .stream()
24319 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
24320 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24321 let (hd, nh, nhkv, nsp) = (
24322 head_dim as i32,
24323 n_head as i32,
24324 n_head_kv as i32,
24325 n_splits as i32,
24326 );
24327 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24328 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
24329 let deep = fa_vec
24332 && head_dim == 256
24333 && fa_v4_at(bucket_max)
24334 && !g
24335 && fa_deep_at(bucket_max)
24336 && !matches!(fa_v4_mode(), "noB3" | "stage");
24337 let (f, cfg) = if fa_vec
24338 && head_dim == 512
24339 && bucket_max >= {
24340 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
24341 *FA512_MIN_DC.get_or_init(|| {
24342 std::env::var("MEMRA_FA512_MIN")
24343 .ok()
24344 .and_then(|v| v.parse().ok())
24345 .unwrap_or(512)
24346 })
24347 } {
24348 let gqa = (n_head / n_head_kv).max(1) as u32;
24350 (
24351 self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
24352 LaunchConfig {
24353 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24354 block_dim: (32, gqa, 1),
24355 shared_mem_bytes: 0,
24356 },
24357 )
24358 } else if fa_vec && head_dim == 512 {
24359 let q_view = q.as_view();
24362 let mut o_view = o.as_view_mut();
24363 return self.fa_decode_scalar_unified(
24364 &q_view,
24365 k,
24366 v,
24367 &mut o_view,
24368 head_dim,
24369 n_head,
24370 n_head_kv,
24371 0,
24372 Some(t_kv_dev),
24373 scale,
24374 n_splits,
24375 sp,
24376 k_tok_bytes,
24377 v_tok_bytes,
24378 g,
24379 &mut *part_o,
24380 &mut *part_m,
24381 &mut *part_l,
24382 q8_out,
24383 );
24384 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
24385 let gqa = (n_head / n_head_kv).max(1) as u32;
24388 let fv = if g {
24389 self.func_g("fa_decode_vec_q_v4_dc")
24390 } else if deep {
24391 self.func("fa_decode_vec_q_v4_deep_dc")
24392 } else {
24393 self.func("fa_decode_vec_q_v4_dc")
24394 };
24395 let shmem =
24396 (if deep { 12160 } else { 11520 } + 32 * head_dim * if g { 1 } else { 2 }) as u32;
24397 use cudarc::driver::sys::CUfunction_attribute_enum as A;
24398 fv.set_attribute(
24399 A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
24400 shmem as i32,
24401 )?;
24402 (
24403 fv,
24404 LaunchConfig {
24405 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24406 block_dim: (32, gqa, 1),
24407 shared_mem_bytes: shmem,
24408 },
24409 )
24410 } else if fa_vec && fa_v3_active(head_dim) {
24411 let gqa = (n_head / n_head_kv).max(1) as u32;
24414 let fv = if g {
24415 self.func_g("fa_decode_vec_q_v3_dc")
24416 } else {
24417 self.func("fa_decode_vec_q_v3_dc")
24418 };
24419 let shmem = (32 * head_dim * 2) as u32; (
24421 fv,
24422 LaunchConfig {
24423 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24424 block_dim: (32, gqa, 1),
24425 shared_mem_bytes: shmem,
24426 },
24427 )
24428 } else if fa_vec && fa_v2_on() {
24429 let gqa = (n_head / n_head_kv).max(1) as u32;
24433 let fv = if g {
24434 self.func_g("fa_decode_vec_q_v2_dc")
24435 } else {
24436 self.func("fa_decode_vec_q_v2_dc")
24437 };
24438 let shmem = (2 * 32 * head_dim * 2) as u32; (
24440 fv,
24441 LaunchConfig {
24442 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24443 block_dim: (32, gqa, 1),
24444 shared_mem_bytes: shmem,
24445 },
24446 )
24447 } else if fa_vec {
24448 let gqa = (n_head / n_head_kv).max(1) as u32;
24449 let fv = if g {
24451 self.func_g("fa_decode_vec_q_dc")
24452 } else {
24453 self.func("fa_decode_vec_q_dc")
24454 };
24455 (
24456 fv,
24457 LaunchConfig {
24458 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24459 block_dim: (32, gqa, 1),
24460 shared_mem_bytes: 0,
24461 },
24462 )
24463 } else {
24464 let q_view = q.as_view();
24465 let mut o_view = o.as_view_mut();
24466 return self.fa_decode_scalar_unified(
24467 &q_view,
24468 k,
24469 v,
24470 &mut o_view,
24471 head_dim,
24472 n_head,
24473 n_head_kv,
24474 0,
24475 Some(t_kv_dev),
24476 scale,
24477 n_splits,
24478 if fa_vec { sp } else { 256 },
24479 k_tok_bytes,
24480 v_tok_bytes,
24481 g,
24482 &mut *part_o,
24483 &mut *part_m,
24484 &mut *part_l,
24485 q8_out,
24486 );
24487 };
24488 let ski = sp as i32; let __s_b = self.gpu.stream();
24490 let mut b = __s_b.launch_builder(&f);
24491 b.arg(q)
24492 .arg(k)
24493 .arg(v)
24494 .arg(&mut *part_o)
24495 .arg(&mut *part_m)
24496 .arg(&mut *part_l)
24497 .arg(&hd)
24498 .arg(&nh)
24499 .arg(&nhkv)
24500 .arg(t_kv_dev)
24501 .arg(&scale)
24502 .arg(&nsp)
24503 .arg(&ski)
24504 .arg(&ktb)
24505 .arg(&vtb);
24506 unsafe {
24507 b.launch(cfg)?;
24508 }
24509 let cfg2 = LaunchConfig {
24510 grid_dim: (n_head as u32, 1, 1),
24511 block_dim: (head_dim as u32, 1, 1),
24512 shared_mem_bytes: 0,
24513 };
24514 if let Some((oq, od)) = q8_out {
24515 let fc = if g {
24516 self.func_g("fa_decode_combine_q8_1")
24517 } else {
24518 self.fa_func("fa_decode_combine_q8_1", head_dim)
24519 };
24520 let __s_b2 = self.gpu.stream();
24521 let mut b2 = __s_b2.launch_builder(&fc);
24522 b2.arg(&*part_o)
24523 .arg(&*part_m)
24524 .arg(&*part_l)
24525 .arg(oq)
24526 .arg(od)
24527 .arg(&hd)
24528 .arg(&nh)
24529 .arg(&nsp);
24530 unsafe {
24531 b2.launch(cfg2)?;
24532 }
24533 return Ok(());
24534 }
24535 let fc = if g {
24536 self.func_g("fa_decode_combine_f32")
24537 } else {
24538 self.fa_func("fa_decode_combine_f32", head_dim)
24539 };
24540 let __s_b2 = self.gpu.stream();
24541 let mut b2 = __s_b2.launch_builder(&fc);
24542 b2.arg(&*part_o)
24543 .arg(&*part_m)
24544 .arg(&*part_l)
24545 .arg(o)
24546 .arg(&hd)
24547 .arg(&nh)
24548 .arg(&nsp);
24549 unsafe {
24550 b2.launch(cfg2)?;
24551 }
24552 Ok(())
24553 }
24554
24555 #[allow(clippy::too_many_arguments)]
24559 pub fn append_kv_quantized_dcw(
24560 &self,
24561 k_row: &CudaSlice<f32>,
24562 v_row: &CudaSlice<f32>,
24563 kc: &mut CudaSlice<u8>,
24564 vc: &mut CudaSlice<u8>,
24565 len_dev: &CudaSlice<i32>,
24566 base_dev: Option<&CudaSlice<i32>>,
24567 kv_dim_k: usize,
24568 kv_dim_v: usize,
24569 k_tok_bytes: usize,
24570 v_tok_bytes: usize,
24571 ) -> Result<(), Box<dyn std::error::Error>> {
24572 let f = self.func("append_quantize_kv_q8_0_q5_1_dcw");
24573 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
24574 let cfg = LaunchConfig {
24575 grid_dim: (nblk, 1, 1),
24576 block_dim: (32, 1, 1),
24577 shared_mem_bytes: 0,
24578 };
24579 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
24580 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24581 let null: u64 = 0;
24582 let __s_b = self.gpu.stream();
24583 let mut b = __s_b.launch_builder(&f);
24584 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(len_dev);
24585 match base_dev {
24586 Some(base) => {
24587 b.arg(base);
24588 }
24589 None => {
24590 b.arg(&null);
24591 }
24592 }
24593 b.arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
24594 unsafe {
24595 b.launch(cfg)?;
24596 }
24597 Ok(())
24598 }
24599
24600 pub fn inc_i32(&self, counter: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
24602 let f = self.func("inc_i32");
24603 let cfg = LaunchConfig {
24604 grid_dim: (1, 1, 1),
24605 block_dim: (1, 1, 1),
24606 shared_mem_bytes: 0,
24607 };
24608 let __s_b = self.gpu.stream();
24609 let mut b = __s_b.launch_builder(&f);
24610 b.arg(counter);
24611 unsafe {
24612 b.launch(cfg)?;
24613 }
24614 Ok(())
24615 }
24616
24617 #[allow(clippy::too_many_arguments)]
24626 fn fa_part_alloc(
24650 &self,
24651 o_len: usize,
24652 ml_len: usize,
24653 co: usize,
24654 cm: usize,
24655 ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
24656 static GROWS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
24657 let n = GROWS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
24658 if n < 64 {
24659 eprintln!(
24660 "[fa-pool] grow #{n} dev={} o_len {co} -> {o_len} ml_len {cm} -> {ml_len} (retired kept, zero={})",
24661 self.ctx().ordinal(),
24662 fa_part_zero_on()
24663 );
24664 }
24665 let mut po = self.alloc_uninit::<f32>(o_len)?;
24666 let mut pm = self.alloc_uninit::<f32>(ml_len)?;
24667 let mut pl = self.alloc_uninit::<f32>(ml_len)?;
24668 if fa_part_zero_on() {
24669 self.gpu.stream().memset_zeros(&mut po)?;
24670 self.gpu.stream().memset_zeros(&mut pm)?;
24671 self.gpu.stream().memset_zeros(&mut pl)?;
24672 }
24673 Ok((po, pm, pl))
24674 }
24675
24676 fn fa_part_pool_grow(
24677 &self,
24678 part_guard: &mut Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>,
24679 o_len: usize,
24680 ml_len: usize,
24681 ) -> Result<(), Box<dyn std::error::Error>> {
24682 if part_guard
24683 .as_ref()
24684 .map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len)
24685 .unwrap_or(true)
24686 {
24687 let old = part_guard.take();
24688 let (co, cm) = old
24689 .as_ref()
24690 .map(|pp| (pp.0.len(), pp.1.len()))
24691 .unwrap_or((0, 0));
24692 if let Some(old) = old {
24693 self.fa_part_retired.lock().unwrap().push(old);
24694 }
24695 *part_guard =
24708 Some(self.fa_part_alloc(o_len.max(2 * co), ml_len.max(2 * cm), co, cm)?);
24709 }
24710 Ok(())
24711 }
24712
24713 pub fn fa_dcw_pool_ensure(
24716 &self,
24717 head_dim: usize,
24718 n_head: usize,
24719 n_head_kv: usize,
24720 bucket_max: usize,
24721 ) -> Result<(), Box<dyn std::error::Error>> {
24722 let sp = fa_split_keys(bucket_max, n_head_kv);
24723 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24724 let o_len = n_head * n_splits * head_dim;
24725 let ml_len = n_head * n_splits;
24726 let mut part_guard = self.fa_part_pool.lock().unwrap();
24727 self.fa_part_pool_grow(&mut part_guard, o_len, ml_len)
24728 }
24729
24730 #[allow(clippy::too_many_arguments)]
24738 pub fn fa_decode_dcw2(
24739 &self,
24740 q2: &CudaSlice<f32>,
24741 k_ring: &cudarc::driver::CudaView<u8>,
24742 v_ring: &cudarc::driver::CudaView<u8>,
24743 o2: &mut CudaSlice<f32>,
24744 head_dim: usize,
24745 n_head: usize,
24746 n_head_kv: usize,
24747 len_dev: &CudaSlice<i32>,
24748 base_dev: Option<&CudaSlice<i32>>,
24749 window: usize,
24750 bucket_max: usize,
24751 scale: f32,
24752 k_tok_bytes: usize,
24753 v_tok_bytes: usize,
24754 gate2: &CudaSlice<f32>,
24755 ) -> Result<(), Box<dyn std::error::Error>> {
24756 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24757 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24758 return Err("fa_decode_dcw2 supports the default v3-vec class only".into());
24759 }
24760 let sp = fa_split_keys(bucket_max, n_head_kv);
24761 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
24762 let o_len = 2 * n_head * n_splits * head_dim;
24764 let ml_len = 2 * n_head * n_splits;
24765 let mut part_guard = self.fa_part_pool.lock().unwrap();
24766 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24767 let pg = part_guard.as_mut().unwrap();
24768 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24769 let (hd, nh, nhkv, nsp) = (
24770 head_dim as i32,
24771 n_head as i32,
24772 n_head_kv as i32,
24773 n_splits as i32,
24774 );
24775 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24776 let (ski, win) = (sp as i32, window as i32);
24777 let gqa = (n_head / n_head_kv).max(1) as u32;
24778 let smem = (32 * head_dim * 2) as u32;
24779 let f = self.func("fa_decode_vec_q_v3_dcw2");
24780 let cfg = LaunchConfig {
24781 grid_dim: (n_head_kv as u32, n_splits as u32, 1),
24782 block_dim: (32, gqa, 1),
24783 shared_mem_bytes: smem,
24784 };
24785 let null: u64 = 0;
24786 {
24787 let __s_b = self.gpu.stream();
24788 let mut b = __s_b.launch_builder(&f);
24789 b.arg(q2)
24790 .arg(k_ring)
24791 .arg(v_ring)
24792 .arg(&mut *part_o)
24793 .arg(&mut *part_m)
24794 .arg(&mut *part_l)
24795 .arg(&hd)
24796 .arg(&nh)
24797 .arg(&nhkv)
24798 .arg(len_dev);
24799 match base_dev {
24800 Some(base) => {
24801 b.arg(base);
24802 }
24803 None => {
24804 b.arg(&null);
24805 }
24806 }
24807 b.arg(&win)
24808 .arg(&scale)
24809 .arg(&nsp)
24810 .arg(&ski)
24811 .arg(&ktb)
24812 .arg(&vtb);
24813 unsafe {
24814 b.launch(cfg)?;
24815 }
24816 }
24817 let fc = {
24821 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24822 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24823 self.func("fa_decode_combine_gate_f32_s")
24824 } else {
24825 self.func("fa_decode_combine_gate_f32")
24826 }
24827 };
24828 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24829 let nh2 = (2 * n_head) as i32;
24830 let cfg2 = LaunchConfig {
24831 grid_dim: ((2 * n_head) as u32, 1, 1),
24832 block_dim: (head_dim as u32, 1, 1),
24833 shared_mem_bytes: if combine_shared {
24834 (2 * n_splits * 4) as u32
24835 } else {
24836 0
24837 },
24838 };
24839 let __s_b2 = self.gpu.stream();
24840 let mut b2 = __s_b2.launch_builder(&fc);
24841 b2.arg(&*part_o)
24842 .arg(&*part_m)
24843 .arg(&*part_l)
24844 .arg(gate2)
24845 .arg(o2)
24846 .arg(&hd)
24847 .arg(&nh2)
24848 .arg(&nsp);
24849 unsafe {
24850 b2.launch(cfg2)?;
24851 }
24852 Ok(())
24853 }
24854
24855 #[allow(clippy::too_many_arguments)]
24864 pub fn fa_decode_dcw_rows(
24865 &self,
24866 q_rows: &CudaSlice<f32>,
24867 tab: &CudaSlice<u64>,
24868 o_rows: &mut CudaSlice<f32>,
24869 t: usize,
24870 head_dim: usize,
24871 n_head: usize,
24872 n_head_kv: usize,
24873 window: usize,
24874 max_ns: usize,
24875 scale: f32,
24876 k_tok_bytes: usize,
24877 v_tok_bytes: usize,
24878 gate_rows: &CudaSlice<f32>,
24879 ) -> Result<(), Box<dyn std::error::Error>> {
24880 if std::env::var("MEMRA_NO_FA_VEC").is_ok()
24881 || head_dim > 256
24882 || head_dim % 32 != 0
24883 || !fa_v3_on()
24884 {
24885 return Err("fa_decode_dcw_rows supports the default v3-vec class only".into());
24886 }
24887 if fa_sm_count() < 128
24888 || std::env::var("MEMRA_FA_SPLIT").is_ok()
24889 || std::env::var("MEMRA_FA_SP_SHORT").is_ok()
24890 || std::env::var("MEMRA_FA_SP16").is_ok()
24891 {
24892 return Err(
24893 "fa_decode_dcw_rows embeds the big-rig split ladder; env split overrides \
24894 (or a <128-SM rig) keep the per-row path"
24895 .into(),
24896 );
24897 }
24898 if t == 0 || t > 32 || max_ns == 0 || tab.len() < t * 6 {
24899 return Err("fa_decode_dcw_rows geometry".into());
24900 }
24901 let o_len = t * n_head * max_ns * head_dim;
24902 let ml_len = t * n_head * max_ns;
24903 let mut part_guard = self.fa_part_pool.lock().unwrap();
24904 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
24905 let pg = part_guard.as_mut().unwrap();
24906 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
24907 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
24908 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
24909 let (win, mns) = (window as i32, max_ns as i32);
24910 let gqa = (n_head / n_head_kv).max(1) as u32;
24911 let smem = (32 * head_dim * 2) as u32;
24912 let f = self.func("fa_decode_vec_q_v3_dcw_rows");
24913 let cfg = LaunchConfig {
24914 grid_dim: (n_head_kv as u32, max_ns as u32, t as u32),
24915 block_dim: (32, gqa, 1),
24916 shared_mem_bytes: smem,
24917 };
24918 {
24919 let __s_b = self.gpu.stream();
24920 let mut b = __s_b.launch_builder(&f);
24921 b.arg(q_rows)
24922 .arg(tab)
24923 .arg(&mut *part_o)
24924 .arg(&mut *part_m)
24925 .arg(&mut *part_l)
24926 .arg(&hd)
24927 .arg(&nh)
24928 .arg(&nhkv)
24929 .arg(&win)
24930 .arg(&scale)
24931 .arg(&mns)
24932 .arg(&ktb)
24933 .arg(&vtb);
24934 unsafe {
24935 b.launch(cfg)?;
24936 }
24937 }
24938 let fc = {
24942 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
24943 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
24944 self.func("fa_decode_combine_gate_f32_s")
24945 } else {
24946 self.func("fa_decode_combine_gate_f32")
24947 }
24948 };
24949 let combine_shared = std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1");
24950 let nht = (t * n_head) as i32;
24951 let cfg2 = LaunchConfig {
24952 grid_dim: ((t * n_head) as u32, 1, 1),
24953 block_dim: (head_dim as u32, 1, 1),
24954 shared_mem_bytes: if combine_shared {
24955 (2 * max_ns * 4) as u32
24956 } else {
24957 0
24958 },
24959 };
24960 let __s_b2 = self.gpu.stream();
24961 let mut b2 = __s_b2.launch_builder(&fc);
24962 b2.arg(&*part_o)
24963 .arg(&*part_m)
24964 .arg(&*part_l)
24965 .arg(gate_rows)
24966 .arg(o_rows)
24967 .arg(&hd)
24968 .arg(&nht)
24969 .arg(&mns);
24970 unsafe {
24971 b2.launch(cfg2)?;
24972 }
24973 Ok(())
24974 }
24975
24976 pub fn fa_decode_dcw(
24977 &self,
24978 q: &CudaSlice<f32>,
24979 k_ring: &cudarc::driver::CudaView<u8>,
24980 v_ring: &cudarc::driver::CudaView<u8>,
24981 o: &mut CudaSlice<f32>,
24982 head_dim: usize,
24983 n_head: usize,
24984 n_head_kv: usize,
24985 len_dev: &CudaSlice<i32>,
24986 base_dev: Option<&CudaSlice<i32>>,
24987 window: usize,
24988 bucket_max: usize,
24989 scale: f32,
24990 k_tok_bytes: usize,
24991 v_tok_bytes: usize,
24992 fused_gate: Option<&CudaSlice<f32>>,
24996 ) -> Result<(), Box<dyn std::error::Error>> {
24997 let fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
24998 if !fa_vec || head_dim > 256 || head_dim % 32 != 0 || !fa_v3_on() {
24999 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"
25000 .into());
25001 }
25002 let sp = fa_split_keys(bucket_max, n_head_kv);
25003 let n_splits = ((bucket_max + sp - 1) / sp).max(1);
25004 let o_len = n_head * n_splits * head_dim;
25005 let ml_len = n_head * n_splits;
25006 let mut part_guard = self.fa_part_pool.lock().unwrap();
25007 Self::fa_part_pool_grow(self, &mut part_guard, o_len, ml_len)?;
25008 let pg = part_guard.as_mut().unwrap();
25009 static MEMSET_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25014 let memset_on = *MEMSET_ON
25019 .get_or_init(|| std::env::var("MEMRA_FA_DCW_MEMSET").as_deref() != Ok("0"))
25020 || crate::tp::token_graph_building();
25021 if memset_on {
25022 self.gpu
25023 .stream()
25024 .memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
25025 self.gpu
25026 .stream()
25027 .memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
25028 self.gpu
25029 .stream()
25030 .memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
25031 }
25032 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
25033 let (hd, nh, nhkv, nsp) = (
25034 head_dim as i32,
25035 n_head as i32,
25036 n_head_kv as i32,
25037 n_splits as i32,
25038 );
25039 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
25040 let (ski, win) = (sp as i32, window as i32);
25041 let gqa = (n_head / n_head_kv).max(1) as u32;
25042 let smem = (32 * head_dim * 2) as u32; static U8: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25046 static HOIST: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
25047 let hoist = *HOIST.get_or_init(|| match std::env::var("MEMRA_FA_HOIST").as_deref() {
25048 Ok("2") => 2,
25049 Ok("1") => 1,
25050 _ => 0,
25051 });
25052 static FPROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25057 let fprof = *FPROF.get_or_init(|| std::env::var("MEMRA_FA_PROF").as_deref() == Ok("1"));
25058 static PROF_BUF: std::sync::Mutex<Option<(usize, CudaSlice<u64>)>> =
25059 std::sync::Mutex::new(None);
25060 static HS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25064 let hs2 = *HS.get_or_init(|| std::env::var("MEMRA_FA_HSPLIT").as_deref() == Ok("2"))
25065 && (n_head / n_head_kv) % 2 == 0
25066 && (n_head / n_head_kv) >= 2;
25067 let f = if fprof {
25068 self.func("fa_decode_vec_q_v3_dcw_prof")
25069 } else if hs2 {
25070 self.func("fa_decode_vec_q_v3_dcw_hs2")
25071 } else if hoist == 2 {
25072 self.func("fa_decode_vec_q_v3_dcw_hc")
25074 } else if hoist == 1 {
25075 self.func("fa_decode_vec_q_v3_dcw_h")
25077 } else if *U8.get_or_init(|| std::env::var("MEMRA_FA_UNROLL").as_deref() == Ok("8")) {
25078 self.func("fa_decode_vec_q_v3_dcw_u8")
25079 } else {
25080 self.func("fa_decode_vec_q_v3_dcw")
25081 };
25082 let cfg = LaunchConfig {
25083 grid_dim: if hs2 {
25084 ((2 * n_head_kv) as u32, n_splits as u32, 1)
25085 } else {
25086 (n_head_kv as u32, n_splits as u32, 1)
25087 },
25088 block_dim: if hs2 { (32, gqa / 2, 1) } else { (32, gqa, 1) },
25089 shared_mem_bytes: smem,
25090 };
25091 let null: u64 = 0;
25092 let __s_b = self.gpu.stream();
25093 let mut b = __s_b.launch_builder(&f);
25094 b.arg(q)
25095 .arg(k_ring)
25096 .arg(v_ring)
25097 .arg(&mut *part_o)
25098 .arg(&mut *part_m)
25099 .arg(&mut *part_l)
25100 .arg(&hd)
25101 .arg(&nh)
25102 .arg(&nhkv)
25103 .arg(len_dev);
25104 match base_dev {
25105 Some(base) => {
25106 b.arg(base);
25107 }
25108 None => {
25109 b.arg(&null);
25110 }
25111 }
25112 b.arg(&win)
25113 .arg(&scale)
25114 .arg(&nsp)
25115 .arg(&ski)
25116 .arg(&ktb)
25117 .arg(&vtb);
25118 if fprof {
25119 let mut guard = PROF_BUF.lock().map_err(|_| "fa prof buffer lock")?;
25120 if guard
25121 .as_ref()
25122 .is_none_or(|(d, _)| *d != self.ctx().ordinal())
25123 {
25124 *guard = Some((self.ctx().ordinal(), self.htod_u64(&vec![0u64; 8])?));
25125 }
25126 let (_, buf) = guard.as_mut().expect("armed above");
25127 b.arg(&*buf);
25128 unsafe {
25129 b.launch(cfg)?;
25130 }
25131 static CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
25132 let n = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
25133 if n % 430 == 0 {
25134 self.stream().synchronize()?;
25135 let h = self.dtoh_u64(buf)?;
25136 let phases = ["setup", "stageV", "b1_klo", "b2_soft", "sync", "b3_vacc"];
25137 let tot: u64 = h[..6].iter().sum();
25138 let mut line = format!("[fa-prof] calls={n} keys={} cycles={tot}", h[6]);
25139 for (i, name) in phases.iter().enumerate() {
25140 let pct = if tot > 0 {
25141 h[i] as f64 / tot as f64 * 100.0
25142 } else {
25143 0.0
25144 };
25145 line.push_str(&format!(" {name}={pct:.1}%"));
25146 }
25147 if h[6] > 0 {
25148 line.push_str(&format!(" cyc/key={:.0}", tot as f64 / h[6] as f64));
25149 }
25150 eprintln!("{line}");
25151 }
25152 } else {
25153 unsafe {
25154 b.launch(cfg)?;
25155 }
25156 }
25157 let mut combine_shared = false;
25158 let fc = if fused_gate.is_some() {
25159 static CS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25162 if *CS.get_or_init(|| std::env::var("MEMRA_FA_COMBINE_S").as_deref() == Ok("1")) {
25163 combine_shared = true;
25164 self.func("fa_decode_combine_gate_f32_s")
25165 } else {
25166 self.func("fa_decode_combine_gate_f32")
25167 }
25168 } else {
25169 self.fa_func("fa_decode_combine_f32", head_dim)
25170 };
25171 let cfg2 = LaunchConfig {
25172 grid_dim: (n_head as u32, 1, 1),
25173 block_dim: (head_dim as u32, 1, 1),
25174 shared_mem_bytes: if combine_shared {
25175 (2 * n_splits * 4) as u32
25176 } else {
25177 0
25178 },
25179 };
25180 let __s_b2 = self.gpu.stream();
25181 let mut b2 = __s_b2.launch_builder(&fc);
25182 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l);
25183 if let Some(gate_row) = fused_gate {
25184 b2.arg(gate_row);
25185 }
25186 b2.arg(o).arg(&hd).arg(&nh).arg(&nsp);
25187 unsafe {
25188 b2.launch(cfg2)?;
25189 }
25190 Ok(())
25191 }
25192
25193 pub fn fa_geom_eager(
25199 &self,
25200 t_kv: usize,
25201 head_dim: usize,
25202 n_head_kv: usize,
25203 g: bool,
25204 ) -> (bool, usize) {
25205 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
25209 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
25215 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
25216 if g && head_dim == 256 && !fa_v4_at(t_kv) {
25222 fa_vec = false;
25223 }
25224 let sp = fa_split_keys(t_kv, n_head_kv);
25225 let n_splits = if fa_vec {
25226 ((t_kv + sp - 1) / sp).max(1)
25227 } else {
25228 ((t_kv + 255) / 256).max(1)
25229 };
25230 (fa_vec, n_splits)
25231 }
25232
25233 pub fn fa_bucket_key(
25239 &self,
25240 t_kv: usize,
25241 head_dim: usize,
25242 n_head_kv: usize,
25243 g: bool,
25244 ) -> (bool, usize) {
25245 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
25246 }
25247
25248 pub fn capture_graph_retained<F>(
25260 &self,
25261 step: F,
25262 ) -> Result<
25263 (
25264 cudarc::driver::CudaGraph,
25265 Vec<Box<dyn std::any::Any + Send>>,
25266 ),
25267 Box<dyn std::error::Error>,
25268 >
25269 where
25270 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25271 {
25272 use cudarc::driver::sys::CUgraphInstantiate_flags;
25273 self.capture_graph_retained_flags(
25274 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25275 step,
25276 )
25277 }
25278
25279 pub fn capture_graph_retained_flags<F>(
25284 &self,
25285 flags: cudarc::driver::sys::CUgraphInstantiate_flags,
25286 mut step: F,
25287 ) -> Result<
25288 (
25289 cudarc::driver::CudaGraph,
25290 Vec<Box<dyn std::any::Any + Send>>,
25291 ),
25292 Box<dyn std::error::Error>,
25293 >
25294 where
25295 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25296 {
25297 use cudarc::driver::sys::CUstreamCaptureMode;
25298 self.capture_keep.lock().unwrap().clear();
25306 let was_tracking = self.gpu.ctx.is_event_tracking();
25307 if was_tracking {
25308 unsafe {
25309 self.gpu.ctx.disable_event_tracking();
25310 }
25311 }
25312 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25313 self.capture_keep_on
25314 .store(true, std::sync::atomic::Ordering::Relaxed);
25315 let w = (|| {
25316 step(self)?;
25317 step(self)
25318 })();
25319 self.capture_keep_on
25320 .store(false, std::sync::atomic::Ordering::Relaxed);
25321 w?;
25322 self.gpu.stream().synchronize()?;
25323 self.gpu
25324 .stream()
25325 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25326 let r = step(self);
25327 let g = self.gpu.stream().end_capture(flags);
25328 r?;
25329 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25330 graph.upload()?;
25331 Ok(graph)
25332 };
25333 let result = run();
25334 self.capture_keep_on
25335 .store(false, std::sync::atomic::Ordering::Relaxed);
25336 if was_tracking {
25337 unsafe {
25338 self.gpu.ctx.enable_event_tracking();
25339 }
25340 }
25341 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
25342 Ok((result?, keeper))
25343 }
25344
25345 pub fn capture_graph_retained_nowarm<F>(
25351 &self,
25352 mut step: F,
25353 ) -> Result<
25354 (
25355 cudarc::driver::CudaGraph,
25356 Vec<Box<dyn std::any::Any + Send>>,
25357 ),
25358 Box<dyn std::error::Error>,
25359 >
25360 where
25361 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25362 {
25363 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
25364 let was_tracking = self.gpu.ctx.is_event_tracking();
25365 if was_tracking {
25366 unsafe {
25367 self.gpu.ctx.disable_event_tracking();
25368 }
25369 }
25370 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25371 self.gpu.stream().synchronize()?;
25372 self.gpu
25373 .stream()
25374 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25375 let r = step(self);
25376 let g = self.gpu.stream().end_capture(
25377 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25378 );
25379 r?;
25380 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25381 graph.upload()?;
25382 Ok(graph)
25383 };
25384 let result = run();
25385 if was_tracking {
25386 unsafe {
25387 self.gpu.ctx.enable_event_tracking();
25388 }
25389 }
25390 Ok((result?, Vec::new()))
25391 }
25392
25393 pub fn capture_graph<F>(
25394 &self,
25395 mut step: F,
25396 ) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
25397 where
25398 F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>,
25399 {
25400 use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
25401 let was_tracking = self.gpu.ctx.is_event_tracking();
25409 if was_tracking {
25410 unsafe {
25411 self.gpu.ctx.disable_event_tracking();
25412 }
25413 }
25414 let iflag = {
25421 static F: std::sync::OnceLock<CUgraphInstantiate_flags> = std::sync::OnceLock::new();
25422 *F.get_or_init(|| match std::env::var("MEMRA_GRAPH_IFLAG").as_deref() {
25423 Ok("upload") => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD,
25426 Ok("priority") => {
25427 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
25428 }
25429 _ => CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
25430 })
25431 };
25432 let ct = {
25439 static T: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
25440 *T.get_or_init(|| std::env::var("MEMRA_GRAPH_CAPTIME").as_deref() == Ok("1"))
25441 };
25442 let warmups = {
25465 static W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
25466 *W.get_or_init(|| {
25467 std::env::var("MEMRA_GRAPH_WARMUPS")
25468 .ok()
25469 .and_then(|v| v.parse().ok())
25470 .filter(|n| *n >= 1)
25471 .unwrap_or(1)
25472 })
25473 };
25474 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
25475 let t_w = std::time::Instant::now();
25476 for _ in 0..warmups {
25478 step(self)?;
25479 }
25480 self.gpu.stream().synchronize()?;
25481 let ms_warm = t_w.elapsed().as_secs_f64() * 1e3;
25482 let t_c = std::time::Instant::now();
25484 self.gpu
25485 .stream()
25486 .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
25487 let r = step(self);
25490 let ms_body = t_c.elapsed().as_secs_f64() * 1e3;
25491 let t_i = std::time::Instant::now();
25492 let g = self.gpu.stream().end_capture(iflag);
25493 let ms_inst = t_i.elapsed().as_secs_f64() * 1e3;
25494 r?;
25495 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
25496 let t_u = std::time::Instant::now();
25497 graph.upload()?;
25498 if ct {
25499 println!(
25500 "[graph-captime] warmup2x {ms_warm:.2} ms capture-body {ms_body:.2} ms \
25501 instantiate {ms_inst:.2} ms upload {:.2} ms",
25502 t_u.elapsed().as_secs_f64() * 1e3
25503 );
25504 }
25505 Ok(graph)
25506 };
25507 let result = run();
25508 if was_tracking {
25509 unsafe {
25510 self.gpu.ctx.enable_event_tracking();
25511 }
25512 }
25513 result
25514 }
25515
25516 pub fn gdn_scan_s128_view(
25518 &self,
25519 q: &CudaSlice<f32>,
25520 k: &CudaSlice<f32>,
25521 v: &CudaSlice<f32>,
25522 g: &CudaSlice<f32>,
25523 beta: &CudaSlice<f32>,
25524 state_in: &cudarc::driver::CudaView<f32>,
25525 state_out: &mut cudarc::driver::CudaViewMut<f32>,
25526 o: &mut CudaSlice<f32>,
25527 n_head: usize,
25528 t: usize,
25529 scale: f32,
25530 ) -> Result<(), Box<dyn std::error::Error>> {
25531 let f = self.func("gdn_scan_s128");
25532 const S_V: u32 = 128;
25533 const WARP: u32 = 32;
25534 const COLS: u32 = 4;
25535 let cfg = LaunchConfig {
25536 grid_dim: (n_head as u32, 1, S_V / COLS),
25537 block_dim: (WARP, COLS, 1),
25538 shared_mem_bytes: 0,
25539 };
25540 let (h, ti) = (n_head as i32, t as i32);
25541 let __s_b = self.gpu.stream();
25542 let mut b = __s_b.launch_builder(&f);
25543 b.arg(q)
25544 .arg(k)
25545 .arg(v)
25546 .arg(g)
25547 .arg(beta)
25548 .arg(state_in)
25549 .arg(state_out)
25550 .arg(o)
25551 .arg(&h)
25552 .arg(&ti)
25553 .arg(&scale);
25554 unsafe {
25555 b.launch(cfg)?;
25556 }
25557 Ok(())
25558 }
25559
25560 pub fn ssm_conv1d_view(
25562 &self,
25563 x: &cudarc::driver::CudaView<f32>,
25564 w: &CudaSlice<f32>,
25565 y: &mut CudaSlice<f32>,
25566 conv_dim: usize,
25567 t: usize,
25568 d_conv: usize,
25569 silu: bool,
25570 ) -> Result<(), Box<dyn std::error::Error>> {
25571 let f = self.func("ssm_conv1d_silu_f32");
25572 let cfg = LaunchConfig {
25574 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25575 block_dim: (256, 1, 1),
25576 shared_mem_bytes: 0,
25577 };
25578 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25579 let __s_b = self.gpu.stream();
25580 let mut b = __s_b.launch_builder(&f);
25581 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25582 unsafe {
25583 b.launch(cfg)?;
25584 }
25585 Ok(())
25586 }
25587
25588 pub fn ssm_conv1d_tm(
25595 &self,
25596 qkv_tm: &CudaSlice<f32>,
25597 w: &CudaSlice<f32>,
25598 y: &mut CudaSlice<f32>,
25599 conv_dim: usize,
25600 t: usize,
25601 d_conv: usize,
25602 ) -> Result<(), Box<dyn std::error::Error>> {
25603 let f = self.func("ssm_conv1d_tm_f32");
25604 let cfg = LaunchConfig {
25605 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25606 block_dim: (256, 1, 1),
25607 shared_mem_bytes: 0,
25608 };
25609 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25610 let __s_b = self.gpu.stream();
25611 let mut b = __s_b.launch_builder(&f);
25612 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
25613 unsafe {
25614 b.launch(cfg)?;
25615 }
25616 Ok(())
25617 }
25618
25619 pub fn ssm_conv1d_tm_state(
25627 &self,
25628 qkv_tm: &CudaSlice<f32>,
25629 conv_state: &mut CudaSlice<f32>,
25630 w: &CudaSlice<f32>,
25631 y: &mut CudaSlice<f32>,
25632 conv_dim: usize,
25633 t: usize,
25634 d_conv: usize,
25635 ) -> Result<(), Box<dyn std::error::Error>> {
25636 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
25637 }
25638
25639 #[allow(clippy::too_many_arguments)]
25642 pub fn ssm_conv1d_tm_state_pad(
25643 &self,
25644 qkv_tm: &CudaSlice<f32>,
25645 conv_state: &mut CudaSlice<f32>,
25646 w: &CudaSlice<f32>,
25647 y: &mut CudaSlice<f32>,
25648 conv_dim: usize,
25649 t: usize,
25650 d_conv: usize,
25651 pad_len: Option<&CudaSlice<i32>>,
25652 ) -> Result<(), Box<dyn std::error::Error>> {
25653 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25654 let ring_old = if t < d_conv - 1 {
25658 Some(self.clone_dtod(conv_state)?)
25659 } else {
25660 None
25661 };
25662 {
25663 let f = self.func("ssm_conv1d_tm_state_f32");
25664 let cfg = LaunchConfig {
25665 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25666 block_dim: (256, 1, 1),
25667 shared_mem_bytes: 0,
25668 };
25669 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25670 let __s_b = self.gpu.stream();
25671 let mut b = __s_b.launch_builder(&f);
25672 b.arg(qkv_tm)
25673 .arg(&*conv_state)
25674 .arg(w)
25675 .arg(y)
25676 .arg(&cd)
25677 .arg(&ti)
25678 .arg(&dc);
25679 unsafe {
25680 b.launch(cfg)?;
25681 }
25682 }
25683 match (ring_old, pad_len) {
25684 (None, Some(len_d)) => {
25685 let f = self.func("ssm_conv_ring_update_dev_f32");
25686 let n = conv_dim * (d_conv - 1);
25687 let cfg = LaunchConfig::for_num_elems(n as u32);
25688 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25689 let __s_b = self.gpu.stream();
25690 let mut b = __s_b.launch_builder(&f);
25691 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25692 unsafe {
25693 b.launch(cfg)?;
25694 }
25695 }
25696 (None, None) => {
25697 let f = self.func("ssm_conv_ring_update_f32");
25698 let n = conv_dim * (d_conv - 1);
25699 let cfg = LaunchConfig::for_num_elems(n as u32);
25700 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25701 let __s_b = self.gpu.stream();
25702 let mut b = __s_b.launch_builder(&f);
25703 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25704 unsafe {
25705 b.launch(cfg)?;
25706 }
25707 }
25708 (Some(old), _) => {
25709 self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?
25710 }
25711 }
25712 Ok(())
25713 }
25714
25715 pub fn ssm_conv1d_tm_state_pad_v(
25717 &self,
25718 qkv_tm: &cudarc::driver::CudaView<f32>,
25719 conv_state: &mut CudaSlice<f32>,
25720 w: &CudaSlice<f32>,
25721 y: &mut CudaSlice<f32>,
25722 conv_dim: usize,
25723 t: usize,
25724 d_conv: usize,
25725 pad_len: Option<&CudaSlice<i32>>,
25726 ) -> Result<(), Box<dyn std::error::Error>> {
25727 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
25728 let ring_old = if t < d_conv - 1 {
25732 Some(self.clone_dtod(conv_state)?)
25733 } else {
25734 None
25735 };
25736 {
25737 let f = self.func("ssm_conv1d_tm_state_f32");
25738 let cfg = LaunchConfig {
25739 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25740 block_dim: (256, 1, 1),
25741 shared_mem_bytes: 0,
25742 };
25743 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25744 let __s_b = self.gpu.stream();
25745 let mut b = __s_b.launch_builder(&f);
25746 b.arg(qkv_tm)
25747 .arg(&*conv_state)
25748 .arg(w)
25749 .arg(y)
25750 .arg(&cd)
25751 .arg(&ti)
25752 .arg(&dc);
25753 unsafe {
25754 b.launch(cfg)?;
25755 }
25756 }
25757 match (ring_old, pad_len) {
25758 (None, Some(len_d)) => {
25759 let f = self.func("ssm_conv_ring_update_dev_f32");
25760 let n = conv_dim * (d_conv - 1);
25761 let cfg = LaunchConfig::for_num_elems(n as u32);
25762 let (cd, dc) = (conv_dim as i32, d_conv as i32);
25763 let __s_b = self.gpu.stream();
25764 let mut b = __s_b.launch_builder(&f);
25765 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
25766 unsafe {
25767 b.launch(cfg)?;
25768 }
25769 }
25770 (None, None) => {
25771 let f = self.func("ssm_conv_ring_update_f32");
25772 let n = conv_dim * (d_conv - 1);
25773 let cfg = LaunchConfig::for_num_elems(n as u32);
25774 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25775 let __s_b = self.gpu.stream();
25776 let mut b = __s_b.launch_builder(&f);
25777 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
25778 unsafe {
25779 b.launch(cfg)?;
25780 }
25781 }
25782 (Some(_), _) => unreachable!(
25783 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"
25784 ),
25785 }
25786 Ok(())
25787 }
25788
25789 pub fn ssm_conv_ring_rebuild(
25794 &self,
25795 qkv_tm: &CudaSlice<f32>,
25796 ring_old: &CudaSlice<f32>,
25797 conv_state: &mut CudaSlice<f32>,
25798 conv_dim: usize,
25799 tc: usize,
25800 d_conv: usize,
25801 ) -> Result<(), Box<dyn std::error::Error>> {
25802 let f = self.func("ssm_conv_ring_rebuild_f32");
25803 let n = conv_dim * (d_conv - 1);
25804 let cfg = LaunchConfig::for_num_elems(n as u32);
25805 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
25806 let __s_b = self.gpu.stream();
25807 let mut b = __s_b.launch_builder(&f);
25808 b.arg(qkv_tm)
25809 .arg(ring_old)
25810 .arg(conv_state)
25811 .arg(&cd)
25812 .arg(&ti)
25813 .arg(&dc);
25814 unsafe {
25815 b.launch(cfg)?;
25816 }
25817 Ok(())
25818 }
25819
25820 #[allow(clippy::too_many_arguments)]
25825 pub fn gdn_prep_decode(
25826 &self,
25827 conv_out: &CudaSlice<f32>,
25828 beta_raw: &CudaSlice<f32>,
25829 alpha: &CudaSlice<f32>,
25830 dt_bias: &CudaSlice<f32>,
25831 a: &CudaSlice<f32>,
25832 q_l2: &mut CudaSlice<f32>,
25833 k_l2: &mut CudaSlice<f32>,
25834 v_g: &mut CudaSlice<f32>,
25835 beta: &mut CudaSlice<f32>,
25836 g_log: &mut CudaSlice<f32>,
25837 d_state: usize,
25838 num_v: usize,
25839 num_k: usize,
25840 key_dim: usize,
25841 eps: f32,
25842 ) -> Result<(), Box<dyn std::error::Error>> {
25843 let f = self.func("gdn_prep_decode_f32");
25844 let cfg = LaunchConfig {
25845 grid_dim: (num_v as u32, 1, 1),
25846 block_dim: (32, 4, 1),
25847 shared_mem_bytes: 0,
25848 };
25849 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25850 let __s_b = self.gpu.stream();
25851 let mut b = __s_b.launch_builder(&f);
25852 b.arg(conv_out)
25853 .arg(beta_raw)
25854 .arg(alpha)
25855 .arg(dt_bias)
25856 .arg(a)
25857 .arg(q_l2)
25858 .arg(k_l2)
25859 .arg(v_g)
25860 .arg(beta)
25861 .arg(g_log)
25862 .arg(&ds)
25863 .arg(&nv)
25864 .arg(&nk)
25865 .arg(&kd)
25866 .arg(&eps);
25867 unsafe {
25868 b.launch(cfg)?;
25869 }
25870 Ok(())
25871 }
25872
25873 #[allow(clippy::too_many_arguments)]
25877 pub fn ssm_conv1d_gdn(
25878 &self,
25879 qkv_tm: &CudaSlice<f32>,
25880 w: &CudaSlice<f32>,
25881 q_g: &mut CudaSlice<f32>,
25882 k_g: &mut CudaSlice<f32>,
25883 v_g: &mut CudaSlice<f32>,
25884 conv_dim: usize,
25885 t: usize,
25886 d_conv: usize,
25887 d_state: usize,
25888 num_v: usize,
25889 num_k: usize,
25890 key_dim: usize,
25891 ) -> Result<(), Box<dyn std::error::Error>> {
25892 let f = self.func("ssm_conv1d_gdn_f32");
25893 let cfg = LaunchConfig {
25894 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
25895 block_dim: (256, 1, 1),
25896 shared_mem_bytes: 0,
25897 };
25898 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
25899 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
25900 let __s_b = self.gpu.stream();
25901 let mut b = __s_b.launch_builder(&f);
25902 b.arg(qkv_tm)
25903 .arg(w)
25904 .arg(q_g)
25905 .arg(k_g)
25906 .arg(v_g)
25907 .arg(&cd)
25908 .arg(&ti)
25909 .arg(&dc)
25910 .arg(&ds)
25911 .arg(&nv)
25912 .arg(&nk)
25913 .arg(&kd);
25914 unsafe {
25915 b.launch(cfg)?;
25916 }
25917 Ok(())
25918 }
25919
25920 pub fn ssm_conv1d(
25921 &self,
25922 x: &CudaSlice<f32>,
25923 w: &CudaSlice<f32>,
25924 y: &mut CudaSlice<f32>,
25925 conv_dim: usize,
25926 t: usize,
25927 d_conv: usize,
25928 silu: bool,
25929 ) -> Result<(), Box<dyn std::error::Error>> {
25930 let f = self.func("ssm_conv1d_silu_f32");
25931 let cfg = LaunchConfig {
25932 grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
25933 block_dim: (256, 1, 1),
25934 shared_mem_bytes: 0,
25935 };
25936 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
25937 let __s_b = self.gpu.stream();
25938 let mut b = __s_b.launch_builder(&f);
25939 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
25940 unsafe {
25941 b.launch(cfg)?;
25942 }
25943 Ok(())
25944 }
25945
25946 pub fn gdn_scan_s128(
25949 &self,
25950 q: &CudaSlice<f32>,
25951 k: &CudaSlice<f32>,
25952 v: &CudaSlice<f32>,
25953 g: &CudaSlice<f32>,
25954 beta: &CudaSlice<f32>,
25955 state_in: &CudaSlice<f32>,
25956 state_out: &mut CudaSlice<f32>,
25957 o: &mut CudaSlice<f32>,
25958 n_head: usize,
25959 t: usize,
25960 scale: f32,
25961 ) -> Result<(), Box<dyn std::error::Error>> {
25962 let f = self.func("gdn_scan_s128");
25963 const S_V: u32 = 128;
25964 const WARP: u32 = 32;
25965 const COLS_PER_BLOCK: u32 = 4;
25966 let cfg = LaunchConfig {
25967 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
25968 block_dim: (WARP, COLS_PER_BLOCK, 1),
25969 shared_mem_bytes: 0,
25970 };
25971 let (h, ti) = (n_head as i32, t as i32);
25972 let __s_b = self.gpu.stream();
25973 let mut b = __s_b.launch_builder(&f);
25974 b.arg(q)
25975 .arg(k)
25976 .arg(v)
25977 .arg(g)
25978 .arg(beta)
25979 .arg(state_in)
25980 .arg(state_out)
25981 .arg(o)
25982 .arg(&h)
25983 .arg(&ti)
25984 .arg(&scale);
25985 unsafe {
25986 b.launch(cfg)?;
25987 }
25988 Ok(())
25989 }
25990
25991 #[allow(clippy::too_many_arguments)]
25996 pub fn ssm_conv1d_fused_decode_b(
25997 &self,
25998 qkv_cols: &CudaSlice<f32>,
25999 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
26000 w: &CudaSlice<f32>,
26001 conv_outs: &mut CudaSlice<f32>,
26002 conv_dim: usize,
26003 d_conv: usize,
26004 b_n: usize,
26005 ) -> Result<(), Box<dyn std::error::Error>> {
26006 let f = self.func("ssm_conv1d_fused_decode_b_f32");
26007 let cfg = LaunchConfig {
26008 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
26009 block_dim: (256, 1, 1),
26010 shared_mem_bytes: 0,
26011 };
26012 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26013 let __s_b = self.gpu.stream();
26014 let mut b = __s_b.launch_builder(&f);
26015 b.arg(qkv_cols)
26016 .arg(conv_state_ptrs)
26017 .arg(w)
26018 .arg(conv_outs)
26019 .arg(&cd)
26020 .arg(&dc);
26021 unsafe {
26022 b.launch(cfg)?;
26023 }
26024 Ok(())
26025 }
26026
26027 #[allow(clippy::too_many_arguments)]
26028 pub fn gdn_prep_decode_b(
26029 &self,
26030 conv_outs: &CudaSlice<f32>,
26031 beta_raws: &CudaSlice<f32>,
26032 alphas: &CudaSlice<f32>,
26033 dt_bias: &CudaSlice<f32>,
26034 a: &CudaSlice<f32>,
26035 q_l2: &mut CudaSlice<f32>,
26036 k_l2: &mut CudaSlice<f32>,
26037 v_g: &mut CudaSlice<f32>,
26038 beta: &mut CudaSlice<f32>,
26039 g_log: &mut CudaSlice<f32>,
26040 d_state: usize,
26041 num_v: usize,
26042 num_k: usize,
26043 key_dim: usize,
26044 eps: f32,
26045 conv_dim: usize,
26046 b_n: usize,
26047 ) -> Result<(), Box<dyn std::error::Error>> {
26048 let f = self.func("gdn_prep_decode_b_f32");
26049 let cfg = LaunchConfig {
26050 grid_dim: (num_v as u32, 1, b_n as u32),
26051 block_dim: (32, 4, 1),
26052 shared_mem_bytes: 0,
26053 };
26054 let (ds, nv, nk, kd, cd) = (
26055 d_state as i32,
26056 num_v as i32,
26057 num_k as i32,
26058 key_dim as i32,
26059 conv_dim as i32,
26060 );
26061 let __s_b = self.gpu.stream();
26062 let mut b = __s_b.launch_builder(&f);
26063 b.arg(conv_outs)
26064 .arg(beta_raws)
26065 .arg(alphas)
26066 .arg(dt_bias)
26067 .arg(a)
26068 .arg(q_l2)
26069 .arg(k_l2)
26070 .arg(v_g)
26071 .arg(beta)
26072 .arg(g_log)
26073 .arg(&ds)
26074 .arg(&nv)
26075 .arg(&nk)
26076 .arg(&kd)
26077 .arg(&eps)
26078 .arg(&cd);
26079 unsafe {
26080 b.launch(cfg)?;
26081 }
26082 Ok(())
26083 }
26084
26085 #[allow(clippy::too_many_arguments)]
26086 pub fn gdn_scan_s128_batched(
26087 &self,
26088 q: &CudaSlice<f32>,
26089 k: &CudaSlice<f32>,
26090 v: &CudaSlice<f32>,
26091 g: &CudaSlice<f32>,
26092 beta: &CudaSlice<f32>,
26093 state_in_ptrs: &cudarc::driver::CudaView<u64>,
26094 state_out_ptrs: &cudarc::driver::CudaView<u64>,
26095 o: &mut CudaSlice<f32>,
26096 n_head: usize,
26097 b_n: usize,
26098 scale: f32,
26099 ) -> Result<(), Box<dyn std::error::Error>> {
26100 let f = self.func("gdn_scan_s128_b");
26101 const S_V: u32 = 128;
26102 const WARP: u32 = 32;
26103 const COLS_PER_BLOCK: u32 = 4;
26104 let cfg = LaunchConfig {
26105 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
26106 block_dim: (WARP, COLS_PER_BLOCK, 1),
26107 shared_mem_bytes: 0,
26108 };
26109 let h = n_head as i32;
26110 let __s_b = self.gpu.stream();
26111 let mut b = __s_b.launch_builder(&f);
26112 b.arg(q)
26113 .arg(k)
26114 .arg(v)
26115 .arg(g)
26116 .arg(beta)
26117 .arg(state_in_ptrs)
26118 .arg(state_out_ptrs)
26119 .arg(o)
26120 .arg(&h)
26121 .arg(&scale);
26122 unsafe {
26123 b.launch(cfg)?;
26124 }
26125 Ok(())
26126 }
26127
26128 #[allow(clippy::too_many_arguments)]
26134 pub fn ssm_conv1d_fused_decode_b_view(
26135 &self,
26136 qkv_cols: &cudarc::driver::CudaView<f32>,
26137 conv_state_ptrs: &cudarc::driver::CudaView<u64>,
26138 w: &CudaSlice<f32>,
26139 conv_outs: &mut CudaSlice<f32>,
26140 conv_dim: usize,
26141 d_conv: usize,
26142 b_n: usize,
26143 ) -> Result<(), Box<dyn std::error::Error>> {
26144 let f = self.func("ssm_conv1d_fused_decode_b_f32");
26145 let cfg = LaunchConfig {
26146 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
26147 block_dim: (256, 1, 1),
26148 shared_mem_bytes: 0,
26149 };
26150 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26151 let __s_b = self.gpu.stream();
26152 let mut b = __s_b.launch_builder(&f);
26153 b.arg(qkv_cols)
26154 .arg(conv_state_ptrs)
26155 .arg(w)
26156 .arg(conv_outs)
26157 .arg(&cd)
26158 .arg(&dc);
26159 unsafe {
26160 b.launch(cfg)?;
26161 }
26162 Ok(())
26163 }
26164
26165 #[allow(clippy::too_many_arguments)]
26166 pub fn gdn_prep_decode_b_view(
26167 &self,
26168 conv_outs: &CudaSlice<f32>,
26169 beta_raws: &cudarc::driver::CudaView<f32>,
26170 alphas: &cudarc::driver::CudaView<f32>,
26171 dt_bias: &CudaSlice<f32>,
26172 a: &CudaSlice<f32>,
26173 q_l2: &mut CudaSlice<f32>,
26174 k_l2: &mut CudaSlice<f32>,
26175 v_g: &mut CudaSlice<f32>,
26176 beta: &mut CudaSlice<f32>,
26177 g_log: &mut CudaSlice<f32>,
26178 d_state: usize,
26179 num_v: usize,
26180 num_k: usize,
26181 key_dim: usize,
26182 eps: f32,
26183 conv_dim: usize,
26184 b_n: usize,
26185 ) -> Result<(), Box<dyn std::error::Error>> {
26186 let f = self.func("gdn_prep_decode_b_f32");
26187 let cfg = LaunchConfig {
26188 grid_dim: (num_v as u32, 1, b_n as u32),
26189 block_dim: (32, 4, 1),
26190 shared_mem_bytes: 0,
26191 };
26192 let (ds, nv, nk, kd, cd) = (
26193 d_state as i32,
26194 num_v as i32,
26195 num_k as i32,
26196 key_dim as i32,
26197 conv_dim as i32,
26198 );
26199 let __s_b = self.gpu.stream();
26200 let mut b = __s_b.launch_builder(&f);
26201 b.arg(conv_outs)
26202 .arg(beta_raws)
26203 .arg(alphas)
26204 .arg(dt_bias)
26205 .arg(a)
26206 .arg(q_l2)
26207 .arg(k_l2)
26208 .arg(v_g)
26209 .arg(beta)
26210 .arg(g_log)
26211 .arg(&ds)
26212 .arg(&nv)
26213 .arg(&nk)
26214 .arg(&kd)
26215 .arg(&eps)
26216 .arg(&cd);
26217 unsafe {
26218 b.launch(cfg)?;
26219 }
26220 Ok(())
26221 }
26222
26223 #[allow(clippy::too_many_arguments)]
26224 pub fn gdn_scan_s128_batched_view(
26225 &self,
26226 q: &CudaSlice<f32>,
26227 k: &CudaSlice<f32>,
26228 v: &CudaSlice<f32>,
26229 g: &CudaSlice<f32>,
26230 beta: &CudaSlice<f32>,
26231 state_in_ptrs: &cudarc::driver::CudaView<u64>,
26232 state_out_ptrs: &cudarc::driver::CudaView<u64>,
26233 o: &mut cudarc::driver::CudaViewMut<f32>,
26234 n_head: usize,
26235 b_n: usize,
26236 scale: f32,
26237 ) -> Result<(), Box<dyn std::error::Error>> {
26238 let f = self.func("gdn_scan_s128_b");
26239 const S_V: u32 = 128;
26240 const WARP: u32 = 32;
26241 const COLS_PER_BLOCK: u32 = 4;
26242 let cfg = LaunchConfig {
26243 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
26244 block_dim: (WARP, COLS_PER_BLOCK, 1),
26245 shared_mem_bytes: 0,
26246 };
26247 let h = n_head as i32;
26248 let __s_b = self.gpu.stream();
26249 let mut b = __s_b.launch_builder(&f);
26250 b.arg(q)
26251 .arg(k)
26252 .arg(v)
26253 .arg(g)
26254 .arg(beta)
26255 .arg(state_in_ptrs)
26256 .arg(state_out_ptrs)
26257 .arg(o)
26258 .arg(&h)
26259 .arg(&scale);
26260 unsafe {
26261 b.launch(cfg)?;
26262 }
26263 Ok(())
26264 }
26265
26266 pub fn gdn_chunked_enabled() -> bool {
26275 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
26276 *E.get_or_init(|| {
26277 std::env::var("MEMRA_GDN_CHUNKED")
26278 .map(|v| v != "0")
26279 .unwrap_or(true)
26280 })
26281 }
26282
26283 pub fn gdn_chunk_size() -> usize {
26288 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
26289 *C.get_or_init(|| {
26290 let c: usize = std::env::var("MEMRA_GDN_CHUNK")
26291 .ok()
26292 .and_then(|v| v.parse().ok())
26293 .unwrap_or(32);
26294 c.clamp(32, 128) / 32 * 32
26295 })
26296 }
26297
26298 #[allow(clippy::too_many_arguments)]
26303 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
26306 #[allow(clippy::too_many_arguments)]
26307 pub fn gdn_chunk_k123(
26308 &self,
26309 q: &CudaSlice<f32>,
26310 k: &CudaSlice<f32>,
26311 v: &CudaSlice<f32>,
26312 g: &CudaSlice<f32>,
26313 beta: &CudaSlice<f32>,
26314 wb16: Option<&mut CudaSlice<u8>>,
26315 n_head: usize,
26316 t: usize,
26317 c: usize,
26318 hk: usize,
26319 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>,
26320 ) -> Result<
26321 (
26322 CudaSlice<f32>,
26323 CudaSlice<f32>,
26324 CudaSlice<f32>,
26325 CudaSlice<f32>,
26326 ),
26327 Box<dyn std::error::Error>,
26328 > {
26329 const D: usize = 128;
26330 let h = n_head;
26331 let nc = (t + c - 1) / c;
26332 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
26333 let mut gcum = self.uninit(t * h)?;
26334 let mut a = self.uninit(nc * h * c * c)?;
26335 let mut p = self.uninit(nc * h * c * c)?;
26336 let mut u = self.uninit(nc * h * c * D)?;
26337 let mut w = self.uninit(nc * h * c * D)?;
26338 {
26339 let f = self.func("gdn_chunk_cumgate_f32");
26341 let cfg = LaunchConfig {
26342 grid_dim: (nc as u32, h as u32, 1),
26343 block_dim: (32, 1, 1),
26344 shared_mem_bytes: 0,
26345 };
26346 let __s_b = self.gpu.stream();
26347 let mut b = __s_b.launch_builder(&f);
26348 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
26349 unsafe {
26350 b.launch(cfg)?;
26351 }
26352 }
26353 if let Some((qb, kb, pb)) = k2w {
26354 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
26357 let f = self.func("gdn_k2_wgmma");
26358 let cfg = LaunchConfig {
26359 grid_dim: (nc as u32, h as u32, 1),
26360 block_dim: (128, 1, 1),
26361 shared_mem_bytes: 0,
26362 };
26363 let hki = hk as i32;
26364 let __s_b = self.gpu.stream();
26365 let mut b = __s_b.launch_builder(&f);
26366 b.arg(qb)
26367 .arg(kb)
26368 .arg(&gcum)
26369 .arg(beta)
26370 .arg(&mut a)
26371 .arg(&mut *pb)
26372 .arg(&hi)
26373 .arg(&ti)
26374 .arg(&ci)
26375 .arg(&hki);
26376 unsafe {
26377 b.launch(cfg)?;
26378 }
26379 } else if c <= 64 && !portable_mma_gated() {
26380 let f = self.func("gdn_chunk_attn_f32");
26382 f.set_attribute(
26383 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26384 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26385 )?;
26386 let jt = ((c + 31) / 32) as u32;
26387 let cfg = LaunchConfig {
26388 grid_dim: (nc as u32, h as u32, jt),
26389 block_dim: (256, 1, 1),
26390 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26391 };
26392 let hki = hk as i32;
26393 let __s_b = self.gpu.stream();
26394 let mut b = __s_b.launch_builder(&f);
26395 b.arg(q)
26396 .arg(k)
26397 .arg(&gcum)
26398 .arg(beta)
26399 .arg(&mut a)
26400 .arg(&mut p)
26401 .arg(&hi)
26402 .arg(&ti)
26403 .arg(&ci)
26404 .arg(&hki);
26405 unsafe {
26406 b.launch(cfg)?;
26407 }
26408 } else {
26409 assert!(
26411 hk == h,
26412 "generic K2 is broadcast-only (de-broadcast rides C==32)"
26413 );
26414 let f = self.func("gdn_chunk_attn_g_f32");
26415 let cfg = LaunchConfig {
26416 grid_dim: (nc as u32, h as u32, 1),
26417 block_dim: (32, 8, 1),
26418 shared_mem_bytes: 0,
26419 };
26420 let __s_b = self.gpu.stream();
26421 let mut b = __s_b.launch_builder(&f);
26422 b.arg(q)
26423 .arg(k)
26424 .arg(&gcum)
26425 .arg(beta)
26426 .arg(&mut a)
26427 .arg(&mut p)
26428 .arg(&hi)
26429 .arg(&ti)
26430 .arg(&ci);
26431 unsafe {
26432 b.launch(cfg)?;
26433 }
26434 }
26435 {
26436 let cfg = LaunchConfig {
26438 grid_dim: (nc as u32, h as u32, 1),
26439 block_dim: (256, 1, 1),
26440 shared_mem_bytes: 0,
26441 };
26442 match c {
26443 32 | 64 => {
26444 let f = self.func(if c == 32 {
26445 "gdn_chunk_solve32_f32"
26446 } else {
26447 "gdn_chunk_solve64_f32"
26448 });
26449 let wb: u64 = match wb16 {
26451 Some(d) => self.addr_u8(d),
26452 None => 0,
26453 };
26454 let hki = hk as i32;
26455 let __s_b = self.gpu.stream();
26456 let mut b = __s_b.launch_builder(&f);
26457 b.arg(v)
26458 .arg(k)
26459 .arg(&a)
26460 .arg(&gcum)
26461 .arg(&mut u)
26462 .arg(&mut w)
26463 .arg(&wb)
26464 .arg(&hi)
26465 .arg(&ti)
26466 .arg(&hki);
26467 unsafe {
26468 b.launch(cfg)?;
26469 }
26470 }
26471 _ => {
26472 assert!(hk == h, "generic K3 is broadcast-only");
26473 let f = self.func("gdn_chunk_solve_f32");
26474 let __s_b = self.gpu.stream();
26475 let mut b = __s_b.launch_builder(&f);
26476 b.arg(v)
26477 .arg(k)
26478 .arg(&a)
26479 .arg(&gcum)
26480 .arg(&mut u)
26481 .arg(&mut w)
26482 .arg(&hi)
26483 .arg(&ti)
26484 .arg(&ci);
26485 unsafe {
26486 b.launch(cfg)?;
26487 }
26488 }
26489 }
26490 }
26491 Ok((gcum, p, u, w))
26492 }
26493
26494 pub fn gdn_db_on() -> bool {
26498 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
26499 }
26500
26501 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
26510 !portable_mma_gated()
26511 && c == 32
26512 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
26513 Ok("1") => true,
26514 Ok("0") => false,
26515 _ => gdn_mma_default_on(),
26516 }
26517 }
26518
26519 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
26526 cfg!(memra_hopper_mma)
26527 && self.gdn_mma_enabled(c)
26528 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0")
26529 }
26530
26531 #[allow(clippy::too_many_arguments)]
26536 pub fn ssm_conv1d_gdn_state_pad(
26537 &self,
26538 qkv_tm: &cudarc::driver::CudaView<f32>,
26539 conv_state: &mut CudaSlice<f32>,
26540 w: &CudaSlice<f32>,
26541 q_g: &mut CudaSlice<f32>,
26542 k_g: &mut CudaSlice<f32>,
26543 v_g: &mut CudaSlice<f32>,
26544 conv_dim: usize,
26545 t: usize,
26546 d_conv: usize,
26547 d_state: usize,
26548 num_v: usize,
26549 num_k: usize,
26550 key_dim: usize,
26551 hk: usize,
26552 pad_len: Option<&CudaSlice<i32>>,
26553 ) -> Result<(), Box<dyn std::error::Error>> {
26554 assert!(
26555 t >= d_conv - 1,
26556 "fused state conv requires T >= pad (PRIME_MIN_T gates)"
26557 );
26558 {
26559 let f = self.func("ssm_conv1d_gdn_state_f32");
26560 let cfg = LaunchConfig {
26561 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
26562 block_dim: (256, 1, 1),
26563 shared_mem_bytes: 0,
26564 };
26565 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
26566 let (ds, nv, nk, kd, hki) = (
26567 d_state as i32,
26568 num_v as i32,
26569 num_k as i32,
26570 key_dim as i32,
26571 hk as i32,
26572 );
26573 let __s_b = self.gpu.stream();
26574 let mut b = __s_b.launch_builder(&f);
26575 b.arg(qkv_tm)
26576 .arg(&*conv_state)
26577 .arg(w)
26578 .arg(q_g)
26579 .arg(k_g)
26580 .arg(v_g)
26581 .arg(&cd)
26582 .arg(&ti)
26583 .arg(&dc)
26584 .arg(&ds)
26585 .arg(&nv)
26586 .arg(&nk)
26587 .arg(&kd)
26588 .arg(&hki);
26589 unsafe {
26590 b.launch(cfg)?;
26591 }
26592 }
26593 match pad_len {
26594 Some(len_d) => {
26595 let f = self.func("ssm_conv_ring_update_dev_f32");
26596 let n = conv_dim * (d_conv - 1);
26597 let cfg = LaunchConfig::for_num_elems(n as u32);
26598 let (cd, dc) = (conv_dim as i32, d_conv as i32);
26599 let __s_b = self.gpu.stream();
26600 let mut b = __s_b.launch_builder(&f);
26601 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
26602 unsafe {
26603 b.launch(cfg)?;
26604 }
26605 }
26606 None => {
26607 let f = self.func("ssm_conv_ring_update_f32");
26608 let n = conv_dim * (d_conv - 1);
26609 let cfg = LaunchConfig::for_num_elems(n as u32);
26610 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
26611 let __s_b = self.gpu.stream();
26612 let mut b = __s_b.launch_builder(&f);
26613 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
26614 unsafe {
26615 b.launch(cfg)?;
26616 }
26617 }
26618 }
26619 Ok(())
26620 }
26621
26622 pub fn gdn_chunk_alloc(
26626 &self,
26627 n_head: usize,
26628 t: usize,
26629 c: usize,
26630 hk: usize,
26631 ) -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
26632 const D: usize = 128;
26633 assert!(
26634 c == 32,
26635 "gdn_chunk_alloc: varlen chain is the C==32 mma pair"
26636 );
26637 let h = n_head;
26638 let nc = (t + c - 1) / c;
26639 Ok(GdnChunkBufs {
26640 gcum: self.uninit(t * h)?,
26641 a: self.uninit(nc * h * c * c)?,
26642 p: self.uninit(nc * h * c * c)?,
26643 u: self.uninit(nc * h * c * D)?,
26644 w: self.uninit(nc * h * c * D)?,
26645 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
26646 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
26647 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
26648 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
26649 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
26650 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
26651 o: self.uninit(D * h * t)?,
26652 t,
26653 nc,
26654 })
26655 }
26656
26657 pub fn f32_to_bf16_v(
26659 &self,
26660 x: &cudarc::driver::CudaView<f32>,
26661 dst: &mut CudaSlice<u8>,
26662 n: usize,
26663 ) -> Result<(), Box<dyn std::error::Error>> {
26664 let f = self.func("f32_to_bf16_bulk");
26665 let ni = n as i64;
26666 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26667 let __s_b = self.gpu.stream();
26668 let mut b = __s_b.launch_builder(&f);
26669 b.arg(x).arg(dst).arg(&ni);
26670 unsafe {
26671 b.launch(cfg)?;
26672 }
26673 Ok(())
26674 }
26675
26676 pub fn f32_to_bf16_into(
26678 &self,
26679 x: &CudaSlice<f32>,
26680 dst: &mut CudaSlice<u8>,
26681 n: usize,
26682 ) -> Result<(), Box<dyn std::error::Error>> {
26683 let f = self.func("f32_to_bf16_bulk");
26684 let ni = n as i64;
26685 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
26686 let __s_b = self.gpu.stream();
26687 let mut b = __s_b.launch_builder(&f);
26688 b.arg(x).arg(dst).arg(&ni);
26689 unsafe {
26690 b.launch(cfg)?;
26691 }
26692 Ok(())
26693 }
26694
26695 pub fn gdn_chunk_k123_vl8(
26698 &self,
26699 seqs: &[GdnSeqVl],
26700 n_head: usize,
26701 hk: usize,
26702 wq: Option<&GdnWVl8>,
26703 ) -> Result<(), Box<dyn std::error::Error>> {
26704 let b = seqs.len();
26705 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
26706 let mut packed = [GdnSeqVl::default(); 8];
26707 packed[..b].copy_from_slice(seqs);
26708 let v = GdnVl8(packed);
26709 let (hi, ci) = (n_head as i32, 32i32);
26710 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
26711 {
26712 let f = self.func("gdn_chunk_cumgate_vl");
26713 let cfg = LaunchConfig {
26714 grid_dim: (max_nc, n_head as u32, b as u32),
26715 block_dim: (32, 1, 1),
26716 shared_mem_bytes: 0,
26717 };
26718 let __s_lb = self.gpu.stream();
26719 let mut lb = __s_lb.launch_builder(&f);
26720 lb.arg(&v).arg(&hi).arg(&ci);
26721 unsafe {
26722 lb.launch(cfg)?;
26723 }
26724 }
26725 let hki = hk as i32;
26726 if let Some(w) = wq {
26727 let f = self.func("gdn_k2_wgmma_vl");
26729 let cfg = LaunchConfig {
26730 grid_dim: (max_nc, n_head as u32, b as u32),
26731 block_dim: (128, 1, 1),
26732 shared_mem_bytes: 0,
26733 };
26734 let __s_lb = self.gpu.stream();
26735 let mut lb = __s_lb.launch_builder(&f);
26736 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
26737 unsafe {
26738 lb.launch(cfg)?;
26739 }
26740 } else {
26741 let f = self.func("gdn_chunk_attn_vl");
26742 f.set_attribute(
26743 CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
26744 GDN_K2_DYNAMIC_SHARED_BYTES as i32,
26745 )?;
26746 let cfg = LaunchConfig {
26747 grid_dim: (max_nc, n_head as u32, b as u32),
26748 block_dim: (256, 1, 1),
26749 shared_mem_bytes: GDN_K2_DYNAMIC_SHARED_BYTES,
26750 };
26751 let __s_lb = self.gpu.stream();
26752 let mut lb = __s_lb.launch_builder(&f);
26753 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26754 unsafe {
26755 lb.launch(cfg)?;
26756 }
26757 }
26758 {
26759 let f = self.func("gdn_chunk_solve32_vl");
26760 let cfg = LaunchConfig {
26761 grid_dim: (max_nc, n_head as u32, b as u32),
26762 block_dim: (256, 1, 1),
26763 shared_mem_bytes: 0,
26764 };
26765 let __s_lb = self.gpu.stream();
26766 let mut lb = __s_lb.launch_builder(&f);
26767 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
26768 unsafe {
26769 lb.launch(cfg)?;
26770 }
26771 }
26772 Ok(())
26773 }
26774
26775 #[allow(clippy::too_many_arguments)]
26779 pub fn gdn_prep_vl8(
26780 &self,
26781 seqs: &[GdnPrepVl],
26782 conv_w: &CudaSlice<f32>,
26783 dt_bias: &CudaSlice<f32>,
26784 a: &CudaSlice<f32>,
26785 conv_dim: usize,
26786 d_conv: usize,
26787 d_state: usize,
26788 num_v: usize,
26789 num_k: usize,
26790 key_dim: usize,
26791 hk: usize,
26792 eps: f32,
26793 ) -> Result<(), Box<dyn std::error::Error>> {
26794 let b = seqs.len();
26795 assert!(b >= 1 && b <= 8);
26796 let mut packed = [GdnPrepVl::default(); 8];
26797 packed[..b].copy_from_slice(seqs);
26798 let v = GdnPrepVl8(packed);
26799 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26800 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
26801 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
26802 assert!(
26803 conv_fuse || hk == num_v,
26804 "de-broadcast requires the fused conv"
26805 );
26806 if conv_fuse {
26807 let f = self.func("ssm_conv1d_gdn_state_vl");
26808 let cfg = LaunchConfig {
26809 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26810 block_dim: (256, 1, 1),
26811 shared_mem_bytes: 0,
26812 };
26813 let (dsi, nvi, nki, kdi, hki) = (
26814 d_state as i32,
26815 num_v as i32,
26816 num_k as i32,
26817 key_dim as i32,
26818 hk as i32,
26819 );
26820 let __s_lb = self.gpu.stream();
26821 let mut lb = __s_lb.launch_builder(&f);
26822 lb.arg(&v)
26823 .arg(conv_w)
26824 .arg(&cdi)
26825 .arg(&dci)
26826 .arg(&dsi)
26827 .arg(&nvi)
26828 .arg(&nki)
26829 .arg(&kdi)
26830 .arg(&hki);
26831 unsafe {
26832 lb.launch(cfg)?;
26833 }
26834 } else {
26835 let f = self.func("ssm_conv1d_tm_state_vl");
26836 let cfg = LaunchConfig {
26837 grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32),
26838 block_dim: (256, 1, 1),
26839 shared_mem_bytes: 0,
26840 };
26841 let __s_lb = self.gpu.stream();
26842 let mut lb = __s_lb.launch_builder(&f);
26843 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
26844 unsafe {
26845 lb.launch(cfg)?;
26846 }
26847 }
26848 {
26849 let f = self.func("ssm_conv_ring_update_vl");
26850 let n = (conv_dim * (d_conv - 1)) as u32;
26851 let cfg = LaunchConfig {
26852 grid_dim: (n.div_ceil(256), 1, b as u32),
26853 block_dim: (256, 1, 1),
26854 shared_mem_bytes: 0,
26855 };
26856 let __s_lb = self.gpu.stream();
26857 let mut lb = __s_lb.launch_builder(&f);
26858 lb.arg(&v).arg(&cdi).arg(&dci);
26859 unsafe {
26860 lb.launch(cfg)?;
26861 }
26862 }
26863 if !conv_fuse {
26864 let f = self.func("qkv_to_gdn_repack_vl");
26865 let n = max_t * (num_v * d_state) as u32;
26866 let cfg = LaunchConfig {
26867 grid_dim: (n.div_ceil(256), 1, b as u32),
26868 block_dim: (256, 1, 1),
26869 shared_mem_bytes: 0,
26870 };
26871 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
26872 let __s_lb = self.gpu.stream();
26873 let mut lb = __s_lb.launch_builder(&f);
26874 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
26875 unsafe {
26876 lb.launch(cfg)?;
26877 }
26878 }
26879 if Self::l2_v2_on(d_state) {
26880 let f = self.func("gdn_l2_v2_vl");
26881 let cfg = LaunchConfig {
26882 grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32),
26883 block_dim: (256, 1, 1),
26884 shared_mem_bytes: 0,
26885 };
26886 let (dsi, nvi) = (d_state as i32, hk as i32);
26887 let __s_lb = self.gpu.stream();
26888 let mut lb = __s_lb.launch_builder(&f);
26889 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26890 unsafe {
26891 lb.launch(cfg)?;
26892 }
26893 } else {
26894 let f = self.func("gdn_l2_vl");
26895 let cfg = LaunchConfig {
26896 grid_dim: (max_t * hk as u32, 2, b as u32),
26897 block_dim: (256, 1, 1),
26898 shared_mem_bytes: 0,
26899 };
26900 let (dsi, nvi) = (d_state as i32, hk as i32);
26901 let __s_lb = self.gpu.stream();
26902 let mut lb = __s_lb.launch_builder(&f);
26903 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
26904 unsafe {
26905 lb.launch(cfg)?;
26906 }
26907 }
26908 {
26909 let f = self.func("gdn_gate_prep_vl");
26910 let n = max_t * num_v as u32;
26911 let cfg = LaunchConfig {
26912 grid_dim: (n.div_ceil(256), 1, b as u32),
26913 block_dim: (256, 1, 1),
26914 shared_mem_bytes: 0,
26915 };
26916 let nvi = num_v as i32;
26917 let __s_lb = self.gpu.stream();
26918 let mut lb = __s_lb.launch_builder(&f);
26919 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
26920 unsafe {
26921 lb.launch(cfg)?;
26922 }
26923 }
26924 Ok(())
26925 }
26926
26927 pub fn gdn_mirror_vl8(
26929 &self,
26930 seqs: &[GdnSeqVl],
26931 n_head: usize,
26932 which: i32,
26933 hk: usize,
26934 ) -> Result<(), Box<dyn std::error::Error>> {
26935 let b = seqs.len();
26936 assert!(b >= 1 && b <= 8);
26937 let mut packed = [GdnSeqVl::default(); 8];
26938 packed[..b].copy_from_slice(seqs);
26939 let v = GdnVl8(packed);
26940 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
26941 let max_n = seqs
26942 .iter()
26943 .map(|s| {
26944 if which == 0 {
26945 s.t as i64 * ept as i64
26946 } else {
26947 s.nc as i64 * ept as i64 * 32
26948 }
26949 })
26950 .max()
26951 .unwrap();
26952 let f = self.func("gdn_mirror_vl");
26953 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
26954 let cfg = LaunchConfig {
26955 grid_dim: (blocks, 1, b as u32),
26956 block_dim: (256, 1, 1),
26957 shared_mem_bytes: 0,
26958 };
26959 let __s_lb = self.gpu.stream();
26960 let mut lb = __s_lb.launch_builder(&f);
26961 lb.arg(&v).arg(&ept).arg(&which);
26962 unsafe {
26963 lb.launch(cfg)?;
26964 }
26965 Ok(())
26966 }
26967
26968 pub fn gdn_tail_vl8(
26970 &self,
26971 seqs: &[GdnPrepVl],
26972 norm_w: &CudaSlice<f32>,
26973 d_state: usize,
26974 num_v: usize,
26975 eps: f32,
26976 ) -> Result<(), Box<dyn std::error::Error>> {
26977 let b = seqs.len();
26978 assert!(b >= 1 && b <= 8);
26979 let mut packed = [GdnPrepVl::default(); 8];
26980 packed[..b].copy_from_slice(seqs);
26981 let v = GdnPrepVl8(packed);
26982 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
26983 let f = self.func("gated_rmsnorm_f16out_vl");
26984 let cfg = LaunchConfig {
26986 grid_dim: (max_t * num_v as u32, 1, b as u32),
26987 block_dim: (128, 1, 1),
26988 shared_mem_bytes: 0,
26989 };
26990 let (dsi, nvi) = (d_state as i32, num_v as i32);
26991 let __s_lb = self.gpu.stream();
26992 let mut lb = __s_lb.launch_builder(&f);
26993 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
26994 unsafe {
26995 lb.launch(cfg)?;
26996 }
26997 Ok(())
26998 }
26999
27000 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
27003 use cudarc::driver::DevicePtr;
27004 let s = self.gpu.stream();
27005 let (p, _g) = x.device_ptr(&s);
27006 p as u64
27007 }
27008 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
27009 use cudarc::driver::DevicePtrMut;
27010 let s = self.gpu.stream();
27011 let (p, _g) = x.device_ptr_mut(&s);
27012 p as u64
27013 }
27014 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
27015 use cudarc::driver::DevicePtr;
27016 let s = self.gpu.stream();
27017 let (p, _g) = x.device_ptr(&s);
27018 p as u64
27019 }
27020 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
27021 use cudarc::driver::DevicePtr;
27022 let s = self.gpu.stream();
27023 let (p, _g) = x.device_ptr(&s);
27024 p as u64
27025 }
27026
27027 pub fn gdn_chunk_vl8(
27031 &self,
27032 seqs: &[GdnSeqVl],
27033 n_head: usize,
27034 scale: f32,
27035 hk: usize,
27036 wq: Option<&GdnWVl8>,
27037 ) -> Result<(), Box<dyn std::error::Error>> {
27038 const NSPLIT: u32 = 4;
27039 let b = seqs.len();
27040 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
27041 let mut packed = [GdnSeqVl::default(); 8];
27042 packed[..b].copy_from_slice(seqs);
27043 let v = GdnVl8(packed);
27044 let (hi, ci) = (n_head as i32, 32i32);
27045 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
27046 let hki = hk as i32;
27047 if let Some(w) = wq {
27048 let f = self.func("gdn_k45_wgmma_vl");
27050 let cfg = LaunchConfig {
27051 grid_dim: (n_head as u32, NSPLIT, b as u32),
27052 block_dim: (256, 1, 1),
27053 shared_mem_bytes: 0,
27054 };
27055 let __s_lb = self.gpu.stream();
27056 let mut lb = __s_lb.launch_builder(&f);
27057 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
27058 unsafe {
27059 lb.launch(cfg)?;
27060 }
27061 let _ = max_nc;
27062 return Ok(());
27063 }
27064 {
27065 let f = self.func("gdn_chunk_state_mma_vl");
27066 let cfg = LaunchConfig {
27067 grid_dim: (n_head as u32, NSPLIT, b as u32),
27068 block_dim: (256, 1, 1),
27069 shared_mem_bytes: 0,
27070 };
27071 let __s_lb = self.gpu.stream();
27072 let mut lb = __s_lb.launch_builder(&f);
27073 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
27074 unsafe {
27075 lb.launch(cfg)?;
27076 }
27077 }
27078 {
27079 let f = self.func("gdn_chunk_output_mma_vl");
27080 let cfg = LaunchConfig {
27081 grid_dim: (max_nc, n_head as u32, b as u32),
27082 block_dim: (256, 1, 1),
27083 shared_mem_bytes: 0,
27084 };
27085 let __s_lb = self.gpu.stream();
27086 let mut lb = __s_lb.launch_builder(&f);
27087 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
27088 unsafe {
27089 lb.launch(cfg)?;
27090 }
27091 }
27092 Ok(())
27093 }
27094 pub fn gdn_scan_chunked(
27095 &self,
27096 q: &CudaSlice<f32>,
27097 k: &CudaSlice<f32>,
27098 v: &CudaSlice<f32>,
27099 g: &CudaSlice<f32>,
27100 beta: &CudaSlice<f32>,
27101 kb16_pre: Option<&CudaSlice<u8>>,
27102 qb16_pre: Option<&CudaSlice<u8>>,
27103 state_in: &CudaSlice<f32>,
27104 state_out: &mut CudaSlice<f32>,
27105 o: &mut CudaSlice<f32>,
27106 n_head: usize,
27107 t: usize,
27108 scale: f32,
27109 c: usize,
27110 hk: usize,
27111 ) -> Result<(), Box<dyn std::error::Error>> {
27112 const D: usize = 128;
27113 const NSPLIT: u32 = 4;
27114 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
27115 let h = n_head;
27116 let nc = (t + c - 1) / c;
27117 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
27118 let gdn_mma_pre = !portable_mma_gated()
27123 && c == 32
27124 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
27125 Ok("1") => true,
27126 Ok("0") => false,
27127 _ => gdn_mma_default_on(),
27128 };
27129 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
27130 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
27131 } else {
27132 None
27133 };
27134 let gdn_wgmma_pre = cfg!(memra_hopper_mma)
27139 && gdn_mma_pre
27140 && std::env::var("MEMRA_GDN_WGMMA").as_deref() != Ok("0");
27141 let nk = t * hk * D;
27142 let mut kb16_local: Option<CudaSlice<u8>> = None;
27143 if gdn_mma_pre && kb16_pre.is_none() {
27144 let mut kb = self.alloc_u8_uninit(nk * 2)?;
27145 let f = self.func("f32_to_bf16_bulk");
27146 let n2 = nk as i64;
27147 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
27148 let __s_b = self.gpu.stream();
27149 let mut b = __s_b.launch_builder(&f);
27150 b.arg(k).arg(&mut kb).arg(&n2);
27151 unsafe {
27152 b.launch(cfg2)?;
27153 }
27154 kb16_local = Some(kb);
27155 }
27156 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
27157 if let Some(kb) = kb16_pre {
27158 assert!(kb.len() >= nk * 2, "kb16_pre too small");
27159 }
27160 let mut qb16: Option<CudaSlice<u8>> = None;
27161 let mut pb16: Option<CudaSlice<u8>> = None;
27162 if gdn_wgmma_pre {
27163 if qb16_pre.is_none() {
27166 let mut qb = self.alloc_u8_uninit(nk * 2)?;
27167 let f = self.func("f32_to_bf16_bulk");
27168 let n2 = nk as i64;
27169 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
27170 let __s_b = self.gpu.stream();
27171 let mut b = __s_b.launch_builder(&f);
27172 b.arg(q).arg(&mut qb).arg(&n2);
27173 unsafe {
27174 b.launch(cfg2)?;
27175 }
27176 qb16 = Some(qb);
27177 } else if let Some(qb) = qb16_pre {
27178 assert!(qb.len() >= nk * 2, "qb16_pre too small");
27179 }
27180 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
27181 }
27182 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
27183 let k2w = if gdn_wgmma_pre {
27184 Some((
27185 *qb16_ref0.as_ref().unwrap(),
27186 *kb16_ref0.as_ref().unwrap(),
27187 pb16.as_mut().unwrap(),
27188 ))
27189 } else {
27190 None
27191 };
27192 let (gcum, p, u, w) =
27193 self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
27194 let _ = &w;
27195 let mut y = self.uninit(nc * h * c * D)?;
27196 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated()
27210 && c == 32
27211 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
27212 Ok("1") => true,
27213 Ok("0") => false,
27214 _ => gdn_mma_default_on(),
27215 };
27216 if gdn_mma {
27217 let wb16 = wb16_pre
27218 .take()
27219 .expect("mma path pre-allocates wb16 (K3 store fold)");
27220 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
27221 if gdn_wgmma_pre {
27233 let qb16 = qb16_ref0.unwrap();
27235 let pb16 = pb16.as_ref().unwrap();
27236 {
27237 let f = self.func("gdn_k45_wgmma");
27238 let cfg = LaunchConfig {
27239 grid_dim: (h as u32, 4, 1),
27240 block_dim: (256, 1, 1),
27241 shared_mem_bytes: 0,
27242 };
27243 let hki = hk as i32;
27244 let __s_b = self.gpu.stream();
27245 let mut b = __s_b.launch_builder(&f);
27246 b.arg(kb16_ref)
27247 .arg(&gcum)
27248 .arg(beta)
27249 .arg(&u)
27250 .arg(&wb16)
27251 .arg(qb16)
27252 .arg(pb16)
27253 .arg(o)
27254 .arg(&scale)
27255 .arg(state_in)
27256 .arg(&mut *state_out)
27257 .arg(&hi)
27258 .arg(&ti)
27259 .arg(&ci)
27260 .arg(&hki);
27261 unsafe {
27262 b.launch(cfg)?;
27263 }
27264 }
27265 return Ok(());
27266 }
27267 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
27271 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
27272 {
27273 let f = self.func("gdn_chunk_state_mma");
27274 let cfg = LaunchConfig {
27275 grid_dim: (h as u32, NSPLIT, 1),
27276 block_dim: (256, 1, 1),
27277 shared_mem_bytes: 0,
27278 };
27279 let hki = hk as i32;
27280 let __s_b = self.gpu.stream();
27281 let mut b = __s_b.launch_builder(&f);
27282 b.arg(kb16_ref)
27283 .arg(&gcum)
27284 .arg(beta)
27285 .arg(&u)
27286 .arg(&wb16)
27287 .arg(&mut y16)
27288 .arg(&mut ssnap16)
27289 .arg(state_in)
27290 .arg(&mut *state_out)
27291 .arg(&hi)
27292 .arg(&ti)
27293 .arg(&ci)
27294 .arg(&hki);
27295 unsafe {
27296 b.launch(cfg)?;
27297 }
27298 }
27299 {
27300 let f = self.func("gdn_chunk_output_mma");
27302 let jt = ((c + 31) / 32) as u32;
27303 let cfg = LaunchConfig {
27304 grid_dim: (nc as u32, h as u32, jt),
27305 block_dim: (256, 1, 1),
27306 shared_mem_bytes: 0,
27307 };
27308 let hki = hk as i32;
27309 let __s_b = self.gpu.stream();
27310 let mut b = __s_b.launch_builder(&f);
27311 b.arg(q)
27312 .arg(&gcum)
27313 .arg(&p)
27314 .arg(&y16)
27315 .arg(&ssnap16)
27316 .arg(o)
27317 .arg(&hi)
27318 .arg(&ti)
27319 .arg(&ci)
27320 .arg(&scale)
27321 .arg(&hki);
27322 unsafe {
27323 b.launch(cfg)?;
27324 }
27325 }
27326 return Ok(());
27327 }
27328 {
27329 let f = self.func("gdn_chunk_state_f32");
27331 let cfg = LaunchConfig {
27332 grid_dim: (h as u32, NSPLIT, 1),
27333 block_dim: (256, 1, 1),
27334 shared_mem_bytes: 0,
27335 };
27336 let __s_b = self.gpu.stream();
27337 let mut b = __s_b.launch_builder(&f);
27338 b.arg(k)
27339 .arg(&gcum)
27340 .arg(beta)
27341 .arg(&u)
27342 .arg(&w)
27343 .arg(&mut y)
27344 .arg(&mut ssnap)
27345 .arg(state_in)
27346 .arg(&mut *state_out)
27347 .arg(&hi)
27348 .arg(&ti)
27349 .arg(&ci);
27350 unsafe {
27351 b.launch(cfg)?;
27352 }
27353 }
27354 {
27355 let f = self.func("gdn_chunk_output_f32");
27357 let jt = ((c + 31) / 32) as u32;
27358 let cfg = LaunchConfig {
27359 grid_dim: (nc as u32, h as u32, jt),
27360 block_dim: (256, 1, 1),
27361 shared_mem_bytes: 0,
27362 };
27363 let __s_b = self.gpu.stream();
27364 let mut b = __s_b.launch_builder(&f);
27365 b.arg(q)
27366 .arg(&gcum)
27367 .arg(&p)
27368 .arg(&y)
27369 .arg(&ssnap)
27370 .arg(o)
27371 .arg(&hi)
27372 .arg(&ti)
27373 .arg(&ci)
27374 .arg(&scale);
27375 unsafe {
27376 b.launch(cfg)?;
27377 }
27378 }
27379 Ok(())
27380 }
27381
27382 #[allow(clippy::too_many_arguments)]
27391 #[allow(clippy::too_many_arguments)]
27392 pub fn gdn_scan_prefill(
27393 &self,
27394 q: &CudaSlice<f32>,
27395 k: &CudaSlice<f32>,
27396 v: &CudaSlice<f32>,
27397 g: &CudaSlice<f32>,
27398 beta: &CudaSlice<f32>,
27399 kb16_pre: Option<&CudaSlice<u8>>,
27400 qb16_pre: Option<&CudaSlice<u8>>,
27401 state_in: &CudaSlice<f32>,
27402 state_out: &mut CudaSlice<f32>,
27403 o: &mut CudaSlice<f32>,
27404 n_head: usize,
27405 t: usize,
27406 scale: f32,
27407 hk: usize,
27408 ) -> Result<(), Box<dyn std::error::Error>> {
27409 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
27410 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
27411 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
27412 }
27413 if Self::gdn_chunked_enabled() && t >= 16 {
27414 self.gdn_scan_chunked(
27415 q,
27416 k,
27417 v,
27418 g,
27419 beta,
27420 kb16_pre,
27421 qb16_pre,
27422 state_in,
27423 state_out,
27424 o,
27425 n_head,
27426 t,
27427 scale,
27428 Self::gdn_chunk_size(),
27429 hk,
27430 )
27431 } else {
27432 assert!(
27433 hk == n_head,
27434 "s128 scan is broadcast-only (prep guarantees by predicate)"
27435 );
27436 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
27437 }
27438 }
27439
27440 #[allow(clippy::too_many_arguments)]
27442 fn gdn_scan_diff(
27443 &self,
27444 q: &CudaSlice<f32>,
27445 k: &CudaSlice<f32>,
27446 v: &CudaSlice<f32>,
27447 g: &CudaSlice<f32>,
27448 beta: &CudaSlice<f32>,
27449 state_in: &CudaSlice<f32>,
27450 state_out: &mut CudaSlice<f32>,
27451 o: &mut CudaSlice<f32>,
27452 n_head: usize,
27453 t: usize,
27454 scale: f32,
27455 ) -> Result<(), Box<dyn std::error::Error>> {
27456 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
27457 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
27458 let mut o_c = self.uninit(o.len())?;
27459 let mut st_c = self.uninit(state_out.len())?;
27460 self.gdn_scan_chunked(
27461 q,
27462 k,
27463 v,
27464 g,
27465 beta,
27466 None,
27467 None,
27468 state_in,
27469 &mut st_c,
27470 &mut o_c,
27471 n_head,
27472 t,
27473 scale,
27474 Self::gdn_chunk_size(),
27475 n_head,
27476 )?;
27477 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
27478 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
27479 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
27480 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
27481 let mut max_abs = 0f32;
27482 let mut max_rel = 0f32;
27483 let mut sum_rel = 0f64;
27484 for (x, y) in a.iter().zip(b) {
27485 let ad = (x - y).abs();
27486 let rel = ad / x.abs().max(y.abs()).max(1e-3);
27487 if ad > max_abs {
27488 max_abs = ad;
27489 }
27490 if rel > max_rel {
27491 max_rel = rel;
27492 }
27493 sum_rel += rel as f64;
27494 }
27495 (max_abs, max_rel, sum_rel / a.len() as f64)
27496 };
27497 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
27498 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
27499 println!(
27500 "[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
27501 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
27502 Self::gdn_chunk_size()
27503 );
27504 Ok(())
27505 }
27506
27507 pub fn gdn_glog(
27509 &self,
27510 alpha: &CudaSlice<f32>,
27511 dt_bias: &CudaSlice<f32>,
27512 a: &CudaSlice<f32>,
27513 g_log: &mut CudaSlice<f32>,
27514 n_head: usize,
27515 t: usize,
27516 ) -> Result<(), Box<dyn std::error::Error>> {
27517 let f = self.func("gdn_glog_f32");
27518 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
27519 let (h, ti) = (n_head as i32, t as i32);
27520 let __s_b = self.gpu.stream();
27521 let mut b = __s_b.launch_builder(&f);
27522 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
27523 unsafe {
27524 b.launch(cfg)?;
27525 }
27526 Ok(())
27527 }
27528
27529 pub fn sigmoid_v(
27532 &self,
27533 x: &cudarc::driver::CudaView<f32>,
27534 y: &mut CudaSlice<f32>,
27535 n: usize,
27536 ) -> Result<(), Box<dyn std::error::Error>> {
27537 let f = self.func("sigmoid_f32");
27538 let cfg = LaunchConfig::for_num_elems(n as u32);
27539 let ni = n as i32;
27540 let __s_b = self.gpu.stream();
27541 let mut b = __s_b.launch_builder(&f);
27542 b.arg(x).arg(y).arg(&ni);
27543 unsafe {
27544 b.launch(cfg)?;
27545 }
27546 Ok(())
27547 }
27548
27549 pub fn gdn_glog_v(
27550 &self,
27551 alpha: &cudarc::driver::CudaView<f32>,
27552 dt_bias: &CudaSlice<f32>,
27553 a: &CudaSlice<f32>,
27554 g_log: &mut CudaSlice<f32>,
27555 n_head: usize,
27556 t: usize,
27557 ) -> Result<(), Box<dyn std::error::Error>> {
27558 let f = self.func("gdn_glog_f32");
27559 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
27560 let (h, ti) = (n_head as i32, t as i32);
27561 let __s_b = self.gpu.stream();
27562 let mut b = __s_b.launch_builder(&f);
27563 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
27564 unsafe {
27565 b.launch(cfg)?;
27566 }
27567 Ok(())
27568 }
27569
27570 pub fn sigmoid(
27571 &self,
27572 x: &CudaSlice<f32>,
27573 y: &mut CudaSlice<f32>,
27574 n: usize,
27575 ) -> Result<(), Box<dyn std::error::Error>> {
27576 let f = self.func("sigmoid_f32");
27577 let cfg = LaunchConfig::for_num_elems(n as u32);
27578 let ni = n as i32;
27579 let __s_b = self.gpu.stream();
27580 let mut b = __s_b.launch_builder(&f);
27581 b.arg(x).arg(y).arg(&ni);
27582 unsafe {
27583 b.launch(cfg)?;
27584 }
27585 Ok(())
27586 }
27587
27588 pub fn sig_mul_f16out(
27591 &self,
27592 a: &CudaSlice<f32>,
27593 g: &CudaSlice<f32>,
27594 dst: &mut CudaSlice<f32>,
27595 dst16: &mut CudaSlice<u8>,
27596 n: usize,
27597 ) -> Result<(), Box<dyn std::error::Error>> {
27598 let f = self.func("sig_mul_f16out_f32");
27599 let cfg = LaunchConfig::for_num_elems(n as u32);
27600 let ni = n as i32;
27601 let __s_b = self.gpu.stream();
27602 let mut b = __s_b.launch_builder(&f);
27603 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
27604 unsafe {
27605 b.launch(cfg)?;
27606 }
27607 Ok(())
27608 }
27609
27610 #[allow(clippy::too_many_arguments)]
27619 pub fn attn_head_gate(
27620 &self,
27621 a: &CudaSlice<f32>,
27622 g: &CudaSlice<f32>,
27623 dst: &mut CudaSlice<f32>,
27624 dst16: Option<&mut CudaSlice<u8>>,
27625 head_dim: usize,
27626 n_head: usize,
27627 t: usize,
27628 ) -> Result<(), Box<dyn std::error::Error>> {
27629 let f = self.func("attn_head_gate_f32");
27630 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27631 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27632 let d16: u64 = match dst16 {
27634 Some(d) => self.addr_u8(d),
27635 None => 0,
27636 };
27637 let __s_b = self.gpu.stream();
27638 let mut b = __s_b.launch_builder(&f);
27639 b.arg(a)
27640 .arg(g)
27641 .arg(dst)
27642 .arg(&d16)
27643 .arg(&hd)
27644 .arg(&nh)
27645 .arg(&ti);
27646 unsafe {
27647 b.launch(cfg)?;
27648 }
27649 Ok(())
27650 }
27651
27652 #[allow(clippy::too_many_arguments)]
27661 pub fn swiglu_clamped_mul_scaled(
27662 &self,
27663 gate: &CudaSlice<f32>,
27664 up: &CudaSlice<f32>,
27665 gs: f32,
27666 us: f32,
27667 limit: f32,
27668 dst: &mut CudaSlice<f32>,
27669 n: usize,
27670 ) -> Result<(), Box<dyn std::error::Error>> {
27671 debug_assert!(
27672 limit > 1e-6,
27673 "swiglu_clamped needs a live limit; use silu_mul_scaled"
27674 );
27675 let f = self.func("swiglu_clamped_mul_scaled_f32");
27676 let cfg = LaunchConfig::for_num_elems(n as u32);
27677 let ni = n as i32;
27678 let __s_b = self.gpu.stream();
27679 let mut b = __s_b.launch_builder(&f);
27680 b.arg(gate)
27681 .arg(up)
27682 .arg(&gs)
27683 .arg(&us)
27684 .arg(&limit)
27685 .arg(dst)
27686 .arg(&ni);
27687 unsafe {
27688 b.launch(cfg)?;
27689 }
27690 Ok(())
27691 }
27692
27693 pub fn gated_rmsnorm(
27695 &self,
27696 o: &CudaSlice<f32>,
27697 w: &CudaSlice<f32>,
27698 z: &CudaSlice<f32>,
27699 dst: &mut CudaSlice<f32>,
27700 ncols: usize,
27701 nrows: usize,
27702 eps: f32,
27703 ) -> Result<(), Box<dyn std::error::Error>> {
27704 let f = self.func("gated_rmsnorm_f32");
27705 let cfg = LaunchConfig {
27706 grid_dim: (nrows as u32, 1, 1),
27707 block_dim: (128, 1, 1),
27708 shared_mem_bytes: 0,
27709 };
27710 let (nc, e) = (ncols as i32, eps);
27711 let __s_b = self.gpu.stream();
27712 let mut b = __s_b.launch_builder(&f);
27713 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27714 unsafe {
27715 b.launch(cfg)?;
27716 }
27717 Ok(())
27718 }
27719
27720 pub fn gated_rmsnorm_f16out(
27723 &self,
27724 o: &CudaSlice<f32>,
27725 w: &CudaSlice<f32>,
27726 z: &CudaSlice<f32>,
27727 dst: &mut CudaSlice<f32>,
27728 dst16: &mut CudaSlice<u8>,
27729 ncols: usize,
27730 nrows: usize,
27731 eps: f32,
27732 ) -> Result<(), Box<dyn std::error::Error>> {
27733 let f = self.func("gated_rmsnorm_f16out_f32");
27734 let cfg = LaunchConfig {
27736 grid_dim: (nrows as u32, 1, 1),
27737 block_dim: (128, 1, 1),
27738 shared_mem_bytes: 0,
27739 };
27740 let (nc, e) = (ncols as i32, eps);
27741 let __s_b = self.gpu.stream();
27742 let mut b = __s_b.launch_builder(&f);
27743 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27744 unsafe {
27745 b.launch(cfg)?;
27746 }
27747 Ok(())
27748 }
27749
27750 #[allow(clippy::too_many_arguments)]
27754 pub fn add_rms_norm_zq8(
27755 &self,
27756 a: &CudaSlice<f32>,
27757 b_in: &CudaSlice<f32>,
27758 w: &CudaSlice<f32>,
27759 res: &mut CudaSlice<f32>,
27760 z: &mut CudaSlice<f32>,
27761 ncols: usize,
27762 nrows: usize,
27763 eps: f32,
27764 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27765 assert!(ncols % 32 == 0);
27766 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
27767 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27768 let f = self.func("add_rms_norm_zq8");
27769 let cfg = LaunchConfig {
27770 grid_dim: (nrows as u32, 1, 1),
27771 block_dim: (1024, 1, 1),
27772 shared_mem_bytes: 0,
27773 };
27774 let (nc, ep) = (ncols as i32, eps);
27775 let __s_b = self.gpu.stream();
27776 let mut b = __s_b.launch_builder(&f);
27777 b.arg(a)
27778 .arg(b_in)
27779 .arg(w)
27780 .arg(res)
27781 .arg(z)
27782 .arg(&mut q)
27783 .arg(&mut d)
27784 .arg(&nc)
27785 .arg(&ep);
27786 unsafe {
27787 b.launch(cfg)?;
27788 }
27789 Ok((q, d))
27790 }
27791
27792 pub fn gated_rmsnorm_zv(
27797 &self,
27798 o: &CudaSlice<f32>,
27799 w: &CudaSlice<f32>,
27800 z: &cudarc::driver::CudaView<f32>,
27801 dst: &mut CudaSlice<f32>,
27802 ncols: usize,
27803 nrows: usize,
27804 eps: f32,
27805 ) -> Result<(), Box<dyn std::error::Error>> {
27806 let f = self.func("gated_rmsnorm_f32");
27807 let cfg = LaunchConfig {
27808 grid_dim: (nrows as u32, 1, 1),
27809 block_dim: (128, 1, 1),
27810 shared_mem_bytes: 0,
27811 };
27812 let (nc, e) = (ncols as i32, eps);
27813 let __s_b = self.gpu.stream();
27814 let mut b = __s_b.launch_builder(&f);
27815 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
27816 unsafe {
27817 b.launch(cfg)?;
27818 }
27819 Ok(())
27820 }
27821
27822 pub fn gated_rmsnorm_f16out_zv(
27823 &self,
27824 o: &CudaSlice<f32>,
27825 w: &CudaSlice<f32>,
27826 z: &cudarc::driver::CudaView<f32>,
27827 dst: &mut CudaSlice<f32>,
27828 dst16: &mut CudaSlice<u8>,
27829 ncols: usize,
27830 nrows: usize,
27831 eps: f32,
27832 ) -> Result<(), Box<dyn std::error::Error>> {
27833 let f = self.func("gated_rmsnorm_f16out_f32");
27834 let cfg = LaunchConfig {
27836 grid_dim: (nrows as u32, 1, 1),
27837 block_dim: (128, 1, 1),
27838 shared_mem_bytes: 0,
27839 };
27840 let (nc, e) = (ncols as i32, eps);
27841 let __s_b = self.gpu.stream();
27842 let mut b = __s_b.launch_builder(&f);
27843 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
27844 unsafe {
27845 b.launch(cfg)?;
27846 }
27847 Ok(())
27848 }
27849
27850 pub fn gated_rmsnorm_q8_1(
27851 &self,
27852 o: &CudaSlice<f32>,
27853 w: &CudaSlice<f32>,
27854 z: &CudaSlice<f32>,
27855 ncols: usize,
27856 nrows: usize,
27857 eps: f32,
27858 ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
27859 assert!(ncols % 32 == 0);
27860 let f = self.func("gated_rmsnorm_q8_1");
27861 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
27862 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
27863 let cfg = LaunchConfig {
27864 grid_dim: (nrows as u32, 1, 1),
27865 block_dim: (128, 1, 1),
27866 shared_mem_bytes: 0,
27867 };
27868 let (nc, ep) = (ncols as i32, eps);
27869 let __s_b = self.gpu.stream();
27870 let mut b = __s_b.launch_builder(&f);
27871 b.arg(o)
27872 .arg(w)
27873 .arg(z)
27874 .arg(&mut out_q)
27875 .arg(&mut out_d)
27876 .arg(&nc)
27877 .arg(&ep);
27878 unsafe {
27879 b.launch(cfg)?;
27880 }
27881 Ok((out_q, out_d))
27882 }
27883
27884 pub fn transpose(
27886 &self,
27887 inp: &CudaSlice<f32>,
27888 rows: usize,
27889 cols: usize,
27890 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
27891 let f = self.func("transpose_f32");
27892 let mut out = self.zeros(rows * cols)?;
27893 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
27894 let (r, c) = (rows as i32, cols as i32);
27895 let __s_b = self.gpu.stream();
27896 let mut b = __s_b.launch_builder(&f);
27897 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
27898 unsafe {
27899 b.launch(cfg)?;
27900 }
27901 Ok(out)
27902 }
27903
27904 pub fn repeat_heads(
27906 &self,
27907 inp: &CudaSlice<f32>,
27908 out: &mut CudaSlice<f32>,
27909 head_dim: usize,
27910 n_in: usize,
27911 n_out: usize,
27912 t: usize,
27913 ) -> Result<(), Box<dyn std::error::Error>> {
27914 let f = self.func("repeat_heads_f32");
27915 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
27916 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
27917 let __s_b = self.gpu.stream();
27918 let mut b = __s_b.launch_builder(&f);
27919 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
27920 unsafe {
27921 b.launch(cfg)?;
27922 }
27923 Ok(())
27924 }
27925
27926 pub fn q_gate_split(
27933 &self,
27934 qf: &CudaSlice<f32>,
27935 q_out: &mut CudaSlice<f32>,
27936 gate_out: &mut CudaSlice<f32>,
27937 head_dim: usize,
27938 n_head: usize,
27939 t: usize,
27940 ) -> Result<(), Box<dyn std::error::Error>> {
27941 memra_gguf::config::check_fused_q_gate_extent(qf.len(), head_dim, n_head, t)?;
27942 let out_need = head_dim * n_head * t;
27943 if q_out.len() < out_need || gate_out.len() < out_need {
27944 return Err(format!(
27945 "q_gate_split destinations too small: need {out_need} each, have q={} gate={}",
27946 q_out.len(),
27947 gate_out.len()
27948 )
27949 .into());
27950 }
27951 let f = self.func("q_gate_split_f32");
27952 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
27953 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
27954 let __s_b = self.gpu.stream();
27955 let mut b = __s_b.launch_builder(&f);
27956 b.arg(qf)
27957 .arg(q_out)
27958 .arg(gate_out)
27959 .arg(&hd)
27960 .arg(&nh)
27961 .arg(&ti);
27962 unsafe {
27963 b.launch(cfg)?;
27964 }
27965 Ok(())
27966 }
27967
27968 pub fn qkv_to_gdn_repack(
27972 &self,
27973 conv_out: &CudaSlice<f32>,
27974 q_g: &mut CudaSlice<f32>,
27975 k_g: &mut CudaSlice<f32>,
27976 v_g: &mut CudaSlice<f32>,
27977 d_state: usize,
27978 num_v: usize,
27979 num_k: usize,
27980 key_dim: usize,
27981 t: usize,
27982 ) -> Result<(), Box<dyn std::error::Error>> {
27983 let f = self.func("qkv_to_gdn_repack_f32");
27984 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
27985 let (ds, nv, nk, kd, ti) = (
27986 d_state as i32,
27987 num_v as i32,
27988 num_k as i32,
27989 key_dim as i32,
27990 t as i32,
27991 );
27992 let __s_b = self.gpu.stream();
27993 let mut b = __s_b.launch_builder(&f);
27994 b.arg(conv_out)
27995 .arg(q_g)
27996 .arg(k_g)
27997 .arg(v_g)
27998 .arg(&ds)
27999 .arg(&nv)
28000 .arg(&nk)
28001 .arg(&kd)
28002 .arg(&ti);
28003 unsafe {
28004 b.launch(cfg)?;
28005 }
28006 Ok(())
28007 }
28008
28009 pub fn conv_left_pad(
28012 &self,
28013 src: &CudaSlice<f32>,
28014 dst: &mut CudaSlice<f32>,
28015 conv_dim: usize,
28016 t: usize,
28017 pad: usize,
28018 ) -> Result<(), Box<dyn std::error::Error>> {
28019 let f = self.func("conv_left_pad_f32");
28020 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
28021 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
28022 let __s_b = self.gpu.stream();
28023 let mut b = __s_b.launch_builder(&f);
28024 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
28025 unsafe {
28026 b.launch(cfg)?;
28027 }
28028 Ok(())
28029 }
28030
28031 pub fn conv_assemble_and_roll(
28035 &self,
28036 qkv_col: &CudaSlice<f32>,
28037 conv_state: &mut CudaSlice<f32>,
28038 conv_in: &mut CudaSlice<f32>,
28039 conv_dim: usize,
28040 pad: usize,
28041 ) -> Result<(), Box<dyn std::error::Error>> {
28042 let f = self.func("conv_assemble_and_roll_f32");
28043 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
28044 let (cd, p) = (conv_dim as i32, pad as i32);
28045 let __s_b = self.gpu.stream();
28046 let mut b = __s_b.launch_builder(&f);
28047 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
28048 unsafe {
28049 b.launch(cfg)?;
28050 }
28051 Ok(())
28052 }
28053
28054 pub fn ssm_conv1d_fused_decode(
28060 &self,
28061 qkv_col: &CudaSlice<f32>,
28062 conv_state: &mut CudaSlice<f32>,
28063 w: &CudaSlice<f32>,
28064 conv_out: &mut CudaSlice<f32>,
28065 conv_dim: usize,
28066 d_conv: usize,
28067 ) -> Result<(), Box<dyn std::error::Error>> {
28068 let f = self.func("ssm_conv1d_fused_decode_f32");
28069 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
28070 let (cd, dc) = (conv_dim as i32, d_conv as i32);
28071 let __s_b = self.gpu.stream();
28072 let mut b = __s_b.launch_builder(&f);
28073 b.arg(qkv_col)
28074 .arg(conv_state)
28075 .arg(w)
28076 .arg(conv_out)
28077 .arg(&cd)
28078 .arg(&dc);
28079 unsafe {
28080 b.launch(cfg)?;
28081 }
28082 Ok(())
28083 }
28084
28085 pub fn slice_range(
28088 &self,
28089 src: &CudaSlice<f32>,
28090 start: usize,
28091 len: usize,
28092 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28093 let host = self.gpu.stream().clone_dtoh(src)?;
28094 self.gpu.stream().synchronize()?;
28095 Ok(self.htod(&host[start..start + len])?)
28096 }
28097}
28098
28099#[cfg(test)]
28100mod target_dispatch_tests {
28101 use super::legacy_quant_gemm_allowed;
28102
28103 #[test]
28104 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
28105 assert!(legacy_quant_gemm_allowed(false, false, false));
28107 assert!(!legacy_quant_gemm_allowed(false, false, true));
28108 assert!(!legacy_quant_gemm_allowed(true, false, false));
28110 assert!(!legacy_quant_gemm_allowed(true, false, true));
28111 assert!(legacy_quant_gemm_allowed(true, true, false));
28113 assert!(!legacy_quant_gemm_allowed(true, true, true));
28114 }
28115
28116 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
28117 #[test]
28118 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
28119 assert!(!legacy_quant_gemm_allowed(
28120 cfg!(memra_portable_cuda),
28121 cfg!(memra_hopper_mma),
28122 false
28123 ));
28124 }
28125
28126 #[cfg(memra_hopper_mma)]
28127 #[test]
28128 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
28129 assert!(legacy_quant_gemm_allowed(
28130 cfg!(memra_portable_cuda),
28131 cfg!(memra_hopper_mma),
28132 false
28133 ));
28134 assert!(super::portable_mma_gated() == false);
28135 }
28136}
28137
28138impl memra_kv::KvDev for Engine {
28141 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28142 Engine::zeros(self, n)
28143 }
28144 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28145 Engine::uninit(self, n)
28146 }
28147 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
28148 Engine::alloc_u8(self, n)
28149 }
28150 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
28151 Engine::htod_i32(self, v)
28152 }
28153 fn clone_dtod(
28154 &self,
28155 src: &CudaSlice<f32>,
28156 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
28157 Engine::clone_dtod(self, src)
28158 }
28159 fn copy_into(
28160 &self,
28161 dst: &mut CudaSlice<f32>,
28162 off: usize,
28163 src: &CudaSlice<f32>,
28164 len: usize,
28165 ) -> Result<(), Box<dyn std::error::Error>> {
28166 Engine::copy_into(self, dst, off, src, len)
28167 }
28168 fn set_i32_one(
28169 &self,
28170 d: &mut CudaSlice<i32>,
28171 v: i32,
28172 ) -> Result<(), Box<dyn std::error::Error>> {
28173 Engine::set_i32_one(self, d, v)
28174 }
28175}
28176
28177#[cfg(test)]
28178mod fused_gate_bounds_tests {
28179 use super::*;
28180
28181 #[test]
28194 #[ignore = "requires a CUDA GPU"]
28195 fn q_gate_split_refuses_a_separate_gate_wq_instead_of_reading_past_it() {
28196 let e = Engine::new(0).unwrap();
28197 let (head_dim, n_head, t) = (8usize, 4usize, 2usize);
28198 let fused = 2 * head_dim * n_head * t;
28199 let out_n = head_dim * n_head * t;
28200
28201 let narrow = e.htod(&vec![1.0f32; out_n]).unwrap();
28203 let mut q = e.uninit(out_n).unwrap();
28204 let mut gate = e.uninit(out_n).unwrap();
28205 let err = e
28206 .q_gate_split(&narrow, &mut q, &mut gate, head_dim, n_head, t)
28207 .expect_err("half-width wq must be refused, not read past")
28208 .to_string();
28209 assert!(err.contains("NO fused gate"), "{err}");
28210 assert!(err.contains(&format!("{fused}")), "{err}");
28211
28212 let host: Vec<f32> = (0..fused).map(|i| i as f32).collect();
28215 let wide = e.htod(&host).unwrap();
28216 e.q_gate_split(&wide, &mut q, &mut gate, head_dim, n_head, t)
28217 .expect("full-width wq splits");
28218 let (qh, gh) = (e.dtoh(&q).unwrap(), e.dtoh(&gate).unwrap());
28219 for tok in 0..t {
28220 for hh in 0..n_head {
28221 for d in 0..head_dim {
28222 let base = tok * (n_head * 2 * head_dim) + hh * (2 * head_dim);
28223 let idx = tok * (n_head * head_dim) + hh * head_dim + d;
28224 assert_eq!(qh[idx], host[base + d], "q t{tok} h{hh} d{d}");
28225 assert_eq!(gh[idx], host[base + head_dim + d], "gate t{tok} h{hh} d{d}");
28226 }
28227 }
28228 }
28229
28230 let mut small = e.uninit(out_n - 1).unwrap();
28232 assert!(
28233 e.q_gate_split(&wide, &mut small, &mut gate, head_dim, n_head, t)
28234 .is_err()
28235 );
28236 }
28237}
28238
28239#[cfg(test)]
28243mod fused_rope_width_tests {
28244 use super::Engine;
28245
28246 #[test]
28249 fn full_width_is_accepted() {
28250 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 256).is_ok());
28251 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_cat", 512, 512).is_ok());
28252 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append", 128, 128).is_ok());
28253 }
28254
28255 #[test]
28268 fn gemma4_official_artifact_widths_pass() {
28269 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 512, 512).is_ok());
28270 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 256, 256).is_ok());
28271 }
28272
28273 #[test]
28276 fn partial_rotary_is_refused_with_the_geometry_named() {
28277 let err = Engine::full_width_rope_only("rms_norm_qkv_rope", 64, 256)
28279 .expect_err("partial rotary must refuse");
28280 let msg = err.to_string();
28281 assert!(msg.contains("PARTIAL ROTARY REFUSED"), "{msg}");
28282 assert!(msg.contains("n_rot 64"), "{msg}");
28283 assert!(msg.contains("head_dim 256"), "{msg}");
28284 assert!(
28285 msg.contains("64..256"),
28286 "names the band it would corrupt: {msg}"
28287 );
28288 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope_append_dc", 64, 128).is_err());
28290 assert!(Engine::full_width_rope_only("rms_norm_qkv_rope", 256, 128).is_err());
28292 }
28293}