1use std::sync::{Arc, Mutex};
4use cudarc::driver::{CudaContext, CudaStream, CudaModule, CudaFunction, CudaSlice, LaunchConfig, PushKernelArg};
5use cudarc::nvrtc::Ptx;
6
7pub use memra_gguf;
8pub use memra_runtime;
9
10pub mod model;
11pub mod forward;
12pub mod hybrid;
13pub mod hybrid_forward;
14pub mod cache {
17 pub use memra_kv::*;
18}
19pub mod decode;
20pub mod decode_batch;
21pub mod mla;
25pub mod pp;
26pub mod spec;
27pub mod gemma_spec;
28pub mod round_stream;
29pub mod graph_update;
30pub mod dflash;
31pub mod eagle;
32pub use memra_sampling as sampler;
33
34pub fn moe_f16g_mode() -> u8 {
78 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
79 *M.get_or_init(|| match std::env::var("MEMRA_MOE_F16G").as_deref() {
80 Ok("0") => 0,
81 Ok("2") => 2,
82 Ok("3") => 3,
83 Ok(_) => 1,
84 Err(_) => 2,
87 })
88}
89pub fn moe_f16g_sk_params() -> (i32, i32) {
103 static P: std::sync::OnceLock<(i32, i32)> = std::sync::OnceLock::new();
104 *P.get_or_init(|| match std::env::var("MEMRA_F16G_SK").as_deref() {
105 Ok("0") => (-1, 0),
106 Ok("32") => (0, i32::MAX),
107 Ok("128") => (0, 1),
108 _ => {
109 let cross = std::env::var("MEMRA_F16G_SK_CROSS").ok()
110 .and_then(|v| v.parse().ok()).unwrap_or(64);
111 (0, cross)
112 }
113 })
114}
115pub fn moe_f16g_direct_on(qtype: i32) -> bool {
126 static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
127 let m = *M.get_or_init(|| match std::env::var("MEMRA_F16G_DIRECT").as_deref() {
128 Ok("0") => 0,
129 Ok("kq") => 1,
130 _ => 2,
131 });
132 match m {
133 0 => false,
134 1 => qtype == QT_Q4_K || qtype == QT_Q6_K,
135 _ => true,
136 }
137}
138pub fn moe_f16g_tail_on() -> bool {
147 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
148 *ON.get_or_init(|| std::env::var("MEMRA_F16G_TAIL").as_deref() != Ok("0"))
149}
150
151pub fn moe_f16g_gemma_on() -> bool {
158 static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
159 *M.get_or_init(|| !matches!(std::env::var("MEMRA_MOE_F16G").as_deref(), Ok("0") | Err(_)))
160}
161
162pub fn moe_fuse_actq_on() -> bool {
166 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
167 *ON.get_or_init(|| std::env::var("MEMRA_MOE_FUSE_ACTQ").as_deref() != Ok("0"))
168}
169
170pub fn router_prefill_exact_on() -> bool {
180 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
181 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_PREFILL_EXACT").as_deref() != Ok("0"))
182}
183
184pub fn router_kernel_on() -> bool {
185 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
186 *ON.get_or_init(|| {
187 let on = std::env::var("MEMRA_ROUTER_KERNEL").as_deref() != Ok("0");
188 if !on { eprintln!("[memra] router kernel OFF (rollback: per-column cuBLAS gemv)"); }
189 on
190 })
191}
192
193pub const ROUTER_BATCH_MIN_T: usize = 8;
208pub fn router_batch_on() -> bool {
209 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
210 *ON.get_or_init(|| std::env::var("MEMRA_ROUTER_BATCH").as_deref() != Ok("0"))
211}
212mod cpu_experts;
213pub mod moe_cache;
214pub mod spill;
215mod spill_pread;
216#[cfg(memra_cutlass)]
217pub mod cutlass_ffi;
218pub mod mmq_ffi;
219pub mod f16_ffi;
220pub mod prime_graph;
221pub mod fp8_ffi;
222
223const FATBIN: &[u8] = include_bytes!(env!("MEMRA_ENGINE_FATBIN"));
230const HYBRID_FATBIN: &[u8] = include_bytes!(env!("MEMRA_HYBRID_FATBIN"));
231const QMATVEC_FATBIN: &[u8] = include_bytes!(env!("MEMRA_QMATVEC_FATBIN"));
232const FLASH_FATBIN: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN"));
233const GEMM_FATBIN: &[u8] = include_bytes!(env!("MEMRA_GEMM_FATBIN"));
234const ROUTER_FATBIN: &[u8] = include_bytes!(env!("MEMRA_ROUTER_FATBIN"));
235const SAMPLE_FATBIN: &[u8] = include_bytes!(env!("MEMRA_SAMPLE_FATBIN"));
237
238fn gemm_fatbin_bytes() -> std::borrow::Cow<'static, [u8]> {
244 assert!(!(portable_mma_gated() && std::env::var_os("MEMRA_GEMM_FATBIN").is_some()),
245 "MEMRA_GEMM_FATBIN overrides are not allowed in the portable CUDA lane");
246 match std::env::var("MEMRA_GEMM_FATBIN") {
247 Ok(path) => std::borrow::Cow::Owned(
248 std::fs::read(&path).unwrap_or_else(|e| panic!("MEMRA_GEMM_FATBIN read {path}: {e}"))),
249 Err(_) => std::borrow::Cow::Borrowed(GEMM_FATBIN),
250 }
251}
252
253pub(crate) const fn portable_mma_gated() -> bool {
260 cfg!(memra_portable_cuda) && !cfg!(memra_hopper_mma)
261}
262
263const fn legacy_quant_gemm_allowed(portable_cuda: bool, hopper_mma: bool, no_gemm: bool) -> bool {
268 (!portable_cuda || hopper_mma) && !no_gemm
269}
270
271const FLASH_FATBIN_VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VQ4"));
279const FLASH_FATBIN_VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_VF8"));
280const FLASH_FATBIN_KF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8"));
281const FLASH_FATBIN_KF8VQ4: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VQ4"));
282const FLASH_FATBIN_KF8VF8: &[u8] = include_bytes!(env!("MEMRA_FLASH_FATBIN_KF8VF8"));
283
284pub use memra_kv::{kv_blk_bytes, kv_cache_formats};
287
288fn flash_fatbin_bytes() -> &'static [u8] {
290 match kv_cache_formats() {
291 ("q8_0", "q5_1") => FLASH_FATBIN,
292 ("q8_0", "q4_0") => FLASH_FATBIN_VQ4,
293 ("q8_0", "fp8") => FLASH_FATBIN_VF8,
294 ("fp8", "q5_1") => FLASH_FATBIN_KF8,
295 ("fp8", "q4_0") => FLASH_FATBIN_KF8VQ4,
296 ("fp8", "fp8") => FLASH_FATBIN_KF8VF8,
297 other => unreachable!("kv_cache_formats returned {other:?}"),
298 }
299}
300
301fn k1_launch_override() -> Option<(u32, u32, u32)> {
308 static K1: std::sync::OnceLock<Option<(u32, u32, u32)>> = std::sync::OnceLock::new();
309 *K1.get_or_init(|| {
310 let v = std::env::var("MEMRA_GEMM_K1_LAUNCH").ok()?;
311 let p: Vec<u32> = v.split(',').filter_map(|s| s.trim().parse().ok()).collect();
312 match p.as_slice() { [bm, bn, w] => Some((*bm, *bn, *w)), _ => None }
313 })
314}
315
316pub(crate) fn wgmma_gemm_enabled() -> bool {
323 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
324 *V.get_or_init(|| std::env::var("MEMRA_WGMMA").as_deref() == Ok("1"))
325}
326
327pub const FA_VEC_MIN_TKV: usize = 96;
342pub fn fa_vec_min_tkv() -> usize {
346 static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
347 *V.get_or_init(|| std::env::var("MEMRA_FA_VEC_MIN").ok()
348 .and_then(|v| v.parse().ok())
349 .unwrap_or_else(|| FA_VEC_MIN_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)))
350}
351
352pub fn fa_f16pv_on() -> bool {
363 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
364 *ON.get_or_init(|| std::env::var("MEMRA_FA_F16PV").map(|v| v != "0")
365 .unwrap_or_else(|_| std::env::var("MEMRA_DRAFT").is_err()))
366}
367
368pub fn fa512_hp_on() -> bool {
372 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
373 *ON.get_or_init(|| std::env::var("MEMRA_FA512_HP").as_deref() != Ok("0"))
374}
375
376pub fn faw_hp_on() -> bool {
380 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
381 *ON.get_or_init(|| std::env::var("MEMRA_FAW_HP").as_deref() != Ok("0"))
382}
383
384pub fn fa512_wide_warps() -> usize {
388 static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389 *N.get_or_init(|| match std::env::var("MEMRA_FA512_W4").as_deref() {
390 Ok("1") => 4, _ => 2,
391 })
392}
393
394pub fn fa512_min_tkv() -> usize {
397 static FA512_MIN: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
398 *FA512_MIN.get_or_init(|| std::env::var("MEMRA_FA512_MIN").ok()
399 .and_then(|v| v.parse().ok()).unwrap_or(512))
400}
401pub static FA_VEC_MIN_DEFAULT: std::sync::atomic::AtomicUsize =
405 std::sync::atomic::AtomicUsize::new(FA_VEC_MIN_TKV);
406pub static FA_SPW_DEFAULT: std::sync::atomic::AtomicUsize =
410 std::sync::atomic::AtomicUsize::new(32);
411pub static FUSED_MR1_DEFAULT: std::sync::atomic::AtomicBool =
417 std::sync::atomic::AtomicBool::new(false);
418pub static ROUTER_W8_DEFAULT: std::sync::atomic::AtomicBool =
425 std::sync::atomic::AtomicBool::new(true);
426pub static FA_SP512_DEFAULT: std::sync::atomic::AtomicUsize =
427 std::sync::atomic::AtomicUsize::new(16);
428pub static RMS_BLOCK_DEFAULT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(256);
433pub static FA_SP_GEMMA: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
435pub static MMQ_SK_FORCE: std::sync::atomic::AtomicI8 = std::sync::atomic::AtomicI8::new(-1);
440pub use memra_kv::KV_FP8_FORCE;
443pub(crate) fn rms_block() -> u32 {
444 static V: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
445 *V.get_or_init(|| std::env::var("MEMRA_RMS_BLOCK").ok()
446 .and_then(|v| v.parse().ok())
447 .unwrap_or_else(|| RMS_BLOCK_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)))
448}
449
450pub(crate) fn fa_split_keys(t_kv: usize, n_head_kv: usize) -> usize {
451 static S: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
452 if let Some(forced) = *S.get_or_init(|| {
453 std::env::var("MEMRA_FA_SPLIT").ok().and_then(|v| v.parse().ok())
454 .filter(|&s: &usize| s >= 8 && s % 8 == 0)
455 }) { return forced; }
456 if FA_SP_GEMMA.load(std::sync::atomic::Ordering::Relaxed)
474 && std::env::var("MEMRA_FA_SP16").as_deref() == Ok("1") {
475 return if t_kv <= 8192 { 16 } else if t_kv <= 16384 { 64 } else { 128 };
476 }
477 let big_rig = fa_sm_count() >= 128;
478 if big_rig {
479 let _ = n_head_kv;
480 if t_kv <= 2048 { 16 } else if t_kv <= 16384 { 64 } else { 128 }
481 } else if n_head_kv <= 4 {
482 if t_kv <= 512 { 8 } else if t_kv <= 16384 { 64 } else { 128 }
503 } else {
504 if t_kv <= 8192 { 32 } else if t_kv <= 16384 { 64 } else { 128 }
505 }
506}
507
508fn fa_sm_count() -> i32 {
511 static N: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
512 *N.get_or_init(|| {
513 cudarc::driver::result::init().ok();
514 cudarc::driver::result::device::get(0)
515 .and_then(|d| unsafe { cudarc::driver::result::device::get_attribute(
516 d, cudarc::driver::sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT) })
517 .unwrap_or(82)
518 })
519}
520
521fn fa_hd_suffix(head_dim: usize) -> Result<&'static str, Box<dyn std::error::Error>> {
525 match head_dim {
526 256 => Ok(""),
527 128 => Ok("_hd128"),
528 d => Err(format!("fa_prefill: no kernel stamped for head_dim={d} (only 256/128); \
529 callers must gate to sdpa_naive").into()),
530 }
531}
532
533pub const QT_Q8_0: i32 = 0;
535pub const QT_Q4_K: i32 = 1;
536pub const QT_Q6_K: i32 = 2;
537pub const QT_Q5_K: i32 = 3;
538pub const QT_Q3_K: i32 = 4;
539pub const QT_IQ4_XS: i32 = 5;
540pub const QT_IQ3_S: i32 = 6;
541pub const QT_NVFP4: i32 = 7;
542pub const QT_F8_E4M3: i32 = 10;
548pub const QT_NVFP4_RP: i32 = 9;
551pub const QT_F32: i32 = 8;
553pub const QT_BF16: i32 = 11;
554pub const QT_Q4_0: i32 = 12; pub const QT_Q2_K: i32 = 13;
559
560pub struct Engine {
562 pub gpu: memra_runtime::Gpu,
563 module: Arc<CudaModule>,
564 hybrid: Arc<CudaModule>,
565 qmatvec: Arc<CudaModule>,
566 flash: Arc<CudaModule>,
567 flash_g: std::sync::OnceLock<Arc<CudaModule>>,
571 gemm: Arc<CudaModule>,
572 router: Arc<CudaModule>,
573 sample: Arc<CudaModule>,
575 moe_cache: Mutex<Option<crate::moe_cache::MoeSlotCache>>,
579 moe_cache_layout: Mutex<Option<Vec<usize>>>,
583 capture_keep_on: std::sync::atomic::AtomicBool,
589 verify_exact: std::sync::atomic::AtomicBool,
594 capture_keep: Mutex<Vec<Box<dyn std::any::Any + Send>>>,
595 pub copy_stream: Arc<CudaStream>,
597 #[cfg(memra_cutlass)]
604 cutlass_scratch: Mutex<Option<crate::cutlass_ffi::CutlassScratch>>,
605 fp8_scratch: Mutex<Option<crate::fp8_ffi::Fp8Scratch>>,
609 fa_vf16_scratch: Mutex<Option<CudaSlice<u8>>>,
612 fa_part_pool: Mutex<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
616 fa_part_retired: Mutex<Vec<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>>,
620 fn_cache: Mutex<std::collections::HashMap<String, CudaFunction>>,
622 f16_scratch: Mutex<Option<crate::f16_ffi::F16Scratch>>,
623 argmax_partials: Mutex<Option<(CudaSlice<f32>, CudaSlice<i32>)>>,
628 prime_deqw_ws: Mutex<Option<(CudaSlice<u8>, CudaSlice<u8>)>>,
633 router_stage: Mutex<Option<PinnedStage>>,
637}
638
639fn fa_v2_on() -> bool {
649 std::env::var("MEMRA_FA_V2").map(|v| v != "0").unwrap_or(true)
655}
656
657fn fa_v3_on() -> bool {
665 std::env::var("MEMRA_FA_V3").map(|v| v != "0").unwrap_or(true)
669}
670
671fn fa_v4_mode() -> &'static str {
676 static M: std::sync::OnceLock<String> = std::sync::OnceLock::new();
677 M.get_or_init(|| std::env::var("MEMRA_FA_V4").unwrap_or_default())
678}
679fn fa_v4_on() -> bool { fa_v4_mode() != "0" } pub static FA_SMEM_TKV_DEFAULT: std::sync::atomic::AtomicUsize =
688 std::sync::atomic::AtomicUsize::new(1024);
689pub static FA_V4_MAX_DEFAULT: std::sync::atomic::AtomicUsize =
690 std::sync::atomic::AtomicUsize::new(usize::MAX);
691pub fn fa_v4_at_pub(t_kv: usize) -> bool { fa_v4_at(t_kv) }
692fn fa_v4_at(t_kv: usize) -> bool {
693 static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
694 let mx = *M.get_or_init(|| std::env::var("MEMRA_FA_V4_MAX").ok()
695 .and_then(|v| v.parse().ok())
696 .unwrap_or_else(|| FA_V4_MAX_DEFAULT.load(std::sync::atomic::Ordering::Relaxed)));
697 fa_v4_on() && t_kv < mx
698}
699pub const FA_DEEP_MIN_DEFAULT: usize = 0;
713fn fa_deep_at(t_kv: usize) -> bool {
714 if std::env::var("MEMRA_FA_DEEP").as_deref() == Ok("0") { return false; }
715 let min = std::env::var("MEMRA_FA_DEEP_MIN").ok().and_then(|v| v.parse().ok())
716 .unwrap_or(FA_DEEP_MIN_DEFAULT);
717 t_kv >= min
718}
719pub fn fa_deep_at_pub(t_kv: usize) -> bool { fa_deep_at(t_kv) }
721
722fn fa_v3_active(head_dim: usize) -> bool {
723 fa_v3_on() && head_dim % 128 == 0 && kv_cache_formats() == ("q8_0", "q5_1")
726 && !Engine::kv_fp8_on()
727}
728
729pub fn fa_seqs_eligible(t_kv: usize, head_dim: usize) -> bool {
737 std::env::var("MEMRA_NO_FA_VEC").is_err()
738 && t_kv >= fa_vec_min_tkv()
739 && head_dim == 256
740 && fa_v4_at(t_kv)
741 && !matches!(fa_v4_mode(), "noB3" | "stage")
742 && !Engine::kv_fp8_on()
743}
744pub fn fa_split_keys_pub(t_kv: usize, n_head_kv: usize) -> usize { fa_split_keys(t_kv, n_head_kv) }
746
747struct PinnedStage {
752 ptr: *mut u8,
753 cap: usize,
754}
755unsafe impl Send for PinnedStage {}
756impl PinnedStage {
757 fn new(cap: usize) -> Result<Self, Box<dyn std::error::Error>> {
758 let ptr = unsafe { cudarc::driver::result::malloc_host(cap, 0)? } as *mut u8;
759 Ok(PinnedStage { ptr, cap })
760 }
761}
762impl Drop for PinnedStage {
763 fn drop(&mut self) {
764 let _ = unsafe { cudarc::driver::result::free_host(self.ptr as _) };
765 }
766}
767
768pub const ARGMAX_NB: usize = 256;
771
772pub(crate) use memra_fa3_vl as fa3_vl_raw;
774
775unsafe extern "C" {
776 fn memra_fa3_prefill(q16: *const core::ffi::c_void, k16: *const core::ffi::c_void,
778 v16: *const core::ffi::c_void, o: *mut f32,
779 t: i32, h: i32, hkv: i32, d: i32, scale: f32,
780 stream: *mut core::ffi::c_void) -> i32;
781 pub(crate) fn memra_fa3_vl(q16s: *const *const core::ffi::c_void, k16s: *const *const core::ffi::c_void,
783 v16s: *const *const core::ffi::c_void, os: *const *mut f32,
784 ts: *const i32, b: i32, h: i32, hkv: i32, d: i32, scale: f32,
785 stream: *mut core::ffi::c_void) -> i32;
786}
787
788#[repr(C)]
793#[derive(Clone, Copy)]
794pub struct WPtr8(pub [u64; 8]);
795unsafe impl cudarc::driver::DeviceRepr for WPtr8 {}
796
797#[repr(C)]
802#[derive(Clone, Copy, Default)]
803pub struct GdnSeqVl {
804 pub kb16: u64, pub gcum: u64, pub beta: u64, pub u: u64, pub wb16: u64,
805 pub y: u64, pub ssnap: u64, pub state_in: u64, pub state_out: u64,
806 pub q: u64, pub p: u64, pub o: u64,
807 pub k: u64, pub v: u64, pub g: u64, pub a: u64, pub w: u64,
808 pub t: i32, pub nc: i32,
809}
810unsafe impl cudarc::driver::DeviceRepr for GdnSeqVl {}
811#[repr(C)]
812#[derive(Clone, Copy)]
813pub struct GdnVl8(pub [GdnSeqVl; 8]);
814unsafe impl cudarc::driver::DeviceRepr for GdnVl8 {}
815
816#[repr(C)]
819#[derive(Clone, Copy, Default)]
820pub struct GdnWVl { pub qb16: u64, pub pb16: u64 }
821unsafe impl cudarc::driver::DeviceRepr for GdnWVl {}
822#[repr(C)]
823#[derive(Clone, Copy)]
824pub struct GdnWVl8(pub [GdnWVl; 8]);
825unsafe impl cudarc::driver::DeviceRepr for GdnWVl8 {}
826
827#[repr(C)]
829#[derive(Clone, Copy, Default)]
830pub struct GdnPrepVl {
831 pub qkv: u64, pub conv_state: u64, pub conv_out: u64,
832 pub q_g: u64, pub k_g: u64, pub v_g: u64,
833 pub q_l2: u64, pub k_l2: u64,
834 pub beta_raw: u64, pub alpha: u64, pub beta: u64, pub g_log: u64,
835 pub o: u64, pub z: u64, pub gn: u64, pub gn16: u64,
836 pub kb16: u64,
837 pub qb16: u64,
838 pub t: i32, pub pad: i32,
839}
840unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl {}
841#[repr(C)]
842#[derive(Clone, Copy)]
843pub struct GdnPrepVl8(pub [GdnPrepVl; 8]);
844unsafe impl cudarc::driver::DeviceRepr for GdnPrepVl8 {}
845
846#[repr(C)]
848#[derive(Clone, Copy, Default)]
849pub struct FaSeqVl {
850 pub q: u64, pub k16: u64, pub v16: u64, pub o: u64, pub kf: u64, pub vf: u64,
851 pub t: i32, pub pad: i32,
852}
853unsafe impl cudarc::driver::DeviceRepr for FaSeqVl {}
854#[repr(C)]
855#[derive(Clone, Copy)]
856pub struct FaVl8(pub [FaSeqVl; 8]);
857unsafe impl cudarc::driver::DeviceRepr for FaVl8 {}
858
859#[repr(C)]
861#[derive(Clone, Copy, Default)]
862pub struct AttnPreVl {
863 pub qf: u64, pub kf: u64, pub vf: u64,
864 pub q: u64, pub gate: u64, pub qn: u64, pub kn: u64,
865 pub kc: u64, pub vc: u64,
866 pub t: i32, pub pad: i32,
867}
868unsafe impl cudarc::driver::DeviceRepr for AttnPreVl {}
869#[repr(C)]
870#[derive(Clone, Copy)]
871pub struct AttnPreVl8(pub [AttnPreVl; 8]);
872unsafe impl cudarc::driver::DeviceRepr for AttnPreVl8 {}
873
874pub struct GdnChunkBufs {
877 pub gcum: CudaSlice<f32>,
878 pub a: CudaSlice<f32>,
879 pub p: CudaSlice<f32>,
880 pub u: CudaSlice<f32>,
881 pub w: CudaSlice<f32>,
882 pub kb16: CudaSlice<u8>,
883 pub wb16: CudaSlice<u8>,
884 pub y16: CudaSlice<u8>,
885 pub ssnap16: CudaSlice<u8>,
886 pub qb16: CudaSlice<u8>,
887 pub pb16: CudaSlice<u8>,
888 pub o: CudaSlice<f32>,
889 pub t: usize,
890 pub nc: usize,
891}
892
893#[repr(C)]
895#[derive(Clone, Copy)]
896pub struct F32x8(pub [f32; 8]);
897unsafe impl cudarc::driver::DeviceRepr for F32x8 {}
898
899pub static PRIME_NANOS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
903
904impl Engine {
905 pub fn new(ordinal: usize) -> Result<Self, Box<dyn std::error::Error>> {
906 let gpu = memra_runtime::Gpu::new(ordinal)?;
907 if std::env::var("MEMRA_ARCH_CHECK").as_deref() != Ok("0") {
911 use cudarc::driver::sys::CUdevice_attribute_enum as A;
912 let (maj, min) = cudarc::driver::result::device::get(ordinal as i32)
913 .and_then(|d| unsafe { Ok((
914 cudarc::driver::result::device::get_attribute(d, A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)?,
915 cudarc::driver::result::device::get_attribute(d, A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)?)) })
916 .unwrap_or((0, 0));
917 let built = env!("MEMRA_BUILT_CUDA_ARCH");
918 let ok = matches!((built, maj, min),
919 ("120a", 12, 0) | ("120a", 12, 1) | ("100a", 10, 0) | ("90a", 9, 0) | ("89", 8, 9));
920 if !ok {
921 return Err(format!(
922 "memra was built for sm_{built} but device {ordinal} reports compute \
923 capability {maj}.{min}. Rebuild on this machine (MEMRA_CUDA_ARCH \
924 auto-detects the GPU) or set MEMRA_ARCH_CHECK=0 to bypass.").into());
925 }
926 }
927 unsafe {
932 use cudarc::driver::sys;
933 let dev: sys::CUdevice = ordinal as sys::CUdevice;
934 let mut pool: sys::CUmemoryPool = std::ptr::null_mut();
935 if sys::cuDeviceGetDefaultMemPool(&mut pool, dev) == sys::CUresult::CUDA_SUCCESS {
936 let mut thresh: u64 = u64::MAX;
937 let _ = sys::cuMemPoolSetAttribute(
938 pool,
939 sys::CUmemPool_attribute_enum::CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
940 &mut thresh as *mut u64 as *mut core::ffi::c_void,
941 );
942 }
943 }
944 let module = gpu.ctx.load_module(Ptx::from_binary(FATBIN.to_vec()))?;
945 let hybrid = gpu.ctx.load_module(Ptx::from_binary(HYBRID_FATBIN.to_vec()))?;
946 let qmatvec = gpu.ctx.load_module(Ptx::from_binary(QMATVEC_FATBIN.to_vec()))?;
947 let flash = gpu.ctx.load_module(Ptx::from_binary(flash_fatbin_bytes().to_vec()))?;
948 let gemm = gpu.ctx.load_module(Ptx::from_binary(gemm_fatbin_bytes().into_owned()))?;
949 let router = gpu.ctx.load_module(Ptx::from_binary(ROUTER_FATBIN.to_vec()))?;
950 let sample = gpu.ctx.load_module(Ptx::from_binary(SAMPLE_FATBIN.to_vec()))?;
951 let copy_stream = gpu.ctx.new_stream()?;
952 if std::env::var("MEMRA_EVT").map(|v| v == "1").unwrap_or(false) {
968 } else {
970 unsafe { gpu.ctx.disable_event_tracking(); }
971 }
972 Ok(Self { gpu, module, hybrid, qmatvec, flash, flash_g: std::sync::OnceLock::new(), gemm, router, sample,
973 moe_cache: Mutex::new(None),
974 moe_cache_layout: Mutex::new(None),
975 copy_stream,
976 capture_keep_on: std::sync::atomic::AtomicBool::new(false),
977 verify_exact: std::sync::atomic::AtomicBool::new(false),
978 capture_keep: Mutex::new(Vec::new()),
979 argmax_partials: Mutex::new(None),
980 prime_deqw_ws: Mutex::new(None),
981 router_stage: Mutex::new(None),
982 fp8_scratch: Mutex::new(None),
983 fa_vf16_scratch: Mutex::new(None),
984 fa_part_pool: Mutex::new(None),
985 fa_part_retired: Mutex::new(Vec::new()),
986 fn_cache: Mutex::new(Default::default()),
987 f16_scratch: Mutex::new(None),
988 #[cfg(memra_cutlass)]
989 cutlass_scratch: Mutex::new(None) })
990 }
991
992 pub fn ctx(&self) -> &Arc<CudaContext> { &self.gpu.ctx }
993 pub fn stream(&self) -> Arc<CudaStream> { self.gpu.stream() }
996 pub fn gkv_on() -> bool {
999 memra_kv::gkv_on()
1000 }
1001
1002 pub fn wkv_on() -> bool {
1014 memra_kv::wkv_on()
1015 }
1016
1017 pub fn kv_fp8_on() -> bool {
1023 memra_kv::kv_fp8_on()
1024 }
1025
1026 fn fa_func(&self, name: &str, head_dim: usize) -> CudaFunction {
1029 if head_dim == 512 && Self::gkv_on() { self.func_g(name) } else { self.func(name) }
1030 }
1031
1032 fn func_g(&self, name: &str) -> CudaFunction {
1036 let m = self.flash_g.get_or_init(|| {
1037 self.gpu.ctx.load_module(cudarc::nvrtc::Ptx::from_binary(FLASH_FATBIN_KF8VF8.to_vec()))
1038 .expect("load kf8vf8 flash fatbin (fp8-globals arm)")
1039 });
1040 let key = format!("g:{name}");
1041 if let Some(f) = self.fn_cache.lock().unwrap().get(&key) { return f.clone(); }
1042 let f = match m.load_function(name) {
1043 Ok(f) => f,
1044 Err(_) => self.func(name),
1045 };
1046 self.fn_cache.lock().unwrap().insert(key, f.clone());
1047 f
1048 }
1049
1050 fn func(&self, name: &str) -> CudaFunction {
1051 if let Some(f) = self.fn_cache.lock().unwrap().get(name) { return f.clone(); }
1054 let f = self.module.load_function(name)
1055 .or_else(|_| self.hybrid.load_function(name))
1056 .or_else(|_| self.qmatvec.load_function(name))
1057 .or_else(|_| self.flash.load_function(name))
1058 .or_else(|_| self.gemm.load_function(name))
1059 .or_else(|_| self.router.load_function(name))
1060 .or_else(|_| self.sample.load_function(name))
1061 .unwrap_or_else(|_| panic!("kernel {name} not in any fatbin"));
1062 self.fn_cache.lock().unwrap().insert(name.to_string(), f.clone());
1063 f
1064 }
1065
1066 pub fn scatter_trim_logits(&self, src: &CudaSlice<f32>, d2t: &CudaSlice<u32>,
1069 dst: &mut CudaSlice<f32>, d_vocab: usize, n_vocab: usize)
1070 -> Result<(), Box<dyn std::error::Error>> {
1071 let f1 = self.func("scatter_trim_logits_f32");
1072 let f2 = self.func("scatter_trim_logits_pass2_f32");
1073 let (dv, nv) = (d_vocab as i32, n_vocab as i32);
1074 let cfg1 = LaunchConfig { grid_dim: (256, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1075 let __s_b1 = self.gpu.stream();
1076 let mut b1 = __s_b1.launch_builder(&f1);
1077 b1.arg(src).arg(d2t).arg(&mut *dst).arg(&dv).arg(&nv);
1078 unsafe { b1.launch(cfg1)?; }
1079 let cfg2 = LaunchConfig { grid_dim: (d_vocab.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1080 let __s_b2 = self.gpu.stream();
1081 let mut b2 = __s_b2.launch_builder(&f2);
1082 b2.arg(src).arg(d2t).arg(&mut *dst).arg(&dv);
1083 unsafe { b2.launch(cfg2)?; }
1084 Ok(())
1085 }
1086
1087 #[allow(clippy::too_many_arguments)]
1093 pub fn filter_stats(&self, x: &CudaSlice<f32>, row_stride: usize, rows: &CudaSlice<i32>,
1094 out_th: &mut CudaSlice<f32>, out_z: &mut CudaSlice<f32>,
1095 out_max: &mut CudaSlice<f32>, n: usize, nrow: usize,
1096 temp: f32, top_k: i32, top_p: f32, min_p: f32)
1097 -> Result<(), Box<dyn std::error::Error>> {
1098 let f = self.func("filter_stats_f32");
1099 let (ni, nr, rs) = (n as i32, nrow as i32, row_stride as i64);
1100 let cfg = LaunchConfig { grid_dim: (nrow as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
1101 let __s_b = self.gpu.stream();
1102 let mut b = __s_b.launch_builder(&f);
1103 b.arg(x).arg(&rs).arg(rows).arg(&mut *out_th).arg(&mut *out_z).arg(&mut *out_max)
1104 .arg(&ni).arg(&nr).arg(&temp).arg(&top_k).arg(&top_p).arg(&min_p);
1105 unsafe { b.launch(cfg)?; }
1106 Ok(())
1107 }
1108
1109 #[allow(clippy::too_many_arguments)]
1111 pub fn softmax_gather_filtered(&self, x: &CudaSlice<f32>, row_stride: usize,
1112 ids: &CudaSlice<u32>, rows: &CudaSlice<i32>,
1113 th: &CudaSlice<f32>, z: &CudaSlice<f32>,
1114 out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32)
1115 -> Result<(), Box<dyn std::error::Error>> {
1116 let f = self.func("softmax_gather_filtered_f32");
1117 let (ni, np, rs) = (n as i32, npair as i32, row_stride as i64);
1118 let cfg = LaunchConfig { grid_dim: (npair as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1119 let __s_b = self.gpu.stream();
1120 let mut b = __s_b.launch_builder(&f);
1121 b.arg(x).arg(&rs).arg(ids).arg(rows).arg(th).arg(z).arg(&mut *out).arg(&ni).arg(&np).arg(&temp);
1122 unsafe { b.launch(cfg)?; }
1123 Ok(())
1124 }
1125
1126 #[allow(clippy::too_many_arguments)]
1128 pub fn residual_sample_filtered(&self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize,
1129 temp: f32, seed: u64, stream_pos: u32,
1130 p_stats: (f32, f32, f32), q_stats: (f32, f32, f32),
1131 out_tok: &mut CudaSlice<u32>)
1132 -> Result<(), Box<dyn std::error::Error>> {
1133 let f = self.func("residual_sample_filtered_f32");
1134 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1135 let has_q: i32 = q.is_some() as i32;
1136 let qbuf = q.unwrap_or(p);
1137 let (pm, pth, pz) = p_stats; let (qm, qth, qz) = q_stats;
1138 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
1139 let __s_b = self.gpu.stream();
1140 let mut b = __s_b.launch_builder(&f);
1141 b.arg(p).arg(qbuf).arg(&has_q).arg(&ni).arg(&temp).arg(&slo).arg(&shi).arg(&stream_pos)
1142 .arg(&pm).arg(&pth).arg(&pz).arg(&qm).arg(&qth).arg(&qz).arg(&mut *out_tok);
1143 unsafe { b.launch(cfg)?; }
1144 Ok(())
1145 }
1146
1147 #[allow(clippy::too_many_arguments)]
1149 pub fn gumbel_perturb_filtered(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1150 seed: u64, stream_pos: u32, temp: f32, row_max: f32, th: f32)
1151 -> Result<(), Box<dyn std::error::Error>> {
1152 let f = self.func("gumbel_perturb_filtered_f32");
1153 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1154 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1155 let __s_b = self.gpu.stream();
1156 let mut b = __s_b.launch_builder(&f);
1157 b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp).arg(&row_max).arg(&th);
1158 unsafe { b.launch(cfg)?; }
1159 Ok(())
1160 }
1161
1162 #[allow(clippy::too_many_arguments)]
1166 pub fn penalize_logits(&self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize,
1167 rep: f32, freq: f32, present: f32, n: usize)
1168 -> Result<(), Box<dyn std::error::Error>> {
1169 if n_hist == 0 { return Ok(()); }
1170 let f = self.func("penalize_logits_f32");
1171 let (nh, ni) = (n_hist as i32, n as i32);
1172 let cfg = LaunchConfig { grid_dim: (n_hist.div_ceil(128) as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
1173 let __s_b = self.gpu.stream();
1174 let mut b = __s_b.launch_builder(&f);
1175 b.arg(&mut *x).arg(hist).arg(&nh).arg(&rep).arg(&freq).arg(&present).arg(&ni);
1176 unsafe { b.launch(cfg)?; }
1177 Ok(())
1178 }
1179
1180 #[allow(clippy::too_many_arguments)]
1182 pub fn penalize_logits_rows(&self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize,
1183 rep: f32, freq: f32, present: f32, n: usize, nrow: usize)
1184 -> Result<(), Box<dyn std::error::Error>> {
1185 if n_hist == 0 || nrow == 0 { return Ok(()); }
1186 let f = self.func("penalize_logits_rows_f32");
1187 let (nh, ni, nr) = (n_hist as i32, n as i32, nrow as i32);
1188 let cfg = LaunchConfig { grid_dim: (n_hist.div_ceil(128) as u32, nrow as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
1189 let __s_b = self.gpu.stream();
1190 let mut b = __s_b.launch_builder(&f);
1191 b.arg(&mut *x).arg(hist).arg(&nh).arg(&rep).arg(&freq).arg(&present).arg(&ni).arg(&nr);
1192 unsafe { b.launch(cfg)?; }
1193 Ok(())
1194 }
1195
1196 pub fn wpf_level() -> u32 {
1204 static ON: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
1205 *ON.get_or_init(|| std::env::var("MEMRA_WPF").ok()
1206 .and_then(|v| v.parse().ok()).unwrap_or(1))
1207 }
1208
1209 pub fn set_verify_exact(&self, on: bool) {
1221 self.verify_exact.store(on, std::sync::atomic::Ordering::Relaxed);
1222 }
1223 pub(crate) fn verify_exact_on(&self) -> bool {
1224 self.verify_exact.load(std::sync::atomic::Ordering::Relaxed)
1225 }
1226
1227 pub fn qkv_append_on() -> bool {
1230 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1231 *ON.get_or_init(|| std::env::var("MEMRA_QKV_APPEND").map(|v| v != "0").unwrap_or(true))
1232 }
1233
1234 pub fn pdl_wb_on() -> bool {
1237 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1238 *ON.get_or_init(|| std::env::var("MEMRA_PDL_WB").map(|v| v != "0").unwrap_or(true))
1239 }
1240
1241 pub fn pdl_mmvq_on() -> bool {
1245 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1246 *ON.get_or_init(|| std::env::var("MEMRA_PDL_MMVQ").map(|v| v != "0").unwrap_or(true))
1247 }
1248
1249 pub fn pdl_on() -> bool {
1250 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1251 *ON.get_or_init(|| std::env::var("MEMRA_PDL").map(|v| v != "0").unwrap_or(true))
1252 }
1253
1254 fn q40_mr1_on() -> bool {
1260 static Q40MR: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
1261 match *Q40MR.get_or_init(|| std::env::var("MEMRA_Q40_MR").ok()
1262 .and_then(|v| v.parse().ok())) {
1263 Some(v) => v == 1,
1264 None => crate::FUSED_MR1_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1265 }
1266 }
1267
1268 fn pdl_func_flash(&self, g: bool, name: &'static str)
1273 -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1274 use cudarc::driver::sys as cu;
1275 static MODS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool), usize>>> =
1282 std::sync::Mutex::new(None);
1283 static FNS: std::sync::Mutex<Option<std::collections::HashMap<(usize, bool, &'static str), usize>>> =
1284 std::sync::Mutex::new(None);
1285 let ctx_key = self.ctx().cu_ctx() as usize;
1286 if let Some(&f) = FNS.lock().unwrap().get_or_insert_with(Default::default)
1287 .get(&(ctx_key, g, name)) { return Ok(f as cu::CUfunction); }
1288 let module = {
1289 let mut mods = MODS.lock().unwrap();
1290 let map = mods.get_or_insert_with(Default::default);
1291 match map.get(&(ctx_key, g)) {
1292 Some(&m) => m,
1293 None => {
1294 let m = self.pdl_load_module_in_ctx(
1295 if g { FLASH_FATBIN_KF8VF8 } else { FLASH_FATBIN })?;
1296 map.insert((ctx_key, g), m);
1297 m
1298 }
1299 }
1300 };
1301 let cname = std::ffi::CString::new(name)?;
1302 let mut f: cu::CUfunction = std::ptr::null_mut();
1303 let r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1304 if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("pdl_func_flash {name} (g={g}): {r:?}").into()); }
1305 FNS.lock().unwrap().get_or_insert_with(Default::default)
1306 .insert((ctx_key, g, name), f as usize);
1307 Ok(f)
1308 }
1309
1310 fn pdl_load_module_in_ctx(&self, bytes: &[u8]) -> Result<usize, Box<dyn std::error::Error>> {
1315 use cudarc::driver::sys as cu;
1316 let mut prev: cu::CUcontext = std::ptr::null_mut();
1317 unsafe { cu::cuCtxGetCurrent(&mut prev).result()?; }
1318 self.ctx().bind_to_thread()?;
1319 let mut m: cu::CUmodule = std::ptr::null_mut();
1320 let r = unsafe { cu::cuModuleLoadData(&mut m, bytes.as_ptr() as *const std::ffi::c_void) };
1321 let restore = if prev.is_null() { cu::CUresult::CUDA_SUCCESS }
1322 else { unsafe { cu::cuCtxSetCurrent(prev) } };
1323 if r != cu::CUresult::CUDA_SUCCESS {
1324 return Err(format!("pdl module load: {r:?}").into());
1325 }
1326 if restore != cu::CUresult::CUDA_SUCCESS {
1327 return Err(format!("pdl module load: ctx restore {restore:?}").into());
1328 }
1329 Ok(m as usize)
1330 }
1331
1332 fn pdl_func(&self, name: &'static str) -> Result<cudarc::driver::sys::CUfunction, Box<dyn std::error::Error>> {
1333 use cudarc::driver::sys as cu;
1334 static MODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1337 std::sync::Mutex::new(None);
1338 static QMODULES: std::sync::Mutex<Option<std::collections::HashMap<usize, usize>>> =
1341 std::sync::Mutex::new(None);
1342 static FNS: std::sync::Mutex<Option<std::collections::HashMap<(usize, &'static str), usize>>> =
1343 std::sync::Mutex::new(None);
1344 let ctx_key = self.ctx().cu_ctx() as usize;
1345 if let Some(&f) = FNS.lock().unwrap().get_or_insert_with(Default::default)
1346 .get(&(ctx_key, name)) { return Ok(f as cu::CUfunction); }
1347 let module = {
1348 let mut mods = MODULES.lock().unwrap();
1349 let map = mods.get_or_insert_with(Default::default);
1350 match map.get(&ctx_key) {
1351 Some(&m) => m,
1352 None => {
1353 let m = self.pdl_load_module_in_ctx(FATBIN)?;
1354 map.insert(ctx_key, m);
1355 m
1356 }
1357 }
1358 };
1359 let cname = std::ffi::CString::new(name)?;
1360 let mut f: cu::CUfunction = std::ptr::null_mut();
1361 let mut r = unsafe { cu::cuModuleGetFunction(&mut f, module as cu::CUmodule, cname.as_ptr()) };
1362 if r == cu::CUresult::CUDA_ERROR_NOT_FOUND {
1363 let qmodule = {
1364 let mut mods = QMODULES.lock().unwrap();
1365 let map = mods.get_or_insert_with(Default::default);
1366 match map.get(&ctx_key) {
1367 Some(&m) => m,
1368 None => {
1369 let m = self.pdl_load_module_in_ctx(QMATVEC_FATBIN)?;
1370 map.insert(ctx_key, m);
1371 m
1372 }
1373 }
1374 };
1375 r = unsafe { cu::cuModuleGetFunction(&mut f, qmodule as cu::CUmodule, cname.as_ptr()) };
1376 }
1377 if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("pdl_func {name}: {r:?}").into()); }
1378 FNS.lock().unwrap().get_or_insert_with(Default::default)
1379 .insert((ctx_key, name), f as usize);
1380 Ok(f)
1381 }
1382
1383 unsafe fn launch_pdl_flash(&self, g: bool, name: &'static str, grid: (u32, u32, u32),
1395 block: (u32, u32, u32), smem: u32,
1396 params: &mut [*mut std::ffi::c_void])
1397 -> Result<(), Box<dyn std::error::Error>> {
1398 use cudarc::driver::sys as cu;
1399 let f = self.pdl_func_flash(g, name)?;
1400 if smem > 0 {
1401 let r = unsafe { cu::cuFuncSetAttribute(f,
1403 cu::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
1404 smem as i32) };
1405 if r != cu::CUresult::CUDA_SUCCESS {
1406 return Err(format!("pdl smem attr {name}: {r:?}").into());
1407 }
1408 }
1409 let mut attr = cu::CUlaunchAttribute {
1410 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1411 pad: [0; 4],
1412 value: cu::CUlaunchAttributeValue { programmaticStreamSerializationAllowed: 1 },
1413 };
1414 let cfg = cu::CUlaunchConfig {
1415 gridDimX: grid.0, gridDimY: grid.1, gridDimZ: grid.2,
1416 blockDimX: block.0, blockDimY: block.1, blockDimZ: block.2,
1417 sharedMemBytes: smem, hStream: self.gpu.stream().cu_stream(),
1418 attrs: &mut attr, numAttrs: 1,
1419 };
1420 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1421 if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("launch_pdl_flash {name}: {r:?}").into()); }
1422 Ok(())
1423 }
1424
1425 unsafe fn launch_pdl(&self, name: &'static str, grid: (u32, u32, u32), block: (u32, u32, u32),
1426 params: &mut [*mut std::ffi::c_void])
1427 -> Result<(), Box<dyn std::error::Error>> {
1428 use cudarc::driver::sys as cu;
1429 let f = self.pdl_func(name)?;
1430 let mut attr = cu::CUlaunchAttribute {
1431 id: cu::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
1432 pad: [0; 4],
1433 value: cu::CUlaunchAttributeValue { programmaticStreamSerializationAllowed: 1 },
1434 };
1435 let cfg = cu::CUlaunchConfig {
1436 gridDimX: grid.0, gridDimY: grid.1, gridDimZ: grid.2,
1437 blockDimX: block.0, blockDimY: block.1, blockDimZ: block.2,
1438 sharedMemBytes: 0, hStream: self.gpu.stream().cu_stream(),
1439 attrs: &mut attr, numAttrs: 1,
1440 };
1441 let r = unsafe { cu::cuLaunchKernelEx(&cfg, f, params.as_mut_ptr(), std::ptr::null_mut()) };
1442 if r != cu::CUresult::CUDA_SUCCESS { return Err(format!("launch_pdl {name}: {r:?}").into()); }
1443 Ok(())
1444 }
1445
1446 pub fn prefetch_weight_l2(&self, w: &crate::model::GpuTensor)
1449 -> Result<(), Box<dyn std::error::Error>> {
1450 if let crate::model::GpuTensor::Quant { bytes, rp4, .. } = w {
1451 let p = rp4.as_ref().unwrap_or(bytes);
1452 self.prefetch_l2(p, p.len())?;
1453 }
1454 Ok(())
1455 }
1456
1457 pub fn gather_row_bf16(&self, table: &CudaSlice<u8>, tok: &CudaSlice<u32>, idx: usize,
1460 dst: &mut CudaSlice<f32>, ncols: usize)
1461 -> Result<(), Box<dyn std::error::Error>> {
1462 let f = self.func("gather_row_bf16_f32");
1463 let cfg = LaunchConfig { grid_dim: (ncols.div_ceil(256) as u32, 1, 1),
1464 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1465 let (nc, ix) = (ncols as i32, idx as i32);
1466 let __s_b = self.gpu.stream();
1467 let mut b = __s_b.launch_builder(&f);
1468 b.arg(table).arg(tok).arg(&ix).arg(dst).arg(&nc);
1469 unsafe { b.launch(cfg)?; }
1470 Ok(())
1471 }
1472
1473 pub fn add_row_inplace(&self, logits: &mut CudaSlice<f32>, bias: &CudaSlice<f32>,
1475 n: usize, row_off: usize)
1476 -> Result<(), Box<dyn std::error::Error>> {
1477 let f = self.func("add_row_inplace_f32");
1478 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1),
1479 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1480 let (ni, off) = (n as i32, row_off as i64);
1481 let __s_b = self.gpu.stream();
1482 let mut b = __s_b.launch_builder(&f);
1483 b.arg(logits).arg(bias).arg(&ni).arg(&off);
1484 unsafe { b.launch(cfg)?; }
1485 Ok(())
1486 }
1487
1488 pub fn prefetch_l2(&self, p: &CudaSlice<u8>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1490 let f = self.func("prefetch_l2_bytes");
1491 let lines = n.div_ceil(128);
1492 let ni = n as i64;
1493 let cfg = LaunchConfig { grid_dim: (lines.div_ceil(256) as u32, 1, 1),
1494 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1495 let __s_b = self.gpu.stream();
1496 let mut b = __s_b.launch_builder(&f);
1497 b.arg(p).arg(&ni);
1498 unsafe { b.launch(cfg)?; }
1499 Ok(())
1500 }
1501
1502 pub fn router_gemv(&self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize,
1505 n_experts: usize, t: usize)
1506 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1507 let w8 = match std::env::var("MEMRA_ROUTER_V2").as_deref() {
1513 Ok("0") => false,
1514 Ok(_) => true,
1515 Err(_) => ROUTER_W8_DEFAULT.load(std::sync::atomic::Ordering::Relaxed),
1516 };
1517 let batch = w8 && t >= ROUTER_BATCH_MIN_T && router_batch_on();
1527 self.router_gemv_form(w, x, n_embd, n_experts, t, w8, batch)
1528 }
1529
1530 pub fn router_gemv_form(&self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize,
1533 n_experts: usize, t: usize, w8: bool, batch: bool)
1534 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1535 debug_assert!(!batch || w8, "batch twin exists for the w8 form only");
1536 let mut y = self.alloc_uninit::<f32>(t * n_experts)?;
1537 let f = if batch { self.func("router_gemv_f32_w8_batch") }
1538 else if w8 { self.func("router_gemv_f32_w8") }
1539 else { self.func("router_gemv_f32") };
1540 let (ne, nx, ti) = (n_embd as i32, n_experts as i32, t as i32);
1541 let cfg = if batch {
1542 LaunchConfig { grid_dim: (n_experts.div_ceil(8) as u32, t.div_ceil(8) as u32, 1),
1543 block_dim: (32, 8, 1), shared_mem_bytes: 0 }
1544 } else {
1545 LaunchConfig { grid_dim: (n_experts as u32, t as u32, 1),
1546 block_dim: (32, if w8 { 8 } else { 1 }, 1), shared_mem_bytes: 0 }
1547 };
1548 let __s_b = self.gpu.stream();
1549 let mut b = __s_b.launch_builder(&f);
1550 b.arg(w).arg(x).arg(&mut y).arg(&ne).arg(&nx).arg(&ti);
1551 unsafe { b.launch(cfg)?; }
1552 Ok(y)
1553 }
1554
1555 pub fn rows_permute(&self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>, nrows: usize,
1557 ncols: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1558 let mut dst = self.alloc_uninit::<f32>(nrows * ncols)?;
1559 let f = self.func("rows_permute_f32");
1560 let (nc, nr) = (ncols as i32, nrows as i32);
1561 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (256, 1, 1),
1562 shared_mem_bytes: 0 };
1563 let __s_b = self.gpu.stream();
1564 let mut b = __s_b.launch_builder(&f);
1565 b.arg(src).arg(idx).arg(&mut dst).arg(&nc).arg(&nr);
1566 unsafe { b.launch(cfg)?; }
1567 Ok(dst)
1568 }
1569
1570 pub fn sigmoid_dot_rows(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, n_embd: usize,
1575 t: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1576 static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1579 if *OFF.get_or_init(|| std::env::var("MEMRA_SHEXP_DOT").as_deref() == Ok("0")) {
1580 let gs = self.linear(x, w, t, n_embd, 1)?;
1581 let mut g = self.uninit(t)?;
1582 self.sigmoid(&gs, &mut g, t)?;
1583 return Ok(g);
1584 }
1585 let mut g = self.alloc_uninit::<f32>(t)?;
1591 let f = self.func("sigmoid_dot_rows_f32");
1592 let (ne, ti) = (n_embd as i32, t as i32);
1593 let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (32, 8, 1),
1594 shared_mem_bytes: 0 };
1595 let __s_b = self.gpu.stream();
1596 let mut b = __s_b.launch_builder(&f);
1597 b.arg(x).arg(w).arg(&mut g).arg(&ne).arg(&ti);
1598 unsafe { b.launch(cfg)?; }
1599 Ok(g)
1600 }
1601
1602 pub fn spec_rollback_stream(&self, len_ptrs: &CudaSlice<u64>, pos_start: &CudaSlice<i32>,
1604 acc: &CudaSlice<u32>, base: usize, n_rows: usize)
1605 -> Result<(), Box<dyn std::error::Error>> {
1606 let f = self.func("spec_rollback_stream");
1607 let (b, nr) = (base as i32, n_rows as i32);
1608 let cfg = LaunchConfig { grid_dim: (n_rows.div_ceil(64) as u32, 1, 1),
1609 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
1610 let __s_bl = self.gpu.stream();
1611 let mut bl = __s_bl.launch_builder(&f);
1612 bl.arg(len_ptrs).arg(pos_start).arg(acc).arg(&b).arg(&nr);
1613 unsafe { bl.launch(cfg)?; }
1614 Ok(())
1615 }
1616
1617 pub fn plain_tok_ring(&self, vam: &CudaSlice<u32>, pos_start: &CudaSlice<i32>,
1619 base: usize, ring: &mut CudaSlice<u32>)
1620 -> Result<(), Box<dyn std::error::Error>> {
1621 let f = self.func("plain_tok_ring");
1622 let (b, cap) = (base as i32, ring.len() as i32);
1623 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1624 let __s_bl = self.gpu.stream();
1625 let mut bl = __s_bl.launch_builder(&f);
1626 bl.arg(vam).arg(pos_start).arg(&b).arg(&mut *ring).arg(&cap);
1627 unsafe { bl.launch(cfg)?; }
1628 Ok(())
1629 }
1630
1631 pub fn spec_ring_commit(&self, vtok: &CudaSlice<u32>, acc: &CudaSlice<u32>,
1633 brk: &CudaSlice<u32>, ring: &mut CudaSlice<u32>,
1634 pend: &mut CudaSlice<u32>)
1635 -> Result<(), Box<dyn std::error::Error>> {
1636 let f = self.func("spec_ring_commit");
1637 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1638 let __s_b = self.gpu.stream();
1639 let mut b = __s_b.launch_builder(&f);
1640 b.arg(vtok).arg(acc).arg(brk).arg(ring).arg(pend);
1641 unsafe { b.launch(cfg)?; }
1642 Ok(())
1643 }
1644 pub fn i32_copy_add(&self, src: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, delta: i32)
1645 -> Result<(), Box<dyn std::error::Error>> {
1646 let f = self.func("i32_copy_add");
1647 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1648 let __s_b = self.gpu.stream();
1649 let mut b = __s_b.launch_builder(&f);
1650 b.arg(src).arg(dst).arg(&delta);
1651 unsafe { b.launch(cfg)?; }
1652 Ok(())
1653 }
1654 pub fn u32_copy(&self, src: &CudaSlice<u32>, dst: &mut CudaSlice<u32>)
1655 -> Result<(), Box<dyn std::error::Error>> {
1656 let f = self.func("u32_copy");
1657 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1658 let __s_b = self.gpu.stream();
1659 let mut b = __s_b.launch_builder(&f);
1660 b.arg(src).arg(dst);
1661 unsafe { b.launch(cfg)?; }
1662 Ok(())
1663 }
1664
1665 pub fn spec_adapt_k(&self, acc: &CudaSlice<u32>, brk: &mut CudaSlice<u32>,
1669 floor: usize, cap: usize)
1670 -> Result<(), Box<dyn std::error::Error>> {
1671 let f = self.func("spec_adapt_k");
1672 let (fl, cp) = (floor as i32, cap as i32);
1673 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1674 let __s_b = self.gpu.stream();
1675 let mut b = __s_b.launch_builder(&f);
1676 b.arg(acc).arg(brk).arg(&fl).arg(&cp);
1677 unsafe { b.launch(cfg)?; }
1678 Ok(())
1679 }
1680
1681 pub fn spec_accept_greedy_dc(&self, preds: &CudaSlice<u32>, vtok: &CudaSlice<u32>,
1683 last_pred: &CudaSlice<u32>, brk: &CudaSlice<u32>,
1684 out: &mut CudaSlice<u32>)
1685 -> Result<(), Box<dyn std::error::Error>> {
1686 let f = self.func("spec_accept_greedy_dc");
1687 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1688 let __s_b = self.gpu.stream();
1689 let mut b = __s_b.launch_builder(&f);
1690 b.arg(preds).arg(vtok).arg(last_pred).arg(brk).arg(out);
1691 unsafe { b.launch(cfg)?; }
1692 Ok(())
1693 }
1694
1695 pub fn pos_iota(&self, pos0: &CudaSlice<i32>, out: &mut CudaSlice<i32>, t: usize)
1697 -> Result<(), Box<dyn std::error::Error>> {
1698 let f = self.func("pos_iota_i32");
1699 let ti = t as i32;
1700 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (t.max(1) as u32, 1, 1),
1701 shared_mem_bytes: 0 };
1702 let __s_b = self.gpu.stream();
1703 let mut b = __s_b.launch_builder(&f);
1704 b.arg(pos0).arg(out).arg(&ti);
1705 unsafe { b.launch(cfg)?; }
1706 Ok(())
1707 }
1708 #[allow(clippy::too_many_arguments)]
1709 pub fn append_kv_quantized_rows_dc(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
1710 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
1711 t0_dev: &CudaSlice<i32>, t: usize,
1712 kv_dim_k: usize, kv_dim_v: usize,
1713 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
1714 -> Result<(), Box<dyn std::error::Error>> {
1715 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_rows_dc") }
1716 else { self.func("append_quantize_kv_q8_0_q5_1_rows_dc") };
1717 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
1718 let cfg = LaunchConfig { grid_dim: (nblk, t as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1719 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
1720 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
1721 let __s_b = self.gpu.stream();
1722 let mut b = __s_b.launch_builder(&f);
1723 b.arg(k_rows).arg(v_rows).arg(kc).arg(vc).arg(t0_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
1724 unsafe { b.launch(cfg)?; }
1725 Ok(())
1726 }
1727
1728 #[allow(clippy::too_many_arguments)]
1731 pub fn append_kv_quantized_row_dc_inc(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
1732 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
1733 t0_dev: &mut CudaSlice<i32>,
1734 kv_dim_k: usize, kv_dim_v: usize,
1735 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
1736 -> Result<(), Box<dyn std::error::Error>> {
1737 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_dc_inc") }
1738 else { self.func("append_quantize_kv_q8_0_q5_1_dc_inc") };
1739 let nthreads = ((kv_dim_k.max(kv_dim_v) / 32) * 32).min(1024) as u32;
1740 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (nthreads, 1, 1),
1741 shared_mem_bytes: 0 };
1742 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
1743 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
1744 let __s_b = self.gpu.stream();
1745 let mut b = __s_b.launch_builder(&f);
1746 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(t0_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
1747 unsafe { b.launch(cfg)?; }
1748 Ok(())
1749 }
1750
1751 pub fn pack_tok_p(&self, tok: &CudaSlice<u32>, p: &CudaSlice<f32>, out: &mut CudaSlice<u32>,
1753 slot: usize) -> Result<(), Box<dyn std::error::Error>> {
1754 let f = self.func("pack_tok_p");
1755 let sl = slot as i32;
1756 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1757 let __s_b = self.gpu.stream();
1758 let mut b = __s_b.launch_builder(&f);
1759 b.arg(tok).arg(p).arg(out).arg(&sl);
1760 unsafe { b.launch(cfg)?; }
1761 Ok(())
1762 }
1763 pub fn tok_map_u32(&self, tok: &mut CudaSlice<u32>, map: &CudaSlice<u32>)
1764 -> Result<(), Box<dyn std::error::Error>> {
1765 let f = self.func("tok_map_u32");
1766 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1767 let __s_b = self.gpu.stream();
1768 let mut b = __s_b.launch_builder(&f);
1769 b.arg(tok).arg(map);
1770 unsafe { b.launch(cfg)?; }
1771 Ok(())
1772 }
1773
1774 #[allow(clippy::too_many_arguments)]
1776 pub fn spec_assemble_verify(&self, tokp: &CudaSlice<u32>, pend: &CudaSlice<u32>,
1777 d2t: Option<&CudaSlice<u32>>, vtok: &mut CudaSlice<u32>,
1778 brk: &mut CudaSlice<u32>, p_min: f32, k: usize, pmin0: bool)
1779 -> Result<(), Box<dyn std::error::Error>> {
1780 let f = self.func("spec_assemble_verify");
1781 let (ki, pm) = (k as i32, if pmin0 { 1i32 } else { 0i32 });
1782 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1783 let __s_b = self.gpu.stream();
1784 let mut b = __s_b.launch_builder(&f);
1785 match d2t {
1786 Some(m) => { b.arg(tokp).arg(pend).arg(m).arg(vtok).arg(brk).arg(&p_min).arg(&ki).arg(&pm);
1787 unsafe { b.launch(cfg)?; } }
1788 None => { let null: u64 = 0;
1789 b.arg(tokp).arg(pend).arg(&null).arg(vtok).arg(brk).arg(&p_min).arg(&ki).arg(&pm);
1790 unsafe { b.launch(cfg)?; } }
1791 }
1792 Ok(())
1793 }
1794
1795 #[allow(clippy::too_many_arguments)]
1797 pub fn ssm_conv_ring_rebuild_dc(&self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>,
1798 conv_state: &mut CudaSlice<f32>, conv_dim: usize,
1799 acc: &CudaSlice<u32>, base: usize, t_v: usize, d_conv: usize)
1800 -> Result<(), Box<dyn std::error::Error>> {
1801 let f = self.func("ssm_conv_ring_rebuild_f32_dc");
1802 let n = conv_dim * (d_conv - 1);
1803 let cfg = LaunchConfig::for_num_elems(n as u32);
1804 let (cd, b0, tv, dc) = (conv_dim as i32, base as i32, t_v as i32, d_conv as i32);
1805 let __s_b = self.gpu.stream();
1806 let mut b = __s_b.launch_builder(&f);
1807 b.arg(qkv_tm).arg(ring_old).arg(conv_state).arg(&cd).arg(acc).arg(&b0).arg(&tv).arg(&dc);
1808 unsafe { b.launch(cfg)?; }
1809 Ok(())
1810 }
1811 #[allow(clippy::too_many_arguments)]
1812 pub fn gdn_scan_s128_dc(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
1813 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
1814 state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
1815 n_head: usize, acc: &CudaSlice<u32>, base: usize, t_v: usize,
1816 scale: f32)
1817 -> Result<(), Box<dyn std::error::Error>> {
1818 let f = self.func("gdn_scan_s128_dc");
1819 const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
1820 let cfg = LaunchConfig {
1821 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
1822 block_dim: (WARP, COLS_PER_BLOCK, 1),
1823 shared_mem_bytes: 0,
1824 };
1825 let (h, b0, tv) = (n_head as i32, base as i32, t_v as i32);
1826 let __s_b = self.gpu.stream();
1827 let mut b = __s_b.launch_builder(&f);
1828 b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o)
1829 .arg(&h).arg(acc).arg(&b0).arg(&tv).arg(&scale);
1830 unsafe { b.launch(cfg)?; }
1831 Ok(())
1832 }
1833
1834 pub fn spec_rollback_kv(&self, len_ptrs: &CudaSlice<u64>, saved: &CudaSlice<i32>,
1836 acc: &CudaSlice<u32>, base: usize, n_layer: usize)
1837 -> Result<(), Box<dyn std::error::Error>> {
1838 let f = self.func("spec_rollback_kv");
1839 let (b, nl) = (base as i32, n_layer as i32);
1840 let cfg = LaunchConfig { grid_dim: (n_layer.div_ceil(64) as u32, 1, 1),
1841 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
1842 let __s_bl = self.gpu.stream();
1843 let mut bl = __s_bl.launch_builder(&f);
1844 bl.arg(len_ptrs).arg(saved).arg(acc).arg(&b).arg(&nl);
1845 unsafe { bl.launch(cfg)?; }
1846 Ok(())
1847 }
1848
1849 pub fn spec_seed_gather(&self, vx: &CudaSlice<f32>, fill_prev: &CudaSlice<f32>,
1852 acc: &CudaSlice<u32>, h_seed: &mut CudaSlice<f32>,
1853 base: usize, n_embd: usize)
1854 -> Result<(), Box<dyn std::error::Error>> {
1855 let f = self.func("spec_seed_gather");
1856 let (b, ne) = (base as i32, n_embd as i32);
1857 let cfg = LaunchConfig { grid_dim: (n_embd.div_ceil(256) as u32, 1, 1),
1858 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1859 let __s_bl = self.gpu.stream();
1860 let mut bl = __s_bl.launch_builder(&f);
1861 bl.arg(vx).arg(fill_prev).arg(acc).arg(h_seed).arg(&b).arg(&ne);
1862 unsafe { bl.launch(cfg)?; }
1863 Ok(())
1864 }
1865
1866
1867 pub fn spec_accept_greedy(&self, preds: &CudaSlice<u32>, draft: &CudaSlice<u32>,
1869 last_pred: u32, base: usize, k_round: usize,
1870 out: &mut CudaSlice<u32>)
1871 -> Result<(), Box<dyn std::error::Error>> {
1872 let f = self.func("spec_accept_greedy");
1873 let (b, k) = (base as i32, k_round as i32);
1874 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
1875 let __s_bl = self.gpu.stream();
1876 let mut bl = __s_bl.launch_builder(&f);
1877 bl.arg(preds).arg(draft).arg(&last_pred).arg(&b).arg(&k).arg(out);
1878 unsafe { bl.launch(cfg)?; }
1879 Ok(())
1880 }
1881
1882 pub fn gumbel_perturb(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1889 seed: u64, stream_pos: u32, temp: f32)
1890 -> Result<(), Box<dyn std::error::Error>> {
1891 let f = self.func("gumbel_perturb_f32");
1892 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1893 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1894 let __s_b = self.gpu.stream();
1895 let mut b = __s_b.launch_builder(&f);
1896 b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp);
1897 unsafe { b.launch(cfg)?; }
1898 Ok(())
1899 }
1900
1901 pub fn mask_logits_col(&self, logits: &mut CudaSlice<f32>, mask: &CudaSlice<u32>,
1909 col: usize, n: usize, mask_words: usize)
1910 -> Result<(), Box<dyn std::error::Error>> {
1911 let f = self.func("mask_logits_f32");
1912 let (ci, ni, mw) = (col as i32, n as i32, mask_words as i32);
1913 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256).min(1024) as u32, 1, 1),
1914 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1915 let __s_b = self.gpu.stream();
1916 let mut b = __s_b.launch_builder(&f);
1917 b.arg(&mut *logits).arg(mask).arg(&ci).arg(&ni).arg(&mw);
1918 unsafe { b.launch(cfg)?; }
1919 Ok(())
1920 }
1921
1922 pub fn gumbel_perturb_col(&self, x: &CudaSlice<f32>, col: usize, y: &mut CudaSlice<f32>,
1929 n: usize, seed: u64, stream_pos: u32, temp: f32)
1930 -> Result<(), Box<dyn std::error::Error>> {
1931 let f = self.func("gumbel_perturb_f32");
1932 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1933 let col_view = x.slice(col * n..(col + 1) * n);
1934 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1935 let __s_b = self.gpu.stream();
1936 let mut b = __s_b.launch_builder(&f);
1937 b.arg(&col_view).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(&stream_pos).arg(&temp);
1938 unsafe { b.launch(cfg)?; }
1939 Ok(())
1940 }
1941
1942 pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn std::error::Error>> {
1947 let f = self.func("memra_sctr_inc");
1948 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
1949 let __s_b = self.gpu.stream();
1950 let mut b = __s_b.launch_builder(&f);
1951 b.arg(&mut *ctr);
1952 unsafe { b.launch(cfg)?; }
1953 Ok(())
1954 }
1955
1956 pub fn gumbel_perturb_ctr(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize,
1961 seed: u64, ctr: &CudaSlice<u32>, temp: f32)
1962 -> Result<(), Box<dyn std::error::Error>> {
1963 let f = self.func("gumbel_perturb_ctr_f32");
1964 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1965 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1966 let __s_b = self.gpu.stream();
1967 let mut b = __s_b.launch_builder(&f);
1968 b.arg(x).arg(&mut *y).arg(&ni).arg(&slo).arg(&shi).arg(ctr).arg(&temp);
1969 unsafe { b.launch(cfg)?; }
1970 Ok(())
1971 }
1972
1973 pub fn softmax_gather(&self, x: &CudaSlice<f32>, row_stride: usize,
1977 ids: &CudaSlice<u32>, rows: &CudaSlice<i32>,
1978 out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32)
1979 -> Result<(), Box<dyn std::error::Error>> {
1980 let f = self.func("softmax_gather_f32");
1981 let (ni, rs) = (n as i32, row_stride as i64);
1982 let np = npair as i32;
1983 let cfg = LaunchConfig { grid_dim: (npair as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
1984 let __s_b = self.gpu.stream();
1985 let mut b = __s_b.launch_builder(&f);
1986 b.arg(x).arg(&rs).arg(ids).arg(rows).arg(&mut *out).arg(&ni).arg(&np).arg(&temp);
1987 unsafe { b.launch(cfg)?; }
1988 Ok(())
1989 }
1990
1991 pub fn residual_sample(&self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize,
1995 temp: f32, seed: u64, stream_pos: u32,
1996 out_tok: &mut CudaSlice<u32>)
1997 -> Result<(), Box<dyn std::error::Error>> {
1998 let f = self.func("residual_sample_f32");
1999 let (ni, slo, shi) = (n as i32, (seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
2000 let nth = 1024u32;
2001 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (nth, 1, 1), shared_mem_bytes: 0 };
2002 let has_q: i32 = q.is_some() as i32;
2003 let qbuf = q.unwrap_or(p); let __s_b = self.gpu.stream();
2005 let mut b = __s_b.launch_builder(&f);
2006 b.arg(p).arg(qbuf).arg(&has_q).arg(&ni).arg(&temp).arg(&slo).arg(&shi).arg(&stream_pos)
2007 .arg(&mut *out_tok);
2008 unsafe { b.launch(cfg)?; }
2009 Ok(())
2010 }
2011
2012 pub fn with_moe_cache<R>(&self, max_block_bytes: usize,
2017 f: impl FnOnce(&mut crate::moe_cache::MoeSlotCache, &Engine) -> Result<R, Box<dyn std::error::Error>>)
2018 -> Result<R, Box<dyn std::error::Error>> {
2019 let mut guard = self.moe_cache.lock().unwrap();
2020 if guard.is_none() {
2021 *guard = Some(crate::moe_cache::MoeSlotCache::new(self, max_block_bytes)?);
2022 }
2023 let cache = guard.as_mut().unwrap();
2024 f(cache, self)
2025 }
2026
2027 pub fn freeze_moe_cache(&self) {
2030 if let Some(cache) = self.moe_cache.lock().unwrap().as_mut() {
2031 cache.freeze();
2032 }
2033 }
2034
2035 pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>> {
2038 self.moe_cache
2039 .lock()
2040 .unwrap()
2041 .as_ref()
2042 .map(crate::moe_cache::MoeSlotCache::export_residency)
2043 }
2044
2045 pub(crate) fn moe_cache_frozen(&self) -> bool {
2046 self.moe_cache
2047 .lock()
2048 .unwrap()
2049 .as_ref()
2050 .is_some_and(crate::moe_cache::MoeSlotCache::is_frozen)
2051 }
2052
2053 pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool {
2060 crate::cpu_experts::configured()
2061 && self.moe_cache_frozen()
2062 && std::env::var("MEMRA_CPU_EXPERT_BATCHED_PRIME").as_deref() != Ok("1")
2063 }
2064
2065 pub(crate) fn configure_moe_cache_layout(&self, block_bytes: Vec<usize>) {
2067 assert!(
2068 self.moe_cache.lock().unwrap().is_none(),
2069 "MoE cache layout configured after cache construction"
2070 );
2071 *self.moe_cache_layout.lock().unwrap() = Some(block_bytes);
2072 }
2073
2074 pub(crate) fn moe_cache_layout(&self) -> Option<Vec<usize>> {
2075 self.moe_cache_layout.lock().unwrap().clone()
2076 }
2077
2078 pub fn moe_cache_enabled() -> bool {
2080 std::env::var("MEMRA_MOE_CACHE").as_deref() != Ok("0")
2081 }
2082
2083 pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)> {
2086 let guard = self.moe_cache.lock().unwrap();
2087 guard.as_ref() .map(|c| (c.hits, c.misses, c.staged_bytes, c.n_slots()))
2088 }
2089
2090 pub fn cpu_expert_stats(
2094 &self,
2095 ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)> {
2096 crate::cpu_experts::configured().then(crate::cpu_experts::stats)
2097 }
2098
2099 pub fn cpu_expert_predictor_stats(&self) -> (u64, u64) {
2102 crate::cpu_experts::predictor_stats()
2103 }
2104
2105 pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64> {
2106 crate::cpu_experts::configured().then(crate::cpu_experts::exposed_wait_ns)
2107 }
2108
2109 pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)> {
2112 crate::cpu_experts::configured().then(crate::cpu_experts::incomplete_gpu_residency_stats)
2113 }
2114
2115 pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)> {
2118
2119 let guard = self.moe_cache.lock().unwrap();
2120 guard.as_ref().and_then(|cache| cache.pread_stats()).map(|stats| (
2121 stats.reads,
2122 stats.bytes,
2123 stats.read_errors,
2124 stats.short_reads,
2125 stats.fallbacks,
2126 stats.buffer_waits,
2127 stats.ring_full,
2128 ))
2129 }
2130
2131 pub fn moe_cache_reset_counters(&self) {
2133 if let Some(c) = self.moe_cache.lock().unwrap().as_mut() { c.reset_counters(); }
2134 }
2135
2136 pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2137 Ok(self.gpu.stream().clone_htod(v)?)
2138 }
2139
2140 pub fn htod_bytes_padded(&self, v: &[u8], pad: usize)
2144 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2145 let mut d = self.alloc_u8_uninit(v.len() + pad)?;
2146 {
2147 let mut view = d.slice_mut(0..v.len());
2148 self.gpu.stream().memcpy_htod(v, &mut view)?;
2149 }
2150 Ok(d)
2151 }
2152
2153 pub fn copy_into(&self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize)
2155 -> Result<(), Box<dyn std::error::Error>> {
2156 let mut view = dst.slice_mut(off..off + len);
2157 self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2158 Ok(())
2159 }
2160
2161 pub fn copy_u8_into(&self, dst: &mut CudaSlice<u8>, off: usize, src: &CudaSlice<u8>, len: usize)
2164 -> Result<(), Box<dyn std::error::Error>> {
2165 let mut view = dst.slice_mut(off..off + len);
2166 self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2167 Ok(())
2168 }
2169
2170 pub fn htod_u8_into(&self, dst: &mut CudaSlice<u8>, off: usize, src: &[u8])
2173 -> Result<(), Box<dyn std::error::Error>> {
2174 let mut view = dst.slice_mut(off..off + src.len());
2175 self.gpu.stream().memcpy_htod(src, &mut view)?;
2176 Ok(())
2177 }
2178
2179 pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> cudarc::driver::CudaView<'a, f32> {
2180 b.slice(0..len)
2181 }
2182
2183 pub fn view_u8_range<'a>(&self, b: &'a CudaSlice<u8>, start: usize, end: usize)
2186 -> cudarc::driver::CudaView<'a, u8> {
2187 b.slice(start..end)
2188 }
2189 pub fn view_u8<'a>(&self, b: &'a CudaSlice<u8>, len: usize) -> cudarc::driver::CudaView<'a, u8> {
2190 b.slice(0..len)
2191 }
2192
2193 pub fn append_kv_quantized(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
2197 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize,
2198 kv_dim_k: usize, kv_dim_v: usize,
2199 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2200 -> Result<(), Box<dyn std::error::Error>> {
2201 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1") } else { self.func("append_quantize_kv_q8_0_q5_1") };
2202 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2203 let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2204 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
2205 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2206 let __s_b = self.gpu.stream();
2207 let mut b = __s_b.launch_builder(&f);
2208 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(&ti).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2209 unsafe { b.launch(cfg)?; }
2210 Ok(())
2211 }
2212
2213 pub fn append_kv_quantized_dc(&self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>,
2217 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t_dev: &CudaSlice<i32>,
2218 kv_dim_k: usize, kv_dim_v: usize,
2219 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2220 -> Result<(), Box<dyn std::error::Error>> {
2221 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2222 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
2223 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2224 if Self::pdl_on() && Self::pdl_wb_on() {
2226 use cudarc::driver::{DevicePtr, DevicePtrMut};
2227 let s = &self.gpu.stream();
2228 let (pk, _g0) = k_row.device_ptr(s); let (pv, _g1) = v_row.device_ptr(s);
2229 let (pkc, _g2) = kc.device_ptr_mut(s); let (pvc, _g3) = vc.device_ptr_mut(s);
2230 let (pt, _g4) = t_dev.device_ptr(s);
2231 let mut ps = [
2232 &pk as *const _ as *mut std::ffi::c_void, &pv as *const _ as *mut _,
2233 &pkc as *const _ as *mut _, &pvc as *const _ as *mut _,
2234 &pt as *const _ as *mut _, &kdk as *const _ as *mut _,
2235 &kdv as *const _ as *mut _, &ktb as *const _ as *mut _,
2236 &vtb as *const _ as *mut _,
2237 ];
2238 unsafe { self.launch_pdl_flash(g, "append_quantize_kv_q8_0_q5_1_dc",
2239 (nblk, 1, 1), (32, 1, 1), 0, &mut ps)?; }
2240 return Ok(());
2241 }
2242 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_dc") } else { self.func("append_quantize_kv_q8_0_q5_1_dc") };
2243 let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2244 let __s_b = self.gpu.stream();
2245 let mut b = __s_b.launch_builder(&f);
2246 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(t_dev).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2247 unsafe { b.launch(cfg)?; }
2248 Ok(())
2249 }
2250
2251 #[allow(clippy::too_many_arguments)]
2258 pub fn append_kv_quantized_rows(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
2259 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
2260 t0: usize, t: usize, kv_dim_k: usize, kv_dim_v: usize,
2261 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2262 -> Result<(), Box<dyn std::error::Error>> {
2263 if std::env::var("MEMRA_PRIME_APPEND_LOOP").is_ok() {
2264 for i in 0..t {
2265 let k_row = k_rows.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
2266 let v_row = v_rows.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
2267 self.append_kv_quantized_view(&k_row, &v_row, kc, vc, t0 + i,
2268 kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes, g)?;
2269 }
2270 return Ok(());
2271 }
2272 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1_rows") } else { self.func("append_quantize_kv_q8_0_q5_1_rows") };
2273 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2274 let cfg = LaunchConfig { grid_dim: (nblk, t as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2275 let (t0i, kdk, kdv) = (t0 as i32, kv_dim_k as i32, kv_dim_v as i32);
2276 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2277 let __s_b = self.gpu.stream();
2278 let mut b = __s_b.launch_builder(&f);
2279 b.arg(k_rows).arg(v_rows).arg(kc).arg(vc).arg(&t0i).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2280 unsafe { b.launch(cfg)?; }
2281 Ok(())
2282 }
2283
2284 pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn std::error::Error>> {
2288 let f = self.func("inc_i32");
2289 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
2290 let __s_b = self.gpu.stream();
2291 let mut b = __s_b.launch_builder(&f);
2292 b.arg(p);
2293 unsafe { b.launch(cfg)?; }
2294 Ok(())
2295 }
2296
2297 pub fn append_kv_quantized_view(&self, k_row: &cudarc::driver::CudaView<f32>,
2300 v_row: &cudarc::driver::CudaView<f32>,
2301 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize,
2302 kv_dim_k: usize, kv_dim_v: usize,
2303 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
2304 -> Result<(), Box<dyn std::error::Error>> {
2305 let f = if g { self.func_g("append_quantize_kv_q8_0_q5_1") }
2306 else { self.func("append_quantize_kv_q8_0_q5_1") };
2307 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
2308 let cfg = LaunchConfig { grid_dim: (nblk, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2309 let (ti, kdk, kdv) = (t as i32, kv_dim_k as i32, kv_dim_v as i32);
2310 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
2311 let __s_b = self.gpu.stream();
2312 let mut b = __s_b.launch_builder(&f);
2313 b.arg(k_row).arg(v_row).arg(kc).arg(vc).arg(&ti).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
2314 unsafe { b.launch(cfg)?; }
2315 Ok(())
2316 }
2317
2318 pub fn copy_view_into(&self, dst: &mut CudaSlice<f32>, off: usize,
2321 src: &cudarc::driver::CudaView<f32>, len: usize)
2322 -> Result<(), Box<dyn std::error::Error>> {
2323 let mut view = dst.slice_mut(off..off + len);
2324 self.gpu.stream().memcpy_dtod(&src.slice(0..len), &mut view)?;
2325 Ok(())
2326 }
2327
2328 pub fn clone_dtod(&self, src: &CudaSlice<f32>) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2332 let mut dst = self.gpu.stream().alloc_zeros::<f32>(src.len())?;
2333 self.gpu.stream().memcpy_dtod(src, &mut dst)?;
2334 Ok(dst)
2335 }
2336
2337 pub fn dtod_copy_view(&self, src: &cudarc::driver::CudaView<f32>, dst: &mut CudaSlice<f32>)
2340 -> Result<(), Box<dyn std::error::Error>> {
2341 self.gpu.stream().memcpy_dtod(src, dst)?;
2342 Ok(())
2343 }
2344
2345 pub fn dtod_copy_view_i8(&self, src: &cudarc::driver::CudaView<i8>, dst: &mut CudaSlice<i8>)
2347 -> Result<(), Box<dyn std::error::Error>> {
2348 self.gpu.stream().memcpy_dtod(src, dst)?;
2349 Ok(())
2350 }
2351
2352 pub fn dtod_copy_into(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, offset: usize)
2354 -> Result<(), Box<dyn std::error::Error>> {
2355 let n = src.len();
2356 let mut dv = dst.slice_mut(offset..offset + n);
2357 self.gpu.stream().memcpy_dtod(src, &mut dv)?;
2358 Ok(())
2359 }
2360
2361 pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
2363 self.alloc_uninit::<i8>(n)
2364 }
2365
2366 pub fn qmatvec(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize,
2368 qtype: i32, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2369 let f = self.func("qmatvec_f32");
2370 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2372 let (inf, outf, mi, qt, rb) = (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
2373 let __s_b = self.gpu.stream();
2374 let mut b = __s_b.launch_builder(&f);
2375 b.arg(w).arg(x).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qt).arg(&rb);
2376 unsafe { b.launch(cfg)?; }
2377 Ok(y)
2378 }
2379
2380 pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2382 let s = self.gpu.stream().alloc_zeros::<u8>(n)?;
2383 self.keep_if_capturing(&s);
2384 Ok(s)
2385 }
2386
2387 pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2391 let s = unsafe { self.gpu.stream().alloc::<u8>(n)? };
2392 self.keep_if_capturing(&s);
2393 Ok(s)
2394 }
2395
2396 pub fn memset_zeros_view(&self, dst: &mut cudarc::driver::CudaViewMut<f32>)
2399 -> Result<(), Box<dyn std::error::Error>> {
2400 self.gpu.stream().memset_zeros(dst)?;
2401 Ok(())
2402 }
2403
2404 pub fn stage_expert(&self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize)
2410 -> Result<(), Box<dyn std::error::Error>> {
2411 let mut dst = scratch.slice_mut(off..off + host_bytes.len()); self.gpu.stream().memcpy_htod(host_bytes, &mut dst)?; Ok(())
2414 }
2415
2416 pub fn moe_router_topk(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
2422 -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2423 let f = self.func("moe_router_topk_f32");
2424 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 { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2427 shared_mem_bytes: 0 };
2428 let (ne, nu) = (n_expert as i32, n_used as i32);
2429 let __s_b = self.gpu.stream();
2430 let mut b = __s_b.launch_builder(&f);
2431 b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu);
2432 unsafe { b.launch(cfg)?; }
2433 Ok((sel_idx, sel_w))
2434 }
2435
2436 pub fn moe_router_topk_scaled(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize,
2439 n_used: usize, ex_scale: &CudaSlice<f32>)
2440 -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2441 let f = self.func("moe_router_topk_scaled_f32");
2446 let mut sel_idx = self.alloc_uninit::<i32>(t * n_used)?;
2447 let mut sel_w = self.alloc_uninit::<f32>(t * n_used)?;
2448 let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2449 shared_mem_bytes: 0 };
2450 let (ne, nu) = (n_expert as i32, n_used as i32);
2451 let __s_b = self.gpu.stream();
2452 let mut b = __s_b.launch_builder(&f);
2453 b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu).arg(ex_scale);
2454 unsafe { b.launch(cfg)?; }
2455 Ok((sel_idx, sel_w))
2456 }
2457
2458 pub fn moe_router_topk_host(&self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize)
2466 -> Result<(Vec<u32>, Vec<f32>), Box<dyn std::error::Error>> {
2467 let f = self.func("moe_router_topk_f32");
2468 let n = t * n_used;
2469 let mut sel_idx = self.alloc_uninit::<i32>(n)?;
2470 let mut sel_w = self.alloc_uninit::<f32>(n)?;
2471 let cfg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (n_expert as u32, 1, 1),
2472 shared_mem_bytes: 0 };
2473 let (ne, nu) = (n_expert as i32, n_used as i32);
2474 let __s_b = self.gpu.stream();
2475 let mut b = __s_b.launch_builder(&f);
2476 b.arg(logits).arg(&mut sel_idx).arg(&mut sel_w).arg(&ne).arg(&nu);
2477 unsafe { b.launch(cfg)?; }
2478 let bytes = n * 8;
2480 let mut guard = self.router_stage.lock().unwrap();
2481 if guard.as_ref().map(|p| p.cap < bytes).unwrap_or(true) {
2482 *guard = Some(PinnedStage::new(bytes.max(4096))?);
2483 }
2484 let stage = guard.as_mut().unwrap();
2485 let (si, sw) = unsafe {
2486 (std::slice::from_raw_parts_mut(stage.ptr as *mut i32, n),
2487 std::slice::from_raw_parts_mut(stage.ptr.add(n * 4) as *mut f32, n))
2488 };
2489 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()))
2493 }
2494
2495 pub fn stage_expert_async(&self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize)
2499 -> Result<cudarc::driver::CudaEvent, Box<dyn std::error::Error>> {
2500 let mut dst = scratch.slice_mut(off..off + host_bytes.len());
2501 self.copy_stream.memcpy_htod(host_bytes, &mut dst)?;
2502 Ok(self.copy_stream.record_event(None)?)
2503 }
2504
2505 pub fn compute_wait(&self, ev: &cudarc::driver::CudaEvent) -> Result<(), Box<dyn std::error::Error>> {
2507 self.gpu.stream().wait(ev)?;
2508 Ok(())
2509 }
2510
2511 pub fn qmatvec_view(&self, w: &CudaSlice<u8>, range: std::ops::Range<usize>,
2516 x: &cudarc::driver::CudaView<f32>, m: usize, in_f: usize, out_f: usize,
2517 qtype: i32, row_bytes: usize)
2518 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2519 let f = self.func("qmatvec_f32");
2520 let wv = w.slice(range); let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2523 let (inf, outf, mi, qt, rb) = (in_f as i32, out_f as i32, m as i32, qtype, row_bytes as i64);
2524 let __s_b = self.gpu.stream();
2525 let mut b = __s_b.launch_builder(&f);
2526 b.arg(&wv).arg(x).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qt).arg(&rb);
2527 unsafe { b.launch(cfg)?; }
2528 Ok(y)
2529 }
2530
2531 #[allow(clippy::too_many_arguments)]
2538 pub fn moe_gate_up_silu8_q8(&self, gp: WPtr8, up: WPtr8,
2542 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2543 in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32,
2544 rb_g: usize, rb_u: usize)
2545 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2546 let f = self.func("moe_gate_up_silu8_q8");
2547 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
2548 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2549 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2550 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
2551 let __s_b = self.gpu.stream();
2552 let mut b = __s_b.launch_builder(&f);
2553 b.arg(&gp).arg(&up).arg(aq).arg(ad).arg(&mut act)
2554 .arg(&inf).arg(&nff).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2555 unsafe { b.launch(cfg)?; }
2556 Ok(act)
2557 }
2558
2559 #[allow(clippy::too_many_arguments)]
2560 pub fn moe_down8_fma_q8(&self, dp: WPtr8, w: F32x8,
2561 aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
2562 dst: &mut cudarc::driver::CudaViewMut<f32>,
2563 in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize)
2564 -> Result<(), Box<dyn std::error::Error>> {
2565 let f = self.func("moe_down8_fma_q8");
2566 let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
2567 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2568 let (inf, outf, nu, rbi) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
2569 let __s_b = self.gpu.stream();
2570 let mut b = __s_b.launch_builder(&f);
2571 b.arg(&dp).arg(&w).arg(aq2).arg(ad2).arg(dst)
2572 .arg(&inf).arg(&outf).arg(&nu).arg(&qt).arg(&rbi);
2573 unsafe { b.launch(cfg)?; }
2574 Ok(())
2575 }
2576
2577 pub fn qmatvec_expert_q8(&self, w: &CudaSlice<u8>, range: std::ops::Range<usize>,
2579 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
2580 in_f: usize, out_f: usize, qtype: i32, row_bytes: usize)
2581 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2582 let f = self.func("qmatvec_expert_q8");
2583 let wv = w.slice(range);
2584 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
2585 const ROWS: u32 = 4; let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, m as u32, 1),
2587 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2588 let (inf, outf, mi, rbi) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
2589 let __s_b = self.gpu.stream();
2590 let mut b = __s_b.launch_builder(&f);
2591 b.arg(&wv).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&qtype).arg(&rbi);
2592 unsafe { b.launch(cfg)?; }
2593 Ok(y)
2594 }
2595
2596 pub fn moe_gate_up_silu8(&self, gp: WPtr8, up: WPtr8, x: &cudarc::driver::CudaView<f32>,
2597 in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32,
2598 rb_g: usize, rb_u: usize)
2599 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2600 let f = self.func("moe_gate_up_silu8_f32");
2601 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2603 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2604 let (inf, nff, rbg, rbu) = (in_f as i32, n_ff as i32, rb_g as i64, rb_u as i64);
2605 let __s_b = self.gpu.stream();
2606 let mut b = __s_b.launch_builder(&f);
2607 b.arg(&gp).arg(&up).arg(x).arg(&mut act)
2608 .arg(&inf).arg(&nff).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2609 unsafe { b.launch(cfg)?; }
2610 Ok(act)
2611 }
2612
2613 #[allow(clippy::too_many_arguments)]
2619 pub fn moe_down8_fma_into(&self, dp: WPtr8, w: F32x8, act: &CudaSlice<f32>,
2620 dst: &mut cudarc::driver::CudaViewMut<f32>,
2621 in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize)
2622 -> Result<(), Box<dyn std::error::Error>> {
2623 let f = self.func("moe_down8_fma_f32");
2624 let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
2625 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2626 let (inf, outf, nu, rbv) = (in_f as i32, out_f as i32, n_used as i32, rb as i64);
2627 let __s_b = self.gpu.stream();
2628 let mut b = __s_b.launch_builder(&f);
2629 b.arg(&dp).arg(&w).arg(act).arg(dst).arg(&inf).arg(&outf).arg(&nu).arg(&qt).arg(&rbv);
2630 unsafe { b.launch(cfg)?; }
2631 Ok(())
2632 }
2633
2634 #[allow(clippy::too_many_arguments)]
2639 #[allow(clippy::too_many_arguments)]
2654 #[allow(clippy::too_many_arguments)]
2656 pub fn moe_pairs_matvec_q8(&self, table: &CudaSlice<u64>, proj: i32,
2657 pair_tok: &CudaSlice<i32>, pair_ex: &CudaSlice<i32>,
2658 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2659 in_f: usize, out_f: usize, n_expert: usize, n_pairs: usize,
2660 qtype: i32, row_bytes: usize)
2661 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2662 let f = self.func("moe_pairs_matvec_q8");
2663 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2664 const ROWS: u32 = 4;
2665 let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_pairs as u32, 1),
2666 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2667 let (inf, outf, ne, np, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2668 n_pairs as i32, row_bytes as i64);
2669 let __s_b = self.gpu.stream();
2670 let mut b = __s_b.launch_builder(&f);
2671 b.arg(table).arg(&proj).arg(pair_tok).arg(pair_ex).arg(aq).arg(ad).arg(&mut y)
2672 .arg(&inf).arg(&outf).arg(&ne).arg(&np).arg(&qtype).arg(&rbi);
2673 unsafe { b.launch(cfg)?; }
2674 Ok(y)
2675 }
2676
2677 #[allow(clippy::too_many_arguments)]
2679 pub fn moe_pairs_matvec_q8_em(&self, table: &CudaSlice<u64>, proj: i32,
2680 ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>,
2681 ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>,
2682 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2683 in_f: usize, out_f: usize, n_expert: usize, n_active: usize,
2684 n_pairs: usize, qtype: i32, row_bytes: usize)
2685 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2686 let f = self.func("moe_pairs_matvec_q8_em");
2687 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2688 const ROWS: u32 = 4;
2689 let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
2690 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2691 let (inf, outf, ne, na, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2692 n_active as i32, row_bytes as i64);
2693 let __s_b = self.gpu.stream();
2694 let mut b = __s_b.launch_builder(&f);
2695 b.arg(table).arg(&proj).arg(ex_ids).arg(ex_off).arg(ex_pairs).arg(pair_tok)
2696 .arg(aq).arg(ad).arg(&mut y)
2697 .arg(&inf).arg(&outf).arg(&ne).arg(&na).arg(&qtype).arg(&rbi);
2698 unsafe { b.launch(cfg)?; }
2699 Ok(y)
2700 }
2701
2702 #[allow(clippy::too_many_arguments)]
2705 pub fn moe_pairs_matvec_q8_dec(&self, table: &CudaSlice<u64>, proj: i32,
2706 ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>,
2707 ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>,
2708 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2709 in_f: usize, out_f: usize, n_expert: usize, n_active: usize,
2710 n_pairs: usize, qtype: i32, row_bytes: usize)
2711 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2712 let f = self.func("moe_pairs_matvec_q8_dec");
2713 let mut y = self.alloc_uninit::<f32>(n_pairs * out_f)?;
2714 const ROWS: u32 = 4;
2715 let cfg = LaunchConfig { grid_dim: ((out_f as u32 + ROWS - 1) / ROWS, n_active as u32, 1),
2716 block_dim: (32, ROWS, 1), shared_mem_bytes: 0 };
2717 let (inf, outf, ne, na, rbi) = (in_f as i32, out_f as i32, n_expert as i32,
2718 n_active as i32, row_bytes as i64);
2719 let __s_b = self.gpu.stream();
2720 let mut b = __s_b.launch_builder(&f);
2721 b.arg(table).arg(&proj).arg(ex_ids).arg(ex_off).arg(ex_pairs).arg(pair_tok)
2722 .arg(aq).arg(ad).arg(&mut y)
2723 .arg(&inf).arg(&outf).arg(&ne).arg(&na).arg(&qtype).arg(&rbi);
2724 unsafe { b.launch(cfg)?; }
2725 Ok(y)
2726 }
2727
2728 pub fn moe_pairs_gelu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize)
2729 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2730 let f = self.func("moe_pairs_gelu_mul");
2731 let mut act = self.alloc_uninit::<f32>(n)?;
2732 let cfg = LaunchConfig::for_num_elems(n as u32);
2733 let nl = n as i64;
2734 let __s_b = self.gpu.stream();
2735 let mut b = __s_b.launch_builder(&f);
2736 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
2737 unsafe { b.launch(cfg)?; }
2738 Ok(act)
2739 }
2740
2741 pub fn moe_pairs_silu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize)
2742 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2743 let f = self.func("moe_pairs_silu_mul");
2744 let mut act = self.alloc_uninit::<f32>(n)?;
2745 let cfg = LaunchConfig::for_num_elems(n as u32);
2746 let nl = n as i64;
2747 let __s_b = self.gpu.stream();
2748 let mut b = __s_b.launch_builder(&f);
2749 b.arg(gate).arg(up).arg(&mut act).arg(&nl);
2750 unsafe { b.launch(cfg)?; }
2751 Ok(act)
2752 }
2753
2754 #[allow(clippy::too_many_arguments)]
2755 pub fn moe_pairs_scatter(&self, y_down: &CudaSlice<f32>, pair_w: &CudaSlice<f32>,
2756 tok_pair_off: &CudaSlice<i32>, tok_pair_ids: &CudaSlice<i32>,
2757 moe_out: &mut CudaSlice<f32>, t: usize, n_embd: usize)
2758 -> Result<(), Box<dyn std::error::Error>> {
2759 let f = self.func("moe_pairs_scatter");
2760 let cfg = LaunchConfig { grid_dim: (((n_embd + 255) / 256) as u32, t as u32, 1),
2761 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2762 let ne = n_embd as i32;
2763 let __s_b = self.gpu.stream();
2764 let mut b = __s_b.launch_builder(&f);
2765 b.arg(y_down).arg(pair_w).arg(tok_pair_off).arg(tok_pair_ids).arg(moe_out).arg(&ne);
2766 unsafe { b.launch(cfg)?; }
2767 Ok(())
2768 }
2769
2770 #[allow(clippy::too_many_arguments)]
2774 pub fn moe_gate_up_gelu8_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
2775 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
2776 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2777 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2778 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2779 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
2780 let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
2781 rb_g as i64, rb_u as i64);
2782 let f = self.func("moe_gate_up_gelu8_dev_q8");
2783 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
2784 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2785 let __s_b = self.gpu.stream();
2786 let mut b = __s_b.launch_builder(&f);
2787 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2788 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
2789 unsafe { b.launch(cfg)?; }
2790 Ok(act)
2791 }
2792
2793 #[allow(clippy::too_many_arguments)]
2795 pub fn moe_gate_up_gelu8_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2796 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize,
2797 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2798 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2799 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2800 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
2801 let (inf, nff, ne, rbg, rbu, nu) = (in_f as i32, n_ff as i32, n_expert as i32,
2802 rb_g as i64, rb_u as i64, n_used as i32);
2803 let f = self.func("moe_gate_up_gelu8_dev_q8_rows");
2804 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, t as u32),
2805 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2806 let __s_b = self.gpu.stream();
2807 let mut b = __s_b.launch_builder(&f);
2808 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2809 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu);
2810 unsafe { b.launch(cfg)?; }
2811 Ok(act)
2812 }
2813
2814 #[allow(clippy::too_many_arguments)]
2816 pub fn moe_gate_up_gelu8_dev_q8_csr(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2817 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, n_pairs: usize,
2818 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
2819 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
2820 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2821 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
2822 let (inf, nff, ne, rbg, rbu, nu, npi) = (in_f as i32, n_ff as i32, n_expert as i32,
2823 rb_g as i64, rb_u as i64, n_used as i32,
2824 n_pairs as i32);
2825 let f = self.func("moe_gate_up_gelu8_dev_q8_csr");
2826 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_pairs as u32, 1),
2827 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2828 let __s_b = self.gpu.stream();
2829 let mut b = __s_b.launch_builder(&f);
2830 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
2831 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(&npi);
2832 unsafe { b.launch(cfg)?; }
2833 Ok(act)
2834 }
2835
2836 #[allow(clippy::too_many_arguments)]
2838 pub fn moe_down8_fma_dev_q8_rows_g(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
2839 w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
2840 dst: &mut CudaSlice<f32>, t: usize,
2841 in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
2842 qt: i32, rb: usize)
2843 -> Result<(), Box<dyn std::error::Error>> {
2844 let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
2845 n_expert as i32, rb as i64);
2846 let f = self.func("moe_down8_fma_dev_q8_rows_g");
2847 let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, t as u32),
2848 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
2849 let __s_b = self.gpu.stream();
2850 let mut b = __s_b.launch_builder(&f);
2851 b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
2852 .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
2853 unsafe { b.launch(cfg)?; }
2854 Ok(())
2855 }
2856
2857 pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn std::error::Error>> {
2861 let (out_f, in_f) = (2048usize, 2816usize);
2862 let nblk = in_f / 32;
2863 let mut seed = 0x9E3779B97F4A7C15u64;
2864 let mut rng = move || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); (seed >> 33) as u8 };
2865 let mut w = vec![0u8; out_f * nblk * 18];
2866 for b in w.iter_mut() { *b = rng(); }
2867 for r in 0..out_f {
2868 for g in 0..nblk {
2869 let off = (r * nblk + g) * 18;
2870 w[off] = 0x00; w[off + 1] = 0x2C; }
2872 }
2873 let qplane = out_f * nblk * 16;
2874 let mut wrp = vec![0u8; w.len()];
2875 for r in 0..out_f {
2876 for g in 0..nblk {
2877 let src = &w[(r * nblk + g) * 18..(r * nblk + g) * 18 + 18];
2878 wrp[qplane + (r * nblk + g) * 2..qplane + (r * nblk + g) * 2 + 2]
2879 .copy_from_slice(&src[0..2]);
2880 wrp[(r * nblk + g) * 16..(r * nblk + g) * 16 + 16].copy_from_slice(&src[2..18]);
2881 }
2882 }
2883 let w_d = self.htod_bytes(&w)?;
2884 let wrp_d = self.htod_bytes(&wrp)?;
2885 let mut aq = vec![0i8; m * in_f];
2886 for v in aq.iter_mut() { *v = rng() as i8; }
2887 let aq_d = self.htod_i8(&aq)?;
2888 let ad_d = self.htod(&vec![0.03125f32; m * nblk])?;
2889 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
2890 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
2891 const RPB: u32 = 4;
2892 let cfg = LaunchConfig { grid_dim: ((out_f as u32).div_ceil(RPB), 1, 1),
2893 block_dim: (32, RPB, 1), shared_mem_bytes: 0 };
2894 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
2895 let (rb, qp) = ((nblk * 18) as i64, qplane as i64);
2896 let fb = self.func("qmatvec_q4_0_mmvq_b4");
2897 let fr = self.func("qmatvec_q4_0_mmvq_b4_rp");
2898 {
2899 let __s_b = self.gpu.stream();
2900 let mut b = __s_b.launch_builder(&fb);
2901 b.arg(&w_d).arg(&aq_d).arg(&ad_d).arg(&mut y0).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
2902 unsafe { b.launch(cfg)?; }
2903 let __s_b = self.gpu.stream();
2904 let mut b = __s_b.launch_builder(&fr);
2905 b.arg(&wrp_d).arg(&aq_d).arg(&ad_d).arg(&mut y1).arg(&inf).arg(&outf).arg(&mi).arg(&qp);
2906 unsafe { b.launch(cfg)?; }
2907 }
2908 self.gpu.stream().synchronize()?;
2909 let (h0, h1) = (self.dtoh(&y0)?, self.dtoh(&y1)?);
2910 let nd = h0.iter().zip(&h1).filter(|(a, b)| a.to_bits() != b.to_bits()).count();
2911 if nd != 0 { return Err(format!("rp twin not bitwise: {nd}/{} diffs", h0.len()).into()); }
2912 let mut time = |rp: bool| -> Result<f64, Box<dyn std::error::Error>> {
2913 self.gpu.stream().synchronize()?;
2914 let t0 = std::time::Instant::now();
2915 for _ in 0..500 {
2916 if rp {
2917 let __s_b = self.gpu.stream();
2918 let mut b = __s_b.launch_builder(&fr);
2919 b.arg(&wrp_d).arg(&aq_d).arg(&ad_d).arg(&mut y1)
2920 .arg(&inf).arg(&outf).arg(&mi).arg(&qp);
2921 unsafe { b.launch(cfg)?; }
2922 } else {
2923 let __s_b = self.gpu.stream();
2924 let mut b = __s_b.launch_builder(&fb);
2925 b.arg(&w_d).arg(&aq_d).arg(&ad_d).arg(&mut y0)
2926 .arg(&inf).arg(&outf).arg(&mi).arg(&rb);
2927 unsafe { b.launch(cfg)?; }
2928 }
2929 }
2930 self.gpu.stream().synchronize()?;
2931 Ok(t0.elapsed().as_secs_f64() * 1e6 / 500.0)
2932 };
2933 let _ = time(false)?; let _ = time(true)?; Ok((time(false)?, time(true)?))
2935 }
2936
2937 pub fn build_q4_rp4(&self, t: &mut crate::model::GpuTensor)
2942 -> Result<(), Box<dyn std::error::Error>> {
2943 use crate::model::GpuTensor;
2944 let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
2945 if *qtype != QT_Q4_0 || rp4.is_some() || ne.len() != 2 { return Ok(()); }
2946 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
2947 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 18 { return Ok(()); }
2948 let nblk = in_f / 32;
2949 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 18)?;
2950 let f = self.func("q4_0_split_rp_build");
2951 let n = (out_f * nblk) as i32;
2952 let cfg = LaunchConfig { grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
2953 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2954 let (of, nb) = (out_f as i32, nblk as i32);
2955 let _ = n;
2956 let __s_b = self.gpu.stream();
2957 let mut b = __s_b.launch_builder(&f);
2958 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
2959 unsafe { b.launch(cfg)?; }
2960 *rp4 = Some(dst);
2961 Ok(())
2962 }
2963
2964 pub fn build_q8_rp4(&self, t: &mut crate::model::GpuTensor)
2969 -> Result<(), Box<dyn std::error::Error>> {
2970 use crate::model::GpuTensor;
2971 let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
2972 if *qtype != QT_Q8_0 || rp4.is_some() || ne.len() != 2 { return Ok(()); }
2973 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
2974 if in_f % 32 != 0 || *row_bytes != (in_f / 32) * 34 { return Ok(()); }
2975 *rp4 = Some(self.build_q8_rp4_raw(bytes, in_f, out_f)?);
2976 Ok(())
2977 }
2978
2979 pub fn build_q8_rp4_raw(&self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize)
2982 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
2983 assert!(in_f % 32 == 0);
2984 let nblk = in_f / 32;
2985 let mut dst = self.alloc_uninit::<u8>(out_f * nblk * 34)?;
2986 let f = self.func("q8_0_split_rp_build");
2987 let cfg = LaunchConfig { grid_dim: (((out_f * nblk) as u32).div_ceil(256), 1, 1),
2988 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
2989 let (of, nb) = (out_f as i32, nblk as i32);
2990 let __s_b = self.gpu.stream();
2991 let mut b = __s_b.launch_builder(&f);
2992 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
2993 unsafe { b.launch(cfg)?; }
2994 Ok(dst)
2995 }
2996
2997 pub fn build_q4k_rp4(&self, t: &mut crate::model::GpuTensor)
3005 -> Result<(), Box<dyn std::error::Error>> {
3006 use crate::model::GpuTensor;
3007 let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
3008 if *qtype != QT_Q4_K || rp4.is_some() || ne.len() != 2 { return Ok(()); }
3009 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
3010 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 144 { return Ok(()); }
3011 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q4_K)?);
3012 Ok(())
3013 }
3014
3015 pub fn build_q6k_rp4(&self, t: &mut crate::model::GpuTensor)
3016 -> Result<(), Box<dyn std::error::Error>> {
3017 use crate::model::GpuTensor;
3018 let GpuTensor::Quant { bytes, qtype, row_bytes, ne, rp4, .. } = t else { return Ok(()) };
3019 if *qtype != QT_Q6_K || rp4.is_some() || ne.len() != 2 { return Ok(()); }
3020 let (in_f, out_f) = (ne[0] as usize, ne[1] as usize);
3021 if in_f % 256 != 0 || *row_bytes != (in_f / 256) * 210 { return Ok(()); }
3022 *rp4 = Some(self.build_kq_rp4_raw(bytes, in_f, out_f, QT_Q6_K)?);
3023 Ok(())
3024 }
3025
3026 pub fn build_kq_rp4_raw(&self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, qtype: i32)
3028 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
3029 assert!(in_f % 256 == 0);
3030 let nsbk = in_f / 256;
3031 let (sb_bytes, kname) = match qtype {
3032 QT_Q4_K => (144usize, "q4_K_split_rp_build"),
3033 QT_Q6_K => (210usize, "q6_K_split_rp_build"),
3034 _ => return Err(format!("build_kq_rp4_raw: qtype {qtype} has no rp mirror").into()),
3035 };
3036 let mut dst = self.alloc_uninit::<u8>(out_f * nsbk * sb_bytes)?;
3037 let f = self.func(kname);
3038 let cfg = LaunchConfig { grid_dim: (((out_f * nsbk) as u32).div_ceil(256), 1, 1),
3039 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3040 let (of, nb) = (out_f as i32, nsbk as i32);
3041 let __s_b = self.gpu.stream();
3042 let mut b = __s_b.launch_builder(&f);
3043 b.arg(&*bytes).arg(&mut dst).arg(&of).arg(&nb);
3044 unsafe { b.launch(cfg)?; }
3045 Ok(dst)
3046 }
3047
3048 pub fn kqrp_enabled() -> bool {
3052 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3053 *ON.get_or_init(|| match std::env::var("MEMRA_KQRP").as_deref() {
3054 Ok("0") => false,
3055 Ok(_) => true,
3056 Err(_) => cfg!(memra_hopper_mma),
3057 })
3058 }
3059
3060 pub fn build_q4_rp_swap(&self, t: &mut crate::model::GpuTensor)
3066 -> Result<bool, Box<dyn std::error::Error>> {
3067 self.build_q4_rp4(t)?;
3068 self.gpu.stream().synchronize()?; use crate::model::GpuTensor;
3070 let GpuTensor::Quant { bytes, rp4, rp, .. } = t else { return Ok(false) };
3071 match rp4.take() {
3072 Some(split) => {
3073 *bytes = split; *rp = true;
3075 Ok(true)
3076 }
3077 None => Ok(false),
3078 }
3079 }
3080
3081 pub fn q4rp_enabled() -> bool {
3083 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3084 *ON.get_or_init(|| std::env::var("MEMRA_Q4RP").map(|v| v != "0").unwrap_or(true))
3085 }
3086
3087 pub fn copy_rows_strided(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
3090 row_elems: usize, n_rows: usize, src_stride: usize, src_off: usize)
3091 -> Result<(), Box<dyn std::error::Error>> {
3092 let f = self.func("copy_rows_strided_f32");
3093 let cfg = LaunchConfig { grid_dim: (((row_elems as u32 + 255) / 256).max(1), n_rows as u32, 1),
3094 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3095 let (re, nr) = (row_elems as i32, n_rows as i32);
3096 let (st, off) = (src_stride as i64, src_off as i64);
3097 let __s_b = self.gpu.stream();
3098 let mut b = __s_b.launch_builder(&f);
3099 b.arg(src).arg(&mut *dst).arg(&re).arg(&nr).arg(&st).arg(&off);
3100 unsafe { b.launch(cfg)?; }
3101 Ok(())
3102 }
3103
3104 pub fn u32_set_k(&self, dst: &mut CudaSlice<u32>, v: u32, idx: usize)
3106 -> Result<(), Box<dyn std::error::Error>> {
3107 let f = self.func("u32_set_k");
3108 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
3109 let ii = idx as i32;
3110 let __s_b = self.gpu.stream();
3111 let mut b = __s_b.launch_builder(&f);
3112 b.arg(dst).arg(&v).arg(&ii);
3113 unsafe { b.launch(cfg)?; }
3114 Ok(())
3115 }
3116
3117 pub fn i32_add_k(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
3119 let f = self.func("i32_add_k");
3120 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3121 let __s_b = self.gpu.stream();
3122 let mut b = __s_b.launch_builder(&f);
3123 b.arg(d).arg(&v);
3124 unsafe { b.launch(cfg)?; }
3125 Ok(())
3126 }
3127
3128 pub fn i32_iota_from(&self, ctr: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, n: usize)
3130 -> Result<(), Box<dyn std::error::Error>> {
3131 let f = self.func("i32_iota_from");
3132 let cfg = LaunchConfig::for_num_elems(n as u32);
3133 let ni = n as i32;
3134 let __s_b = self.gpu.stream();
3135 let mut b = __s_b.launch_builder(&f);
3136 b.arg(ctr).arg(dst).arg(&ni);
3137 unsafe { b.launch(cfg)?; }
3138 Ok(())
3139 }
3140
3141 pub fn u32_map_k(&self, buf: &mut CudaSlice<u32>, map: &CudaSlice<u32>, idx: usize)
3143 -> Result<(), Box<dyn std::error::Error>> {
3144 let f = self.func("u32_map_k");
3145 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
3146 let ii = idx as i32;
3147 let __s_b = self.gpu.stream();
3148 let mut b = __s_b.launch_builder(&f);
3149 b.arg(buf).arg(map).arg(&ii);
3150 unsafe { b.launch(cfg)?; }
3151 Ok(())
3152 }
3153
3154 #[allow(clippy::too_many_arguments)]
3156 pub fn u32_pack2(&self, a: &CudaSlice<u32>, off_a: usize, n1: usize,
3157 b_in: &CudaSlice<u32>, n2: usize, out: &mut CudaSlice<u32>)
3158 -> Result<(), Box<dyn std::error::Error>> {
3159 let f = self.func("u32_pack2");
3160 let cfg = LaunchConfig::for_num_elems((n1 + n2) as u32);
3161 let (oa, i1, i2) = (off_a as i32, n1 as i32, n2 as i32);
3162 let __s_b = self.gpu.stream();
3163 let mut b = __s_b.launch_builder(&f);
3164 b.arg(a).arg(&oa).arg(&i1).arg(b_in).arg(&i2).arg(out);
3165 unsafe { b.launch(cfg)?; }
3166 Ok(())
3167 }
3168
3169 pub fn moe_w_exscale(&self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>,
3171 s: &CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
3172 let f = self.func("moe_w_exscale");
3173 let cfg = LaunchConfig::for_num_elems(n as u32);
3174 let ni = n as i32;
3175 let __s_b = self.gpu.stream();
3176 let mut b = __s_b.launch_builder(&f);
3177 b.arg(w).arg(sel).arg(s).arg(&ni);
3178 unsafe { b.launch(cfg)?; }
3179 Ok(())
3180 }
3181
3182 pub fn moe_w_scale_by_expert(&self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>,
3185 macros: &CudaSlice<f32>, n_expert: usize, n: usize)
3186 -> Result<(), Box<dyn std::error::Error>> {
3187 let f = self.func("moe_w_scale_by_expert");
3188 let cfg = LaunchConfig { grid_dim: (n.div_ceil(64) as u32, 1, 1),
3189 block_dim: (64, 1, 1), shared_mem_bytes: 0 };
3190 let (ne, nn) = (n_expert as i32, n as i32);
3191 let __s_b = self.gpu.stream();
3192 let mut b = __s_b.launch_builder(&f);
3193 b.arg(w).arg(sel).arg(macros).arg(&ne).arg(&nn);
3194 unsafe { b.launch(cfg)?; }
3195 Ok(())
3196 }
3197
3198 pub fn moe_gate_up_silu8_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3199 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3200 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3201 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3202 macros: &CudaSlice<f32>)
3203 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3204 static GU: std::sync::OnceLock<(String, u32)> = std::sync::OnceLock::new();
3205 let (mode, wpb) = GU.get_or_init(|| {
3206 let mode = std::env::var("MEMRA_MOE_DEVQ8_GU").unwrap_or_default();
3207 let wpb = std::env::var("MEMRA_MOE_DEVQ8_WPB").ok()
3208 .and_then(|v| v.parse().ok()).unwrap_or(4u32).clamp(1, 16);
3209 (mode, wpb)
3210 });
3211 let (mode, wpb) = (mode.as_str(), *wpb);
3212 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
3213 let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3214 rb_g as i64, rb_u as i64);
3215 let (f, cfg) = match mode {
3216 "1" | "2" | "4" => {
3217 let rpw: u32 = mode.parse().unwrap();
3218 let f = self.func(match rpw { 1 => "moe_gate_up_silu8_dev_q8_r1",
3219 2 => "moe_gate_up_silu8_dev_q8_r2",
3220 _ => "moe_gate_up_silu8_dev_q8_r4" });
3221 let rows_per_block = (rpw * wpb) as usize;
3222 let gx = n_ff.div_ceil(rows_per_block) as u32;
3223 (f, LaunchConfig { grid_dim: (gx, n_used as u32, 1),
3224 block_dim: (32, wpb, 1), shared_mem_bytes: 0 })
3225 }
3226 "j8" if n_used <= 32 => (self.func("moe_gate_up_silu8_dev_q8_j8"),
3227 LaunchConfig { grid_dim: (n_ff as u32, 1, 1),
3228 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3229 "vsm2" => {
3231 let f = self.func("moe_gate_up_silu8_dev_q8_vsm2");
3232 let sh = (rb_g + rb_u) as u32;
3233 use cudarc::driver::sys::CUfunction_attribute_enum as A;
3234 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
3235 (f, LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3236 block_dim: (32, 1, 1), shared_mem_bytes: sh })
3237 }
3238 "vsm" => {
3239 let f = self.func("moe_gate_up_silu8_dev_q8_vsm");
3240 let sh = (rb_g + rb_u) as u32;
3241 use cudarc::driver::sys::CUfunction_attribute_enum as A;
3242 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
3243 (f, LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3244 block_dim: (32, 1, 1), shared_mem_bytes: sh })
3245 }
3246 "sg" => (self.func("moe_gate_up_silu8_dev_q8_sg"),
3247 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3248 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3249 "j8sg" if n_used <= 32 => (self.func("moe_gate_up_silu8_dev_q8_j8sg"),
3250 LaunchConfig { grid_dim: (n_ff as u32, 1, 1),
3251 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3252 "u64" if in_f == 2048 => (self.func("moe_gate_up_silu8_dev_q8_u64"),
3253 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3254 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3255 "gs4" if in_f == 2048 => (self.func("moe_gate_up_silu8_dev_q8_gs4"),
3256 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3257 block_dim: (32, 4, 1), shared_mem_bytes: 0 }),
3258 "v" | "" => (self.func("moe_gate_up_silu8_dev_q8_v"),
3260 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3261 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3262 "s2" => (self.func("moe_gate_up_silu8_dev_q8_s2"),
3263 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3264 block_dim: (32, 2, 1), shared_mem_bytes: 0 }),
3265 "s2z" => {
3266 let rz = wpb.min(16); (self.func("moe_gate_up_silu8_dev_q8_s2z"),
3268 LaunchConfig { grid_dim: (n_ff.div_ceil(rz as usize) as u32, n_used as u32, 1),
3269 block_dim: (32, 2, rz), shared_mem_bytes: 0 })
3270 }
3271 _ => (self.func("moe_gate_up_silu8_dev_q8"),
3272 LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3273 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3274 };
3275 let __s_b = self.gpu.stream();
3276 let mut b = __s_b.launch_builder(&f);
3277 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3278 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(macros);
3279 unsafe { b.launch(cfg)?; }
3280 Ok(act)
3281 }
3282
3283 #[allow(clippy::too_many_arguments)]
3284 pub fn moe_down8_fma_dev_q8(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3285 w: &cudarc::driver::CudaView<f32>,
3286 aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3287 dst: &mut cudarc::driver::CudaViewMut<f32>,
3288 in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3289 qt: i32, rb: usize)
3290 -> Result<(), Box<dyn std::error::Error>> {
3291 static DOWN: std::sync::OnceLock<String> = std::sync::OnceLock::new();
3292 let mode = DOWN.get_or_init(|| std::env::var("MEMRA_MOE_DEVQ8_DOWN").unwrap_or_default());
3293 let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3294 n_expert as i32, rb as i64);
3295 let (f, cfg) = match mode.as_str() {
3298 m @ ("1" | "2" | "4") if n_used <= 8 => {
3299 let rpw: usize = m.parse().unwrap();
3300 let f = self.func(match rpw { 1 => "moe_down8_fma_dev_q8_w8r1",
3301 2 => "moe_down8_fma_dev_q8_w8r2",
3302 _ => "moe_down8_fma_dev_q8_w8r4" });
3303 (f, LaunchConfig { grid_dim: (out_f.div_ceil(rpw) as u32, 1, 1),
3304 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3305 }
3306 "h2" if in_f == 512 => (self.func("moe_down8_fma_dev_q8_h2"),
3307 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3308 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3309 "" if in_f == 704 && n_used <= 8 =>
3312 (self.func("moe_down8_fma_dev_q8_w8r2"),
3313 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3314 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3315 "w8h2v" | "" if in_f == 512 && n_used <= 8 =>
3319 (self.func("moe_down8_fma_dev_q8_w8h2v"),
3320 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3321 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3322 "w8h2r2v" if in_f == 512 && n_used <= 8 =>
3323 (self.func("moe_down8_fma_dev_q8_w8h2r2v"),
3324 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3325 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3326 "w8h2r2" if in_f == 512 && n_used <= 8 =>
3327 (self.func("moe_down8_fma_dev_q8_w8h2r2"),
3328 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3329 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3330 "w8h2" if in_f == 512 && n_used <= 8 =>
3331 (self.func("moe_down8_fma_dev_q8_w8h2"),
3332 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3333 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 }),
3334 _ => (self.func("moe_down8_fma_dev_q8"),
3335 LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3336 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3337 };
3338 let __s_b = self.gpu.stream();
3339 let mut b = __s_b.launch_builder(&f);
3340 b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3341 .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3342 unsafe { b.launch(cfg)?; }
3343 Ok(())
3344 }
3345
3346 #[allow(clippy::too_many_arguments)]
3353 pub fn moe_gate_up_silu8_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3354 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize,
3355 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3356 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3357 macros: &CudaSlice<f32>)
3358 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3359 let f = self.func("moe_gate_up_silu8_dev_q8_v_rows");
3360 let mut act = self.alloc_uninit::<f32>(t * n_used * n_ff)?;
3361 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, t as u32),
3362 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3363 let (inf, nff, ne, nu, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3364 n_used as i32, rb_g as i64, rb_u as i64);
3365 let __s_b = self.gpu.stream();
3366 let mut b = __s_b.launch_builder(&f);
3367 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3368 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(macros);
3369 unsafe { b.launch(cfg)?; }
3370 Ok(act)
3371 }
3372
3373 #[allow(clippy::too_many_arguments)]
3378 pub fn moe_down8_fma_dev_q8_rows(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3379 w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3380 dst: &mut CudaSlice<f32>, t: usize,
3381 in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3382 qt: i32, rb: usize)
3383 -> Result<(), Box<dyn std::error::Error>> {
3384 assert!(in_f == 512 && n_used <= 8, "down rows twin is w8h2v shape-gated");
3385 let f = self.func("moe_down8_fma_dev_q8_w8h2v_rows");
3386 let cfg = LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, t as u32),
3387 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 };
3388 let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3389 n_expert as i32, rb as i64);
3390 let __s_b = self.gpu.stream();
3391 let mut b = __s_b.launch_builder(&f);
3392 b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3393 .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3394 unsafe { b.launch(cfg)?; }
3395 Ok(())
3396 }
3397
3398 #[allow(clippy::too_many_arguments)]
3402 pub fn moe_gate_up_silu8_dev_q8_csr(&self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>,
3403 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3404 n_pairs: usize, in_f: usize, n_ff: usize, n_used: usize,
3405 n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize)
3406 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3407 let f = self.func("moe_gate_up_silu8_dev_q8_csr_iq4");
3408 let mut act = self.alloc_uninit::<f32>(n_pairs * n_ff)?;
3409 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_pairs as u32, 1),
3410 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3411 let (inf, nff, ne, nu, npi, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3412 n_used as i32, n_pairs as i32, rb_g as i64, rb_u as i64);
3413 let __s_b = self.gpu.stream();
3414 let mut b = __s_b.launch_builder(&f);
3415 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3416 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(&nu).arg(&npi);
3417 unsafe { b.launch(cfg)?; }
3418 Ok(act)
3419 }
3420
3421
3422 #[allow(clippy::too_many_arguments)]
3426 pub fn moe_down8_fma_dev_q8_variant(&self, variant: &str, table: &CudaSlice<u64>,
3427 sel: &cudarc::driver::CudaView<i32>,
3428 w: &cudarc::driver::CudaView<f32>,
3429 aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>,
3430 dst: &mut cudarc::driver::CudaViewMut<f32>,
3431 in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3432 qt: i32, rb: usize)
3433 -> Result<(), Box<dyn std::error::Error>> {
3434 let (inf, outf, nu, ne, rbi) = (in_f as i32, out_f as i32, n_used as i32,
3435 n_expert as i32, rb as i64);
3436 let (f, cfg) = match variant {
3437 "w8h2" | "w8h2v" => {
3438 (self.func(if variant == "w8h2" { "moe_down8_fma_dev_q8_w8h2" }
3439 else { "moe_down8_fma_dev_q8_w8h2v" }),
3440 LaunchConfig { grid_dim: (out_f.div_ceil(2) as u32, 1, 1),
3441 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3442 }
3443 "w8h2r2" | "w8h2r2v" => {
3444 (self.func(if variant == "w8h2r2" { "moe_down8_fma_dev_q8_w8h2r2" }
3445 else { "moe_down8_fma_dev_q8_w8h2r2v" }),
3446 LaunchConfig { grid_dim: (out_f.div_ceil(4) as u32, 1, 1),
3447 block_dim: (32, n_used as u32, 1), shared_mem_bytes: 0 })
3448 }
3449 _ => (self.func("moe_down8_fma_dev_q8"),
3450 LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3451 block_dim: (32, 1, 1), shared_mem_bytes: 0 }),
3452 };
3453 let __s_b = self.gpu.stream();
3454 let mut b = __s_b.launch_builder(&f);
3455 b.arg(table).arg(sel).arg(w).arg(aq2).arg(ad2).arg(dst)
3456 .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbi);
3457 unsafe { b.launch(cfg)?; }
3458 Ok(())
3459 }
3460
3461 #[allow(clippy::too_many_arguments)]
3463 pub fn moe_gate_up_silu8_dev_q8_variant(&self, variant: &str, table: &CudaSlice<u64>,
3464 sel: &cudarc::driver::CudaView<i32>,
3465 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
3466 in_f: usize, n_ff: usize, n_used: usize,
3467 n_expert: usize, qt_g: i32, qt_u: i32,
3468 rb_g: usize, rb_u: usize)
3469 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3470 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?;
3471 let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3472 rb_g as i64, rb_u as i64);
3473 let f = self.func(if variant == "v" { "moe_gate_up_silu8_dev_q8_v" }
3474 else { "moe_gate_up_silu8_dev_q8" });
3475 let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3476 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
3477 let __s_b = self.gpu.stream();
3478 let mut b = __s_b.launch_builder(&f);
3479 b.arg(table).arg(sel).arg(aq).arg(ad).arg(&mut act)
3480 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu);
3481 unsafe { b.launch(cfg)?; }
3482 Ok(act)
3483 }
3484
3485 pub fn moe_gate_up_silu8_dev(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3486 x: &cudarc::driver::CudaView<f32>,
3487 in_f: usize, n_ff: usize, n_used: usize, n_expert: usize,
3488 qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize,
3489 macros: &CudaSlice<f32>)
3490 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3491 let f = self.func("moe_gate_up_silu8_dev");
3492 let mut act = self.alloc_uninit::<f32>(n_used * n_ff)?; let cfg = LaunchConfig { grid_dim: (n_ff as u32, n_used as u32, 1),
3494 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3495 let (inf, nff, ne, rbg, rbu) = (in_f as i32, n_ff as i32, n_expert as i32,
3496 rb_g as i64, rb_u as i64);
3497 let __s_b = self.gpu.stream();
3498 let mut b = __s_b.launch_builder(&f);
3499 b.arg(table).arg(sel).arg(x).arg(&mut act)
3500 .arg(&inf).arg(&nff).arg(&ne).arg(&qt_g).arg(&qt_u).arg(&rbg).arg(&rbu).arg(macros);
3501 unsafe { b.launch(cfg)?; }
3502 Ok(act)
3503 }
3504
3505 #[allow(clippy::too_many_arguments)]
3508 pub fn moe_down8_fma_dev(&self, table: &CudaSlice<u64>, sel: &cudarc::driver::CudaView<i32>,
3509 w: &cudarc::driver::CudaView<f32>, act: &CudaSlice<f32>,
3510 dst: &mut cudarc::driver::CudaViewMut<f32>,
3511 in_f: usize, out_f: usize, n_used: usize, n_expert: usize,
3512 qt: i32, rb: usize)
3513 -> Result<(), Box<dyn std::error::Error>> {
3514 let f = self.func("moe_down8_fma_dev");
3515 let cfg = LaunchConfig { grid_dim: (out_f as u32, 1, 1),
3516 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3517 let (inf, outf, nu, ne, rbv) = (in_f as i32, out_f as i32, n_used as i32,
3518 n_expert as i32, rb as i64);
3519 let __s_b = self.gpu.stream();
3520 let mut b = __s_b.launch_builder(&f);
3521 b.arg(table).arg(sel).arg(w).arg(act).arg(dst)
3522 .arg(&inf).arg(&outf).arg(&nu).arg(&ne).arg(&qt).arg(&rbv);
3523 unsafe { b.launch(cfg)?; }
3524 Ok(())
3525 }
3526
3527 pub fn axpy_into(&self, src: &CudaSlice<f32>, alpha: f32,
3529 dst: &mut cudarc::driver::CudaViewMut<f32>, n: usize)
3530 -> Result<(), Box<dyn std::error::Error>> {
3531 let f = self.func("axpy_f32");
3532 let cfg = LaunchConfig::for_num_elems(n as u32);
3533 let (a, ni) = (alpha, n as i32);
3534 let __s_b = self.gpu.stream();
3535 let mut b = __s_b.launch_builder(&f);
3536 b.arg(src).arg(dst).arg(&a).arg(&ni);
3537 unsafe { b.launch(cfg)?; }
3538 Ok(())
3539 }
3540
3541 pub fn add_scaled_rows(&self, src: &CudaSlice<f32>, scale: &CudaSlice<f32>,
3543 dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize)
3544 -> Result<(), Box<dyn std::error::Error>> {
3545 let f = self.func("add_scaled_rows_f32");
3546 let cfg = LaunchConfig::for_num_elems((ncols * nrows) as u32);
3547 let (nc, nr) = (ncols as i32, nrows as i32);
3548 let __s_b = self.gpu.stream();
3549 let mut b = __s_b.launch_builder(&f);
3550 b.arg(src).arg(scale).arg(dst).arg(&nc).arg(&nr);
3551 unsafe { b.launch(cfg)?; }
3552 Ok(())
3553 }
3554
3555 pub fn gather_rows(&self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>,
3559 dst: &mut CudaSlice<f32>, ncols: usize, m_e: usize)
3560 -> Result<(), Box<dyn std::error::Error>> {
3561 let f = self.func("gather_rows_f32");
3562 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
3563 let (nc, me) = (ncols as i32, m_e as i32);
3564 let __s_b = self.gpu.stream();
3565 let mut b = __s_b.launch_builder(&f);
3566 b.arg(src).arg(idx).arg(dst).arg(&nc).arg(&me);
3567 unsafe { b.launch(cfg)?; }
3568 Ok(())
3569 }
3570
3571 pub fn scatter_slot(&self, src: &CudaSlice<f32>, tok_idx: &CudaSlice<i32>,
3576 slot_idx: &CudaSlice<i32>, weight: &CudaSlice<f32>,
3577 dst: &mut CudaSlice<f32>, wbuf: &mut CudaSlice<f32>,
3578 ncols: usize, n_used: usize, m_e: usize)
3579 -> Result<(), Box<dyn std::error::Error>> {
3580 let f = self.func("scatter_add_slot_f32");
3581 let cfg = LaunchConfig::for_num_elems((m_e * ncols) as u32);
3582 let (nc, nu, me) = (ncols as i32, n_used as i32, m_e as i32);
3583 let __s_b = self.gpu.stream();
3584 let mut b = __s_b.launch_builder(&f);
3585 b.arg(src).arg(tok_idx).arg(slot_idx).arg(weight).arg(dst).arg(wbuf).arg(&nc).arg(&nu).arg(&me);
3586 unsafe { b.launch(cfg)?; }
3587 Ok(())
3588 }
3589
3590 pub fn reduce_slots(&self, slots: &CudaSlice<f32>, wbuf: &CudaSlice<f32>,
3594 dst: &mut CudaSlice<f32>, ncols: usize, n_used: usize, t: usize)
3595 -> Result<(), Box<dyn std::error::Error>> {
3596 let f = self.func("reduce_slots_f32");
3597 let cfg = LaunchConfig::for_num_elems((t * ncols) as u32);
3598 let (nc, nu, ti) = (ncols as i32, n_used as i32, t as i32);
3599 let __s_b = self.gpu.stream();
3600 let mut b = __s_b.launch_builder(&f);
3601 b.arg(slots).arg(wbuf).arg(dst).arg(&nc).arg(&nu).arg(&ti);
3602 unsafe { b.launch(cfg)?; }
3603 Ok(())
3604 }
3605
3606 pub fn quantize_q8_1_view(&self, x: &cudarc::driver::CudaView<f32>, m: usize, in_f: usize)
3613 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3614 let f = self.func("quantize_q8_1");
3615 let nblk = in_f / 32;
3616 let mut q = self.alloc_uninit::<i8>(m * in_f)?;
3617 let mut d = self.alloc_uninit::<f32>(m * nblk)?;
3618 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
3619 let (inf, mi) = (in_f as i32, m as i32);
3620 let __s_b = self.gpu.stream();
3621 let mut b = __s_b.launch_builder(&f);
3622 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
3623 unsafe { b.launch(cfg)?; }
3624 Ok((q, d))
3625 }
3626
3627 pub fn quantize_q8_1(&self, x: &CudaSlice<f32>, m: usize, in_f: usize)
3628 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3629 let nblk = in_f / 32;
3630 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);
3634 let (inf, mi) = (in_f as i32, m as i32);
3635 if Self::pdl_on() && Self::pdl_wb_on() {
3636 {
3637 use cudarc::driver::{DevicePtr, DevicePtrMut};
3638 let s = &self.gpu.stream();
3639 let (px, _g0) = x.device_ptr(s);
3640 let (pq, _g1) = q.device_ptr_mut(s); let (pd, _g2) = d.device_ptr_mut(s);
3641 let mut ps = [
3642 &px as *const _ as *mut std::ffi::c_void, &pq as *const _ as *mut _,
3643 &pd as *const _ as *mut _, &inf as *const _ as *mut _,
3644 &mi as *const _ as *mut _,
3645 ];
3646 unsafe { self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?; }
3647 }
3648 return Ok((q, d));
3649 }
3650 let f = self.func("quantize_q8_1");
3651 let __s_b = self.gpu.stream();
3652 let mut b = __s_b.launch_builder(&f);
3653 b.arg(x).arg(&mut q).arg(&mut d).arg(&inf).arg(&mi);
3654 unsafe { b.launch(cfg)?; }
3655 Ok((q, d))
3656 }
3657
3658 pub fn quantize_fp4_act(&self, x: &CudaSlice<f32>, m: usize, in_f: usize)
3662 -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn std::error::Error>> {
3663 let f = self.func("quantize_fp4_act");
3664 let nb16 = in_f / 16;
3665 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);
3668 let (inf, mi) = (in_f as i32, m as i32);
3669 let __s_b = self.gpu.stream();
3670 let mut b = __s_b.launch_builder(&f);
3671 b.arg(x).arg(&mut aq4).arg(&mut ad4).arg(&inf).arg(&mi);
3672 unsafe { b.launch(cfg)?; }
3673 Ok((aq4, ad4))
3674 }
3675
3676 pub fn qmatvec_gemm_nvfp4_fp4(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3681 in_f: usize, out_f: usize, row_bytes: usize, scale: f32)
3682 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3683 assert!(in_f % 64 == 0, "FP4 GEMM requires in_f % 64 == 0, got {in_f}");
3684 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
3685 let mut y = self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)?;
3686 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
3687 Ok(y)
3688 }
3689
3690 fn fp4_gemm_launch(&self, bytes: &CudaSlice<u8>, aq4: &CudaSlice<u32>, ad4: &CudaSlice<u8>,
3693 m: usize, in_f: usize, out_f: usize, row_bytes: usize)
3694 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3695 let f = self.func("qmatvec_gemm_nvfp4_fp4");
3696 let mut y = self.alloc_uninit::<f32>(m * out_f)?; const BM: u32 = 64; const BN: u32 = 256;
3698 let cfg = LaunchConfig {
3699 grid_dim: ((out_f as u32 + BM - 1) / BM, (m as u32 + BN - 1) / BN, 1),
3700 block_dim: (32, 4, 1), shared_mem_bytes: 0,
3701 };
3702 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3703 let __s_b = self.gpu.stream();
3704 let mut b = __s_b.launch_builder(&f);
3705 b.arg(bytes).arg(aq4).arg(ad4).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3706 unsafe { b.launch(cfg)?; }
3707 Ok(y)
3708 }
3709
3710 pub fn qmatvec_gemm_nvfp4_fp4_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3712 in_f: usize, out_f: usize, row_bytes: usize)
3713 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3714 assert!(in_f % 64 == 0, "FP4 GEMM requires in_f % 64 == 0, got {in_f}");
3715 let (aq4, ad4) = self.quantize_fp4_act(x, m, in_f)?;
3716 self.fp4_gemm_launch(bytes, &aq4, &ad4, m, in_f, out_f, row_bytes)
3717 }
3718
3719 pub fn qmatvec_q8_0_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3721 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3722 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3723 let f = self.func("qmatvec_q8_0_dp4a");
3724 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3726 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3727 let __s_b = self.gpu.stream();
3728 let mut b = __s_b.launch_builder(&f);
3729 b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3730 unsafe { b.launch(cfg)?; }
3731 Ok(y)
3732 }
3733
3734 #[allow(non_snake_case)] pub fn qmatvec_q4_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3737 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3738 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3739 let f = self.func("qmatvec_q4_K_dp4a");
3740 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3742 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3743 let __s_b = self.gpu.stream();
3744 let mut b = __s_b.launch_builder(&f);
3745 b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3746 unsafe { b.launch(cfg)?; }
3747 Ok(y)
3748 }
3749
3750 #[allow(non_snake_case)] pub fn qmatvec_q6_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3753 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3754 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3755 let f = self.func("qmatvec_q6_K_dp4a");
3756 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3758 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3759 let __s_b = self.gpu.stream();
3760 let mut b = __s_b.launch_builder(&f);
3761 b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3762 unsafe { b.launch(cfg)?; }
3763 Ok(y)
3764 }
3765
3766 #[allow(non_snake_case)] pub fn qmatvec_q5_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3769 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3770 self.qmatvec_dp4a_named("qmatvec_q5_K_dp4a", w, x, m, in_f, out_f, row_bytes)
3771 }
3772 #[allow(non_snake_case)] pub fn qmatvec_q3_K_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3775 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3776 self.qmatvec_dp4a_named("qmatvec_q3_K_dp4a", w, x, m, in_f, out_f, row_bytes)
3777 }
3778 pub fn qmatvec_nvfp4_fast_rp(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3780 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3781 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}");
3782 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a_rp", w, x, m, in_f, out_f, row_bytes)
3783 }
3784 pub fn qmatvec_nvfp4_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3786 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3787 assert!(in_f % 64 == 0, "NVFP4 dp4a requires in_f % 64 == 0, got {in_f}");
3790 self.qmatvec_dp4a_named("qmatvec_nvfp4_dp4a", w, x, m, in_f, out_f, row_bytes)
3791 }
3792 #[allow(non_snake_case)] pub fn qmatvec_iq4_XS_fast(&self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
3795 out_f: usize, row_bytes: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3796 self.qmatvec_dp4a_named("qmatvec_iq4_XS_dp4a", w, x, m, in_f, out_f, row_bytes)
3797 }
3798
3799 fn qmatvec_dp4a_named(&self, name: &str, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
3801 in_f: usize, out_f: usize, row_bytes: usize)
3802 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3803 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
3804 let f = self.func(name);
3805 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
3807 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
3808 let __s_b = self.gpu.stream();
3809 let mut b = __s_b.launch_builder(&f);
3810 b.arg(w).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
3811 unsafe { b.launch(cfg)?; }
3812 Ok(y)
3813 }
3814
3815 pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3816 Ok(self.gpu.stream().clone_htod(v)?)
3817 }
3818 pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3819 Ok(self.gpu.stream().clone_htod(v)?)
3820 }
3821 pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
3823 Ok(self.gpu.stream().clone_htod(v)?)
3824 }
3825 pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
3826 Ok(self.gpu.stream().clone_htod(v)?)
3827 }
3828 pub fn dtoh_view(&self, d: &cudarc::driver::CudaView<f32>)
3830 -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3831 let v = self.gpu.stream().clone_dtoh(d)?;
3832 self.gpu.stream().synchronize()?;
3833 Ok(v)
3834 }
3835 pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3836 let v = self.gpu.stream().clone_dtoh(d)?;
3837 self.gpu.stream().synchronize()?;
3838 Ok(v)
3839 }
3840 pub fn dtoh_pair(
3844 &self,
3845 a: &CudaSlice<f32>,
3846 b: &CudaSlice<f32>,
3847 ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn std::error::Error>> {
3848 let av = self.gpu.stream().clone_dtoh(a)?;
3849 let bv = self.gpu.stream().clone_dtoh(b)?;
3850 self.gpu.stream().synchronize()?;
3851 Ok((av, bv))
3852 }
3853 pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn std::error::Error>> {
3855 let v = self.gpu.stream().clone_dtoh(d)?;
3856 self.gpu.stream().synchronize()?;
3857 Ok(v)
3858 }
3859 pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
3861 let v = self.gpu.stream().clone_dtoh(d)?;
3862 self.gpu.stream().synchronize()?;
3863 Ok(v)
3864 }
3865 pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3866 let s = self.gpu.stream().alloc_zeros::<f32>(n)?;
3867 self.keep_if_capturing(&s);
3868 Ok(s)
3869 }
3870
3871 pub fn prob_of_token_device(&self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>, n_vocab: usize)
3880 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3881 let nb = ARGMAX_NB;
3882 let mut part = self.alloc_uninit::<f32>(nb)?;
3883 let mut p = self.alloc_uninit::<f32>(1)?;
3884 let f1 = self.func("prob_of_token_partial_f32");
3885 let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3886 let nv = n_vocab as i32;
3887 let __s_b1 = self.gpu.stream();
3888 let mut b1 = __s_b1.launch_builder(&f1);
3889 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
3890 unsafe { b1.launch(cfg1)?; }
3891 let f2 = self.func("prob_of_token_final_f32");
3892 let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3893 let nbi = nb as i32;
3894 let __s_b2 = self.gpu.stream();
3895 let mut b2 = __s_b2.launch_builder(&f2);
3896 b2.arg(&part).arg(&mut p).arg(&nbi);
3897 unsafe { b2.launch(cfg2)?; }
3898 Ok(p)
3899 }
3900
3901 pub fn prob_of_token_device_col(&self, logits: &CudaSlice<f32>,
3908 tok_all: &CudaSlice<u32>, tok_idx: usize,
3909 p_out: &mut CudaSlice<f32>, p_idx: usize, n_vocab: usize)
3910 -> Result<(), Box<dyn std::error::Error>> {
3911 let tok_v = tok_all.slice(tok_idx..tok_idx + 1);
3912 let mut p_v = p_out.slice_mut(p_idx..p_idx + 1);
3913 let nb = ARGMAX_NB;
3914 let mut part = self.alloc_uninit::<f32>(nb)?;
3915 let f1 = self.func("prob_of_token_partial_f32");
3916 let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3917 let nv = n_vocab as i32;
3918 let __s_b1 = self.gpu.stream();
3919 let mut b1 = __s_b1.launch_builder(&f1);
3920 b1.arg(logits).arg(&tok_v).arg(&mut part).arg(&nv);
3921 unsafe { b1.launch(cfg1)?; }
3922 let f2 = self.func("prob_of_token_final_f32");
3923 let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3924 let nbi = nb as i32;
3925 let __s_b2 = self.gpu.stream();
3926 let mut b2 = __s_b2.launch_builder(&f2);
3927 b2.arg(&part).arg(&mut p_v).arg(&nbi);
3928 unsafe { b2.launch(cfg2)?; }
3929 Ok(())
3930 }
3931
3932 pub fn prob_of_token_device_into(&self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>,
3933 p_out: &mut CudaSlice<f32>, n_vocab: usize)
3934 -> Result<(), Box<dyn std::error::Error>> {
3935 let nb = ARGMAX_NB;
3936 let mut part = self.alloc_uninit::<f32>(nb)?;
3937 let f1 = self.func("prob_of_token_partial_f32");
3938 let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3939 let nv = n_vocab as i32;
3940 let __s_b1 = self.gpu.stream();
3941 let mut b1 = __s_b1.launch_builder(&f1);
3942 b1.arg(logits).arg(tok).arg(&mut part).arg(&nv);
3943 unsafe { b1.launch(cfg1)?; }
3944 let f2 = self.func("prob_of_token_final_f32");
3945 let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3946 let nbi = nb as i32;
3947 let __s_b2 = self.gpu.stream();
3948 let mut b2 = __s_b2.launch_builder(&f2);
3949 b2.arg(&part).arg(p_out).arg(&nbi);
3950 unsafe { b2.launch(cfg2)?; }
3951 Ok(())
3952 }
3953
3954 pub fn argmax_token_device(&self, logits: &CudaSlice<f32>, n_vocab: usize)
3955 -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
3956 let mut tok = unsafe { self.gpu.stream().alloc::<u32>(1)? };
3957 self.argmax_token_device_into(logits, &mut tok, n_vocab)?;
3958 Ok(tok)
3959 }
3960 pub fn argmax_token_device_into(&self, logits: &CudaSlice<f32>, tok: &mut CudaSlice<u32>,
3967 n_vocab: usize) -> Result<(), Box<dyn std::error::Error>> {
3968 let nb = ARGMAX_NB;
3969 let f1 = self.func("argmax_partial_f32");
3970 let f2 = self.func("argmax_final_f32");
3971 let mut guard = self.argmax_partials.lock().unwrap();
3972 if guard.is_none() {
3973 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
3976 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
3977 *guard = Some((pv, pi));
3978 }
3979 let (part_v, part_i) = guard.as_mut().unwrap();
3980 let nv = n_vocab as i32;
3981 let nbi = nb as i32;
3982 let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3984 let __s_b1 = self.gpu.stream();
3985 let mut b1 = __s_b1.launch_builder(&f1);
3986 b1.arg(logits).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
3987 unsafe { b1.launch(cfg1)?; }
3988 let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
3990 let __s_b2 = self.gpu.stream();
3991 let mut b2 = __s_b2.launch_builder(&f2);
3992 b2.arg(&*part_v).arg(&*part_i).arg(tok).arg(&nbi);
3993 unsafe { b2.launch(cfg2)?; }
3994 Ok(())
3995 }
3996 pub fn argmax_token_device_col(&self, logits: &CudaSlice<f32>, col: usize, n_vocab: usize,
4002 toks: &mut CudaSlice<u32>, out_idx: usize)
4003 -> Result<(), Box<dyn std::error::Error>> {
4004 let nb = ARGMAX_NB;
4005 let f1 = self.func("argmax_partial_f32");
4006 let f2 = self.func("argmax_final_f32");
4007 let mut guard = self.argmax_partials.lock().unwrap();
4008 if guard.is_none() {
4009 let pv = self.gpu.stream().alloc_zeros::<f32>(nb)?;
4010 let pi = self.gpu.stream().alloc_zeros::<i32>(nb)?;
4011 *guard = Some((pv, pi));
4012 }
4013 let (part_v, part_i) = guard.as_mut().unwrap();
4014 let col_view = logits.slice(col * n_vocab..(col + 1) * n_vocab);
4015 let nv = n_vocab as i32;
4016 let nbi = nb as i32;
4017 let cfg1 = LaunchConfig { grid_dim: (nb as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4018 let __s_b1 = self.gpu.stream();
4019 let mut b1 = __s_b1.launch_builder(&f1);
4020 b1.arg(&col_view).arg(&mut *part_v).arg(&mut *part_i).arg(&nv);
4021 unsafe { b1.launch(cfg1)?; }
4022 let mut tok_view = toks.slice_mut(out_idx..out_idx + 1);
4023 let cfg2 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4024 let __s_b2 = self.gpu.stream();
4025 let mut b2 = __s_b2.launch_builder(&f2);
4026 b2.arg(&*part_v).arg(&*part_i).arg(&mut tok_view).arg(&nbi);
4027 unsafe { b2.launch(cfg2)?; }
4028 Ok(())
4029 }
4030 pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
4032 Ok(self.gpu.stream().clone_htod(v)?)
4033 }
4034 pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4035 let v = self.gpu.stream().clone_dtoh(d)?;
4036 self.gpu.stream().synchronize()?;
4037 Ok(v)
4038 }
4039 pub fn htod_u32_into(&self, dst: &mut CudaSlice<u32>, src: &[u32])
4043 -> Result<(), Box<dyn std::error::Error>> {
4044 let mut view = dst.slice_mut(0..src.len());
4045 self.gpu.stream().memcpy_htod(src, &mut view)?;
4046 Ok(())
4047 }
4048
4049 pub fn alloc_u32_zeroed(&self, n: usize) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
4050 let s = self.gpu.stream().alloc_zeros::<u32>(n)?;
4051 self.keep_if_capturing(&s);
4052 Ok(s)
4053 }
4054 pub fn embed_gather_device_into(&self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>,
4057 x_out: &mut CudaSlice<f32>, n_embd: usize, qtype: i32,
4058 row_bytes: usize) -> Result<(), Box<dyn std::error::Error>> {
4059 let f = self.func("embed_gather_u32");
4060 let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
4061 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4062 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
4063 let __s_b = self.gpu.stream();
4064 let mut b = __s_b.launch_builder(&f);
4065 b.arg(embd).arg(token_d).arg(x_out).arg(&ne).arg(&qt).arg(&rb);
4066 unsafe { b.launch(cfg)?; }
4067 Ok(())
4068 }
4069 pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn std::error::Error>> {
4071 let v = self.gpu.stream().clone_dtoh(d)?;
4072 self.gpu.stream().synchronize()?;
4073 Ok(v[0])
4074 }
4075 pub fn i32_set_k(&self, dst: &mut CudaSlice<i32>, v: i32)
4082 -> Result<(), Box<dyn std::error::Error>> {
4083 let f = self.func("i32_set_k");
4084 let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
4085 let idx = 0i32;
4086 let __s_b = self.gpu.stream();
4087 let mut b = __s_b.launch_builder(&f);
4088 b.arg(dst).arg(&v).arg(&idx);
4089 unsafe { b.launch(cfg)?; }
4090 Ok(())
4091 }
4092
4093 pub fn set_i32_one(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
4094 self.gpu.stream().memcpy_htod(&[v], d)?;
4095 Ok(())
4096 }
4097 pub fn set_u32_one(&self, d: &mut CudaSlice<u32>, v: u32) -> Result<(), Box<dyn std::error::Error>> {
4100 self.gpu.stream().memcpy_htod(&[v], d)?;
4101 Ok(())
4102 }
4103 pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn std::error::Error>> {
4105 let v = self.gpu.stream().clone_dtoh(d)?;
4106 self.gpu.stream().synchronize()?;
4107 Ok(v[0])
4108 }
4109 pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
4111 Ok(self.gpu.stream().clone_htod(bytes)?)
4112 }
4113 pub fn embed_gather_device(&self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>,
4117 n_embd: usize, qtype: i32, row_bytes: usize)
4118 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4119 let f = self.func("embed_gather_u32");
4120 let mut x = self.alloc_uninit::<f32>(n_embd)?;
4121 let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), 1, 1),
4122 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4123 let (ne, qt, rb) = (n_embd as i32, qtype, row_bytes as i64);
4124 let __s_b = self.gpu.stream();
4125 let mut b = __s_b.launch_builder(&f);
4126 b.arg(embd).arg(token_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb);
4127 unsafe { b.launch(cfg)?; }
4128 Ok(x)
4129 }
4130
4131
4132 pub fn embed_gather_device_t(&self, embd: &CudaSlice<u8>, tokens: &[u32],
4136 n_embd: usize, qtype: i32, row_bytes: usize)
4137 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4138 let t = tokens.len();
4139 let tok_d = self.gpu.stream().clone_htod(tokens)?;
4140 let f = self.func("embed_gather_u32_t");
4141 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4142 let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4143 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4144 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4145 let __s_b = self.gpu.stream();
4146 let mut b = __s_b.launch_builder(&f);
4147 b.arg(embd).arg(&tok_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4148 unsafe { b.launch(cfg)?; }
4149 Ok(x)
4150 }
4151
4152 pub fn embed_gather_device_tv(&self, embd: &CudaSlice<u8>, tok_v: &cudarc::driver::CudaView<u32>,
4157 t: usize, n_embd: usize, qtype: i32, row_bytes: usize)
4158 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4159 let f = self.func("embed_gather_u32_t");
4160 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4161 let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4162 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4163 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4164 let __s_b = self.gpu.stream();
4165 let mut b = __s_b.launch_builder(&f);
4166 b.arg(embd).arg(tok_v).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4167 unsafe { b.launch(cfg)?; }
4168 Ok(x)
4169 }
4170
4171 pub fn embed_gather_device_td(&self, embd: &CudaSlice<u8>, tok_d: &CudaSlice<u32>, t: usize,
4172 n_embd: usize, qtype: i32, row_bytes: usize)
4173 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4174 let f = self.func("embed_gather_u32_t");
4175 let mut x = self.alloc_uninit::<f32>(t * n_embd)?;
4176 let cfg = LaunchConfig { grid_dim: (((n_embd as u32 + 255) / 256).max(1), t as u32, 1),
4177 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
4178 let (ne, qt, rb, ti) = (n_embd as i32, qtype, row_bytes as i64, t as i32);
4179 let __s_b = self.gpu.stream();
4180 let mut b = __s_b.launch_builder(&f);
4181 b.arg(embd).arg(tok_d).arg(&mut x).arg(&ne).arg(&qt).arg(&rb).arg(&ti);
4182 unsafe { b.launch(cfg)?; }
4183 Ok(x)
4184 }
4185
4186 #[inline]
4192 fn keep_if_capturing<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, s: &CudaSlice<T>) {
4194 if self.capture_keep_on.load(std::sync::atomic::Ordering::Relaxed) {
4195 self.capture_keep.lock().unwrap().push(Box::new(s.clone()));
4196 }
4197 }
4198
4199 fn alloc_uninit<T: cudarc::driver::DeviceRepr + Send + 'static>(&self, n: usize)
4200 -> Result<CudaSlice<T>, Box<dyn std::error::Error>> {
4201 let mut s = unsafe { self.gpu.stream().alloc::<T>(n)? };
4202 {
4206 static Z: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4207 if *Z.get_or_init(|| std::env::var("MEMRA_DEBUG_ZERO_ALLOCS").as_deref() == Ok("1")) {
4208 use cudarc::driver::DevicePtrMut;
4210 let n_bytes = s.len() * std::mem::size_of::<T>();
4211 let stream = self.gpu.stream();
4212 let (p_, _g) = s.device_ptr_mut(&stream);
4213 unsafe {
4214 cudarc::driver::sys::cuMemsetD8Async(p_, 0, n_bytes, stream.cu_stream())
4215 .result()?;
4216 }
4217 }
4218 }
4219 self.keep_if_capturing(&s);
4220 Ok(s)
4221 }
4222
4223 pub fn uninit_q8_pair(&self, n: usize)
4228 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4229 Ok((self.alloc_uninit::<i8>(n)?, self.alloc_uninit::<f32>(n / 32)?))
4230 }
4231
4232 pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4233 self.alloc_uninit::<f32>(n)
4234 }
4235
4236 pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn std::error::Error>> {
4238 self.alloc_uninit::<i8>(n)
4239 }
4240
4241 #[allow(clippy::too_many_arguments)]
4245 pub fn rms_norm3(&self, x: &CudaSlice<f32>, w0: &CudaSlice<f32>, w1: &CudaSlice<f32>,
4246 w2: &CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>,
4247 d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4248 -> Result<(), Box<dyn std::error::Error>> {
4249 let f = self.func("rms_norm3_f32");
4250 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4251 let (nc, e) = (ncols as i32, eps);
4252 let __s_b = self.gpu.stream();
4253 let mut b = __s_b.launch_builder(&f);
4254 b.arg(x).arg(w0).arg(w1).arg(w2).arg(d0).arg(d1).arg(d2).arg(&nc).arg(&e);
4255 unsafe { b.launch(cfg)?; }
4256 Ok(())
4257 }
4258
4259 #[allow(clippy::too_many_arguments)]
4261 pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool {
4264 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4265 *WARP_ON.get_or_init(|| {
4266 std::env::var("MEMRA_QKVNORM_W").map(|v| v != "0").unwrap_or(true)
4267 }) && ncols % 4 == 0 && rows >= 64
4268 }
4269
4270 #[allow(clippy::too_many_arguments)]
4273 pub fn rms_norm_qkv_w4b(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
4274 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4275 dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>,
4276 dvb: &mut CudaSlice<u8>,
4277 ncols: usize, rq: usize, rk: usize, eps: f32, vf16: bool)
4278 -> Result<(), Box<dyn std::error::Error>> {
4279 assert!(ncols % 4 == 0 && rq + 2 * rk >= 64);
4280 let f = self.func("rms_norm_qkv_w4b_f32");
4281 let rows = (rq + 2 * rk) as u32;
4282 let cfg = LaunchConfig {
4283 grid_dim: (rows.div_ceil(8), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0,
4284 };
4285 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
4286 let vf = vf16 as i32;
4287 let __s_b = self.gpu.stream();
4288 let mut b = __s_b.launch_builder(&f);
4289 b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv).arg(&mut *dvb)
4290 .arg(&nc).arg(&rqi).arg(&rki).arg(&rvi).arg(&e).arg(&vf);
4291 unsafe { b.launch(cfg)?; }
4292 Ok(())
4293 }
4294
4295 pub fn rms_norm_qkv(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
4296 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4297 dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>,
4298 ncols: usize, rq: usize, rk: usize, eps: f32)
4299 -> Result<(), Box<dyn std::error::Error>> {
4300 static WARP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4304 let warp_on = *WARP_ON.get_or_init(|| {
4305 std::env::var("MEMRA_QKVNORM_W").map(|v| v != "0").unwrap_or(true)
4306 });
4307 if warp_on && ncols % 4 == 0 && rq + 2 * rk >= 64 {
4310 let f = self.func("rms_norm_qkv_w4_f32");
4311 let rows = (rq + 2 * rk) as u32;
4312 let cfg = LaunchConfig {
4313 grid_dim: (rows.div_ceil(8), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0,
4314 };
4315 let (nc, rqi, rki, rvi, e) = (ncols as i32, rq as i32, rk as i32, rk as i32, eps);
4316 let __s_b = self.gpu.stream();
4317 let mut b = __s_b.launch_builder(&f);
4318 b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv)
4319 .arg(&nc).arg(&rqi).arg(&rki).arg(&rvi).arg(&e);
4320 unsafe { b.launch(cfg)?; }
4321 return Ok(());
4322 }
4323 let f = self.func("rms_norm_qkv_f32");
4324 let grid = (rq + 2 * rk) as u32;
4325 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4326 let (nc, rqi, rki, e) = (ncols as i32, rq as i32, rk as i32, eps);
4327 let __s_b = self.gpu.stream();
4328 let mut b = __s_b.launch_builder(&f);
4329 b.arg(q).arg(k).arg(v).arg(wq).arg(wk).arg(wv).arg(dq).arg(dk).arg(dv)
4330 .arg(&nc).arg(&rqi).arg(&rki).arg(&e);
4331 unsafe { b.launch(cfg)?; }
4332 Ok(())
4333 }
4334
4335 #[allow(clippy::too_many_arguments)]
4337 pub fn rms_norm2x(&self, a: &CudaSlice<f32>, bb: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4338 wb: &CudaSlice<f32>, da: &mut CudaSlice<f32>, db: &mut CudaSlice<f32>,
4339 ncols: usize, nrows: usize, eps: f32)
4340 -> Result<(), Box<dyn std::error::Error>> {
4341 let f = self.func("rms_norm2x_f32");
4342 let cfg = LaunchConfig { grid_dim: (2 * nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4343 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
4344 let __s_b = self.gpu.stream();
4345 let mut b = __s_b.launch_builder(&f);
4346 b.arg(a).arg(bb).arg(wa).arg(wb).arg(da).arg(db).arg(&nc).arg(&nr).arg(&e);
4347 unsafe { b.launch(cfg)?; }
4348 Ok(())
4349 }
4350
4351 pub fn softcap(&self, y: &mut CudaSlice<f32>, cap: f32, n: usize)
4353 -> Result<(), Box<dyn std::error::Error>> {
4354 let f = self.func("softcap_f32");
4355 let cfg = LaunchConfig::for_num_elems(n as u32);
4356 let ni = n as i32;
4357 let __s_b = self.gpu.stream();
4358 let mut b = __s_b.launch_builder(&f);
4359 b.arg(y).arg(&cap).arg(&ni);
4360 unsafe { b.launch(cfg)?; }
4361 Ok(())
4362 }
4363
4364 pub fn mask_ids_rows(&self, y: &mut CudaSlice<f32>, ids: &CudaSlice<i32>, n_ids: usize,
4367 n_vocab: usize, t: usize)
4368 -> Result<(), Box<dyn std::error::Error>> {
4369 let f = self.func("mask_ids_rows_f32");
4370 let cfg = LaunchConfig::for_num_elems((n_ids * t) as u32);
4371 let (ni, nv, ti) = (n_ids as i32, n_vocab as i32, t as i32);
4372 let __s_b = self.gpu.stream();
4373 let mut b = __s_b.launch_builder(&f);
4374 b.arg(y).arg(ids).arg(&ni).arg(&nv).arg(&ti);
4375 unsafe { b.launch(cfg)?; }
4376 Ok(())
4377 }
4378
4379 #[allow(clippy::too_many_arguments)]
4381 pub fn add_scale_rms_norm(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4382 w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4383 ncols: usize, nrows: usize, eps: f32)
4384 -> Result<(), Box<dyn std::error::Error>> {
4385 let f = self.func("add_scale_rms_norm_f32");
4386 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4387 let (nc, e2) = (ncols as i32, eps);
4388 let __s_b = self.gpu.stream();
4389 let mut b = __s_b.launch_builder(&f);
4390 b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(dst).arg(&nc).arg(&e2);
4391 unsafe { b.launch(cfg)?; }
4392 Ok(())
4393 }
4394
4395 #[allow(clippy::too_many_arguments)]
4398 pub fn add_scale_rms_norm_q8_1(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4399 w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4400 ncols: usize, nrows: usize, eps: f32)
4401 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4402 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4403 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4404 let (nc, e2) = (ncols as i32, eps);
4405 if Self::pdl_on() && Self::pdl_wb_on() {
4406 {
4407 use cudarc::driver::{DevicePtr, DevicePtrMut};
4408 let s = &self.gpu.stream();
4409 let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b_in.device_ptr(s);
4410 let (pw, _g2) = w.device_ptr(s); let (pr, _g3) = res.device_ptr_mut(s);
4411 let (pq, _g4) = out_q.device_ptr_mut(s); let (pd, _g5) = out_d.device_ptr_mut(s);
4412 let mut ps = [
4413 &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4414 &c as *const _ as *mut _, &pw as *const _ as *mut _,
4415 &pr as *const _ as *mut _, &pq as *const _ as *mut _,
4416 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4417 &e2 as *const _ as *mut _,
4418 ];
4419 unsafe { self.launch_pdl("add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4420 (rms_block(), 1, 1), &mut ps)?; }
4421 }
4422 return Ok((out_q, out_d));
4423 }
4424 let f = self.func("add_scale_rms_norm_q8_1");
4425 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4426 let __s_b = self.gpu.stream();
4427 let mut b = __s_b.launch_builder(&f);
4428 b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e2);
4429 unsafe { b.launch(cfg)?; }
4430 Ok((out_q, out_d))
4431 }
4432
4433 #[allow(clippy::too_many_arguments)]
4435 pub fn add_scale_rms_norm_q8_1_into(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4436 w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4437 ncols: usize, nrows: usize, eps: f32,
4438 out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>)
4439 -> Result<(), Box<dyn std::error::Error>> {
4440 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
4441 let (nc, e2) = (ncols as i32, eps);
4442 if Self::pdl_on() && Self::pdl_wb_on() {
4443 use cudarc::driver::{DevicePtr, DevicePtrMut};
4444 let s = &self.gpu.stream();
4445 let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b_in.device_ptr(s);
4446 let (pw, _g2) = w.device_ptr(s); let (pr, _g3) = res.device_ptr_mut(s);
4447 let (pq, _g4) = out_q.device_ptr_mut(s); let (pd, _g5) = out_d.device_ptr_mut(s);
4448 let mut ps = [
4449 &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4450 &c as *const _ as *mut _, &pw as *const _ as *mut _,
4451 &pr as *const _ as *mut _, &pq as *const _ as *mut _,
4452 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4453 &e2 as *const _ as *mut _,
4454 ];
4455 unsafe { self.launch_pdl("add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4456 (rms_block(), 1, 1), &mut ps)?; }
4457 return Ok(());
4458 }
4459 let f = self.func("add_scale_rms_norm_q8_1");
4460 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4461 let __s_b = self.gpu.stream();
4462 let mut b = __s_b.launch_builder(&f);
4463 b.arg(a).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut *out_q).arg(&mut *out_d).arg(&nc).arg(&e2);
4464 unsafe { b.launch(cfg)?; }
4465 Ok(())
4466 }
4467
4468 #[allow(clippy::too_many_arguments)]
4471 pub fn rms_pre_add_scale_rms_norm_q8_1(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4472 b_in: &CudaSlice<f32>, c: f32,
4473 w: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
4474 ncols: usize, nrows: usize, eps: f32)
4475 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4476 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4477 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4478 let (nc, e2) = (ncols as i32, eps);
4479 if Self::pdl_on() {
4480 {
4481 use cudarc::driver::{DevicePtr, DevicePtrMut};
4482 let s = &self.gpu.stream();
4483 let (pa, _g0) = a.device_ptr(s); let (pwa, _g1) = wa.device_ptr(s);
4484 let (pb, _g2) = b_in.device_ptr(s); let (pw, _g3) = w.device_ptr(s);
4485 let (pr, _g4) = res.device_ptr_mut(s);
4486 let (pq, _g5) = out_q.device_ptr_mut(s); let (pd, _g6) = out_d.device_ptr_mut(s);
4487 let mut ps = [
4488 &pa as *const _ as *mut std::ffi::c_void, &pwa as *const _ as *mut _,
4489 &pb as *const _ as *mut _, &c as *const _ as *mut _,
4490 &pw as *const _ as *mut _, &pr as *const _ as *mut _,
4491 &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4492 &nc as *const _ as *mut _, &e2 as *const _ as *mut _,
4493 ];
4494 unsafe { self.launch_pdl("rms_pre_add_scale_rms_norm_q8_1", (nrows as u32, 1, 1),
4495 (rms_block(), 1, 1), &mut ps)?; }
4496 }
4497 return Ok((out_q, out_d));
4498 }
4499 let f = self.func("rms_pre_add_scale_rms_norm_q8_1");
4500 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4501 let __s_b = self.gpu.stream();
4502 let mut b = __s_b.launch_builder(&f);
4503 b.arg(a).arg(wa).arg(b_in).arg(&c).arg(w).arg(res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e2);
4504 unsafe { b.launch(cfg)?; }
4505 Ok((out_q, out_d))
4506 }
4507
4508 pub fn gelu_tanh_mul_q8_1(&self, gate: &CudaSlice<f32>, up: &cudarc::driver::CudaView<f32>,
4511 act: &mut CudaSlice<f32>, ncols: usize, nrows: usize)
4512 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4513 debug_assert!(ncols % 128 == 0);
4514 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4515 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4516 let nc = ncols as i32;
4517 if Self::pdl_on() {
4518 {
4519 use cudarc::driver::{DevicePtr, DevicePtrMut};
4520 let s = &self.gpu.stream();
4521 let (pg, _g0) = gate.device_ptr(s); let (pu, _g1) = up.device_ptr(s);
4522 let (pact, _g2) = act.device_ptr_mut(s);
4523 let (pq, _g3) = out_q.device_ptr_mut(s); let (pd, _g4) = out_d.device_ptr_mut(s);
4524 let mut ps = [
4525 &pg as *const _ as *mut std::ffi::c_void, &pu as *const _ as *mut _,
4526 &pact as *const _ as *mut _, &pq as *const _ as *mut _,
4527 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4528 ];
4529 unsafe { self.launch_pdl("gelu_tanh_mul_q8_1", (nrows as u32, 1, 1),
4530 (rms_block(), 1, 1), &mut ps)?; }
4531 }
4532 return Ok((out_q, out_d));
4533 }
4534 let f = self.func("gelu_tanh_mul_q8_1");
4535 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4536 let __s_b = self.gpu.stream();
4537 let mut b = __s_b.launch_builder(&f);
4538 b.arg(gate).arg(up).arg(act).arg(&mut out_q).arg(&mut out_d).arg(&nc);
4539 unsafe { b.launch(cfg)?; }
4540 Ok((out_q, out_d))
4541 }
4542
4543 #[allow(clippy::too_many_arguments)]
4545 pub fn gelu_tanh_mul_q8_1_into(&self, gate: &CudaSlice<f32>, up: &cudarc::driver::CudaView<f32>,
4546 act: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
4547 out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>)
4548 -> Result<(), Box<dyn std::error::Error>> {
4549 debug_assert!(ncols % 128 == 0);
4550 debug_assert!(out_q.len() >= nrows * ncols && out_d.len() >= nrows * (ncols / 32));
4551 let nc = ncols as i32;
4552 if Self::pdl_on() {
4553 use cudarc::driver::{DevicePtr, DevicePtrMut};
4554 let s = &self.gpu.stream();
4555 let (pg, _g0) = gate.device_ptr(s); let (pu, _g1) = up.device_ptr(s);
4556 let (pact, _g2) = act.device_ptr_mut(s);
4557 let (pq, _g3) = out_q.device_ptr_mut(s); let (pd, _g4) = out_d.device_ptr_mut(s);
4558 let mut ps = [
4559 &pg as *const _ as *mut std::ffi::c_void, &pu as *const _ as *mut _,
4560 &pact as *const _ as *mut _, &pq as *const _ as *mut _,
4561 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4562 ];
4563 unsafe { self.launch_pdl("gelu_tanh_mul_q8_1", (nrows as u32, 1, 1),
4564 (rms_block(), 1, 1), &mut ps)?; }
4565 return Ok(());
4566 }
4567 let f = self.func("gelu_tanh_mul_q8_1");
4568 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4569 let __s_b = self.gpu.stream();
4570 let mut b = __s_b.launch_builder(&f);
4571 b.arg(gate).arg(up).arg(&mut *act).arg(&mut *out_q).arg(&mut *out_d).arg(&nc);
4572 unsafe { b.launch(cfg)?; }
4573 Ok(())
4574 }
4575
4576 #[allow(clippy::too_many_arguments)]
4578 pub fn add_rms_norm3_q8z(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>,
4579 w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>,
4580 res: &mut CudaSlice<f32>, out1: &mut CudaSlice<f32>,
4581 ncols: usize, nrows: usize, eps: f32)
4582 -> Result<((CudaSlice<i8>, CudaSlice<f32>), (CudaSlice<i8>, CudaSlice<f32>)), Box<dyn std::error::Error>> {
4583 let mut q0 = self.alloc_uninit::<i8>(nrows * ncols)?;
4584 let mut d0 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4585 let mut q2 = self.alloc_uninit::<i8>(nrows * ncols)?;
4586 let mut d2 = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4587 let f = self.func("add_rms_norm3_q8z_f32");
4588 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4589 let (nc, e2) = (ncols as i32, eps);
4590 let __s_b = self.gpu.stream();
4591 let mut b = __s_b.launch_builder(&f);
4592 b.arg(a).arg(b_in).arg(w0).arg(w1).arg(w2).arg(res)
4593 .arg(&mut q0).arg(&mut d0).arg(out1).arg(&mut q2).arg(&mut d2).arg(&nc).arg(&e2);
4594 unsafe { b.launch(cfg)?; }
4595 Ok(((q0, d0), (q2, d2)))
4596 }
4597
4598 #[allow(clippy::too_many_arguments)]
4600 pub fn add_rms_norm3(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>,
4601 w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>,
4602 res: &mut CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>,
4603 d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4604 -> Result<(), Box<dyn std::error::Error>> {
4605 let f = self.func("add_rms_norm3_f32");
4606 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4607 let (nc, e2) = (ncols as i32, eps);
4608 let __s_b = self.gpu.stream();
4609 let mut b = __s_b.launch_builder(&f);
4610 b.arg(a).arg(b_in).arg(w0).arg(w1).arg(w2).arg(res).arg(d0).arg(d1).arg(d2).arg(&nc).arg(&e2);
4611 unsafe { b.launch(cfg)?; }
4612 Ok(())
4613 }
4614
4615 pub fn add_scale(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32,
4617 dst: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
4618 let f = self.func("add_scale_f32");
4619 let cfg = LaunchConfig::for_num_elems(n as u32);
4620 let ni = n as i32;
4621 let __s_b = self.gpu.stream();
4622 let mut b = __s_b.launch_builder(&f);
4623 b.arg(a).arg(b_in).arg(&c).arg(dst).arg(&ni);
4624 unsafe { b.launch(cfg)?; }
4625 Ok(())
4626 }
4627
4628 pub fn rms_norm(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4629 ncols: usize, nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4630 let (nc, e) = (ncols as i32, eps);
4631 if Self::pdl_on() && Self::pdl_wb_on() {
4632 use cudarc::driver::{DevicePtr, DevicePtrMut};
4633 let s = &self.gpu.stream();
4634 let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4635 let (pd, _g2) = dst.device_ptr_mut(s);
4636 let mut ps = [
4637 &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4638 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4639 &e as *const _ as *mut _,
4640 ];
4641 unsafe { self.launch_pdl("rms_norm_f32", (nrows as u32, 1, 1),
4642 (rms_block(), 1, 1), &mut ps)?; }
4643 return Ok(());
4644 }
4645 let f = self.func("rms_norm_f32");
4646 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4647 let __s_b = self.gpu.stream();
4648 let mut b = __s_b.launch_builder(&f);
4649 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
4650 unsafe { b.launch(cfg)?; }
4651 Ok(())
4652 }
4653
4654 pub fn rms_norm_decode(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4662 ncols: usize, nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4663 let f = self.func("rms_norm_f32");
4664 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4665 let (nc, e) = (ncols as i32, eps);
4666 let __s_b = self.gpu.stream();
4667 let mut b = __s_b.launch_builder(&f);
4668 b.arg(x).arg(w).arg(dst).arg(&nc).arg(&e);
4669 unsafe { b.launch(cfg)?; }
4670 Ok(())
4671 }
4672
4673 pub fn rms_norm_q8_1(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize, nrows: usize,
4677 eps: f32) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4678 let nblk = ncols / 32;
4679 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
4680 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
4681 let (nc, e) = (ncols as i32, eps);
4682 if Self::pdl_on() {
4683 {
4684 use cudarc::driver::{DevicePtr, DevicePtrMut};
4685 let s = &self.gpu.stream();
4686 let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4687 let (pq, _g2) = q.device_ptr_mut(s); let (pd, _g3) = d.device_ptr_mut(s);
4688 let mut ps = [
4689 &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4690 &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4691 &nc as *const _ as *mut _, &e as *const _ as *mut _,
4692 ];
4693 unsafe { self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1),
4694 &mut ps)?; }
4695 }
4696 return Ok((q, d));
4697 }
4698 let f = self.func("rms_norm_q8_1");
4699 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4702 let __s_b = self.gpu.stream();
4703 let mut b = __s_b.launch_builder(&f);
4704 b.arg(x).arg(w).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
4705 unsafe { b.launch(cfg)?; }
4706 Ok((q, d))
4707 }
4708
4709 pub fn rms_norm_q8_1_into(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize,
4712 nrows: usize, eps: f32,
4713 q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>)
4714 -> Result<(), Box<dyn std::error::Error>> {
4715 let nblk = ncols / 32;
4716 debug_assert!(q.len() >= nrows * ncols && d.len() >= nrows * nblk);
4717 let (nc, e) = (ncols as i32, eps);
4718 if Self::pdl_on() {
4719 use cudarc::driver::{DevicePtr, DevicePtrMut};
4720 let s = &self.gpu.stream();
4721 let (px, _g0) = x.device_ptr(s); let (pw, _g1) = w.device_ptr(s);
4722 let (pq, _g2) = q.device_ptr_mut(s); let (pd, _g3) = d.device_ptr_mut(s);
4723 let mut ps = [
4724 &px as *const _ as *mut std::ffi::c_void, &pw as *const _ as *mut _,
4725 &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4726 &nc as *const _ as *mut _, &e as *const _ as *mut _,
4727 ];
4728 unsafe { self.launch_pdl("rms_norm_q8_1", (nrows as u32, 1, 1), (1024, 1, 1),
4729 &mut ps)?; }
4730 return Ok(());
4731 }
4732 let f = self.func("rms_norm_q8_1");
4733 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4734 let __s_b = self.gpu.stream();
4735 let mut b = __s_b.launch_builder(&f);
4736 b.arg(x).arg(w).arg(&mut *q).arg(&mut *d).arg(&nc).arg(&e);
4737 unsafe { b.launch(cfg)?; }
4738 Ok(())
4739 }
4740
4741 pub fn quantize_q8_1_into(&self, x: &CudaSlice<f32>, m: usize, in_f: usize,
4743 q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>)
4744 -> Result<(), Box<dyn std::error::Error>> {
4745 let nblk = in_f / 32;
4746 debug_assert!(q.len() >= m * in_f && d.len() >= m * nblk);
4747 let cfg = LaunchConfig::for_num_elems((m * in_f) as u32);
4748 let (inf, mi) = (in_f as i32, m as i32);
4749 if Self::pdl_on() && Self::pdl_wb_on() {
4750 use cudarc::driver::{DevicePtr, DevicePtrMut};
4751 let s = &self.gpu.stream();
4752 let (px, _g0) = x.device_ptr(s);
4753 let (pq, _g1) = q.device_ptr_mut(s); let (pd, _g2) = d.device_ptr_mut(s);
4754 let mut ps = [
4755 &px as *const _ as *mut std::ffi::c_void, &pq as *const _ as *mut _,
4756 &pd as *const _ as *mut _, &inf as *const _ as *mut _,
4757 &mi as *const _ as *mut _,
4758 ];
4759 unsafe { self.launch_pdl("quantize_q8_1", cfg.grid_dim, cfg.block_dim, &mut ps)?; }
4760 return Ok(());
4761 }
4762 let f = self.func("quantize_q8_1");
4763 let __s_b = self.gpu.stream();
4764 let mut b = __s_b.launch_builder(&f);
4765 b.arg(x).arg(&mut *q).arg(&mut *d).arg(&inf).arg(&mi);
4766 unsafe { b.launch(cfg)?; }
4767 Ok(())
4768 }
4769
4770 pub fn add_rms_norm_q8_1(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>,
4774 res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
4775 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4776 let nblk = ncols / 32;
4777 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
4778 let mut d = self.alloc_uninit::<f32>(nrows * nblk)?;
4779 let f = self.func("add_rms_norm_q8_1");
4780 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
4782 let (nc, e) = (ncols as i32, eps);
4783 let __s_bld = self.gpu.stream();
4784 let mut bld = __s_bld.launch_builder(&f);
4785 bld.arg(a).arg(b_in).arg(w).arg(res).arg(&mut q).arg(&mut d).arg(&nc).arg(&e);
4786 unsafe { bld.launch(cfg)?; }
4787 Ok((q, d))
4788 }
4789
4790 pub fn add_rms_norm(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4794 res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
4795 eps: f32) -> Result<(), Box<dyn std::error::Error>> {
4796 let (nc, e) = (ncols as i32, eps);
4797 if Self::pdl_on() && Self::pdl_wb_on() {
4798 use cudarc::driver::{DevicePtr, DevicePtrMut};
4799 let s = &self.gpu.stream();
4800 let (pa, _g0) = a.device_ptr(s); let (pb, _g1) = b.device_ptr(s);
4801 let (pw, _g2) = w.device_ptr(s);
4802 let (pr, _g3) = res.device_ptr_mut(s); let (pd, _g4) = dst.device_ptr_mut(s);
4803 let mut ps = [
4804 &pa as *const _ as *mut std::ffi::c_void, &pb as *const _ as *mut _,
4805 &pw as *const _ as *mut _, &pr as *const _ as *mut _,
4806 &pd as *const _ as *mut _, &nc as *const _ as *mut _,
4807 &e as *const _ as *mut _,
4808 ];
4809 unsafe { self.launch_pdl("add_rms_norm_f32", (nrows as u32, 1, 1),
4810 (rms_block(), 1, 1), &mut ps)?; }
4811 return Ok(());
4812 }
4813 let f = self.func("add_rms_norm_f32");
4814 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4815 let __s_b2 = self.gpu.stream();
4816 let mut b2 = __s_b2.launch_builder(&f);
4817 b2.arg(a).arg(b).arg(w).arg(&mut *res).arg(&mut *dst).arg(&nc).arg(&e);
4818 unsafe { b2.launch(cfg)?; }
4819 Ok(())
4820 }
4821
4822 #[allow(clippy::too_many_arguments)]
4825 pub fn rms_pre_add_rms_norm(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4826 b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4827 res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4828 ncols: usize, nrows: usize, eps: f32)
4829 -> Result<(), Box<dyn std::error::Error>> {
4830 let f = self.func("rms_pre_add_rms_norm_f32");
4831 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4832 let (nc, e) = (ncols as i32, eps);
4833 let __s_b2 = self.gpu.stream();
4834 let mut b2 = __s_b2.launch_builder(&f);
4835 b2.arg(a).arg(wa).arg(b).arg(w).arg(&mut *res).arg(&mut *dst).arg(&nc).arg(&e);
4836 unsafe { b2.launch(cfg)?; }
4837 Ok(())
4838 }
4839
4840 #[allow(clippy::too_many_arguments)]
4842 pub fn rms_pre_add_rms_norm_q8z(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>,
4843 b: &CudaSlice<f32>, w: &CudaSlice<f32>,
4844 res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
4845 ncols: usize, nrows: usize, eps: f32)
4846 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4847 debug_assert!(ncols % 128 == 0);
4848 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
4849 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
4850 let (nc, e) = (ncols as i32, eps);
4851 if Self::pdl_on() {
4852 {
4853 use cudarc::driver::{DevicePtr, DevicePtrMut};
4854 let s = &self.gpu.stream();
4855 let (pa, _g0) = a.device_ptr(s); let (pwa, _g1) = wa.device_ptr(s);
4856 let (pb, _g2) = b.device_ptr(s); let (pw, _g3) = w.device_ptr(s);
4857 let (pr, _g4) = res.device_ptr_mut(s); let (pdst, _g5) = dst.device_ptr_mut(s);
4858 let (pq, _g6) = out_q.device_ptr_mut(s); let (pd, _g7) = out_d.device_ptr_mut(s);
4859 let mut ps = [
4860 &pa as *const _ as *mut std::ffi::c_void, &pwa as *const _ as *mut _,
4861 &pb as *const _ as *mut _, &pw as *const _ as *mut _,
4862 &pr as *const _ as *mut _, &pdst as *const _ as *mut _,
4863 &pq as *const _ as *mut _, &pd as *const _ as *mut _,
4864 &nc as *const _ as *mut _, &e as *const _ as *mut _,
4865 ];
4866 unsafe { self.launch_pdl("rms_pre_add_rms_norm_q8z_f32", (nrows as u32, 1, 1),
4867 (rms_block(), 1, 1), &mut ps)?; }
4868 }
4869 return Ok((out_q, out_d));
4870 }
4871 let f = self.func("rms_pre_add_rms_norm_q8z_f32");
4872 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4873 let __s_b2 = self.gpu.stream();
4874 let mut b2 = __s_b2.launch_builder(&f);
4875 b2.arg(a).arg(wa).arg(b).arg(w).arg(&mut *res).arg(&mut *dst)
4876 .arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&e);
4877 unsafe { b2.launch(cfg)?; }
4878 Ok((out_q, out_d))
4879 }
4880
4881 pub fn build_q4_out_concat3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
4885 w2: &crate::model::GpuTensor)
4886 -> Result<Option<crate::model::GpuTensor>, Box<dyn std::error::Error>> {
4887 use crate::model::GpuTensor;
4888 let part = |w: &GpuTensor| -> Option<(usize, usize)> {
4889 match w {
4890 GpuTensor::Quant { qtype, row_bytes, rp, .. }
4891 if *qtype == QT_Q4_0 && !*rp => Some((*row_bytes, w.out_features())),
4892 _ => None,
4893 }
4894 };
4895 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (part(w0), part(w1), part(w2))
4896 else { return Ok(None) };
4897 if rb0 != rb1 || rb0 != rb2
4898 || w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
4899 return Ok(None);
4900 }
4901 fn bytes_of(w: &crate::model::GpuTensor) -> &CudaSlice<u8> {
4902 match w { crate::model::GpuTensor::Quant { bytes, .. } => bytes, _ => unreachable!() }
4903 }
4904 let (b0, b1, b2) = (bytes_of(w0), bytes_of(w1), bytes_of(w2));
4905 let total = rb0 * (o0 + o1 + o2);
4906 let mut cat = self.alloc_u8(total)?;
4907 self.copy_u8_into(&mut cat, 0, b0, rb0 * o0)?;
4908 self.copy_u8_into(&mut cat, rb0 * o0, b1, rb1 * o1)?;
4909 self.copy_u8_into(&mut cat, rb0 * (o0 + o1), b2, rb2 * o2)?;
4910 Ok(Some(GpuTensor::Quant {
4911 bytes: cat, qtype: QT_Q4_0, row_bytes: rb0,
4912 ne: vec![w0.in_features() as u64, (o0 + o1 + o2) as u64], scale: 1.0, rp: false,
4913 #[cfg(memra_cutlass)]
4914 cutlass: None,
4915 fp8: None, rp4: None, f16: None,
4916 }))
4917 }
4918
4919 #[allow(clippy::too_many_arguments)]
4921 pub fn rms_norm_qkv_rope_cat(&self, qkv: &CudaSlice<f32>,
4922 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4923 q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
4924 head_dim: usize, rq: usize, rk: usize,
4925 pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
4926 base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32)
4927 -> Result<(), Box<dyn std::error::Error>> {
4928 let rows = rq + rk + rk;
4929 let theta_scale = base.powf(-2.0 / head_dim as f32);
4930 let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
4931 if Self::pdl_on() {
4932 use cudarc::driver::{DevicePtr, DevicePtrMut};
4933 let s = &self.gpu.stream();
4934 let (pqkv, _g0) = qkv.device_ptr(s);
4935 let (pwq, _g1) = wq.device_ptr(s); let (pwk, _g2) = wk.device_ptr(s);
4936 let (pwv, _g3) = wv.device_ptr(s);
4937 let (pq, _g4) = q.device_ptr_mut(s); let (pk, _g5) = k.device_ptr_mut(s);
4938 let (pv, _g6) = v.device_ptr_mut(s);
4939 let (ppos, _g7) = pos.device_ptr(s);
4940 let (pff, _g8) = match ff {
4941 Some(t) => { let (p, g) = t.device_ptr(s); (p, Some(g)) }
4942 None => (0, None),
4943 };
4944 let mut ps = [
4945 &pqkv as *const _ as *mut std::ffi::c_void,
4946 &pwq as *const _ as *mut _, &pwk as *const _ as *mut _,
4947 &pwv as *const _ as *mut _,
4948 &pq as *const _ as *mut _, &pk as *const _ as *mut _,
4949 &pv as *const _ as *mut _,
4950 &nc as *const _ as *mut _, &rqi as *const _ as *mut _,
4951 &rki as *const _ as *mut _, &ppos as *const _ as *mut _,
4952 &nhq as *const _ as *mut _, &nhk as *const _ as *mut _,
4953 &theta_scale as *const _ as *mut _, &freq_scale as *const _ as *mut _,
4954 &pff as *const _ as *mut _, &eps as *const _ as *mut _,
4955 ];
4956 unsafe { self.launch_pdl("rms_norm_qkv_rope_cat_f32", (rows as u32, 1, 1),
4957 (rms_block(), 1, 1), &mut ps)?; }
4958 return Ok(());
4959 }
4960 let f = self.func("rms_norm_qkv_rope_cat_f32");
4961 let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4962 let __s_b = self.gpu.stream();
4963 let mut b = __s_b.launch_builder(&f);
4964 match ff {
4965 Some(t) => { b.arg(qkv).arg(wq).arg(wk).arg(wv)
4966 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4967 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
4968 .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps);
4969 unsafe { b.launch(cfg)?; } }
4970 None => { let null: u64 = 0;
4971 b.arg(qkv).arg(wq).arg(wk).arg(wv)
4972 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4973 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
4974 .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps);
4975 unsafe { b.launch(cfg)?; } }
4976 }
4977 Ok(())
4978 }
4979
4980 #[allow(clippy::too_many_arguments)]
4982 pub fn rms_norm_qkv_rope(&self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>, v0: &CudaSlice<f32>,
4983 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
4984 q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
4985 head_dim: usize, rq: usize, rk: usize,
4986 pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
4987 base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32)
4988 -> Result<(), Box<dyn std::error::Error>> {
4989 let f = self.func("rms_norm_qkv_rope_f32");
4990 let rows = rq + rk + rk; let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
4992 let theta_scale = base.powf(-2.0 / head_dim as f32);
4993 let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
4994 let __s_b = self.gpu.stream();
4995 let mut b = __s_b.launch_builder(&f);
4996 match ff {
4997 Some(t) => { b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
4998 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
4999 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5000 .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps);
5001 unsafe { b.launch(cfg)?; } }
5002 None => { let null: u64 = 0;
5003 b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5004 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5005 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5006 .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps);
5007 unsafe { b.launch(cfg)?; } }
5008 }
5009 Ok(())
5010 }
5011
5012 #[allow(clippy::too_many_arguments)]
5016 pub fn rms_norm_qkv_rope_append_dc(&self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>,
5017 v0: &CudaSlice<f32>,
5018 wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>,
5019 q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>,
5020 head_dim: usize, rq: usize, rk: usize,
5021 pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize,
5022 base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32,
5023 kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>,
5024 t_dev: &CudaSlice<i32>, k_tok_bytes: usize, v_tok_bytes: usize,
5025 g: bool)
5026 -> Result<(), Box<dyn std::error::Error>> {
5027 let rows = rq + rk + rk;
5028 let theta_scale = base.powf(-2.0 / head_dim as f32);
5029 let (nc, rqi, rki, nhq, nhk) = (head_dim as i32, rq as i32, rk as i32, nh_q as i32, nh_k as i32);
5030 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
5031 if Self::pdl_on() && Self::pdl_wb_on() {
5032 use cudarc::driver::{DevicePtr, DevicePtrMut};
5033 let s = &self.gpu.stream();
5034 let (p0, _a0) = q0.device_ptr(s); let (p1, _a1) = k0.device_ptr(s);
5035 let (p2, _a2) = v0.device_ptr(s);
5036 let (pwq, _a3) = wq.device_ptr(s); let (pwk, _a4) = wk.device_ptr(s);
5037 let (pwv, _a5) = wv.device_ptr(s);
5038 let (pq, _a6) = q.device_ptr_mut(s); let (pk, _a7) = k.device_ptr_mut(s);
5039 let (pv, _a8) = v.device_ptr_mut(s);
5040 let (pp, _a9) = pos.device_ptr(s);
5041 let pff: u64 = match ff { Some(t) => { let (p, _gg) = t.device_ptr(s); p as u64 }
5042 None => 0 };
5043 let (pkc, _a10) = kc.device_ptr_mut(s); let (pvc, _a11) = vc.device_ptr_mut(s);
5044 let (pt, _a12) = t_dev.device_ptr(s);
5045 let mut ps = [
5046 &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
5047 &p2 as *const _ as *mut _, &pwq as *const _ as *mut _,
5048 &pwk as *const _ as *mut _, &pwv as *const _ as *mut _,
5049 &pq as *const _ as *mut _, &pk as *const _ as *mut _,
5050 &pv as *const _ as *mut _, &nc as *const _ as *mut _,
5051 &rqi as *const _ as *mut _, &rki as *const _ as *mut _,
5052 &pp as *const _ as *mut _, &nhq as *const _ as *mut _,
5053 &nhk as *const _ as *mut _, &theta_scale as *const _ as *mut _,
5054 &freq_scale as *const _ as *mut _, &pff as *const _ as *mut _,
5055 &eps as *const _ as *mut _, &pkc as *const _ as *mut _,
5056 &pvc as *const _ as *mut _, &pt as *const _ as *mut _,
5057 &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
5058 ];
5059 unsafe { self.launch_pdl_flash(g, "rms_norm_qkv_rope_append_dc_f32",
5060 (rows as u32, 1, 1), (rms_block(), 1, 1), 0, &mut ps)?; }
5061 return Ok(());
5062 }
5063 let f = if g { self.func_g("rms_norm_qkv_rope_append_dc_f32") }
5064 else { self.func("rms_norm_qkv_rope_append_dc_f32") };
5065 let cfg = LaunchConfig { grid_dim: (rows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
5066 let __s_b = self.gpu.stream();
5067 let mut b = __s_b.launch_builder(&f);
5068 match ff {
5069 Some(t) => { b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5070 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5071 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5072 .arg(&theta_scale).arg(&freq_scale).arg(t).arg(&eps)
5073 .arg(&mut *kc).arg(&mut *vc).arg(t_dev).arg(&ktb).arg(&vtb);
5074 unsafe { b.launch(cfg)?; } }
5075 None => { let null: u64 = 0;
5076 b.arg(q0).arg(k0).arg(v0).arg(wq).arg(wk).arg(wv)
5077 .arg(&mut *q).arg(&mut *k).arg(&mut *v)
5078 .arg(&nc).arg(&rqi).arg(&rki).arg(pos).arg(&nhq).arg(&nhk)
5079 .arg(&theta_scale).arg(&freq_scale).arg(&null).arg(&eps)
5080 .arg(&mut *kc).arg(&mut *vc).arg(t_dev).arg(&ktb).arg(&vtb);
5081 unsafe { b.launch(cfg)?; } }
5082 }
5083 Ok(())
5084 }
5085
5086 pub fn add_q8_1(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, res: &mut CudaSlice<f32>,
5088 ncols: usize, nrows: usize)
5089 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5090 debug_assert!(ncols % 128 == 0);
5091 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
5092 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
5093 let f = self.func("add_q8_1_f32");
5094 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
5095 let nc = ncols as i32;
5096 let __s_b2 = self.gpu.stream();
5097 let mut b2 = __s_b2.launch_builder(&f);
5098 b2.arg(a).arg(b).arg(&mut *res).arg(&mut out_q).arg(&mut out_d).arg(&nc);
5099 unsafe { b2.launch(cfg)?; }
5100 Ok((out_q, out_d))
5101 }
5102
5103 pub fn rms_pre_add_q8_1(&self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b: &CudaSlice<f32>,
5107 res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
5108 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5109 debug_assert!(ncols % 128 == 0);
5110 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
5111 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
5112 let f = self.func("rms_pre_add_q8_1_f32");
5113 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1),
5114 shared_mem_bytes: 0 };
5115 let (nc, ep) = (ncols as i32, eps);
5116 let __s_b2 = self.gpu.stream();
5117 let mut b2 = __s_b2.launch_builder(&f);
5118 b2.arg(a).arg(wa).arg(b).arg(&mut *res).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&ep);
5119 unsafe { b2.launch(cfg)?; }
5120 Ok((out_q, out_d))
5121 }
5122
5123 pub fn l2_v2_on(ncols: usize) -> bool {
5127 ncols == 128 && std::env::var("MEMRA_L2_V2").as_deref() != Ok("0")
5128 }
5129
5130 pub fn l2_norm_pp(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
5131 dst16: Option<&mut CudaSlice<u8>>, ncols: usize, nrows: usize,
5132 eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5133 if Self::l2_v2_on(ncols) {
5134 let f = self.func("l2_norm_pp_v2_f32");
5135 let rows_per_block = 8u32; let cfg = LaunchConfig { grid_dim: ((nrows as u32).div_ceil(rows_per_block), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
5137 let (nc, nr, e) = (ncols as i32, nrows as i32, eps);
5138 let d16: u64 = match dst16 { Some(d) => self.addr_u8(d), None => 0 };
5140 let __s_b = self.gpu.stream();
5141 let mut b = __s_b.launch_builder(&f);
5142 b.arg(x).arg(dst).arg(&d16).arg(&nc).arg(&nr).arg(&e);
5143 unsafe { b.launch(cfg)?; }
5144 return Ok(());
5145 }
5146 self.l2_norm(x, dst, ncols, nrows, eps)
5147 }
5148
5149 pub fn l2_norm(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize,
5150 eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5151 let f = self.func("l2_norm_f32");
5152 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
5153 let (nc, e) = (ncols as i32, eps);
5154 let __s_b = self.gpu.stream();
5155 let mut b = __s_b.launch_builder(&f);
5156 b.arg(x).arg(dst).arg(&nc).arg(&e);
5157 unsafe { b.launch(cfg)?; }
5158 Ok(())
5159 }
5160
5161 pub fn l2_norm_decode(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize,
5167 nrows: usize, eps: f32) -> Result<(), Box<dyn std::error::Error>> {
5168 let f = self.func("l2_norm_f32");
5169 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
5170 let (nc, e) = (ncols as i32, eps);
5171 let __s_b = self.gpu.stream();
5172 let mut b = __s_b.launch_builder(&f);
5173 b.arg(x).arg(dst).arg(&nc).arg(&e);
5174 unsafe { b.launch(cfg)?; }
5175 Ok(())
5176 }
5177
5178 pub fn rope_neox(&self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize,
5180 n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32, freq_scale: f32)
5181 -> Result<(), Box<dyn std::error::Error>> {
5182 let f = self.func("rope_neox_f32");
5183 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5184 let grid = (n_heads * n_tokens) as u32;
5185 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5186 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
5187 let __s_b = self.gpu.stream();
5188 let mut b = __s_b.launch_builder(&f);
5189 b.arg(x).arg(pos).arg(&hd).arg(&nd).arg(&nh).arg(&theta_scale).arg(&freq_scale);
5190 unsafe { b.launch(cfg)?; }
5191 Ok(())
5192 }
5193
5194 pub fn rope_neox_ff(&self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize,
5196 n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32,
5197 freq_scale: f32, ff: &CudaSlice<f32>)
5198 -> Result<(), Box<dyn std::error::Error>> {
5199 let f = self.func("rope_neox_ff_f32");
5200 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5201 let grid = (n_heads * n_tokens) as u32;
5202 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5203 let (hd, nd, nh) = (head_dim as i32, n_dims as i32, n_heads as i32);
5204 let __s_b = self.gpu.stream();
5205 let mut b = __s_b.launch_builder(&f);
5206 b.arg(x).arg(pos).arg(&hd).arg(&nd).arg(&nh).arg(&theta_scale).arg(&freq_scale).arg(ff);
5207 unsafe { b.launch(cfg)?; }
5208 Ok(())
5209 }
5210
5211 #[allow(clippy::too_many_arguments)]
5213 pub fn rope_neox2(&self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>,
5214 pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize,
5215 nh_q: usize, nh_k: usize, n_tokens: usize, freq_base: f32,
5216 freq_scale: f32, ff: Option<&CudaSlice<f32>>)
5217 -> Result<(), Box<dyn std::error::Error>> {
5218 let f = self.func("rope_neox2_f32");
5219 let theta_scale = (freq_base).powf(-2.0 / n_dims as f32);
5220 let grid = ((nh_q + nh_k) * n_tokens) as u32;
5221 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
5222 let (hd, nd, nq, nk, nt) = (head_dim as i32, n_dims as i32, nh_q as i32, nh_k as i32, n_tokens as i32);
5223 let __s_b = self.gpu.stream();
5224 let mut b = __s_b.launch_builder(&f);
5225 b.arg(q).arg(k).arg(pos).arg(&hd).arg(&nd).arg(&nq).arg(&nk).arg(&nt)
5226 .arg(&theta_scale).arg(&freq_scale);
5227 match ff {
5228 Some(ffv) => { b.arg(ffv); unsafe { b.launch(cfg)?; } }
5229 None => {
5230 let null: u64 = 0;
5231 b.arg(&null);
5232 unsafe { b.launch(cfg)?; }
5233 }
5234 }
5235 Ok(())
5236 }
5237
5238 pub fn gelu_tanh_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5240 -> Result<(), Box<dyn std::error::Error>> {
5241 let f = self.func("gelu_tanh_mul_f32");
5242 let cfg = LaunchConfig::for_num_elems(n as u32);
5243 let ni = n as i32;
5244 let __s_b = self.gpu.stream();
5245 let mut b = __s_b.launch_builder(&f);
5246 b.arg(gate).arg(up).arg(dst).arg(&ni);
5247 unsafe { b.launch(cfg)?; }
5248 Ok(())
5249 }
5250
5251 pub fn silu_mul(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5252 -> Result<(), Box<dyn std::error::Error>> {
5253 let f = self.func("silu_mul_f32");
5254 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5256 let ni = n as i32;
5257 let __s_b = self.gpu.stream();
5258 let mut b = __s_b.launch_builder(&f);
5259 b.arg(gate).arg(up).arg(dst).arg(&ni);
5260 unsafe { b.launch(cfg)?; }
5261 Ok(())
5262 }
5263
5264 pub fn silu_mul_f16out(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>,
5267 dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize)
5268 -> Result<(), Box<dyn std::error::Error>> {
5269 let f = self.func("silu_mul_f16out_f32");
5270 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5271 let ni = n as i32;
5272 let __s_b = self.gpu.stream();
5273 let mut b = __s_b.launch_builder(&f);
5274 b.arg(gate).arg(up).arg(dst).arg(dst16).arg(&ni);
5275 unsafe { b.launch(cfg)?; }
5276 Ok(())
5277 }
5278
5279 pub fn silu_mul_scaled(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5286 dst: &mut CudaSlice<f32>, n: usize) -> Result<(), Box<dyn std::error::Error>> {
5287 let f = self.func("silu_mul_scaled_f32");
5288 let cfg = LaunchConfig::for_num_elems(n as u32);
5289 let ni = n as i32;
5290 let (gsf, usf) = (gs, us);
5291 let __s_b = self.gpu.stream();
5292 let mut b = __s_b.launch_builder(&f);
5293 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(dst).arg(&ni);
5294 unsafe { b.launch(cfg)?; }
5295 Ok(())
5296 }
5297
5298 #[allow(clippy::too_many_arguments)]
5302 pub fn swigluoai_mul_scaled(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5303 alpha: f32, limit: f32, dst: &mut CudaSlice<f32>, n: usize)
5304 -> Result<(), Box<dyn std::error::Error>> {
5305 let f = self.func("swigluoai_mul_scaled_f32");
5306 let cfg = LaunchConfig::for_num_elems(n as u32);
5307 let ni = n as i32;
5308 let __s_b = self.gpu.stream();
5309 let mut b = __s_b.launch_builder(&f);
5310 b.arg(gate).arg(up).arg(&gs).arg(&us).arg(&alpha).arg(&limit).arg(dst).arg(&ni);
5311 unsafe { b.launch(cfg)?; }
5312 Ok(())
5313 }
5314
5315 pub fn silu_mul_scaled_q8_1(&self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32,
5323 n: usize)
5324 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5325 let f = self.func("silu_mul_scaled_q8_1");
5326 let nblk = n / 32;
5327 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);
5331 let (gsf, usf, ni) = (gs, us, n as i32);
5332 let __s_b = self.gpu.stream();
5333 let mut b = __s_b.launch_builder(&f);
5334 b.arg(gate).arg(up).arg(&gsf).arg(&usf).arg(&mut aq).arg(&mut ad).arg(&ni);
5335 unsafe { b.launch(cfg)?; }
5336 Ok((aq, ad))
5337 }
5338
5339 pub fn add(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5340 -> Result<(), Box<dyn std::error::Error>> {
5341 let f = self.func("add_f32");
5342 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
5344 let ni = n as i32;
5345 let __s_bld = self.gpu.stream();
5346 let mut bld = __s_bld.launch_builder(&f);
5347 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
5348 unsafe { bld.launch(cfg)?; }
5349 Ok(())
5350 }
5351
5352 pub fn mul(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize)
5353 -> Result<(), Box<dyn std::error::Error>> {
5354 let f = self.func("mul_f32");
5355 let cfg = LaunchConfig::for_num_elems(n as u32);
5356 let ni = n as i32;
5357 let __s_bld = self.gpu.stream();
5358 let mut bld = __s_bld.launch_builder(&f);
5359 bld.arg(a).arg(b_in).arg(dst).arg(&ni);
5360 unsafe { bld.launch(cfg)?; }
5361 Ok(())
5362 }
5363
5364 pub fn matmul(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize)
5367 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5368 use crate::model::GpuTensor;
5369 let in_f = w.in_features();
5370 let out_f = w.out_features();
5371 #[allow(non_snake_case)]
5379 let GEMM_M_THRESHOLD = if self.verify_exact_on() { usize::MAX } else { 16usize };
5382
5383 const GEMM_MIN_OUT_F: usize = 128; if m >= GEMM_M_THRESHOLD {
5408 if let Some(y) = self.try_fp8_gemm(w, x, m)? { return Ok(y); }
5409 if let Some(y) = self.try_fp8_blk_mmq(w, x, m)? { return Ok(y); }
5413 if let Some(y) = self.try_f16_gemm(w, x, m)? { return Ok(y); }
5416 }
5417 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.mmq_supports(w) {
5418 return self.qmatvec_mmq(w, x, m);
5419 }
5420 if m >= GEMM_M_THRESHOLD && out_f >= GEMM_MIN_OUT_F && self.gemm_supports(w) {
5421 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5422 return self.qmatvec_gemm(w, &aq, &ad, m);
5423 }
5424 if m >= GEMM_M_THRESHOLD {
5427 if let Some(y) = self.try_fp4_gemm(w, x, m, in_f, out_f)? { return Ok(y); }
5428 }
5429 let fast = std::env::var("MEMRA_FAST").as_deref() != Ok("0");
5433 if m == 1 && fast {
5438 if let GpuTensor::Quant { bytes, qtype, row_bytes, rp, rp4, scale, .. } = w {
5439 if self.mmvq_supports(*qtype) {
5440 let (bytes, rp) = match rp4 { Some(m4) => (m4, true), None => (bytes, *rp) };
5444 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5445 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, *scale, rp);
5446 }
5447 }
5448 }
5449 if (2..=16).contains(&m) && fast && std::env::var("MEMRA_NO_BATCHED").is_err()
5465 && (m <= 4 || Self::b8_enabled()) {
5466 let m_ok = m <= 8 || matches!(w, GpuTensor::Quant { qtype, rp4, .. }
5470 if *qtype == QT_Q4_0 || *qtype == QT_Q6_K || (*qtype == QT_Q8_0 && rp4.is_some()));
5471 if m_ok {
5472 if let GpuTensor::Quant { bytes, qtype, row_bytes, rp, rp4, .. } = w {
5473 if self.batched_supports(*qtype) && self.mmvq_supports(*qtype) {
5474 let (bytes, rp) = match rp4 { Some(m4) => (m4, true), None => (bytes, *rp) };
5475 let mcols = Self::batched_mcols(m);
5476 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5477 let mut y = self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes, mcols, 1.0, rp)?;
5478 if let GpuTensor::Quant { scale, .. } = w {
5479 if *scale != 1.0 { self.scale_inplace(&mut y, *scale, m * out_f)?; }
5480 }
5481 return Ok(y);
5482 }
5483 }
5484 }
5485 }
5486 if fast {
5492 if let GpuTensor::Quant { bytes, qtype, row_bytes, scale, .. } = w {
5493 if *qtype == QT_F8_E4M3 {
5494 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5495 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, *qtype, *row_bytes,
5496 *scale, false);
5497 }
5498 }
5499 }
5500 let mut y = match w {
5501 GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q8_0 =>
5502 self.qmatvec_q8_0_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5503 GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q4_K =>
5504 self.qmatvec_q4_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5505 GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q6_K =>
5506 self.qmatvec_q6_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5507 GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q5_K =>
5508 self.qmatvec_q5_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5509 GpuTensor::Quant { bytes, qtype, row_bytes, .. } if fast && *qtype == QT_Q3_K =>
5510 self.qmatvec_q3_K_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5511 GpuTensor::Quant { bytes, qtype, row_bytes, rp, .. } if fast && *qtype == QT_NVFP4 =>
5512 self.qmatvec_dp4a_named(
5513 if *rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
5514 bytes, x, m, in_f, out_f, *row_bytes)?,
5515 GpuTensor::Quant { bytes, qtype, row_bytes, .. }
5519 if fast && *qtype == QT_IQ4_XS && Self::iq_fast_enabled() =>
5520 self.qmatvec_iq4_XS_fast(bytes, x, m, in_f, out_f, *row_bytes)?,
5521 GpuTensor::Quant { bytes, qtype, row_bytes, rp, .. } =>
5526 self.qmatvec(bytes, x, m, in_f, out_f,
5529 if *rp && *qtype == QT_NVFP4 { QT_NVFP4_RP } else { *qtype },
5530 *row_bytes)?,
5531 GpuTensor::Float { data, .. } => self.linear(x, data, m, in_f, out_f)?,
5532 GpuTensor::FloatBf16 { data, .. } =>
5535 self.linear_bf16_chunked(x, data, m, in_f, out_f, false)?,
5536 };
5537 if let GpuTensor::Quant { scale, .. } = w {
5539 if *scale != 1.0 { self.scale_inplace(&mut y, *scale, m * out_f)?; }
5540 }
5541 Ok(y)
5542 }
5543
5544 pub fn uses_q8_1_fast(&self, w: &crate::model::GpuTensor) -> bool {
5547 use crate::model::GpuTensor;
5548 if std::env::var("MEMRA_FAST").as_deref() == Ok("0") { return false; }
5549 match w {
5550 GpuTensor::Quant { qtype, .. } => matches!(*qtype,
5551 QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q3_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0)
5552 || (*qtype == QT_IQ4_XS && Self::iq_fast_enabled()),
5553 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
5554 }
5555 }
5556
5557 pub fn matmul_pre(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
5562 x_fallback: &CudaSlice<f32>, m: usize)
5563 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5564 use crate::model::GpuTensor;
5565 let x_raw_ok = x_fallback.len() >= m * w.in_features();
5571 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
5574 if let Some(y) = self.try_fp8_gemm(w, x_fallback, m)? { return Ok(y); }
5575 if let Some(y) = self.try_fp8_blk_mmq(w, x_fallback, m)? { return Ok(y); }
5578 if let Some(y) = self.try_f16_gemm(w, x_fallback, m)? { return Ok(y); }
5580 }
5581 if m >= 16 && w.out_features() >= 128 && self.mmq_supports(w) && !self.verify_exact_on()
5586 && x_raw_ok {
5587 return self.qmatvec_mmq(w, x_fallback, m);
5588 }
5589 if m >= 16 && x_raw_ok && !self.verify_exact_on() {
5592 if let Some(y) = self.try_fp4_gemm(w, x_fallback, m, w.in_features(), w.out_features())? {
5593 return Ok(y);
5594 }
5595 }
5596 if m >= 16 && self.gemm_supports(w) && !self.verify_exact_on() {
5599 return self.qmatvec_gemm(w, aq, ad, m);
5600 }
5601 if !self.uses_q8_1_fast(w) { return self.matmul(w, x_fallback, m); }
5602 let in_f = w.in_features();
5603 let out_f = w.out_features();
5604 let (bytes, qtype, row_bytes, scale, rp) = match w {
5605 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5606 _ => unreachable!("uses_q8_1_fast guaranteed Quant"),
5607 };
5608 let (mbytes, mrp) = match w {
5611 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5612 _ => (bytes, rp),
5613 };
5614 if m == 1 && self.mmvq_supports(qtype) {
5618 return self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, mrp);
5619 }
5620 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5633 && std::env::var("MEMRA_NO_BATCHED").is_err()
5634 && (m <= 4 || Self::b8_enabled())
5635 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && mrp)) {
5639 let mcols = Self::batched_mcols(m);
5640 return self.qmatvec_mmvq_batched(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, mrp);
5641 }
5642 if qtype == QT_F8_E4M3 || qtype == QT_Q4_0 {
5648 let (b2, r2) = if qtype == QT_Q4_0 { (mbytes, mrp) } else { (bytes, rp) };
5649 return self.qmatvec_mmvq(b2, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, r2);
5650 }
5651 let name = match qtype {
5652 QT_Q8_0 => "qmatvec_q8_0_dp4a", QT_Q4_K => "qmatvec_q4_K_dp4a",
5653 QT_Q6_K => "qmatvec_q6_K_dp4a", QT_Q5_K => "qmatvec_q5_K_dp4a",
5654 QT_Q3_K => "qmatvec_q3_K_dp4a",
5655 QT_NVFP4 => if rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
5656 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
5657 _ => unreachable!(),
5658 };
5659 let f = self.func(name);
5660 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
5662 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5663 let __s_b = self.gpu.stream();
5664 let mut b = __s_b.launch_builder(&f);
5665 b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
5666 unsafe { b.launch(cfg)?; }
5667 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
5668 Ok(y)
5669 }
5670
5671 pub fn matmul_decode_exact(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize)
5679 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5680 use crate::model::GpuTensor;
5681 if let GpuTensor::Float { data, .. } = w {
5689 return self.linear_decode_exact(x, data, m, w.in_features(), w.out_features());
5690 }
5691 if let GpuTensor::FloatBf16 { data, .. } = w {
5694 let (in_f, out_f) = (w.in_features(), w.out_features());
5695 return self.linear_bf16_chunked(x, data, m, in_f, out_f, true);
5696 }
5697 if !self.uses_q8_1_fast(w) { return self.matmul(w, x, m); }
5698 let in_f = w.in_features();
5699 let out_f = w.out_features();
5700 let (bytes, qtype, row_bytes, scale, rp) = match w {
5701 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5702 _ => return self.matmul(w, x, m),
5703 };
5704 let (bytes, rp) = match w {
5707 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5708 _ => (bytes, rp),
5709 };
5710 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5711 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5720 && std::env::var("MEMRA_NO_BATCHED").is_err()
5721 && (m <= 4 || Self::b8_enabled())
5722 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && rp)) {
5724 let mcols = Self::batched_mcols(m);
5725 return self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp);
5726 }
5727 if self.mmvq_supports(qtype) {
5728 return self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
5731 }
5732 self.matmul_pre(w, &aq, &ad, x, m)
5735 }
5736
5737 pub fn matmul_decode_exact_pre(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>,
5747 ad: &CudaSlice<f32>, m: usize)
5748 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5749 use crate::model::GpuTensor;
5750 debug_assert!(self.uses_q8_1_fast(w),
5751 "matmul_decode_exact_pre: caller must guarantee q8_1-fast");
5752 let in_f = w.in_features();
5753 let out_f = w.out_features();
5754 let (bytes, qtype, row_bytes, scale, rp) = match w {
5755 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } =>
5756 (bytes, *qtype, *row_bytes, *scale, *rp),
5757 _ => return Err("matmul_decode_exact_pre: Quant tensor required (q8_1-fast contract)".into()),
5758 };
5759 let (bytes, rp) = match w {
5761 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
5762 _ => (bytes, rp),
5763 };
5764 if (2..=16).contains(&m) && self.batched_supports(qtype) && self.mmvq_supports(qtype)
5766 && std::env::var("MEMRA_NO_BATCHED").is_err()
5767 && (m <= 4 || Self::b8_enabled())
5768 && (m <= 8 || qtype == QT_Q4_0 || qtype == QT_Q6_K || (qtype == QT_Q8_0 && rp)) {
5769 let mcols = Self::batched_mcols(m);
5770 return self.qmatvec_mmvq_batched(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, mcols, scale, rp);
5771 }
5772 if self.mmvq_supports(qtype) {
5773 return self.qmatvec_mmvq(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp);
5774 }
5775 let x0 = self.zeros(0)?;
5778 self.matmul_pre(w, aq, ad, &x0, m)
5779 }
5780
5781 pub fn matmul_decode_exact_dual_pre(&self, w0: &crate::model::GpuTensor,
5790 w1: &crate::model::GpuTensor,
5791 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
5792 -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>> {
5793 use crate::model::GpuTensor;
5794 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5795 let on = *ON.get_or_init(|| {
5796 std::env::var("MEMRA_SPEC_DUAL_T").map(|v| v != "0").unwrap_or(true)
5797 });
5798 if !on || !(2..=7).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok()
5799 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
5800 return Ok(None);
5801 }
5802 if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5807 let (in_f, out_f) = (w0.in_features(), w0.out_features());
5808 if w1.in_features() != in_f || w1.out_features() != out_f {
5809 return Ok(None);
5810 }
5811 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
5812 (GpuTensor::Quant { bytes: b0, qtype: q0, row_bytes: rb0, scale: s0, rp: rp0, rp4: None, .. },
5813 GpuTensor::Quant { bytes: b1, qtype: q1, row_bytes: rb1, scale: s1, rp: rp1, rp4: None, .. })
5814 if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 =>
5815 (b0, b1, *rb0, *s0, *s1, *rp0),
5816 _ => return Ok(None),
5817 };
5818 if m > 4 && !(rp && Self::b8_enabled()
5821 && std::env::var("MEMRA_B567").as_deref() != Ok("0")) {
5822 return Ok(None);
5823 }
5824 let (y0, y1) = self.qmatvec_batched_dual_raw(b0, b1, aq, ad, m, in_f, out_f, row_bytes, rp)?;
5825 Ok(Some(((y0, s0), (y1, s1))))
5826 }
5827
5828 pub fn matmul_decode_exact_dual(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
5844 x: &CudaSlice<f32>, m: usize)
5845 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
5846 use crate::model::GpuTensor;
5847 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5848 let on = *ON.get_or_init(|| {
5849 std::env::var("MEMRA_SPEC_DUAL_T").map(|v| v != "0").unwrap_or(true)
5850 });
5851 if !on || !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok()
5852 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) {
5853 return Ok(None);
5854 }
5855 if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5860 let (in_f, out_f) = (w0.in_features(), w0.out_features());
5861 if w1.in_features() != in_f || w1.out_features() != out_f {
5862 return Ok(None);
5863 }
5864 let (b0, b1, row_bytes, s0, s1, rp) = match (w0, w1) {
5865 (GpuTensor::Quant { bytes: b0, qtype: q0, row_bytes: rb0, scale: s0, rp: rp0, rp4: None, .. },
5866 GpuTensor::Quant { bytes: b1, qtype: q1, row_bytes: rb1, scale: s1, rp: rp1, rp4: None, .. })
5867 if *q0 == QT_NVFP4 && *q1 == QT_NVFP4 && rb0 == rb1 && rp0 == rp1 =>
5868 (b0, b1, *rb0, *s0, *s1, *rp0),
5869 _ => return Ok(None),
5870 };
5871 if std::env::var("MEMRA_DEBUG").is_ok() {
5874 static ONCE: std::sync::Once = std::sync::Once::new();
5875 ONCE.call_once(|| eprintln!("[memra] dual gate+up batched ENGAGED (m={m} rp={rp})"));
5876 }
5877 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
5878 let (y0, y1) = self.qmatvec_batched_dual_raw(b0, b1, &aq, &ad, m, in_f, out_f, row_bytes, rp)?;
5879 let mut y0 = y0;
5880 let mut y1 = y1;
5881 if s0 != 1.0 { self.scale_inplace(&mut y0, s0, m * out_f)?; }
5882 if s1 != 1.0 { self.scale_inplace(&mut y1, s1, m * out_f)?; }
5883 Ok(Some((y0, y1)))
5884 }
5885
5886 #[allow(clippy::too_many_arguments)]
5891 pub fn qmatvec_batched_dual_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
5892 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
5893 m: usize, in_f: usize, out_f: usize, row_bytes: usize, rp: bool)
5894 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5895 const ROWS_PER_BLOCK: u32 = 4;
5896 let mcols = Self::batched_mcols(m);
5897 let (name, rows_per_block) = match (mcols, rp, m) {
5900 (2, false, _) => ("qmatvec_nvfp4_mmvq_dual_b2", ROWS_PER_BLOCK),
5901 (4, false, _) => ("qmatvec_nvfp4_mmvq_dual_b4_r2", ROWS_PER_BLOCK * 2),
5902 (2, true, _) => ("qmatvec_nvfp4_mmvq_dual_b2_rp", ROWS_PER_BLOCK),
5903 (4, true, _) => ("qmatvec_nvfp4_mmvq_dual_b4_rpr2", ROWS_PER_BLOCK * 2),
5904 (8, true, 5) => ("qmatvec_nvfp4_mmvq_dual_b5_rpr2", ROWS_PER_BLOCK * 2),
5905 (8, true, 6) => ("qmatvec_nvfp4_mmvq_dual_b6_rpr2", ROWS_PER_BLOCK * 2),
5906 (8, true, 7) => ("qmatvec_nvfp4_mmvq_dual_b7_rpr2", ROWS_PER_BLOCK * 2),
5907 _ => return Err(format!("qmatvec_batched_dual_raw: no dual kernel for m {m}").into()),
5908 };
5909 let f = self.func(name);
5910 let mut y0 = self.alloc_uninit::<f32>(m * out_f)?;
5911 let mut y1 = self.alloc_uninit::<f32>(m * out_f)?;
5912 let cfg = LaunchConfig {
5913 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
5914 block_dim: (32, ROWS_PER_BLOCK, 1),
5915 shared_mem_bytes: 0,
5916 };
5917 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
5918 let __s_b = self.gpu.stream();
5919 let mut b = __s_b.launch_builder(&f);
5920 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
5921 .arg(&inf).arg(&outf).arg(&mi).arg(&rb);
5922 unsafe { b.launch(cfg)?; }
5923 Ok((y0, y1))
5924 }
5925
5926 pub fn matmul_pre_dual_noscale(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
5938 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
5939 -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn std::error::Error>> {
5940 use crate::model::GpuTensor;
5941 if m != 1 || !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) { return Ok(None); }
5942 if !self.mmvq_supports(QT_NVFP4) { return Ok(None); }
5952 let (in_f, out_f) = (w0.in_features(), w0.out_features());
5953 if w1.in_features() != in_f || w1.out_features() != out_f { return Ok(None); }
5954 let no_mirror = |w: &crate::model::GpuTensor| {
5967 !matches!(w, GpuTensor::Quant { rp4: Some(_), .. })
5968 };
5969 if self.q8_ffn_fuse2_on()
5970 && no_mirror(w0) && no_mirror(w1)
5971 && let Some([p0, p1]) = self.q8_fused_params(&[w0, w1])
5972 {
5973 let (y0, y1) = self.q8_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2)?;
5974 return Ok(Some(((y0, 1.0), (y1, 1.0))));
5975 }
5976 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
5986 let (y0, y1) = self.e4m3_fused2_core(p0.0, p1.0, aq, ad, in_f, p0.1, p1.1, p0.2,
5987 1.0, 1.0)?;
5988 return Ok(Some(((y0, p0.3), (y1, p1.3))));
5989 }
5990 let (b0, q0, rb0, s0, rp0) = match w0 {
5991 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5992 _ => return Ok(None),
5993 };
5994 let (b1, q1, rb1, s1, rp1) = match w1 {
5995 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
5996 _ => return Ok(None),
5997 };
5998 if q0 != QT_NVFP4 || q1 != QT_NVFP4 || rb0 != rb1 || rp0 != rp1 { return Ok(None); }
5999 const ROWS_PER_BLOCK: u32 = 4; const RPW: u32 = 2;
6001 let rows_per_block = ROWS_PER_BLOCK * RPW;
6002 let f = self.func(if rp0 { "qmatvec_nvfp4_mmvq_dual_mr2_rp" } else { "qmatvec_nvfp4_mmvq_dual_mr2" });
6003 let mut y0 = self.alloc_uninit::<f32>(out_f)?;
6004 let mut y1 = self.alloc_uninit::<f32>(out_f)?;
6005 let cfg = LaunchConfig {
6006 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 2, 1),
6007 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0,
6008 };
6009 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, rb0 as i64);
6010 let one = 1.0f32;
6013 let __s_b = self.gpu.stream();
6014 let mut b = __s_b.launch_builder(&f);
6015 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6016 .arg(&inf).arg(&outf).arg(&mi).arg(&rb).arg(&one).arg(&one);
6017 unsafe { b.launch(cfg)?; }
6018 Ok(Some(((y0, s0), (y1, s1))))
6019 }
6020
6021 pub fn matmul_q8_fused2(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6029 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6030 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6031 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6037 return Ok(Some(self.e4m3_fused2_core(p0.0, p1.0, aq, ad, w0.in_features(),
6038 p0.1, p1.1, p0.2, p0.3, p1.3)?));
6039 }
6040 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6041 Ok(Some(self.q8_fused2_core(p0.0, p1.0, aq, ad, w0.in_features(), p0.1, p1.1, p0.2)?))
6042 }
6043
6044 #[allow(clippy::too_many_arguments)]
6045 fn q8_fused2_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6046 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6047 in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6048 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6049 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6051 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6052 let f = self.func("qmatvec_q8_0_mmvq_fused2");
6053 let mut y0 = self.alloc_uninit::<f32>(out0)?;
6054 let mut y1 = self.alloc_uninit::<f32>(out1)?;
6055 let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6056 shared_mem_bytes: 0 };
6057 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
6058 let __s_b = self.gpu.stream();
6059 let mut b = __s_b.launch_builder(&f);
6060 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6061 .arg(&inf).arg(&o0).arg(&o1).arg(&rbl);
6062 unsafe { b.launch(cfg)?; }
6063 Ok((y0, y1))
6064 }
6065
6066 pub fn matmul_q8_fused2_x(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6072 x: &CudaSlice<f32>)
6073 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6074 if !self.uses_q8_1_fast(w0) || !self.uses_q8_1_fast(w1) { return Ok(None); }
6075 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6076 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
6077 return Ok(Some(self.e4m3_fused2_core(p0.0, p1.0, &aq, &ad, w0.in_features(),
6078 p0.1, p1.1, p0.2, p0.3, p1.3)?));
6079 }
6080 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6081 let (aq, ad) = self.quantize_q8_1(x, 1, w0.in_features())?;
6082 Ok(Some(self.q8_fused2_core(p0.0, p1.0, &aq, &ad, w0.in_features(), p0.1, p1.1, p0.2)?))
6083 }
6084
6085 #[allow(clippy::too_many_arguments)]
6088 pub fn qmatvec_q8_fused2_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>,
6089 in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6090 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6091 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6092 self.q8_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes)
6093 }
6094
6095 pub fn matmul_q4_fused3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6101 w2: &crate::model::GpuTensor,
6102 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6103 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6104 use crate::model::GpuTensor;
6105 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6106 match w {
6107 GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6108 Some((*row_bytes, w.out_features())),
6109 _ => None,
6110 }
6111 };
6112 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2))
6113 else { return Ok(None) };
6114 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6115 return Ok(None);
6116 }
6117 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6121 match w {
6122 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6123 Some(m) => (m, true),
6124 None => (bytes, *rp),
6125 },
6126 _ => unreachable!(),
6127 }
6128 }
6129 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6130 if rp0 != rp1 || rp1 != rp2 { return Ok(None); }
6131 let rp = rp0;
6132 let rpb: u32 = 4;
6133 let mr1 = rp && Self::q40_mr1_on();
6137 let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6138 else { (o as u32).div_ceil(2).div_ceil(rpb) };
6139 let grid = nb(o0) + nb(o1) + nb(o2);
6140 let mut y0 = self.alloc_uninit::<f32>(o0)?;
6141 let mut y1 = self.alloc_uninit::<f32>(o1)?;
6142 let mut y2 = self.alloc_uninit::<f32>(o2)?;
6143 let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused3_mr1_rp" }
6144 else if rp { "qmatvec_q4_0_mmvq_fused3_rp" }
6145 else { "qmatvec_q4_0_mmvq_fused3" });
6146 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6147 let inf = w0.in_features() as i32;
6148 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
6149 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
6150 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6153 {
6154 use cudarc::driver::{DevicePtr, DevicePtrMut};
6155 let s = &self.gpu.stream();
6156 let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6157 let (p2, _g2) = b2.device_ptr(s); let (paq, _g3) = aq.device_ptr(s);
6158 let (pad, _g4) = ad.device_ptr(s);
6159 let (py0, _g5) = y0.device_ptr_mut(s); let (py1, _g6) = y1.device_ptr_mut(s);
6160 let (py2, _g7) = y2.device_ptr_mut(s);
6161 let mut ps = [
6162 &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6163 &p2 as *const _ as *mut _, &paq as *const _ as *mut _,
6164 &pad as *const _ as *mut _, &py0 as *const _ as *mut _,
6165 &py1 as *const _ as *mut _, &py2 as *const _ as *mut _,
6166 &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6167 &oo1 as *const _ as *mut _, &oo2 as *const _ as *mut _,
6168 &r0 as *const _ as *mut _, &r1 as *const _ as *mut _,
6169 &r2 as *const _ as *mut _,
6170 ];
6171 unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused3_mr1_rp",
6172 (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6173 }
6174 return Ok(Some((y0, y1, y2)));
6175 }
6176 let __s_b = self.gpu.stream();
6177 let mut b = __s_b.launch_builder(&f);
6178 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6179 .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&r0).arg(&r1).arg(&r2);
6180 unsafe { b.launch(cfg)?; }
6181 Ok(Some((y0, y1, y2)))
6182 }
6183
6184 #[allow(clippy::too_many_arguments)]
6187 pub fn matmul_q4_fused3_into(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6188 w2: &crate::model::GpuTensor,
6189 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6190 y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>,
6191 y2: &mut CudaSlice<f32>)
6192 -> Result<bool, Box<dyn std::error::Error>> {
6193 use crate::model::GpuTensor;
6194 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6195 match w {
6196 GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6197 Some((*row_bytes, w.out_features())),
6198 _ => None,
6199 }
6200 };
6201 let (Some((rb0, o0)), Some((rb1, o1)), Some((rb2, o2))) = (q4(w0), q4(w1), q4(w2))
6202 else { return Ok(false) };
6203 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6204 return Ok(false);
6205 }
6206 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6207 match w {
6208 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6209 Some(m) => (m, true),
6210 None => (bytes, *rp),
6211 },
6212 _ => unreachable!(),
6213 }
6214 }
6215 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6216 if rp0 != rp1 || rp1 != rp2 { return Ok(false); }
6217 let rp = rp0;
6218 let rpb: u32 = 4;
6219 let mr1 = rp && Self::q40_mr1_on();
6220 let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6221 else { (o as u32).div_ceil(2).div_ceil(rpb) };
6222 let grid = nb(o0) + nb(o1) + nb(o2);
6223 debug_assert!(y0.len() >= o0 && y1.len() >= o1 && y2.len() >= o2);
6224 let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused3_mr1_rp" }
6225 else if rp { "qmatvec_q4_0_mmvq_fused3_rp" }
6226 else { "qmatvec_q4_0_mmvq_fused3" });
6227 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6228 let inf = w0.in_features() as i32;
6229 let (oo0, oo1, oo2) = (o0 as i32, o1 as i32, o2 as i32);
6230 let (r0, r1, r2) = (rb0 as i64, rb1 as i64, rb2 as i64);
6231 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6233 use cudarc::driver::{DevicePtr, DevicePtrMut};
6234 let s = &self.gpu.stream();
6235 let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6236 let (p2, _g2) = b2.device_ptr(s); let (paq, _g3) = aq.device_ptr(s);
6237 let (pad, _g4) = ad.device_ptr(s);
6238 let (py0, _g5) = y0.device_ptr_mut(s); let (py1, _g6) = y1.device_ptr_mut(s);
6239 let (py2, _g7) = y2.device_ptr_mut(s);
6240 let mut ps = [
6241 &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6242 &p2 as *const _ as *mut _, &paq as *const _ as *mut _,
6243 &pad as *const _ as *mut _, &py0 as *const _ as *mut _,
6244 &py1 as *const _ as *mut _, &py2 as *const _ as *mut _,
6245 &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6246 &oo1 as *const _ as *mut _, &oo2 as *const _ as *mut _,
6247 &r0 as *const _ as *mut _, &r1 as *const _ as *mut _,
6248 &r2 as *const _ as *mut _,
6249 ];
6250 unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused3_mr1_rp",
6251 (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6252 return Ok(true);
6253 }
6254 let __s_b = self.gpu.stream();
6255 let mut b = __s_b.launch_builder(&f);
6256 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut *y0).arg(&mut *y1).arg(&mut *y2)
6257 .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&r0).arg(&r1).arg(&r2);
6258 unsafe { b.launch(cfg)?; }
6259 Ok(true)
6260 }
6261
6262 pub fn matmul_q4_fused2(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6264 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6265 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6266 use crate::model::GpuTensor;
6267 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6268 match w {
6269 GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6270 Some((*row_bytes, w.out_features())),
6271 _ => None,
6272 }
6273 };
6274 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(None) };
6275 if w0.in_features() != w1.in_features() { return Ok(None); }
6276 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6278 match w {
6279 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6280 Some(m) => (m, true),
6281 None => (bytes, *rp),
6282 },
6283 _ => unreachable!(),
6284 }
6285 }
6286 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6287 if rp0 != rp1 { return Ok(None); }
6288 let rp = rp0;
6289 let rpb: u32 = 4;
6290 let mr1 = rp && Self::q40_mr1_on();
6292 let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6293 else { (o as u32).div_ceil(2).div_ceil(rpb) };
6294 let grid = nb(o0) + nb(o1);
6295 let mut y0 = self.alloc_uninit::<f32>(o0)?;
6296 let mut y1 = self.alloc_uninit::<f32>(o1)?;
6297 let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused2_mr1_rp" }
6298 else if rp { "qmatvec_q4_0_mmvq_fused2_rp" }
6299 else { "qmatvec_q4_0_mmvq_fused2" });
6300 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6301 let inf = w0.in_features() as i32;
6302 let (oo0, oo1) = (o0 as i32, o1 as i32);
6303 let (r0, r1) = (rb0 as i64, rb1 as i64);
6304 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6306 {
6307 use cudarc::driver::{DevicePtr, DevicePtrMut};
6308 let s = &self.gpu.stream();
6309 let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6310 let (paq, _g2) = aq.device_ptr(s); let (pad, _g3) = ad.device_ptr(s);
6311 let (py0, _g4) = y0.device_ptr_mut(s); let (py1, _g5) = y1.device_ptr_mut(s);
6312 let mut ps = [
6313 &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6314 &paq as *const _ as *mut _, &pad as *const _ as *mut _,
6315 &py0 as *const _ as *mut _, &py1 as *const _ as *mut _,
6316 &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6317 &oo1 as *const _ as *mut _, &r0 as *const _ as *mut _,
6318 &r1 as *const _ as *mut _,
6319 ];
6320 unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused2_mr1_rp",
6321 (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6322 }
6323 return Ok(Some((y0, y1)));
6324 }
6325 let __s_b = self.gpu.stream();
6326 let mut b = __s_b.launch_builder(&f);
6327 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6328 .arg(&inf).arg(&oo0).arg(&oo1).arg(&r0).arg(&r1);
6329 unsafe { b.launch(cfg)?; }
6330 Ok(Some((y0, y1)))
6331 }
6332
6333 pub fn matmul_q4_fused2_into(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6335 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6336 y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>)
6337 -> Result<bool, Box<dyn std::error::Error>> {
6338 use crate::model::GpuTensor;
6339 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6340 match w {
6341 GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6342 Some((*row_bytes, w.out_features())),
6343 _ => None,
6344 }
6345 };
6346 let (Some((rb0, o0)), Some((rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(false) };
6347 if w0.in_features() != w1.in_features() { return Ok(false); }
6348 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6349 match w {
6350 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6351 Some(m) => (m, true),
6352 None => (bytes, *rp),
6353 },
6354 _ => unreachable!(),
6355 }
6356 }
6357 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6358 if rp0 != rp1 { return Ok(false); }
6359 let rp = rp0;
6360 let rpb: u32 = 4;
6361 let mr1 = rp && Self::q40_mr1_on();
6362 let nb = |o: usize| if mr1 { (o as u32).div_ceil(rpb) }
6363 else { (o as u32).div_ceil(2).div_ceil(rpb) };
6364 let grid = nb(o0) + nb(o1);
6365 debug_assert!(y0.len() >= o0 && y1.len() >= o1);
6366 let f = self.func(if mr1 { "qmatvec_q4_0_mmvq_fused2_mr1_rp" }
6367 else if rp { "qmatvec_q4_0_mmvq_fused2_rp" }
6368 else { "qmatvec_q4_0_mmvq_fused2" });
6369 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1), shared_mem_bytes: 0 };
6370 let inf = w0.in_features() as i32;
6371 let (oo0, oo1) = (o0 as i32, o1 as i32);
6372 let (r0, r1) = (rb0 as i64, rb1 as i64);
6373 if mr1 && Self::pdl_on() && Self::pdl_mmvq_on() {
6375 use cudarc::driver::{DevicePtr, DevicePtrMut};
6376 let s = &self.gpu.stream();
6377 let (p0, _g0) = b0.device_ptr(s); let (p1, _g1) = b1.device_ptr(s);
6378 let (paq, _g2) = aq.device_ptr(s); let (pad, _g3) = ad.device_ptr(s);
6379 let (py0, _g4) = y0.device_ptr_mut(s); let (py1, _g5) = y1.device_ptr_mut(s);
6380 let mut ps = [
6381 &p0 as *const _ as *mut std::ffi::c_void, &p1 as *const _ as *mut _,
6382 &paq as *const _ as *mut _, &pad as *const _ as *mut _,
6383 &py0 as *const _ as *mut _, &py1 as *const _ as *mut _,
6384 &inf as *const _ as *mut _, &oo0 as *const _ as *mut _,
6385 &oo1 as *const _ as *mut _, &r0 as *const _ as *mut _,
6386 &r1 as *const _ as *mut _,
6387 ];
6388 unsafe { self.launch_pdl("qmatvec_q4_0_mmvq_fused2_mr1_rp",
6389 (grid, 1, 1), (32, rpb, 1), &mut ps)?; }
6390 return Ok(true);
6391 }
6392 let __s_b = self.gpu.stream();
6393 let mut b = __s_b.launch_builder(&f);
6394 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut *y0).arg(&mut *y1)
6395 .arg(&inf).arg(&oo0).arg(&oo1).arg(&r0).arg(&r1);
6396 unsafe { b.launch(cfg)?; }
6397 Ok(true)
6398 }
6399
6400 pub fn matmul_q4_fused2_batched(&self, w0: &crate::model::GpuTensor,
6405 w1: &crate::model::GpuTensor,
6406 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6407 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6408 use crate::model::GpuTensor;
6409 if m < 2 || m > 8 { return Ok(None); }
6410 let q4 = |w: &GpuTensor| -> Option<(usize, usize)> {
6411 match w {
6412 GpuTensor::Quant { qtype, row_bytes, .. } if *qtype == QT_Q4_0 =>
6413 Some((*row_bytes, w.out_features())),
6414 _ => None,
6415 }
6416 };
6417 let (Some((rb0, o0)), Some((_rb1, o1))) = (q4(w0), q4(w1)) else { return Ok(None) };
6418 if w0.in_features() != w1.in_features() { return Ok(None); }
6419 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6420 match w {
6421 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6422 Some(mr) => (mr, true),
6423 None => (bytes, *rp),
6424 },
6425 _ => unreachable!(),
6426 }
6427 }
6428 let ((b0, rp0), (b1, rp1)) = (eff(w0), eff(w1));
6429 if !rp0 || !rp1 { return Ok(None); }
6430 let mcols = Self::batched_mcols(m);
6431 let rpb: u32 = 4;
6432 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
6433 let grid = nb(o0) + nb(o1);
6434 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
6435 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
6436 let f = self.func(match mcols { 2 => "qmatvec_q4_0_mmvq_b2_f2_rp",
6437 4 => "qmatvec_q4_0_mmvq_b4_f2_rp",
6438 _ => "qmatvec_q4_0_mmvq_b8_f2_rp" });
6439 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1),
6440 shared_mem_bytes: 0 };
6441 let inf = w0.in_features() as i32;
6442 let (oo0, oo1, mi) = (o0 as i32, o1 as i32, m as i32);
6443 let rb = rb0 as i64;
6444 let __s_b = self.gpu.stream();
6445 let mut b = __s_b.launch_builder(&f);
6446 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6447 .arg(&inf).arg(&oo0).arg(&oo1).arg(&mi).arg(&rb);
6448 unsafe { b.launch(cfg)?; }
6449 Ok(Some((y0, y1)))
6450 }
6451
6452 #[allow(clippy::too_many_arguments)]
6455 pub fn matmul_q4_fused3_batched(&self, w0: &crate::model::GpuTensor,
6456 w1: &crate::model::GpuTensor, w2: &crate::model::GpuTensor,
6457 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6458 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6459 use crate::model::GpuTensor;
6460 if m < 2 || m > 8 { return Ok(None); }
6461 let q4 = |w: &GpuTensor| -> Option<usize> {
6462 match w {
6463 GpuTensor::Quant { qtype, .. } if *qtype == QT_Q4_0 => Some(w.out_features()),
6464 _ => None,
6465 }
6466 };
6467 let (Some(o0), Some(o1), Some(o2)) = (q4(w0), q4(w1), q4(w2)) else { return Ok(None) };
6468 if w0.in_features() != w1.in_features() || w0.in_features() != w2.in_features() {
6469 return Ok(None);
6470 }
6471 fn eff(w: &GpuTensor) -> (&CudaSlice<u8>, bool) {
6472 match w {
6473 GpuTensor::Quant { bytes, rp4, rp, .. } => match rp4 {
6474 Some(mr) => (mr, true),
6475 None => (bytes, *rp),
6476 },
6477 _ => unreachable!(),
6478 }
6479 }
6480 let ((b0, rp0), (b1, rp1), (b2, rp2)) = (eff(w0), eff(w1), eff(w2));
6481 if !rp0 || !rp1 || !rp2 { return Ok(None); }
6482 let mcols = Self::batched_mcols(m);
6483 let rpb: u32 = 4;
6484 let nb = |o: usize| (o as u32).div_ceil(2 * rpb);
6485 let grid = nb(o0) + nb(o1) + nb(o2);
6486 let mut y0 = self.alloc_uninit::<f32>(m * o0)?;
6487 let mut y1 = self.alloc_uninit::<f32>(m * o1)?;
6488 let mut y2 = self.alloc_uninit::<f32>(m * o2)?;
6489 let f = self.func(match mcols { 2 => "qmatvec_q4_0_mmvq_b2_f3_rp",
6490 4 => "qmatvec_q4_0_mmvq_b4_f3_rp",
6491 _ => "qmatvec_q4_0_mmvq_b8_f3_rp" });
6492 let cfg = LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (32, rpb, 1),
6493 shared_mem_bytes: 0 };
6494 let inf = w0.in_features() as i32;
6495 let (oo0, oo1, oo2, mi) = (o0 as i32, o1 as i32, o2 as i32, m as i32);
6496 let rb = 0i64;
6497 let __s_b = self.gpu.stream();
6498 let mut b = __s_b.launch_builder(&f);
6499 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6500 .arg(&inf).arg(&oo0).arg(&oo1).arg(&oo2).arg(&mi).arg(&rb);
6501 unsafe { b.launch(cfg)?; }
6502 Ok(Some((y0, y1, y2)))
6503 }
6504
6505 pub fn matmul_q8_fused3(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6506 w2: &crate::model::GpuTensor,
6507 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>)
6508 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6509 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
6512 return Ok(Some(self.e4m3_fused3_core(p0.0, p1.0, p2.0, aq, ad, w0.in_features(),
6513 p0.1, p1.1, p2.1, p0.2,
6514 p0.3, p1.3, p2.3)?));
6515 }
6516 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else { return Ok(None) };
6517 Ok(Some(self.q8_fused3_core(p0.0, p1.0, p2.0, aq, ad, w0.in_features(),
6518 p0.1, p1.1, p2.1, p0.2)?))
6519 }
6520
6521 #[allow(clippy::too_many_arguments)]
6522 fn q8_fused3_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6523 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6524 in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize)
6525 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6526 const ROWS_PER_BLOCK: u32 = 4;
6527 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6528 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6529 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6530 let f = self.func("qmatvec_q8_0_mmvq_fused3");
6531 let mut y0 = self.alloc_uninit::<f32>(out0)?;
6532 let mut y1 = self.alloc_uninit::<f32>(out1)?;
6533 let mut y2 = self.alloc_uninit::<f32>(out2)?;
6534 let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6535 shared_mem_bytes: 0 };
6536 let (inf, o0, o1, o2, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32, row_bytes as i64);
6537 let __s_b = self.gpu.stream();
6538 let mut b = __s_b.launch_builder(&f);
6539 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6540 .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&rbl);
6541 unsafe { b.launch(cfg)?; }
6542 Ok((y0, y1, y2))
6543 }
6544
6545 #[allow(clippy::too_many_arguments)]
6547 pub fn qmatvec_q8_fused3_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6548 x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize,
6549 out2: usize, row_bytes: usize)
6550 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6551 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6552 self.q8_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes)
6553 }
6554
6555 pub fn matmul_q8_fused2_t(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6566 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6567 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6568 if !(2..=8).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() { return Ok(None); }
6572 if let Some([p0, p1]) = self.e4m3_fused_params(&[w0, w1]) {
6575 if m > 4 && !Self::b8_enabled() { return Ok(None); }
6576 return Ok(Some(self.e4m3_fused2_t_core(p0.0, p1.0, aq, ad, m, w0.in_features(),
6577 p0.1, p1.1, p0.2, p0.3, p1.3)?));
6578 }
6579 let Some([p0, p1]) = self.q8_fused_params(&[w0, w1]) else { return Ok(None) };
6580 Ok(Some(self.q8_fused2_t_core(p0.0, p1.0, aq, ad, m, w0.in_features(), p0.1, p1.1, p0.2)?))
6581 }
6582
6583 #[allow(clippy::too_many_arguments)]
6584 fn q8_fused2_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6585 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6586 in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6587 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6588 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6590 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6591 let f = self.func(match Self::batched_mcols(m) {
6592 2 => "qmatvec_q8_0_mmvq_fused2_b2",
6593 4 => "qmatvec_q8_0_mmvq_fused2_b4",
6594 _ => "qmatvec_q8_0_mmvq_fused2_b8",
6596 });
6597 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6598 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6599 let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6600 shared_mem_bytes: 0 };
6601 let (inf, o0, o1, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, m as i32, row_bytes as i64);
6602 let __s_b = self.gpu.stream();
6603 let mut b = __s_b.launch_builder(&f);
6604 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6605 .arg(&inf).arg(&o0).arg(&o1).arg(&mi).arg(&rbl);
6606 unsafe { b.launch(cfg)?; }
6607 Ok((y0, y1))
6608 }
6609
6610 #[allow(clippy::too_many_arguments)]
6613 pub fn qmatvec_q8_fused2_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6614 x: &CudaSlice<f32>, m: usize,
6615 in_f: usize, out0: usize, out1: usize, row_bytes: usize)
6616 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6617 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6618 self.q8_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes)
6619 }
6620
6621 #[allow(clippy::too_many_arguments)]
6624 pub fn matmul_q8_fused3_t(&self, w0: &crate::model::GpuTensor, w1: &crate::model::GpuTensor,
6625 w2: &crate::model::GpuTensor,
6626 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize)
6627 -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>> {
6628 if !(2..=4).contains(&m) || std::env::var("MEMRA_NO_BATCHED").is_ok() { return Ok(None); }
6629 if let Some([p0, p1, p2]) = self.e4m3_fused_params(&[w0, w1, w2]) {
6630 return Ok(Some(self.e4m3_fused3_t_core(p0.0, p1.0, p2.0, aq, ad, m, w0.in_features(),
6631 p0.1, p1.1, p2.1, p0.2,
6632 p0.3, p1.3, p2.3)?));
6633 }
6634 let Some([p0, p1, p2]) = self.q8_fused_params(&[w0, w1, w2]) else { return Ok(None) };
6635 Ok(Some(self.q8_fused3_t_core(p0.0, p1.0, p2.0, aq, ad, m, w0.in_features(),
6636 p0.1, p1.1, p2.1, p0.2)?))
6637 }
6638
6639 #[allow(clippy::too_many_arguments)]
6640 fn q8_fused3_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6641 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6642 in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize)
6643 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6644 const ROWS_PER_BLOCK: u32 = 4;
6645 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6646 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6647 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6648 let f = self.func(if Self::batched_mcols(m) == 2 { "qmatvec_q8_0_mmvq_fused3_b2" }
6649 else { "qmatvec_q8_0_mmvq_fused3_b4" });
6650 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6651 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6652 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
6653 let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6654 shared_mem_bytes: 0 };
6655 let (inf, o0, o1, o2, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6656 m as i32, row_bytes as i64);
6657 let __s_b = self.gpu.stream();
6658 let mut b = __s_b.launch_builder(&f);
6659 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6660 .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&mi).arg(&rbl);
6661 unsafe { b.launch(cfg)?; }
6662 Ok((y0, y1, y2))
6663 }
6664
6665 #[allow(clippy::too_many_arguments)]
6667 pub fn qmatvec_q8_fused3_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6668 x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize,
6669 out1: usize, out2: usize, row_bytes: usize)
6670 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6671 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6672 self.q8_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes)
6673 }
6674
6675 pub fn q8_ffn_fuse2_on(&self) -> bool {
6679 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6680 *ON.get_or_init(|| std::env::var("MEMRA_Q8_FFN_FUSE2").as_deref() != Ok("0"))
6681 }
6682
6683 #[allow(clippy::type_complexity)]
6689 fn q8_fused_params<'w, const N: usize>(&self, ws: &[&'w crate::model::GpuTensor; N])
6690 -> Option<[(&'w CudaSlice<u8>, usize, usize); N]> {
6691 use crate::model::GpuTensor;
6692 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") { return None; }
6693 if std::env::var("MEMRA_Q8_DUAL").is_ok_and(|v| v == "0") { return None; }
6694 let in_f = ws[0].in_features();
6695 let mut out: [Option<(&CudaSlice<u8>, usize, usize)>; N] = [None; N];
6696 for (i, w) in ws.iter().enumerate() {
6697 match w {
6698 GpuTensor::Quant { bytes, qtype, row_bytes, scale, .. }
6699 if *qtype == QT_Q8_0 && *scale == 1.0 && w.in_features() == in_f =>
6700 out[i] = Some((bytes, w.out_features(), *row_bytes)),
6701 _ => return None,
6702 }
6703 }
6704 Some(out.map(|o| o.unwrap()))
6705 }
6706
6707 pub fn e4m3_dual_on(&self) -> bool {
6710 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6711 *ON.get_or_init(|| std::env::var("MEMRA_E4M3_DUAL").as_deref() != Ok("0"))
6712 }
6713
6714 #[allow(clippy::type_complexity)]
6726 fn e4m3_fused_params<'w, const N: usize>(&self, ws: &[&'w crate::model::GpuTensor; N])
6727 -> Option<[(&'w CudaSlice<u8>, usize, usize, f32); N]> {
6728 use crate::model::GpuTensor;
6729 if !self.e4m3_dual_on() { return None; }
6730 let in_f = ws[0].in_features();
6731 let mut out: [Option<(&CudaSlice<u8>, usize, usize, f32)>; N] = [None; N];
6732 for (i, w) in ws.iter().enumerate() {
6733 match w {
6734 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, rp4, .. }
6735 if *qtype == QT_F8_E4M3 && w.in_features() == in_f
6736 && *row_bytes == in_f && !*rp && rp4.is_none() =>
6737 out[i] = Some((bytes, w.out_features(), *row_bytes, *scale)),
6738 _ => return None,
6739 }
6740 }
6741 Some(out.map(|o| o.unwrap()))
6742 }
6743
6744 #[allow(clippy::too_many_arguments)]
6748 fn e4m3_fused2_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6749 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6750 in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6751 ws0: f32, ws1: f32)
6752 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6753 const ROWS_PER_BLOCK: u32 = 4; let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6755 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6756 let f = self.func("qmatvec_e4m3_mmvq_fused2");
6757 let mut y0 = self.alloc_uninit::<f32>(out0)?;
6758 let mut y1 = self.alloc_uninit::<f32>(out1)?;
6759 let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6760 shared_mem_bytes: 0 };
6761 let (inf, o0, o1, rbl) = (in_f as i32, out0 as i32, out1 as i32, row_bytes as i64);
6762 let __s_b = self.gpu.stream();
6763 let mut b = __s_b.launch_builder(&f);
6764 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6765 .arg(&inf).arg(&o0).arg(&o1).arg(&rbl).arg(&ws0).arg(&ws1);
6766 unsafe { b.launch(cfg)?; }
6767 Ok((y0, y1))
6768 }
6769
6770 #[allow(clippy::too_many_arguments)]
6772 fn e4m3_fused3_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6773 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6774 in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize,
6775 ws0: f32, ws1: f32, ws2: f32)
6776 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6777 const ROWS_PER_BLOCK: u32 = 4;
6778 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6779 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6780 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6781 let f = self.func("qmatvec_e4m3_mmvq_fused3");
6782 let mut y0 = self.alloc_uninit::<f32>(out0)?;
6783 let mut y1 = self.alloc_uninit::<f32>(out1)?;
6784 let mut y2 = self.alloc_uninit::<f32>(out2)?;
6785 let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6786 shared_mem_bytes: 0 };
6787 let (inf, o0, o1, o2, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6788 row_bytes as i64);
6789 let __s_b = self.gpu.stream();
6790 let mut b = __s_b.launch_builder(&f);
6791 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6792 .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&rbl).arg(&ws0).arg(&ws1).arg(&ws2);
6793 unsafe { b.launch(cfg)?; }
6794 Ok((y0, y1, y2))
6795 }
6796
6797 #[allow(clippy::too_many_arguments)]
6801 fn e4m3_fused2_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6802 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6803 in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6804 ws0: f32, ws1: f32)
6805 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6806 const ROWS_PER_BLOCK: u32 = 4;
6807 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6808 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6809 let f = self.func(match Self::batched_mcols(m) {
6810 2 => "qmatvec_e4m3_mmvq_fused2_b2",
6811 4 => "qmatvec_e4m3_mmvq_fused2_b4",
6812 _ => "qmatvec_e4m3_mmvq_fused2_b8",
6813 });
6814 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6815 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6816 let cfg = LaunchConfig { grid_dim: (nb0 + nb1, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6817 shared_mem_bytes: 0 };
6818 let (inf, o0, o1, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, m as i32,
6819 row_bytes as i64);
6820 let __s_b = self.gpu.stream();
6821 let mut b = __s_b.launch_builder(&f);
6822 b.arg(b0).arg(b1).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1)
6823 .arg(&inf).arg(&o0).arg(&o1).arg(&mi).arg(&rbl);
6824 unsafe { b.launch(cfg)?; }
6825 if ws0 != 1.0 { self.scale_inplace(&mut y0, ws0, m * out0)?; }
6826 if ws1 != 1.0 { self.scale_inplace(&mut y1, ws1, m * out1)?; }
6827 Ok((y0, y1))
6828 }
6829
6830 #[allow(clippy::too_many_arguments)]
6832 fn e4m3_fused3_t_core(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6833 aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize,
6834 in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize,
6835 ws0: f32, ws1: f32, ws2: f32)
6836 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6837 const ROWS_PER_BLOCK: u32 = 4;
6838 let nb0 = (out0 as u32).div_ceil(ROWS_PER_BLOCK);
6839 let nb1 = (out1 as u32).div_ceil(ROWS_PER_BLOCK);
6840 let nb2 = (out2 as u32).div_ceil(ROWS_PER_BLOCK);
6841 let f = self.func(if Self::batched_mcols(m) == 2 { "qmatvec_e4m3_mmvq_fused3_b2" }
6842 else { "qmatvec_e4m3_mmvq_fused3_b4" });
6843 let mut y0 = self.alloc_uninit::<f32>(m * out0)?;
6844 let mut y1 = self.alloc_uninit::<f32>(m * out1)?;
6845 let mut y2 = self.alloc_uninit::<f32>(m * out2)?;
6846 let cfg = LaunchConfig { grid_dim: (nb0 + nb1 + nb2, 1, 1), block_dim: (32, ROWS_PER_BLOCK, 1),
6847 shared_mem_bytes: 0 };
6848 let (inf, o0, o1, o2, mi, rbl) = (in_f as i32, out0 as i32, out1 as i32, out2 as i32,
6849 m as i32, row_bytes as i64);
6850 let __s_b = self.gpu.stream();
6851 let mut b = __s_b.launch_builder(&f);
6852 b.arg(b0).arg(b1).arg(b2).arg(aq).arg(ad).arg(&mut y0).arg(&mut y1).arg(&mut y2)
6853 .arg(&inf).arg(&o0).arg(&o1).arg(&o2).arg(&mi).arg(&rbl);
6854 unsafe { b.launch(cfg)?; }
6855 if ws0 != 1.0 { self.scale_inplace(&mut y0, ws0, m * out0)?; }
6856 if ws1 != 1.0 { self.scale_inplace(&mut y1, ws1, m * out1)?; }
6857 if ws2 != 1.0 { self.scale_inplace(&mut y2, ws2, m * out2)?; }
6858 Ok((y0, y1, y2))
6859 }
6860
6861 #[allow(clippy::too_many_arguments)]
6864 pub fn qmatvec_e4m3_fused2_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>,
6865 in_f: usize, out0: usize, out1: usize, row_bytes: usize,
6866 ws0: f32, ws1: f32)
6867 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6868 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6869 self.e4m3_fused2_core(b0, b1, &aq, &ad, in_f, out0, out1, row_bytes, ws0, ws1)
6870 }
6871
6872 #[allow(clippy::too_many_arguments)]
6873 pub fn qmatvec_e4m3_fused3_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>,
6874 x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize,
6875 out2: usize, row_bytes: usize, ws0: f32, ws1: f32, ws2: f32)
6876 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6877 let (aq, ad) = self.quantize_q8_1(x, 1, in_f)?;
6878 self.e4m3_fused3_core(b0, b1, b2, &aq, &ad, in_f, out0, out1, out2, row_bytes,
6879 ws0, ws1, ws2)
6880 }
6881
6882 #[allow(clippy::too_many_arguments)]
6883 pub fn qmatvec_e4m3_fused2_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6884 x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize,
6885 out1: usize, row_bytes: usize, ws0: f32, ws1: f32)
6886 -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6887 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6888 self.e4m3_fused2_t_core(b0, b1, &aq, &ad, m, in_f, out0, out1, row_bytes, ws0, ws1)
6889 }
6890
6891 #[allow(clippy::too_many_arguments)]
6892 pub fn qmatvec_e4m3_fused3_t_raw(&self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>,
6893 b2: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
6894 in_f: usize, out0: usize, out1: usize, out2: usize,
6895 row_bytes: usize, ws0: f32, ws1: f32, ws2: f32)
6896 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6897 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
6898 self.e4m3_fused3_t_core(b0, b1, b2, &aq, &ad, m, in_f, out0, out1, out2, row_bytes,
6899 ws0, ws1, ws2)
6900 }
6901
6902 pub fn matmul_pre_noscale(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6903 m: usize) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn std::error::Error>> {
6904 use crate::model::GpuTensor;
6905 if m != 1 || !self.uses_q8_1_fast(w) { return Ok(None); }
6907 let in_f = w.in_features();
6908 let out_f = w.out_features();
6909 let (bytes, qtype, row_bytes, scale, rp) = match w {
6910 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
6911 _ => return Ok(None),
6912 };
6913 if self.mmvq_supports(qtype) {
6915 let (mbytes, mrp) = match w {
6917 GpuTensor::Quant { rp4: Some(m4), .. } => (m4, true),
6918 _ => (bytes, rp),
6919 };
6920 let y = self.qmatvec_mmvq(mbytes, aq, ad, m, in_f, out_f, qtype, row_bytes, 1.0, mrp)?;
6921 return Ok(Some((y, scale)));
6922 }
6923 let name = match qtype {
6925 QT_Q8_0 => "qmatvec_q8_0_dp4a", QT_Q4_K => "qmatvec_q4_K_dp4a",
6926 QT_Q6_K => "qmatvec_q6_K_dp4a", QT_Q5_K => "qmatvec_q5_K_dp4a",
6927 QT_Q3_K => "qmatvec_q3_K_dp4a",
6928 QT_NVFP4 => if rp { "qmatvec_nvfp4_dp4a_rp" } else { "qmatvec_nvfp4_dp4a" },
6929 QT_IQ4_XS => "qmatvec_iq4_XS_dp4a",
6930 _ => return Ok(None),
6931 };
6932 let f = self.func(name);
6933 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
6934 let cfg = LaunchConfig { grid_dim: (out_f as u32, m as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
6935 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
6936 let __s_b = self.gpu.stream();
6937 let mut b = __s_b.launch_builder(&f);
6938 b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
6939 unsafe { b.launch(cfg)?; }
6940 Ok(Some((y, scale)))
6941 }
6942
6943 pub fn mmvq_supports(&self, qtype: i32) -> bool {
6946 if qtype == QT_F8_E4M3 { return true; }
6951 if std::env::var("MEMRA_MMVQ").as_deref() == Ok("0") { return false; }
6952 matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_Q4_0)
6953 }
6954
6955 pub fn qmatvec_mmvq(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6960 m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32,
6961 rp: bool)
6962 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6963 let mut y = self.alloc_uninit::<f32>(m * out_f)?; self.qmatvec_mmvq_into(bytes, aq, ad, m, in_f, out_f, qtype, row_bytes, scale, rp, &mut y)?;
6965 Ok(y)
6966 }
6967
6968 #[allow(clippy::too_many_arguments)]
6970 pub fn qmatvec_mmvq_into(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
6971 m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32,
6972 rp: bool, y: &mut CudaSlice<f32>)
6973 -> Result<(), Box<dyn std::error::Error>> {
6974 debug_assert!(y.len() >= m * out_f);
6975 const ROWS_PER_BLOCK: u32 = 4; if qtype == QT_Q8_0 && rp && m == 1 && out_f >= 64
6981 && (out_f as u32).div_ceil(ROWS_PER_BLOCK) < 4 * self.sm_count() as u32
6982 && {
6983 static G2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6984 *G2.get_or_init(|| std::env::var("MEMRA_Q80_G2").as_deref() != Ok("0"))
6985 }
6986 {
6987 let f = self.func("qmatvec_q8_0_mmvq_rp_g2");
6988 let cfg = LaunchConfig {
6989 grid_dim: ((out_f as u32).div_ceil(2), 1, 1),
6990 block_dim: (32, 2, 1),
6991 shared_mem_bytes: 0,
6992 };
6993 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, 1i32, row_bytes as i64);
6994 let __s_b = self.gpu.stream();
6995 let mut b = __s_b.launch_builder(&f);
6996 b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
6997 unsafe { b.launch(cfg)?; }
6998 if scale != 1.0 { self.scale_inplace(y, scale, out_f)?; }
6999 return Ok(());
7000 }
7001 let mut mr: u32 = if m == 1 && (qtype == QT_NVFP4 || qtype == QT_Q5_K) { 2 } else { 1 };
7010 if m == 1 && qtype == QT_Q4_0 {
7015 static Q40MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
7016 mr = *Q40MR.get_or_init(|| std::env::var("MEMRA_Q40_MR").ok()
7019 .and_then(|v| v.parse().ok()).unwrap_or(1));
7020 }
7021 let q5_mode = std::env::var("MEMRA_Q5K_ISSUE").ok();
7032 let q5_force = q5_mode.as_deref() == Some("2");
7033 let q5_il = qtype == QT_Q5_K && m == 1
7036 && (q5_force || q5_mode.as_deref().map(|v| v != "0").unwrap_or(true));
7037 if q5_il && !q5_force && out_f > 65536 { mr = 1; }
7038 if qtype == QT_Q4_0 && rp && mr != 1 { mr = 2; }
7041 if qtype == QT_Q8_0 && rp {
7045 static Q80MR: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
7046 mr = *Q80MR.get_or_init(|| std::env::var("MEMRA_Q80_MR").ok()
7047 .and_then(|v| v.parse().ok()).unwrap_or(1));
7048 }
7049 let name = match (qtype, mr, rp) {
7050 (QT_NVFP4, 2, false) => "qmatvec_nvfp4_mmvq_mr2",
7051 (QT_NVFP4, 2, true) => "qmatvec_nvfp4_mmvq_mr2_rp",
7052 (QT_NVFP4, _, true) => "qmatvec_nvfp4_mmvq_rp",
7053 (QT_Q4_0, 1, true) => "qmatvec_q4_0_mmvq_rp",
7054 (QT_Q4_0, _, true) => "qmatvec_q4_0_mmvq_mr2_rp",
7055 (QT_Q5_K, 2, _) => if q5_il { "qmatvec_q5_K_mmvq_mr2_il" } else { "qmatvec_q5_K_mmvq_mr2" },
7056 (QT_Q8_0, 2, true) => "qmatvec_q8_0_mmvq_mr2_rp",
7057 (QT_Q8_0, _, true) if in_f % 1024 == 0 && {
7062 static CA: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7063 *CA.get_or_init(|| std::env::var("MEMRA_Q80_CA").as_deref() == Ok("1"))
7064 } => "qmatvec_q8_0_mmvq_rpca",
7065 (QT_Q8_0, _, true) => "qmatvec_q8_0_mmvq_rp",
7066 (QT_Q8_0, _, _) => "qmatvec_q8_0_mmvq",
7067 (QT_Q4_K, _, true) => "qmatvec_q4_K_mmvq_rp",
7071 (QT_Q6_K, _, true) => "qmatvec_q6_K_mmvq_rp",
7072 (QT_Q4_K, _, _) => "qmatvec_q4_K_mmvq",
7073 (QT_Q4_0, 2, false) => "qmatvec_q4_0_mmvq_mr2",
7074 (QT_Q4_0, _, false) => "qmatvec_q4_0_mmvq",
7075 (QT_Q5_K, _, _) => if q5_il { "qmatvec_q5_K_mmvq_il" } else { "qmatvec_q5_K_mmvq" },
7076 (QT_Q6_K, _, _) => "qmatvec_q6_K_mmvq",
7077 (QT_NVFP4, _, false) => "qmatvec_nvfp4_mmvq",
7078 (QT_F8_E4M3, _, _) => "qmatvec_e4m3_mmvq",
7079 _ => panic!("qmatvec_mmvq: qtype {qtype} has no MMVQ kernel"),
7080 };
7081 let f = self.func(name);
7082 let rows_per_block = ROWS_PER_BLOCK * mr;
7084 let cfg = LaunchConfig {
7085 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, m as u32, 1),
7086 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0, };
7089 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7090 let __s_b = self.gpu.stream();
7091 let mut b = __s_b.launch_builder(&f);
7092 if qtype == QT_NVFP4 || qtype == QT_F8_E4M3 {
7097 b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb).arg(&scale);
7098 unsafe { b.launch(cfg)?; }
7099 } else if Self::pdl_on() && Self::pdl_mmvq_on()
7100 && matches!(name, "qmatvec_q4_0_mmvq_rp" | "qmatvec_q6_K_mmvq"
7101 | "qmatvec_q6_K_mmvq_rp") {
7102 {
7106 use cudarc::driver::{DevicePtr, DevicePtrMut};
7107 let s = &self.gpu.stream();
7108 let (pw, _g0) = bytes.device_ptr(s); let (paq, _g1) = aq.device_ptr(s);
7109 let (pad, _g2) = ad.device_ptr(s); let (py, _g3) = y.device_ptr_mut(s);
7110 let mut ps = [
7111 &pw as *const _ as *mut std::ffi::c_void, &paq as *const _ as *mut _,
7112 &pad as *const _ as *mut _, &py as *const _ as *mut _,
7113 &inf as *const _ as *mut _, &outf as *const _ as *mut _,
7114 &mi as *const _ as *mut _, &rb as *const _ as *mut _,
7115 ];
7116 unsafe { self.launch_pdl(name, cfg.grid_dim, cfg.block_dim, &mut ps)?; }
7117 }
7118 if scale != 1.0 { self.scale_inplace(y, scale, m * out_f)?; }
7119 } else {
7120 b.arg(bytes).arg(aq).arg(ad).arg(&mut *y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7121 unsafe { b.launch(cfg)?; }
7122 if scale != 1.0 { self.scale_inplace(y, scale, m * out_f)?; }
7123 }
7124 Ok(())
7125 }
7126
7127 pub fn qmatvec_mmvq_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
7131 out_f: usize, qtype: i32, row_bytes: usize, rp: bool)
7132 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7133 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7134 self.qmatvec_mmvq(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, 1.0, rp)
7135 }
7136
7137 pub fn batched_supports(&self, qtype: i32) -> bool {
7141 matches!(qtype, QT_Q8_0 | QT_Q4_K | QT_Q5_K | QT_Q6_K | QT_NVFP4 | QT_F8_E4M3 | QT_Q4_0)
7142 }
7143
7144 pub fn iq_fast_enabled() -> bool {
7152 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7153 *ON.get_or_init(|| std::env::var("MEMRA_IQ_FAST").map(|v| v != "0").unwrap_or(true))
7154 }
7155
7156 pub fn b8_enabled() -> bool {
7159 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7160 *ON.get_or_init(|| std::env::var("MEMRA_B8").map(|v| v != "0").unwrap_or(true))
7161 }
7162
7163 pub fn batched_mcols(m: usize) -> usize {
7165 if m == 2 { 2 } else if m <= 4 { 4 } else if m <= 8 { 8 } else { 16 }
7166 }
7167
7168 fn batched_kernel_name(qtype: i32, mcols: usize) -> Option<&'static str> {
7173 Some(match (qtype, mcols) {
7174 (QT_Q8_0, 2) => "qmatvec_q8_0_mmvq_b2", (QT_Q8_0, 4) => "qmatvec_q8_0_mmvq_b4",
7175 (QT_Q8_0, 8) => "qmatvec_q8_0_mmvq_b8",
7176 (QT_Q8_0, 16) => "qmatvec_q8_0_mmvq_b16",
7179 (QT_Q4_K, 2) => "qmatvec_q4_K_mmvq_b2", (QT_Q4_K, 4) => "qmatvec_q4_K_mmvq_b4",
7180 (QT_Q4_K, 8) => "qmatvec_q4_K_mmvq_b8",
7181 (QT_Q5_K, 2) => "qmatvec_q5_K_mmvq_b2", (QT_Q5_K, 4) => "qmatvec_q5_K_mmvq_b4",
7182 (QT_Q5_K, 8) => "qmatvec_q5_K_mmvq_b8",
7183 (QT_Q6_K, 2) => "qmatvec_q6_K_mmvq_b2", (QT_Q6_K, 4) => "qmatvec_q6_K_mmvq_b4",
7184 (QT_Q6_K, 8) => "qmatvec_q6_K_mmvq_b8", (QT_Q6_K, 16) => "qmatvec_q6_K_mmvq_b16",
7185 (QT_NVFP4, 2) => "qmatvec_nvfp4_mmvq_b2", (QT_NVFP4, 4) => "qmatvec_nvfp4_mmvq_b4",
7186 (QT_NVFP4, 8) => "qmatvec_nvfp4_mmvq_b8",
7187 (QT_F8_E4M3, 2) => "qmatvec_e4m3_mmvq_b2", (QT_F8_E4M3, 4) => "qmatvec_e4m3_mmvq_b4",
7188 (QT_F8_E4M3, 8) => "qmatvec_e4m3_mmvq_b8",
7189 (QT_Q4_0, 2) => "qmatvec_q4_0_mmvq_b2", (QT_Q4_0, 4) => "qmatvec_q4_0_mmvq_b4",
7190 (QT_Q4_0, 8) => "qmatvec_q4_0_mmvq_b8", (QT_Q4_0, 16) => "qmatvec_q4_0_mmvq_b16",
7191 _ => return None,
7192 })
7193 }
7194
7195 pub fn sm_count(&self) -> i32 {
7230 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
7231 *SMS.get_or_init(|| {
7232 use cudarc::driver::sys::CUdevice_attribute_enum as A;
7233 self.gpu.ctx.attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT).unwrap_or(82)
7234 })
7235 }
7236
7237 pub fn batched_variant(&self, _m: usize, in_f: usize, out_f: usize, qtype: i32,
7238 row_bytes: usize, mcols: usize, rp: bool) -> &'static str {
7239 if qtype == QT_Q8_0 {
7244 return if rp { "rp" } else { "base" };
7245 }
7246 static BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7247 let bv = *BV.get_or_init(|| match std::env::var("MEMRA_MMVQ_BV").as_deref() {
7248 Ok("base") => "base", Ok("pf") => "pf", Ok("r2") => "r2", Ok("r2w8") => "r2w8",
7249 Ok("pfr2") => "pfr2", Ok("ca") => "ca", Ok("car2") => "car2",
7250 Ok("rp") => "rp", Ok("rpr2") => "rpr2", Ok("rpr2w8") => "rpr2w8",
7253 Ok("rpca") => "rpca", Ok("rpcar2") => "rpcar2",
7256 Ok("rpsc") => "rpsc", Ok("rpms") => "rpms", Ok("rpmsc") => "rpmsc",
7263 Ok("rpks") => "rpks", Ok("rpksc") => "rpksc",
7264 _ => "auto",
7265 });
7266 let ca_ok = qtype == QT_NVFP4 && (row_bytes % 16 == 0) && (in_f % 1024 == 0);
7270 static KS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7275 let ks_on = *KS_ON.get_or_init(|| std::env::var("MEMRA_KS").as_deref() != Ok("0"));
7276 let sc_ok = ks_on && qtype == QT_NVFP4 && (in_f % 256 == 0) && (in_f / 64 <= 272);
7277 let ks_ok = ks_on && qtype == QT_NVFP4 && (in_f % 512 == 0) && (in_f / 64 <= 272);
7278 static SMS: std::sync::OnceLock<i32> = std::sync::OnceLock::new();
7279 let sms = *SMS.get_or_init(|| {
7280 use cudarc::driver::sys::CUdevice_attribute_enum as A;
7281 self.gpu.ctx.attribute(A::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT).unwrap_or(82)
7282 });
7283 let kq_r2 = matches!(qtype, QT_Q4_K | QT_Q5_K | QT_Q6_K);
7303 static KQBV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7306 let kq_bv = *KQBV.get_or_init(|| match std::env::var("MEMRA_KQ_BV").as_deref() {
7307 Ok("base") => "base", Ok("r2") => "r2", Ok("r2w8") => "r2w8",
7308 _ => "auto",
7309 });
7310 let variant: &'static str = if qtype == QT_Q4_0 {
7311 static Q40BV: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
7315 let q40 = *Q40BV.get_or_init(|| match std::env::var("MEMRA_Q40_BV").as_deref() {
7316 Ok("base") => "base", Ok("r2") => "r2", Ok("ms") => "ms", Ok("sm") => "sm",
7322 Ok("la") => "la", _ => "auto",
7323 });
7324 let v = if q40 != "auto" { q40 }
7325 else if (out_f as u32).div_ceil(8) >= 4 * sms as u32 { "r2" } else { "base" };
7326 if rp { match v { "ms" => "r2ms_rp", "sm" => "r2sm_rp", "la" => "r2la_rp",
7331 "r2" => "r2_rp", _ => "rp" } }
7332 else if matches!(v, "ms" | "sm" | "la") { "r2" } else { v }
7333 } else if qtype != QT_NVFP4 && !kq_r2 {
7334 "base"
7335 } else if kq_r2 && rp {
7336 "rp"
7340 } else if kq_r2 {
7341 if kq_bv != "auto" {
7344 if kq_bv == "r2w8" && mcols != 4 { "r2" } else { kq_bv }
7345 } else if bv != "auto" {
7346 match bv {
7347 "r2" | "pfr2" | "rpr2" | "car2" => "r2",
7348 "r2w8" | "rpr2w8" => if mcols != 4 { "r2" } else { "r2w8" },
7349 _ => "base", }
7351 } else {
7352 let blocks = (out_f + 7) / 8;
7353 let waves = blocks as f64 / (7 * sms as usize) as f64;
7354 let filled = blocks >= 4 * sms as usize;
7355 let use_r2 = if qtype == QT_Q4_K { filled } else { waves >= 2.0 };
7356 if use_r2 { "r2" } else { "base" }
7357 }
7358 } else if bv != "auto" {
7359 let v = if bv == "r2w8" && mcols == 2 { "r2" }
7364 else if bv == "ca" && (!ca_ok || mcols == 8) { "pf" }
7365 else if bv == "car2" && (!ca_ok || mcols == 8) { "r2" }
7366 else if bv == "pfr2" && mcols == 8 { "r2" }
7367 else if (bv == "rpr2w8" || bv == "rpr2") && mcols == 2 { "rpr2" }
7368 else if (bv == "rpca" || bv == "rpcar2") && (!ca_ok || mcols == 8) {
7370 if mcols == 8 { "rpr2w8" } else { "rpr2" }
7371 }
7372 else if bv == "rpcar2" && mcols == 2 { "rpca" }
7373 else if (bv == "rpsc" || bv == "rpmsc") && !sc_ok { "rpr2" }
7376 else if (bv == "rpks" || bv == "rpksc") && !ks_ok { "rpr2" }
7377 else { bv };
7378 if rp {
7379 match v {
7380 "base" | "pf" | "ca" | "rp" => "rp",
7381 "r2" | "pfr2" | "car2" | "rpr2" => "rpr2",
7382 "r2w8" | "rpr2w8" => if mcols == 2 { "rpr2" } else { "rpr2w8" },
7383 other => other, }
7385 } else { v }
7386 } else if mcols == 8 {
7387 if rp { if sc_ok { "rpsc" } else { "rpr2w8" } } else { "r2w8" }
7398 } else if mcols >= 4 {
7399 let blocks = (out_f + 7) / 8;
7403 let r7 = 7 * sms as usize;
7404 let r8 = 8 * sms as usize;
7405 let waves = blocks as f64 / r7 as f64;
7406 let filled = blocks >= 4 * sms as usize;
7407 if filled && blocks.div_ceil(r8) < blocks.div_ceil(r7) {
7411 if rp { "rpr2w8" } else { "r2w8" }
7415 } else if waves >= 2.0 || (waves <= 1.0 && filled) {
7416 if rp { "rpr2" } else { "r2" }
7419 } else {
7420 if rp { "rp" } else { "pf" }
7424 }
7425 } else if in_f >= 6144 {
7426 if rp { "rpr2" } else { "r2" }
7430 }
7431 else if rp {
7432 let waves = ((out_f + 7) / 8) as f64 / (7 * sms as usize) as f64;
7437 if sc_ok && waves >= 0.9 && waves <= 1.1 { "rpsc" } else { "rp" }
7438 } else { "base" };
7439 variant
7440 }
7441
7442 pub fn qmatvec_mmvq_batched(&self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
7443 m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize,
7444 mcols: usize, scale: f32, rp: bool)
7445 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7446 const ROWS_PER_BLOCK: u32 = 4;
7447 let forced: Option<&'static str> = {
7452 static V: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
7453 V.get_or_init(|| std::env::var("MEMRA_BVAR").ok())
7454 .as_deref()
7455 .map(|s| Box::leak(s.to_string().into_boxed_str()) as &'static str)
7456 };
7457 let variant = match forced {
7458 Some(v) if !rp || v.contains("rp") => v,
7459 _ => self.batched_variant(m, in_f, out_f, qtype, row_bytes, mcols, rp),
7460 };
7461 let base_name = Self::batched_kernel_name(qtype, mcols)
7462 .ok_or_else(|| format!("qmatvec_mmvq_batched: no kernel for qtype {qtype} mcols {mcols}"))?;
7463 let variant = if mcols == 16 { if rp { "rp" } else { "base" } } else { variant };
7467 static B567: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7474 let b567 = *B567.get_or_init(|| std::env::var("MEMRA_B567").as_deref() != Ok("0"));
7475 if b567 && qtype == QT_NVFP4 && rp && mcols == 8 && (5..=7).contains(&m)
7476 && matches!(variant, "rpsc" | "rpr2w8") {
7477 let f = self.func(&format!("qmatvec_nvfp4_mmvq_b{m}_{variant}"));
7478 let rows_per_block = ROWS_PER_BLOCK * 2; let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7480 let cfg = LaunchConfig {
7481 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
7482 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: 0 };
7483 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7484 let __s_b = self.gpu.stream();
7485 let mut b = __s_b.launch_builder(&f);
7486 b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7487 unsafe { b.launch(cfg)?; }
7488 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7489 return Ok(y);
7490 }
7491 let (name, rows_per_block): (std::borrow::Cow<'static, str>, u32) = match variant {
7492 "base" => (base_name.into(), ROWS_PER_BLOCK),
7493 "pf" => (format!("{base_name}_pf").into(), ROWS_PER_BLOCK),
7494 "ca" => (format!("{base_name}_ca").into(), ROWS_PER_BLOCK),
7495 "rp" => (format!("{base_name}_rp").into(), ROWS_PER_BLOCK),
7496 "rpca" => (format!("{base_name}_rpca").into(), ROWS_PER_BLOCK), "rpks" => (format!("{base_name}_rpks").into(), ROWS_PER_BLOCK),
7500 "rpksc" => (format!("{base_name}_rpksc").into(), ROWS_PER_BLOCK),
7501 "rpms" => (format!("{base_name}_rpms").into(), ROWS_PER_BLOCK),
7502 "rpmsc" => (format!("{base_name}_rpmsc").into(), ROWS_PER_BLOCK),
7503 "r2ms_rp" => (format!("{base_name}_r2ms_rp").into(), ROWS_PER_BLOCK),
7504 "r2sm_rp" => (format!("{base_name}_r2sm_rp").into(), ROWS_PER_BLOCK * 2),
7505 "r2la_rp" => (format!("{base_name}_r2la_rp").into(), ROWS_PER_BLOCK * 2),
7506 v => (format!("{base_name}_{v}").into(), ROWS_PER_BLOCK * 2), };
7508 debug_assert!(!rp || name.contains("_rp"), "rp weight dispatched to a GGUF-layout kernel");
7509 let f = self.func(&name);
7510 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7511 let smem = if name.contains("_r2sm_rp") { (mcols * 32 * 9 * 4 + mcols * 32 * 4) as u32 }
7513 else { 0 };
7514 let cfg = LaunchConfig {
7515 grid_dim: ((out_f as u32 + rows_per_block - 1) / rows_per_block, 1, 1),
7516 block_dim: (32, ROWS_PER_BLOCK, 1), shared_mem_bytes: smem };
7517 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7518 let __s_b = self.gpu.stream();
7519 let mut b = __s_b.launch_builder(&f);
7520 b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7521 unsafe { b.launch(cfg)?; }
7522 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7523 Ok(y)
7524 }
7525
7526 pub fn qmatvec_batched_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
7530 in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize,
7531 rp: bool)
7532 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7533 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7534 self.qmatvec_mmvq_batched(bytes, &aq, &ad, m, in_f, out_f, qtype, row_bytes, mcols, 1.0, rp)
7535 }
7536
7537 pub fn qmatvec_nvfp4_batched_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize,
7539 in_f: usize, out_f: usize, row_bytes: usize, mcols: usize,
7540 rp: bool)
7541 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7542 self.qmatvec_batched_raw(bytes, x, m, in_f, out_f, QT_NVFP4, row_bytes, mcols, rp)
7543 }
7544
7545 fn try_fp4_gemm(&self, w: &crate::model::GpuTensor, x: &CudaSlice<f32>, m: usize,
7549 in_f: usize, out_f: usize)
7550 -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7551 use crate::model::GpuTensor;
7552 if cfg!(memra_portable_cuda) { return Ok(None); }
7553 if std::env::var("MEMRA_FP4").is_err() { return Ok(None); }
7554 #[cfg(memra_cutlass)]
7563 if m >= 128 && std::env::var("MEMRA_FP4_CUTLASS").is_ok() {
7564 if let GpuTensor::Quant { bytes, qtype, scale, row_bytes, cutlass, .. } = w {
7565 if *qtype == QT_NVFP4 && in_f % 64 == 0 {
7566 if let Some(cw) = cutlass {
7567 let y = self.cutlass_fp4_gemm(&cw.b_packed, &cw.sfb_swizzled, x, *scale,
7569 m, out_f, in_f)?;
7570 return Ok(Some(y));
7571 } else if std::env::var("MEMRA_FP4_CUTLASS_OTF").is_ok() {
7572 let (b_packed, sfb_sw) = self.build_cutlass_weight(bytes, out_f, in_f, *row_bytes)?;
7577 let y = self.cutlass_fp4_gemm(&b_packed, &sfb_sw, x, *scale, m, out_f, in_f)?;
7578 return Ok(Some(y));
7579 }
7580 }
7581 }
7582 }
7583 if let GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } = w {
7584 if *qtype == QT_NVFP4 && in_f % 64 == 0 && !*rp {
7587 let y = self.qmatvec_gemm_nvfp4_fp4(bytes, x, m, in_f, out_f, *row_bytes, *scale)?;
7588 return Ok(Some(y));
7589 }
7590 }
7591 Ok(None)
7592 }
7593
7594 pub fn rms_norm_f16out(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>,
7598 dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
7599 ncols: usize, nrows: usize, eps: f32)
7600 -> Result<(), Box<dyn std::error::Error>> {
7601 let f = self.func("rms_norm_f16out_f32");
7602 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
7603 let (nc, e) = (ncols as i32, eps);
7604 let __s_b = self.gpu.stream();
7605 let mut b = __s_b.launch_builder(&f);
7606 b.arg(x).arg(w).arg(dst).arg(dst16).arg(&nc).arg(&e);
7607 unsafe { b.launch(cfg)?; }
7608 Ok(())
7609 }
7610
7611 #[allow(clippy::too_many_arguments)]
7614 pub fn add_rms_norm_f16out(&self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>,
7615 res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>,
7616 dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32)
7617 -> Result<(), Box<dyn std::error::Error>> {
7618 let f = self.func("add_rms_norm_f16out_f32");
7619 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
7620 let (nc, e) = (ncols as i32, eps);
7621 let __s_lb = self.gpu.stream();
7622 let mut lb = __s_lb.launch_builder(&f);
7623 lb.arg(a).arg(b).arg(w).arg(res).arg(dst).arg(dst16).arg(&nc).arg(&e);
7624 unsafe { lb.launch(cfg)?; }
7625 Ok(())
7626 }
7627
7628 pub fn matmul_group_xh(&self, ws: &[&crate::model::GpuTensor], x: &CudaSlice<f32>,
7631 xh: &CudaSlice<u8>, m: usize)
7632 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7633 let mut out = Vec::with_capacity(ws.len());
7634 let in_f = ws[0].in_features();
7635 for w in ws {
7636 if w.in_features() == in_f && m >= 16 && !self.verify_exact_on() {
7637 if let Some(y) = self.try_f16_gemm_pre(w, xh, m)? {
7638 out.push(y);
7639 continue;
7640 }
7641 }
7642 out.push(self.matmul(w, x, m)?);
7643 }
7644 Ok(out)
7645 }
7646
7647 pub fn gdn_pad_mask(&self, beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
7650 len_d: &CudaSlice<i32>, h: usize, t: usize)
7651 -> Result<(), Box<dyn std::error::Error>> {
7652 let f = self.func("gdn_pad_mask_f32");
7653 let cfg = LaunchConfig::for_num_elems((t * h) as u32);
7654 let (hi, ti) = (h as i32, t as i32);
7655 let __s_b = self.gpu.stream();
7656 let mut b = __s_b.launch_builder(&f);
7657 b.arg(beta).arg(g_log).arg(len_d).arg(&hi).arg(&ti);
7658 unsafe { b.launch(cfg)?; }
7659 Ok(())
7660 }
7661
7662 pub fn row_gather_dev(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
7665 len_d: &CudaSlice<i32>, ncols: usize)
7666 -> Result<(), Box<dyn std::error::Error>> {
7667 let f = self.func("row_gather_dev_f32");
7668 let cfg = LaunchConfig::for_num_elems(ncols as u32);
7669 let nc = ncols as i32;
7670 let __s_b = self.gpu.stream();
7671 let mut b = __s_b.launch_builder(&f);
7672 b.arg(src).arg(dst).arg(len_d).arg(&nc);
7673 unsafe { b.launch(cfg)?; }
7674 Ok(())
7675 }
7676
7677 pub fn matmul_group(&self, ws: &[&crate::model::GpuTensor], x: &CudaSlice<f32>, m: usize)
7684 -> Result<Vec<CudaSlice<f32>>, Box<dyn std::error::Error>> {
7685 use crate::model::GpuTensor;
7686 let mut out = Vec::with_capacity(ws.len());
7687 let any_mirror = ws.iter().any(|w| matches!(w, GpuTensor::Quant { f16: Some(_), .. }));
7688 if m >= 16 && any_mirror && !self.verify_exact_on() {
7689 let in_f = ws[0].in_features();
7690 let xh = self.f16_act(x, m * in_f, in_f)?;
7691 for w in ws {
7692 if w.in_features() == in_f {
7693 if let Some(y) = self.try_f16_gemm_pre(w, &xh, m)? {
7694 out.push(y);
7695 continue;
7696 }
7697 }
7698 out.push(self.matmul(w, x, m)?);
7699 }
7700 return Ok(out);
7701 }
7702 for w in ws {
7703 out.push(self.matmul(w, x, m)?);
7704 }
7705 Ok(out)
7706 }
7707
7708 pub fn matmul_group_multi(&self, ws: &[&crate::model::GpuTensor],
7715 xs: &[&CudaSlice<f32>], ms: &[usize])
7716 -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn std::error::Error>> {
7717 assert_eq!(xs.len(), ms.len());
7718 let in_f = ws[0].in_features();
7719 let total: usize = ms.iter().sum();
7720 let mut xcat = self.uninit(total * in_f)?;
7721 let mut off = 0usize;
7722 for (x, &m) in xs.iter().zip(ms) {
7723 self.copy_into(&mut xcat, off * in_f, x, m * in_f)?;
7724 off += m;
7725 }
7726 let ys = self.matmul_group(ws, &xcat, total)?;
7727 let mut out: Vec<Vec<CudaSlice<f32>>> = (0..xs.len()).map(|_| Vec::new()).collect();
7728 for (w, y) in ws.iter().zip(ys) {
7729 let out_f = w.out_features();
7730 let mut off = 0usize;
7731 for (s, &m) in ms.iter().enumerate() {
7732 let mut ys_s = self.uninit(m * out_f)?;
7733 let src = y.slice(off * out_f..(off + m) * out_f);
7734 self.gpu.stream().memcpy_dtod(&src, &mut ys_s)?;
7735 out[s].push(ys_s);
7736 off += m;
7737 }
7738 }
7739 Ok(out)
7740 }
7741
7742 pub fn gemm_supports(&self, w: &crate::model::GpuTensor) -> bool {
7752 use crate::model::GpuTensor;
7753 if !legacy_quant_gemm_allowed(
7754 cfg!(memra_portable_cuda),
7755 cfg!(memra_hopper_mma),
7756 std::env::var_os("MEMRA_NO_GEMM").is_some(),
7757 ) {
7758 return false;
7759 }
7760 match w {
7761 GpuTensor::Quant { qtype, .. } =>
7762 matches!(*qtype, QT_Q8_0 | QT_Q4_K | QT_Q6_K | QT_Q5_K | QT_Q4_0)
7763 || (*qtype == QT_NVFP4 && w.in_features() % 64 == 0),
7764 GpuTensor::Float { .. } | GpuTensor::FloatBf16 { .. } => false,
7765 }
7766 }
7767
7768 pub fn qmatvec_gemm(&self, w: &crate::model::GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>,
7775 m: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7776 use crate::model::GpuTensor;
7777 let in_f = w.in_features();
7778 let out_f = w.out_features();
7779 let (bytes, qtype, row_bytes, scale, rp) = match w {
7780 GpuTensor::Quant { bytes, qtype, row_bytes, scale, rp, .. } => (bytes, *qtype, *row_bytes, *scale, *rp),
7781 _ => unreachable!("gemm_supports guaranteed Quant"),
7782 };
7783 if cfg!(memra_hopper_mma) && qtype == QT_Q8_0 && out_f % 64 == 0 && wgmma_gemm_enabled() {
7789 if let GpuTensor::Quant { rp4: Some(m4), .. } = w {
7790 let mut y = self.qmatvec_gemm_q8_0_wgmma_raw(m4, aq, ad, m, in_f, out_f)?;
7791 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7792 return Ok(y);
7793 }
7794 }
7795 let name = match qtype {
7796 QT_Q8_0 => "qmatvec_gemm_q8_0", QT_Q4_K => "qmatvec_gemm_q4_K",
7797 QT_Q4_0 => if rp { "qmatvec_gemm_q4_0_rp" } else { "qmatvec_gemm_q4_0" },
7798 QT_Q5_K => "qmatvec_gemm_q5_K",
7799 QT_Q6_K => "qmatvec_gemm_q6_K",
7800 QT_NVFP4 => if rp { "qmatvec_gemm_nvfp4_rp" } else { "qmatvec_gemm_nvfp4" },
7801 _ => unreachable!(),
7802 };
7803 let f = self.func(name);
7804 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);
7809 let k1_tile = if is_k1 { k1_launch_override().unwrap_or((128, 128, 8)) } else { (128, 128, 8) };
7811 let (bm, bn): (u32, u32) = if is_k1 { (k1_tile.0, k1_tile.1) } else { (64, 256) };
7812 let warps: u32 = if is_k1 { k1_tile.2 } else {
7813 match qtype { QT_NVFP4 => 8, _ => 4 }
7814 };
7815 let cfg = LaunchConfig {
7816 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
7817 block_dim: (32, warps, 1),
7818 shared_mem_bytes: 0,
7819 };
7820 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7821 let __s_b = self.gpu.stream();
7822 let mut b = __s_b.launch_builder(&f);
7823 b.arg(bytes).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7824 unsafe { b.launch(cfg)?; }
7825 if scale != 1.0 { self.scale_inplace(&mut y, scale, m * out_f)?; }
7826 Ok(y)
7827 }
7828
7829 pub fn qmatvec_gemm_raw(&self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize,
7834 out_f: usize, qtype: i32, row_bytes: usize)
7835 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7836 let (aq, ad) = self.quantize_q8_1(x, m, in_f)?;
7837 let name = match qtype {
7838 QT_Q8_0 => "qmatvec_gemm_q8_0", QT_Q4_K => "qmatvec_gemm_q4_K",
7839 QT_Q4_0 => "qmatvec_gemm_q4_0",
7840 QT_Q5_K => "qmatvec_gemm_q5_K",
7841 QT_Q6_K => "qmatvec_gemm_q6_K", QT_NVFP4 => "qmatvec_gemm_nvfp4",
7842 QT_NVFP4_RP => "qmatvec_gemm_nvfp4_rp",
7843 _ => panic!("qmatvec_gemm_raw: qtype {qtype} has no GEMM kernel"),
7844 };
7845 let f = self.func(name);
7846 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);
7850 let k1_tile = if is_k1 { k1_launch_override().unwrap_or((128, 128, 8)) } else { (128, 128, 8) };
7852 let (bm, bn): (u32, u32) = if is_k1 { (k1_tile.0, k1_tile.1) } else { (64, 256) };
7853 let warps: u32 = if is_k1 { k1_tile.2 } else {
7854 match qtype { QT_NVFP4 | QT_NVFP4_RP => 8, _ => 4 }
7855 };
7856 let cfg = LaunchConfig {
7857 grid_dim: ((out_f as u32 + bm - 1) / bm, (m as u32 + bn - 1) / bn, 1),
7858 block_dim: (32, warps, 1), shared_mem_bytes: 0,
7859 };
7860 let (inf, outf, mi, rb) = (in_f as i32, out_f as i32, m as i32, row_bytes as i64);
7861 let __s_b = self.gpu.stream();
7862 let mut b = __s_b.launch_builder(&f);
7863 b.arg(bytes).arg(&aq).arg(&ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi).arg(&rb);
7864 unsafe { b.launch(cfg)?; }
7865 Ok(y)
7866 }
7867
7868 pub fn qmatvec_gemm_q8_0_wgmma_raw(&self, rp4: &CudaSlice<u8>, aq: &CudaSlice<i8>,
7875 ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize)
7876 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7877 assert!(out_f % 64 == 0 && in_f % 32 == 0, "wgmma GEMM needs out_f%64==0, in_f%32==0");
7878 let f = self.func("qmatvec_gemm_q8_0_wgmma");
7879 let mut y = self.alloc_uninit::<f32>(m * out_f)?; let cfg = LaunchConfig {
7881 grid_dim: ((out_f / 64) as u32, (m as u32).div_ceil(64), 1),
7882 block_dim: (128, 1, 1), shared_mem_bytes: 0,
7883 };
7884 let (inf, outf, mi) = (in_f as i32, out_f as i32, m as i32);
7885 let __s_b = self.gpu.stream();
7886 let mut b = __s_b.launch_builder(&f);
7887 b.arg(rp4).arg(aq).arg(ad).arg(&mut y).arg(&inf).arg(&outf).arg(&mi);
7888 unsafe { b.launch(cfg)?; }
7889 Ok(y)
7890 }
7891
7892 pub fn scale_inplace(&self, y: &mut CudaSlice<f32>, s: f32, n: usize)
7894 -> Result<(), Box<dyn std::error::Error>> {
7895 let f = self.func("scale_f32");
7896 let cfg = LaunchConfig::for_num_elems(n as u32);
7897 let (sf, ni) = (s, n as i32);
7898 let __s_b = self.gpu.stream();
7899 let mut b = __s_b.launch_builder(&f);
7900 b.arg(y).arg(&sf).arg(&ni);
7901 unsafe { b.launch(cfg)?; }
7902 Ok(())
7903 }
7904
7905 pub fn bf16_to_f32(&self, data: &cudarc::driver::CudaView<'_, u8>, n: usize)
7910 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7911 let mut out = self.alloc_uninit::<f32>(n)?;
7912 let f = self.func("bf16_to_f32");
7913 let cfg = LaunchConfig::for_num_elems(n as u32);
7914 let ni = n as i32;
7915 let __s_b = self.gpu.stream();
7916 let mut b = __s_b.launch_builder(&f);
7917 b.arg(data).arg(&mut out).arg(&ni);
7918 unsafe { b.launch(cfg)?; }
7919 Ok(out)
7920 }
7921
7922 fn linear_bf16_chunked(&self, x: &CudaSlice<f32>, data: &CudaSlice<u8>, m: usize,
7929 in_f: usize, out_f: usize, exact: bool)
7930 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7931 const CHUNK_BYTES: usize = 256 << 20;
7932 let chunk_rows = (CHUNK_BYTES / (in_f * 4)).max(1).min(out_f);
7933 if chunk_rows >= out_f {
7934 let wf32 = self.bf16_to_f32(&data.slice(0..in_f * out_f * 2), in_f * out_f)?;
7935 return if exact { self.linear_decode_exact(x, &wf32, m, in_f, out_f) }
7936 else { self.linear(x, &wf32, m, in_f, out_f) };
7937 }
7938 let mut y = self.alloc_uninit::<f32>(m * out_f)?;
7939 let mut r0 = 0usize;
7940 while r0 < out_f {
7941 let rows = chunk_rows.min(out_f - r0);
7942 let wslice = data.slice(r0 * in_f * 2..(r0 + rows) * in_f * 2);
7943 let wf32 = self.bf16_to_f32(&wslice, in_f * rows)?;
7944 let yc = if exact { self.linear_decode_exact(x, &wf32, m, in_f, rows)? }
7945 else { self.linear(x, &wf32, m, in_f, rows)? };
7946 for mi in 0..m {
7948 let src = yc.slice(mi * rows..(mi + 1) * rows);
7949 let mut dst = y.slice_mut(mi * out_f + r0..mi * out_f + r0 + rows);
7950 self.gpu.stream().memcpy_dtod(&src, &mut dst)?;
7951 }
7952 r0 += rows;
7953 }
7954 Ok(y)
7955 }
7956
7957 pub fn linear_decode_exact(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize,
7964 in_f: usize, out_f: usize)
7965 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7966 if m_tokens == 1 { return self.linear(x, w, 1, in_f, out_f); }
7967 let xv = self.view(x, m_tokens * in_f);
7968 let mut y = self.alloc_uninit::<f32>(m_tokens * out_f)?;
7969 for t in 0..m_tokens {
7970 let row = xv.slice(t * in_f..(t + 1) * in_f);
7971 let mut xr = self.alloc_uninit::<f32>(in_f)?;
7972 self.copy_view_into(&mut xr, 0, &row, in_f)?;
7973 let yr = self.linear(&xr, w, 1, in_f, out_f)?;
7974 self.copy_into(&mut y, t * out_f, &yr, out_f)?;
7975 }
7976 Ok(y)
7977 }
7978
7979 pub fn linear(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize, in_f: usize, out_f: usize)
7980 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7981 use cudarc::cublaslt::{Matmul, MatmulConfig};
7982 let mut c = self.alloc_uninit::<f32>(m_tokens * out_f)?; let cfg = MatmulConfig {
7984 transa: true, transb: false, transc: false,
7985 m: out_f as u64, n: m_tokens as u64, k: in_f as u64,
7986 alpha: 1.0, lda: in_f as i64, ldb: in_f as i64, beta: 0.0, ldc: out_f as i64,
7987 stride_a: None, stride_b: None, stride_c: None, stride_bias: None, batch_size: None,
7988 };
7989 unsafe { self.gpu.blas.matmul(cfg, w, x, &mut c, None, None)?; }
7990 Ok(c)
7991 }
7992
7993 pub fn sdpa_naive(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
7995 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
7996 t: usize, t_kv: usize, scale: f32, causal: bool)
7997 -> Result<(), Box<dyn std::error::Error>> {
7998 let f = self.func("sdpa_naive_f32");
7999 let cfg = LaunchConfig {
8000 grid_dim: (n_head as u32, t as u32, 1),
8001 block_dim: (128, 1, 1),
8002 shared_mem_bytes: (t_kv * 4) as u32,
8003 };
8004 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8005 let __s_b = self.gpu.stream();
8006 let mut b = __s_b.launch_builder(&f);
8007 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8008 unsafe { b.launch(cfg)?; }
8009 Ok(())
8010 }
8011
8012 #[allow(clippy::too_many_arguments)]
8014 pub fn sdpa_naive_w(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8015 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8016 t: usize, t_kv: usize, scale: f32, causal: bool, window: usize)
8017 -> Result<(), Box<dyn std::error::Error>> {
8018 let f = self.func("sdpa_naive_w_f32");
8019 let cfg = LaunchConfig {
8020 grid_dim: (n_head as u32, t as u32, 1),
8021 block_dim: (128, 1, 1),
8022 shared_mem_bytes: (t_kv * 4) as u32,
8023 };
8024 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8025 t as i32, t_kv as i32, causal as i32, window as i32);
8026 let __s_b = self.gpu.stream();
8027 let mut b = __s_b.launch_builder(&f);
8028 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8029 .arg(&scale).arg(&cz).arg(&wi);
8030 unsafe { b.launch(cfg)?; }
8031 Ok(())
8032 }
8033
8034 pub fn sdpa_naive_view(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<f32>,
8036 v: &cudarc::driver::CudaView<f32>, o: &mut CudaSlice<f32>,
8037 head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize,
8038 scale: f32, causal: bool) -> Result<(), Box<dyn std::error::Error>> {
8039 let f = self.func("sdpa_naive_f32");
8040 let cfg = LaunchConfig {
8041 grid_dim: (n_head as u32, t as u32, 1), block_dim: (128, 1, 1),
8042 shared_mem_bytes: (t_kv * 4) as u32,
8043 };
8044 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8045 let __s_b = self.gpu.stream();
8046 let mut b = __s_b.launch_builder(&f);
8047 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8048 unsafe { b.launch(cfg)?; }
8049 Ok(())
8050 }
8051
8052 #[allow(clippy::too_many_arguments)]
8060 pub fn fa_dequant_kv_view_f32(&self, k: &cudarc::driver::CudaView<u8>,
8061 v: &cudarc::driver::CudaView<u8>,
8062 kf: &mut CudaSlice<f32>, vf: &mut CudaSlice<f32>,
8063 kv_dim_k: usize, kv_dim_v: usize, t_kv: usize,
8064 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8065 -> Result<(), Box<dyn std::error::Error>> {
8066 let f = if g { self.func_g("fa_dequant_kv_ws_f32") } else { self.func("fa_dequant_kv_ws_f32") };
8067 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
8068 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8069 let cfg = LaunchConfig { grid_dim: (nblk.max(1), 1, 1), block_dim: (256, 1, 1),
8070 shared_mem_bytes: 0 };
8071 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
8072 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8073 let __s_b = self.gpu.stream();
8074 let mut b = __s_b.launch_builder(&f);
8075 b.arg(k).arg(v).arg(&mut *kf).arg(&mut *vf).arg(&kdk).arg(&kdv).arg(&tkvi).arg(&ktb).arg(&vtb);
8076 unsafe { b.launch(cfg)?; }
8077 Ok(())
8078 }
8079
8080 #[allow(clippy::too_many_arguments)]
8081 pub fn sdpa_naive_quantized_view(
8082 &self,
8083 q: &CudaSlice<f32>,
8084 k: &cudarc::driver::CudaView<u8>,
8085 v: &cudarc::driver::CudaView<u8>,
8086 o: &mut CudaSlice<f32>,
8087 head_dim: usize,
8088 n_head: usize,
8089 n_head_kv: usize,
8090 t: usize,
8091 t_kv: usize,
8092 scale: f32,
8093 causal: bool,
8094 k_tok_bytes: usize,
8095 v_tok_bytes: usize,
8096 ) -> Result<(), Box<dyn std::error::Error>> {
8097 let kv_dim = n_head_kv * head_dim;
8098 let mut kf = self.uninit(t_kv * kv_dim)?;
8099 let mut vf = self.uninit(t_kv * kv_dim)?;
8100 let f = self.func("fa_dequant_kv_ws_f32");
8101 let total = (2 * t_kv * kv_dim) as u64;
8102 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8103 let cfg = LaunchConfig {
8104 grid_dim: (nblk.max(1), 1, 1),
8105 block_dim: (256, 1, 1),
8106 shared_mem_bytes: 0,
8107 };
8108 let (kv_dim_i, t_kv_i) = (kv_dim as i32, t_kv as i32);
8109 let (k_tok_bytes_i, v_tok_bytes_i) = (k_tok_bytes as i64, v_tok_bytes as i64);
8110 let __s_b = self.gpu.stream();
8111 let mut b = __s_b.launch_builder(&f);
8112 b.arg(k)
8113 .arg(v)
8114 .arg(&mut kf)
8115 .arg(&mut vf)
8116 .arg(&kv_dim_i)
8117 .arg(&kv_dim_i)
8118 .arg(&t_kv_i)
8119 .arg(&k_tok_bytes_i)
8120 .arg(&v_tok_bytes_i);
8121 unsafe { b.launch(cfg)? };
8122 self.sdpa_naive(
8123 q, &kf, &vf, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
8124 )
8125 }
8126
8127 pub fn fa_prefill(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8131 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8132 t: usize, t_kv: usize, scale: f32, causal: bool)
8133 -> Result<(), Box<dyn std::error::Error>> {
8134 if portable_mma_gated() {
8135 return self.sdpa_naive(q, k, v, o, head_dim, n_head, n_head_kv,
8136 t, t_kv, scale, causal);
8137 }
8138 let fa3_on = head_dim == 256 && causal && t == t_kv
8146 && match std::env::var("MEMRA_FA3").as_deref() {
8147 Ok("0") => false,
8148 Ok("1") => true,
8149 _ => cfg!(memra_hopper_mma),
8150 };
8151 if fa3_on {
8152 let n = t * n_head * head_dim;
8153 let nkv = t * n_head_kv * head_dim;
8154 let mut q16 = self.alloc_u8_uninit(n * 2)?;
8155 let mut k16 = self.alloc_u8_uninit(nkv * 2)?;
8156 let mut v16 = self.alloc_u8_uninit(nkv * 2)?;
8157 self.f32_to_bf16_into(q, &mut q16, n)?;
8158 self.f32_to_bf16_into(k, &mut k16, nkv)?;
8159 self.f32_to_bf16_into(v, &mut v16, nkv)?;
8160 let rc = {
8161 use cudarc::driver::{DevicePtr, DevicePtrMut};
8162 let stream = self.gpu.stream();
8163 let (qp, _g1) = q16.device_ptr(&stream);
8164 let (kp, _g2) = k16.device_ptr(&stream);
8165 let (vp, _g3) = v16.device_ptr(&stream);
8166 let (op, _g4) = o.device_ptr_mut(&stream);
8167 unsafe {
8168 memra_fa3_prefill(qp as *const core::ffi::c_void,
8169 kp as *const core::ffi::c_void,
8170 vp as *const core::ffi::c_void,
8171 op as *mut f32,
8172 t as i32, n_head as i32, n_head_kv as i32,
8173 head_dim as i32, scale,
8174 stream.cu_stream() as *mut core::ffi::c_void)
8175 }
8176 };
8177 if rc != 0 {
8178 return Err(format!("memra_fa3_prefill rc={rc}").into());
8179 }
8180 return Ok(());
8181 }
8182 static FA_P1: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8187 let fa_p1 = *FA_P1.get_or_init(|| std::env::var("MEMRA_FA_P1").as_deref() == Ok("1"));
8188 if fa_p1 && head_dim == 256 && !std::env::var("MEMRA_FA_FLOOR").is_ok() {
8189 const BLOCK_Q: usize = 64; const BKX: usize = 32;
8190 let f = self.func("fa_prefill_bf16_p1");
8191 let shmem = (2 * (2 * BKX * head_dim + BLOCK_Q * BKX)
8192 + 4 * (BLOCK_Q * BKX + 2 * BLOCK_Q)) as u32;
8193 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8194 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8195 let cfg = LaunchConfig {
8196 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8197 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8198 };
8199 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32,
8200 n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8201 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8202 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8203 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8204 let __s_b = self.gpu.stream();
8205 let mut b = __s_b.launch_builder(&f);
8206 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti)
8207 .arg(&tkvi).arg(&scale).arg(&cz);
8208 unsafe { b.launch(cfg)?; }
8209 return Ok(());
8210 }
8211 const BK: usize = 32;
8217 let w2 = std::env::var("MEMRA_FA_PP_W2").as_deref() == Ok("1");
8220 let (block_q, warps, w2_sfx): (usize, u32, &str) =
8221 if w2 { (32, 2, "_w2") } else { (64, 4, "") };
8222 let hd_sfx = fa_hd_suffix(head_dim)?;
8226 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
8227 let bf16kv = !floor && !w2
8232 && std::env::var("MEMRA_FA_BF16KV").as_deref() != Ok("0");
8233 let (kb16, vb16) = if bf16kv {
8234 let n = t_kv * n_head_kv * head_dim;
8235 let mut kb = self.alloc_u8_uninit(n * 2)?;
8236 let mut vb = self.alloc_u8_uninit(n * 2)?;
8237 let fcv = self.func("f32_to_bf16_bulk");
8238 let ni = n as i64;
8239 let cfgc = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
8240 let __s_b = self.gpu.stream();
8241 let mut b = __s_b.launch_builder(&fcv);
8242 b.arg(k).arg(&mut kb).arg(&ni);
8243 unsafe { b.launch(cfgc)?; }
8244 let __s_b = self.gpu.stream();
8245 let mut b = __s_b.launch_builder(&fcv);
8246 b.arg(v).arg(&mut vb).arg(&ni);
8247 unsafe { b.launch(cfgc)?; }
8248 (Some(kb), Some(vb))
8249 } else {
8250 (None, None)
8251 };
8252 let f = self.func(&if bf16kv {
8253 format!("fa_prefill_bf16kv_pp{hd_sfx}")
8254 } else {
8255 format!("fa_prefill_f32{}{}{hd_sfx}",
8256 if floor { "" } else { "_pp" },
8257 if floor { "" } else { w2_sfx })
8258 });
8259 let kv_stages = if bf16kv { 2 } else { 1 };
8262 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
8263 + 4 * (block_q * BK + 2 * block_q)) as u32;
8264 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8265 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8266 let cfg = LaunchConfig {
8267 grid_dim: ((t as u32 + block_q as u32 - 1) / block_q as u32, n_head as u32, 1),
8268 block_dim: (32, warps, 1), shared_mem_bytes: shmem,
8269 };
8270 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8271 let __s_b = self.gpu.stream();
8272 let mut b = __s_b.launch_builder(&f);
8273 b.arg(q);
8274 match (&kb16, &vb16) {
8275 (Some(kb), Some(vb)) => { b.arg(kb).arg(vb); }
8276 _ => { b.arg(k).arg(v); }
8277 }
8278 b.arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz);
8279 unsafe { b.launch(cfg)?; }
8280 Ok(())
8281 }
8282
8283 #[allow(clippy::too_many_arguments)]
8287 pub fn fa_prefill_w(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8288 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize,
8289 t: usize, t_kv: usize, scale: f32, causal: bool, window: usize)
8290 -> Result<(), Box<dyn std::error::Error>> {
8291 if portable_mma_gated() {
8294 return self.sdpa_naive_w(q, k, v, o, head_dim, n_head, n_head_kv,
8295 t, t_kv, scale, causal, window);
8296 }
8297 static FAW_F32: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8301 let faw_f32 = *FAW_F32.get_or_init(|| {
8302 std::env::var("MEMRA_FAW_STAGE").as_deref() == Ok("f32")
8303 });
8304 let floor = std::env::var("MEMRA_FA_FLOOR").is_ok();
8305 self.fa_prefill_w_arm(q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale, causal,
8306 window, floor || faw_f32, floor)
8307 }
8308
8309 #[allow(clippy::too_many_arguments)]
8312 pub fn fa_prefill_w_pre(&self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>,
8313 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8314 n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8315 window: usize, v_f16: bool)
8316 -> Result<(), Box<dyn std::error::Error>> {
8317 const BLOCK_Q: usize = 64; const BK: usize = 32;
8318 debug_assert_eq!(head_dim, 256);
8319 let hp = fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0
8320 && (n_head / n_head_kv) % 2 == 0;
8321 debug_assert!(!v_f16 || hp, "f16 V emitted but the SWA hp arm is off");
8322 if hp {
8323 const BLOCK_QH: usize = 32;
8324 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
8327 let vh: &CudaSlice<u8> = if v_f16 { vb } else {
8328 let n = t_kv * n_head_kv * head_dim;
8329 if vguard.as_ref().map(|b| b.len() < n * 2).unwrap_or(true) {
8330 *vguard = Some(self.alloc_uninit::<u8>(n * 2)?);
8331 }
8332 self.bf16_to_f16_into(vb, n, vguard.as_mut().unwrap())?;
8333 vguard.as_ref().unwrap()
8334 };
8335 let f = self.func("fa_prefill_w_bf16_p1h2");
8336 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK)
8337 + 4 * (2 * BLOCK_QH)) as u32;
8338 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8339 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8340 let cfg = LaunchConfig {
8341 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
8342 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8343 };
8344 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8345 n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8346 let __s_b = self.gpu.stream();
8347 let mut b = __s_b.launch_builder(&f);
8348 b.arg(qb).arg(kb).arg(vh).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8349 .arg(&scale).arg(&cz).arg(&wi);
8350 unsafe { b.launch(cfg)?; }
8351 return Ok(());
8352 }
8353 let f = self.func("fa_prefill_w_bf16_p1");
8354 let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8355 + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8356 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8357 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8358 let cfg = LaunchConfig {
8359 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8360 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8361 };
8362 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8363 n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8364 let __s_b = self.gpu.stream();
8365 let mut b = __s_b.launch_builder(&f);
8366 b.arg(qb).arg(kb).arg(vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8367 .arg(&scale).arg(&cz).arg(&wi);
8368 unsafe { b.launch(cfg)?; }
8369 Ok(())
8370 }
8371
8372 #[allow(clippy::too_many_arguments)]
8374 pub fn fa_prefill_w_arm(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8375 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8376 n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8377 window: usize, f32_stage: bool, floor: bool)
8378 -> Result<(), Box<dyn std::error::Error>> {
8379 const BLOCK_Q: usize = 64; const BK: usize = 32;
8380 debug_assert_eq!(head_dim, 256, "fa_prefill_w is stamped hd256 only");
8381 static P1_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8385 let p1 = !floor && !f32_stage
8386 && *P1_ON.get_or_init(|| {
8387 std::env::var("MEMRA_FAW_P1").map(|v| v != "0").unwrap_or(true)
8388 });
8389 let hp = p1 && fa_f16pv_on() && faw_hp_on() && n_head % 2 == 0
8390 && (n_head / n_head_kv) % 2 == 0;
8391 if hp {
8392 const BLOCK_QH: usize = 32;
8393 let f = self.func("fa_prefill_w_bf16_p1h2");
8394 let shmem = (2 * (2 * BK * head_dim + 2 * BLOCK_QH * BK)
8395 + 4 * (2 * BLOCK_QH)) as u32;
8396 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8397 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8398 let cfg = LaunchConfig {
8399 grid_dim: ((t as u32).div_ceil(BLOCK_QH as u32), (n_head / 2) as u32, 1),
8400 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8401 };
8402 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8403 n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8404 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8405 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8406 let vh = self.f32_to_f16(v, t_kv * n_head_kv * head_dim)?;
8407 let __s_b = self.gpu.stream();
8408 let mut b = __s_b.launch_builder(&f);
8409 b.arg(&qb).arg(&kb).arg(&vh).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8410 .arg(&scale).arg(&cz).arg(&wi);
8411 unsafe { b.launch(cfg)?; }
8412 return Ok(());
8413 }
8414 if p1 {
8415 let f = self.func("fa_prefill_w_bf16_p1");
8416 let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8417 + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8418 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8419 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8420 let cfg = LaunchConfig {
8421 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8422 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8423 };
8424 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8425 n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8426 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8427 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8428 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8429 let __s_b = self.gpu.stream();
8430 let mut b = __s_b.launch_builder(&f);
8431 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8432 .arg(&scale).arg(&cz).arg(&wi);
8433 unsafe { b.launch(cfg)?; }
8434 return Ok(());
8435 }
8436 static G4_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8439 let g4 = !floor && !f32_stage && n_head_kv == 1 && n_head % 4 == 0
8440 && *G4_ON.get_or_init(|| {
8441 std::env::var("MEMRA_FAW_G4").map(|v| v != "0").unwrap_or(true)
8442 });
8443 if g4 {
8444 const SP_M: usize = 16;
8445 static O2_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8448 let o2 = *O2_ON.get_or_init(|| {
8449 std::env::var("MEMRA_FAW_O2").map(|v| v != "0").unwrap_or(true)
8450 });
8451 let f = self.func(if o2 { "fa_prefill_w_bf16_g4o2" } else { "fa_prefill_w_bf16_g4" });
8452 let shmem = if o2 {
8453 (2 * (4 * SP_M * head_dim + 4 * SP_M * BK) + 4 * (4 * SP_M)) as u32
8454 } else {
8455 (2 * (2 * BK * head_dim + 4 * SP_M * head_dim + 4 * SP_M * BK)
8456 + 4 * (4 * SP_M)) as u32
8457 };
8458 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8459 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8460 let cfg = LaunchConfig {
8461 grid_dim: ((t as u32).div_ceil(SP_M as u32), (n_head / 4) as u32, 1),
8462 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8463 };
8464 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32,
8465 n_head_kv as i32, t as i32, t_kv as i32, causal as i32, window as i32);
8466 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8467 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8468 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8469 let __s_b = self.gpu.stream();
8470 let mut b = __s_b.launch_builder(&f);
8471 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8472 .arg(&scale).arg(&cz).arg(&wi);
8473 unsafe { b.launch(cfg)?; }
8474 return Ok(());
8475 }
8476 let f = self.func(if floor { "fa_prefill_w_f32" }
8477 else if f32_stage { "fa_prefill_w_f32_pp" }
8478 else { "fa_prefill_w_bf16_pp" });
8479 let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8480 + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8481 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8482 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8483 let cfg = LaunchConfig {
8484 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8485 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8486 };
8487 let (hd, nh, nhkv, ti, tkvi, cz, wi) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8488 t as i32, t_kv as i32, causal as i32, window as i32);
8489 if f32_stage {
8490 let __s_b = self.gpu.stream();
8491 let mut b = __s_b.launch_builder(&f);
8492 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8493 .arg(&scale).arg(&cz).arg(&wi);
8494 unsafe { b.launch(cfg)?; }
8495 } else {
8496 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8497 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8498 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8499 let __s_b = self.gpu.stream();
8500 let mut b = __s_b.launch_builder(&f);
8501 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8502 .arg(&scale).arg(&cz).arg(&wi);
8503 unsafe { b.launch(cfg)?; }
8504 }
8505 Ok(())
8506 }
8507
8508 #[allow(clippy::too_many_arguments)]
8512 pub fn fa_prefill_hd512(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8513 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8514 n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool)
8515 -> Result<(), Box<dyn std::error::Error>> {
8516 if portable_mma_gated() {
8518 return self.sdpa_naive(q, k, v, o, head_dim, n_head, n_head_kv,
8519 t, t_kv, scale, causal);
8520 }
8521 static F32_STAGE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8527 let f32_stage = *F32_STAGE.get_or_init(|| {
8528 std::env::var("MEMRA_FA512_STAGE").as_deref() == Ok("f32")
8529 });
8530 static SP_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8534 let sp = !f32_stage
8535 && *SP_ON.get_or_init(|| {
8536 std::env::var("MEMRA_FA512_SP").map(|v| v != "0").unwrap_or(true)
8537 });
8538 self.fa_prefill_hd512_arm(q, k, v, o, head_dim, n_head, n_head_kv, t, t_kv, scale,
8539 causal, f32_stage, sp, sp && fa_f16pv_on())
8540 }
8541
8542 #[allow(clippy::too_many_arguments)]
8544 pub fn fa_prefill_hd512_pre(&self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>,
8545 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8546 n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8547 v_f16: bool)
8548 -> Result<(), Box<dyn std::error::Error>> {
8549 debug_assert_eq!(head_dim, 512);
8550 const SP_M: usize = 16; const BKS: usize = 32;
8551 let f16pv = fa_f16pv_on();
8555 let nw = if f16pv { fa512_wide_warps() } else { 2 };
8556 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
8557 debug_assert!(!v_f16 || f16pv, "f16 V emitted without the door on");
8558 let mut vguard = self.fa_vf16_scratch.lock().unwrap();
8559 let vref: &CudaSlice<u8> = if f16pv && !v_f16 {
8560 let n = t_kv * n_head_kv * head_dim;
8562 let need = n * 2;
8563 if vguard.as_ref().map(|b| b.len() < need).unwrap_or(true) {
8564 *vguard = Some(self.alloc_uninit::<u8>(need)?);
8565 }
8566 let dst = vguard.as_mut().unwrap();
8567 self.bf16_to_f16_into(vb, n, dst)?;
8568 vguard.as_ref().unwrap()
8569 } else { vb };
8570 let f = self.func(if hp { "fa_prefill_bf16_hd512_sp16h2" }
8571 else { match (f16pv, nw) {
8572 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
8573 (true, _) => "fa_prefill_bf16_hd512_sp16",
8574 _ => "fa_prefill_bf16_hd512_sp",
8575 } });
8576 let (nwarp, npart) = if hp { (4usize, 4usize) } else if nw > 2 { (nw, nw) } else { (2, 1) };
8577 let shmem = if hp {
8579 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
8580 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
8581 } else {
8582 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
8583 + 4 * (npart * SP_M * BKS + SP_M)) as u32
8584 };
8585 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8586 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8587 let grid_y = if hp { (n_head / 2) as u32 } else { n_head as u32 };
8588 let cfg = LaunchConfig {
8589 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
8590 block_dim: (32, nwarp as u32, 1), shared_mem_bytes: shmem,
8591 };
8592 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8593 t as i32, t_kv as i32, causal as i32);
8594 let __s_b = self.gpu.stream();
8595 let mut b = __s_b.launch_builder(&f);
8596 b.arg(qb).arg(kb).arg(vref).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8597 .arg(&scale).arg(&cz);
8598 unsafe { b.launch(cfg)?; }
8599 Ok(())
8600 }
8601
8602 #[allow(clippy::too_many_arguments)]
8605 pub fn fa_prefill_hd512_arm(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
8606 o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize,
8607 n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool,
8608 f32_stage: bool, sp: bool, f16pv: bool)
8609 -> Result<(), Box<dyn std::error::Error>> {
8610 debug_assert_eq!(head_dim, 512, "fa_prefill_hd512 is hd512 only");
8611 if sp && !f32_stage {
8612 const SP_M: usize = 16; const BKS: usize = 32;
8616 let nw = if f16pv { fa512_wide_warps() } else { 2 };
8617 let hp = f16pv && fa512_hp_on() && n_head % 2 == 0 && (n_head / n_head_kv) % 2 == 0;
8618 let f = self.func(if hp { "fa_prefill_bf16_hd512_sp16h2" }
8619 else { match (f16pv, nw) {
8620 (true, 4) => "fa_prefill_bf16_hd512_sp16w4",
8621 (true, _) => "fa_prefill_bf16_hd512_sp16",
8622 _ => "fa_prefill_bf16_hd512_sp",
8623 } });
8624 let (nwarp, npart) = if hp { (4usize, 4usize) } else if nw > 2 { (nw, nw) } else { (2, 1) };
8625 let shmem = if hp {
8626 (2 * (2 * BKS * head_dim + 2 * SP_M * BKS)
8627 + 4 * (2 * npart * SP_M * BKS + 2 * SP_M)) as u32
8628 } else {
8629 (2 * (SP_M * head_dim + 2 * BKS * head_dim + SP_M * BKS)
8630 + 4 * (npart * SP_M * BKS + SP_M)) as u32
8631 };
8632 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8633 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8634 let grid_y = if hp { (n_head / 2) as u32 } else { n_head as u32 };
8635 let cfg = LaunchConfig {
8636 grid_dim: ((t as u32).div_ceil(SP_M as u32), grid_y, 1),
8637 block_dim: (32, nwarp as u32, 1), shared_mem_bytes: shmem,
8638 };
8639 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8640 t as i32, t_kv as i32, causal as i32);
8641 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8642 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8643 let vb = if f16pv { self.f32_to_f16(v, t_kv * n_head_kv * head_dim)? }
8644 else { self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)? };
8645 let __s_b = self.gpu.stream();
8646 let mut b = __s_b.launch_builder(&f);
8647 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8648 .arg(&scale).arg(&cz);
8649 unsafe { b.launch(cfg)?; }
8650 return Ok(());
8651 }
8652 const BLOCK_Q: usize = 32; const BK: usize = 32; const HALF: usize = 256;
8653 let f = self.func(if f32_stage { "fa_prefill_f32_hd512" } else { "fa_prefill_bf16_hd512" });
8654 let shmem = (2 * (BLOCK_Q * head_dim + BK * head_dim + BK * HALF + BLOCK_Q * BK)
8656 + 4 * BLOCK_Q) as u32;
8657 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8658 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8659 let cfg = LaunchConfig {
8660 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 2),
8661 block_dim: (32, 2, 1), shared_mem_bytes: shmem,
8662 };
8663 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32,
8664 t as i32, t_kv as i32, causal as i32);
8665 if f32_stage {
8666 let __s_b = self.gpu.stream();
8667 let mut b = __s_b.launch_builder(&f);
8668 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8669 .arg(&scale).arg(&cz);
8670 unsafe { b.launch(cfg)?; }
8671 } else {
8672 let qb = self.f32_to_bf16(q, t * n_head * head_dim)?;
8673 let kb = self.f32_to_bf16(k, t_kv * n_head_kv * head_dim)?;
8674 let vb = self.f32_to_bf16(v, t_kv * n_head_kv * head_dim)?;
8675 let __s_b = self.gpu.stream();
8676 let mut b = __s_b.launch_builder(&f);
8677 b.arg(&qb).arg(&kb).arg(&vb).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi)
8678 .arg(&scale).arg(&cz);
8679 unsafe { b.launch(cfg)?; }
8680 }
8681 Ok(())
8682 }
8683
8684 #[allow(clippy::too_many_arguments)]
8688 pub fn rope_neox2_bf16e(&self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>,
8689 qb: &mut CudaSlice<u8>, kb: &mut CudaSlice<u8>,
8690 pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize,
8691 nh_q: usize, nh_k: usize, n_tokens: usize, base: f32,
8692 freq_scale: f32, ff: Option<&CudaSlice<f32>>)
8693 -> Result<(), Box<dyn std::error::Error>> {
8694 let f = self.func("rope_neox2_bf16e_f32");
8695 let rows = ((nh_q + nh_k) * n_tokens) as u32;
8696 let cfg = LaunchConfig { grid_dim: (rows, 1, 1),
8697 block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
8698 let theta_scale = base.powf(-2.0 / n_dims as f32);
8699 let (hd, nd, nhq, nhk, nt) = (head_dim as i32, n_dims as i32, nh_q as i32,
8700 nh_k as i32, n_tokens as i32);
8701 let __s_b = self.gpu.stream();
8702 let mut b = __s_b.launch_builder(&f);
8703 match ff {
8704 Some(t) => { b.arg(&mut *q).arg(&mut *k).arg(&mut *qb).arg(&mut *kb).arg(pos)
8705 .arg(&hd).arg(&nd).arg(&nhq).arg(&nhk).arg(&nt)
8706 .arg(&theta_scale).arg(&freq_scale).arg(t);
8707 unsafe { b.launch(cfg)?; } }
8708 None => { let null: u64 = 0;
8709 b.arg(&mut *q).arg(&mut *k).arg(&mut *qb).arg(&mut *kb).arg(pos)
8710 .arg(&hd).arg(&nd).arg(&nhq).arg(&nhk).arg(&nt)
8711 .arg(&theta_scale).arg(&freq_scale).arg(&null);
8712 unsafe { b.launch(cfg)?; } }
8713 }
8714 Ok(())
8715 }
8716
8717 pub fn f32_to_bf16(&self, x: &CudaSlice<f32>, n: usize)
8720 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8721 assert!(n % 4 == 0, "f32_to_bf16 requires n % 4 == 0, got {n}");
8722 let mut y = self.alloc_uninit::<u8>(n * 2)?;
8723 let f = self.func("f32_to_bf16_flat");
8724 let n_i = n as i64;
8725 let cfg = LaunchConfig {
8726 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
8727 block_dim: (256, 1, 1), shared_mem_bytes: 0,
8728 };
8729 let __s_b = self.gpu.stream();
8730 let mut b = __s_b.launch_builder(&f);
8731 b.arg(x).arg(&mut y).arg(&n_i);
8732 unsafe { b.launch(cfg)?; }
8733 Ok(y)
8734 }
8735
8736 pub fn f32_to_f16(&self, x: &CudaSlice<f32>, n: usize)
8737 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8738 assert!(n % 4 == 0, "f32_to_f16 requires n % 4 == 0, got {n}");
8739 let mut y = self.alloc_uninit::<u8>(n * 2)?;
8740 let f = self.func("f32_to_f16_flat");
8741 let n_i = n as i64;
8742 let cfg = LaunchConfig {
8743 grid_dim: (((n / 4) as u32).div_ceil(256), 1, 1),
8744 block_dim: (256, 1, 1), shared_mem_bytes: 0,
8745 };
8746 let __s_b = self.gpu.stream();
8747 let mut b = __s_b.launch_builder(&f);
8748 b.arg(x).arg(&mut y).arg(&n_i);
8749 unsafe { b.launch(cfg)?; }
8750 Ok(y)
8751 }
8752
8753 pub fn bf16_to_f16(&self, xb: &CudaSlice<u8>, n: usize)
8755 -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
8756 let mut y = self.alloc_uninit::<u8>(n * 2)?;
8757 self.bf16_to_f16_into(xb, n, &mut y)?;
8758 Ok(y)
8759 }
8760
8761 pub fn bf16_to_f16_into(&self, xb: &CudaSlice<u8>, n: usize, y: &mut CudaSlice<u8>)
8763 -> Result<(), Box<dyn std::error::Error>> {
8764 assert!(n % 2 == 0, "bf16_to_f16 requires n % 2 == 0, got {n}");
8765 assert!(y.len() >= n * 2);
8766 let f = self.func("bf16_to_f16_flat");
8767 let n2 = (n / 2) as i64;
8768 let cfg = LaunchConfig {
8769 grid_dim: (((n / 2) as u32).div_ceil(256), 1, 1),
8770 block_dim: (256, 1, 1), shared_mem_bytes: 0,
8771 };
8772 let __s_b = self.gpu.stream();
8773 let mut b = __s_b.launch_builder(&f);
8774 b.arg(xb).arg(y).arg(&n2);
8775 unsafe { b.launch(cfg)?; }
8776 Ok(())
8777 }
8778
8779 #[allow(clippy::too_many_arguments)]
8784 pub fn fa_prefill_vl8(&self, seqs: &[FaSeqVl], head_dim: usize, n_head: usize,
8785 n_head_kv: usize, scale: f32)
8786 -> Result<(), Box<dyn std::error::Error>> {
8787 const BK: usize = 32;
8788 let b = seqs.len();
8789 assert!(b >= 1 && b <= 8);
8790 let mut packed = [FaSeqVl::default(); 8];
8791 packed[..b].copy_from_slice(seqs);
8792 let v = FaVl8(packed);
8793 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
8794 let ept = (n_head_kv * head_dim) as i32;
8795 {
8796 let f = self.func("fa_mirror_vl");
8797 let max_n = (max_t as i64) * ept as i64;
8798 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
8799 for which in 0..2i32 {
8800 let cfg = LaunchConfig { grid_dim: (blocks, 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8801 let __s_lb = self.gpu.stream();
8802 let mut lb = __s_lb.launch_builder(&f);
8803 lb.arg(&v).arg(&ept).arg(&which);
8804 unsafe { lb.launch(cfg)?; }
8805 }
8806 }
8807 let hd_sfx = fa_hd_suffix(head_dim)?;
8808 let f = self.func(&format!("fa_prefill_bf16kv_vl{hd_sfx}"));
8809 let block_q = 64usize;
8810 let kv_stages = 2usize;
8811 let shmem = (2 * (kv_stages * 2 * BK * head_dim + block_q * BK)
8812 + 4 * (block_q * BK + 2 * block_q)) as u32;
8813 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8814 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8815 let cfg = LaunchConfig {
8816 grid_dim: (max_t.div_ceil(block_q as u32), n_head as u32, b as u32),
8817 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8818 };
8819 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
8820 let __s_lb = self.gpu.stream();
8821 let mut lb = __s_lb.launch_builder(&f);
8822 lb.arg(&v).arg(&hd).arg(&nh).arg(&nhkv).arg(&scale);
8823 unsafe { lb.launch(cfg)?; }
8824 Ok(())
8825 }
8826
8827 #[allow(clippy::too_many_arguments)]
8831 pub fn attn_pre_vl8(&self, seqs: &[AttnPreVl], wq: &CudaSlice<f32>, wk: &CudaSlice<f32>,
8832 head_dim: usize, rope_dims: usize, n_head: usize, n_head_kv: usize,
8833 eps: f32, freq_base: f32, freq_scale: f32,
8834 kv_dim_k: usize, kv_dim_v: usize,
8835 k_tok_bytes: usize, v_tok_bytes: usize)
8836 -> Result<(), Box<dyn std::error::Error>> {
8837 let b = seqs.len();
8838 assert!(b >= 1 && b <= 8);
8839 let mut packed = [AttnPreVl::default(); 8];
8840 packed[..b].copy_from_slice(seqs);
8841 let v = AttnPreVl8(packed);
8842 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
8843 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
8844 {
8845 let f = self.func("q_gate_split_vl");
8846 let n = max_t * (n_head * head_dim) as u32;
8847 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8848 let __s_lb = self.gpu.stream();
8849 let mut lb = __s_lb.launch_builder(&f);
8850 lb.arg(&v).arg(&hd).arg(&nh);
8851 unsafe { lb.launch(cfg)?; }
8852 }
8853 {
8854 let f = self.func("attn_rms_vl");
8855 let cfg = LaunchConfig { grid_dim: (max_t * n_head as u32, 2, b as u32), block_dim: (rms_block(), 1, 1), shared_mem_bytes: 0 };
8856 let __s_lb = self.gpu.stream();
8857 let mut lb = __s_lb.launch_builder(&f);
8858 lb.arg(&v).arg(wq).arg(wk).arg(&hd).arg(&nh).arg(&nhkv).arg(&eps);
8859 unsafe { lb.launch(cfg)?; }
8860 }
8861 {
8862 let f = self.func("attn_rope_vl");
8863 let theta_scale = freq_base.powf(-2.0 / rope_dims as f32);
8864 let nd = rope_dims as i32;
8865 let cfg = LaunchConfig { grid_dim: (max_t * n_head as u32, 2, b as u32), block_dim: ((head_dim / 2) as u32, 1, 1), shared_mem_bytes: 0 };
8866 let __s_lb = self.gpu.stream();
8867 let mut lb = __s_lb.launch_builder(&f);
8868 lb.arg(&v).arg(&hd).arg(&nd).arg(&nh).arg(&nhkv).arg(&theta_scale).arg(&freq_scale);
8869 unsafe { lb.launch(cfg)?; }
8870 }
8871 {
8872 let f = self.func("append_kv_vl");
8873 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
8874 let cfg = LaunchConfig { grid_dim: (nblk, max_t, b as u32), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
8875 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
8876 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8877 let __s_lb = self.gpu.stream();
8878 let mut lb = __s_lb.launch_builder(&f);
8879 lb.arg(&v).arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
8880 unsafe { lb.launch(cfg)?; }
8881 }
8882 Ok(())
8883 }
8884
8885 pub fn fa_prefill_view(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
8890 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
8891 head_dim: usize, n_head: usize, n_head_kv: usize,
8892 t: usize, t_kv: usize, scale: f32, causal: bool,
8893 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8894 -> Result<(), Box<dyn std::error::Error>> {
8895 if portable_mma_gated() {
8896 return self.sdpa_naive_quantized_view(q, k, v, o, head_dim, n_head, n_head_kv,
8897 t, t_kv, scale, causal,
8898 k_tok_bytes, v_tok_bytes);
8899 }
8900 const BLOCK_Q: usize = 64; const BK: usize = 32;
8901 let name = format!("fa_prefill_q{}", fa_hd_suffix(head_dim)?);
8904 let f = if g { self.func_g(&name) } else { self.func(&name) };
8905 let shmem = (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8906 + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32;
8907 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8908 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8909 let cfg = LaunchConfig {
8910 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8911 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8912 };
8913 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
8914 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8915 let __s_b = self.gpu.stream();
8916 let mut b = __s_b.launch_builder(&f);
8917 b.arg(q).arg(k).arg(v).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz)
8918 .arg(&ktb).arg(&vtb);
8919 unsafe { b.launch(cfg)?; }
8920 Ok(())
8921 }
8922
8923 #[allow(clippy::too_many_arguments)]
8933 pub fn fa_prefill_view_ws(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
8934 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
8935 head_dim: usize, n_head: usize, n_head_kv: usize,
8936 t: usize, t_kv: usize, scale: f32, causal: bool,
8937 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
8938 -> Result<(), Box<dyn std::error::Error>> {
8939 if portable_mma_gated() {
8940 return self.sdpa_naive_quantized_view(q, k, v, o, head_dim, n_head, n_head_kv,
8941 t, t_kv, scale, causal,
8942 k_tok_bytes, v_tok_bytes);
8943 }
8944 const BLOCK_Q: usize = 64; const BK: usize = 32;
8945 let kv_dim_k = n_head_kv * head_dim;
8946 let kv_dim_v = n_head_kv * head_dim;
8947 let k_ws_bytes = t_kv * kv_dim_k * 2; let v_ws_bytes = t_kv * kv_dim_v * 2;
8949 let mut guard = self.prime_deqw_ws.lock().unwrap();
8951 let need_grow = match guard.as_ref() {
8952 Some((kw, vw)) => kw.len() < k_ws_bytes || vw.len() < v_ws_bytes,
8953 None => true,
8954 };
8955 if need_grow {
8956 let grow = |cur: usize, need: usize| if cur >= need { cur } else { need };
8957 let (ck, cv) = guard.as_ref().map(|(a, b)| (a.len(), b.len())).unwrap_or((0, 0));
8958 *guard = Some((self.alloc_u8(grow(ck, k_ws_bytes))?, self.alloc_u8(grow(cv, v_ws_bytes))?));
8959 }
8960 let (kw, vw) = guard.as_mut().unwrap();
8961 {
8963 let f = if g { self.func_g("fa_dequant_kv_ws_bf16") } else { self.func("fa_dequant_kv_ws_bf16") };
8965 let total = (t_kv * (kv_dim_k + kv_dim_v)) as u64;
8966 let nblk = ((total + 255) / 256).min(65535 * 16) as u32;
8967 let cfg = LaunchConfig { grid_dim: (nblk.max(1), 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
8968 let (kdk, kdv, tkvi) = (kv_dim_k as i32, kv_dim_v as i32, t_kv as i32);
8969 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
8970 let __s_b = self.gpu.stream();
8971 let mut b = __s_b.launch_builder(&f);
8972 b.arg(k).arg(v).arg(&mut *kw).arg(&mut *vw).arg(&kdk).arg(&kdv).arg(&tkvi).arg(&ktb).arg(&vtb);
8973 unsafe { b.launch(cfg)?; }
8974 }
8975 let db = std::env::var("MEMRA_PRIME_DEQW_DB").map(|v| v != "0").unwrap_or(true);
8983 {
8984 let hd_sfx = fa_hd_suffix(head_dim)?;
8985 let f = self.func(&format!("fa_prefill_qw{}{hd_sfx}", if db { "_db" } else { "" }));
8986 let shmem = if db {
8987 (2 * (4 * BK * head_dim + BLOCK_Q * BK) + 4 * BLOCK_Q) as u32
8989 } else {
8990 (2 * (2 * BK * head_dim + BLOCK_Q * BK)
8991 + 4 * (BLOCK_Q * BK + 2 * BLOCK_Q)) as u32
8992 };
8993 use cudarc::driver::sys::CUfunction_attribute_enum as A;
8994 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
8995 let cfg = LaunchConfig {
8996 grid_dim: ((t as u32 + BLOCK_Q as u32 - 1) / BLOCK_Q as u32, n_head as u32, 1),
8997 block_dim: (32, 4, 1), shared_mem_bytes: shmem,
8998 };
8999 let (hd, nh, nhkv, ti, tkvi, cz) = (head_dim as i32, n_head as i32, n_head_kv as i32, t as i32, t_kv as i32, causal as i32);
9000 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
9001 let __s_b = self.gpu.stream();
9002 let mut b = __s_b.launch_builder(&f);
9003 b.arg(q).arg(&*kw).arg(&*vw).arg(o).arg(&hd).arg(&nh).arg(&nhkv).arg(&ti).arg(&tkvi).arg(&scale).arg(&cz)
9004 .arg(&kdk).arg(&kdv);
9005 unsafe { b.launch(cfg)?; }
9006 }
9007 Ok(())
9008 }
9009
9010 pub fn fa_decode(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9014 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9015 head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32,
9016 k_tok_bytes: usize, v_tok_bytes: usize)
9017 -> Result<(), Box<dyn std::error::Error>> {
9018 self.fa_decode_kvmod(q, k, v, o, head_dim, n_head, n_head_kv, t_kv, scale,
9019 k_tok_bytes, v_tok_bytes, false)
9020 }
9021
9022 #[allow(clippy::too_many_arguments)]
9026 #[allow(clippy::too_many_arguments)]
9030 #[allow(clippy::too_many_arguments)]
9031 fn fa_decode_scalar_unified(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9032 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9033 head_dim: usize, n_head: usize, n_head_kv: usize,
9034 t_kv_host: usize, t_kv_dev: Option<&CudaSlice<i32>>,
9035 scale: f32, n_splits: usize, split_keys: usize,
9036 k_tok_bytes: usize, v_tok_bytes: usize, g: bool,
9037 part_o: &mut CudaSlice<f32>, part_m: &mut CudaSlice<f32>,
9038 part_l: &mut CudaSlice<f32>,
9039 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9040 -> Result<(), Box<dyn std::error::Error>> {
9041 let f = if g { self.func_g("fa_decode_f32") } else { self.fa_func("fa_decode_f32", head_dim) };
9042 let cfg = LaunchConfig { grid_dim: (n_head as u32, n_splits as u32, 1),
9043 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: (4 * (head_dim + 32)) as u32 };
9044 let (hd, nh, nhkv, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, n_splits as i32);
9045 let (ktb, vtb, tkvi, ski) = (k_tok_bytes as i64, v_tok_bytes as i64, t_kv_host as i32,
9046 split_keys as i32);
9047 let __s_b = self.gpu.stream();
9048 let mut b = __s_b.launch_builder(&f);
9049 match t_kv_dev {
9050 Some(d) => { b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9051 .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(d).arg(&scale).arg(&nsp)
9052 .arg(&ski).arg(&ktb).arg(&vtb);
9053 unsafe { b.launch(cfg)?; } }
9054 None => { let null: u64 = 0;
9055 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9056 .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(&null).arg(&scale).arg(&nsp)
9057 .arg(&ski).arg(&ktb).arg(&vtb);
9058 unsafe { b.launch(cfg)?; } }
9059 }
9060 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, 1, 1),
9061 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9062 if let Some((oq, od)) = q8_out {
9063 let fc = if g { self.func_g("fa_decode_combine_q8_1") }
9065 else { self.fa_func("fa_decode_combine_q8_1", head_dim) };
9066 let __s_b2 = self.gpu.stream();
9067 let mut b2 = __s_b2.launch_builder(&fc);
9068 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh).arg(&nsp);
9069 unsafe { b2.launch(cfg2)?; }
9070 return Ok(());
9071 }
9072 let fc = if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) };
9073 let __s_b2 = self.gpu.stream();
9074 let mut b2 = __s_b2.launch_builder(&fc);
9075 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
9076 unsafe { b2.launch(cfg2)?; }
9077 Ok(())
9078 }
9079
9080 pub fn fa_decode_kvmod(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9081 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9082 head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32,
9083 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
9084 -> Result<(), Box<dyn std::error::Error>> {
9085 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
9106 if g && head_dim == 256 && !fa_v4_at(t_kv) { fa_vec = false; }
9110 let sp = fa_split_keys(t_kv, n_head_kv);
9111 let n_splits = if fa_vec { ((t_kv + sp - 1) / sp).max(1) } else { ((t_kv + 255) / 256).max(1) };
9112 let o_len = n_head * n_splits * head_dim;
9113 let ml_len = n_head * n_splits;
9114 let mut part_guard = self.fa_part_pool.lock().unwrap();
9115 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9116 let old = part_guard.take();
9127 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9128 if let Some(old) = old {
9129 self.fa_part_retired.lock().unwrap().push(old);
9130 }
9131 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9132 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9133 }
9134 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9135 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9136 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9137 }
9138 let pg = part_guard.as_mut().unwrap();
9139 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9140 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9141 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9142 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9143 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
9144 let (hd, nh, nhkv, tkvi, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, t_kv as i32, n_splits as i32);
9145 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9146 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
9150 let fa512_min = fa512_min_tkv();
9155 let deep = fa_vec && head_dim == 256 && fa_v4_at(t_kv) && !g
9158 && fa_deep_at(t_kv) && !matches!(fa_v4_mode(), "noB3" | "stage");
9159 let (f, cfg) = if fa_vec && head_dim == 512 && t_kv >= fa512_min {
9160 let gqa = (n_head / n_head_kv).max(1) as u32;
9163 let fv = self.fa_func("fa_decode_vec_q_dpl16", head_dim);
9164 (fv, LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9165 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
9166 } else if fa_vec && head_dim <= 256 {
9167 let gqa = (n_head / n_head_kv).max(1) as u32;
9168 static SMEM_TKV: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9179 let smem_tkv = *SMEM_TKV.get_or_init(|| {
9180 std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9181 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9182 });
9183 if fa_v4_at(t_kv) && head_dim == 256 {
9184 let v4name = match fa_v4_mode() {
9188 "noB3" => "fa_decode_vec_q_v4_noB3", "stage" => "fa_decode_vec_q_v4_stage", _ if deep => "fa_decode_vec_q_v4_deep",
9191 _ => "fa_decode_vec_q_v4",
9192 };
9193 let fv = if g { self.func_g(v4name) } else { self.func(v4name) };
9194 let shmem = (if deep { 12160 } else { 11520 }
9197 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
9198 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9199 fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9200 (fv,
9201 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9202 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9203 } else if fa_v3_active(head_dim) {
9204 let fv = if g { self.func_g("fa_decode_vec_q_v3") } else { self.func("fa_decode_vec_q_v3") };
9207 let shmem = (32 * head_dim * 2) as u32; (fv,
9209 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9210 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9211 } else if fa_v2_on() {
9212 let fv = if g { self.func_g("fa_decode_vec_q_v2") } else { self.func("fa_decode_vec_q_v2") };
9216 let shmem = (2 * 32 * head_dim * 2) as u32; (fv,
9218 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9219 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9220 } else if smem_tkv > 0 && t_kv >= smem_tkv && !g
9221 && !(head_dim == 512 && Self::gkv_on()) {
9222 let fv = if g { self.func_g("fa_decode_vec_q_smem") } else { self.func("fa_decode_vec_q_smem") };
9226 let shmem = (2 * 32 * head_dim * 2) as u32; use cudarc::driver::sys::CUfunction_attribute_enum as A;
9228 fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9229 (fv,
9230 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9231 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
9232 } else {
9233 let fv = if g { self.func_g("fa_decode_vec_q") } else { self.func("fa_decode_vec_q") };
9236 (fv,
9237 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
9238 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
9239 }
9240 } else {
9241 return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
9244 t_kv, None, scale, n_splits,
9245 if fa_vec { sp } else { 256 },
9246 k_tok_bytes, v_tok_bytes, g,
9247 part_o, part_m, part_l, None);
9248 };
9249 let __s_b = self.gpu.stream();
9250 let mut b = __s_b.launch_builder(&f);
9251 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9252 .arg(&hd).arg(&nh).arg(&nhkv).arg(&tkvi).arg(&scale).arg(&nsp).arg(&ktb).arg(&vtb);
9253 unsafe { b.launch(cfg)?; }
9254 let (fc, cfg2) = (if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) },
9257 LaunchConfig { grid_dim: (n_head as u32, 1, 1), block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 });
9258 let __s_b2 = self.gpu.stream();
9259 let mut b2 = __s_b2.launch_builder(&fc);
9260 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
9261 unsafe { b2.launch(cfg2)?; }
9262 Ok(())
9263 }
9264
9265 #[allow(clippy::too_many_arguments)]
9276 pub fn fa_decode_batch_seqs_v4(&self, q: &CudaSlice<f32>,
9277 kv_ptrs: &cudarc::driver::CudaView<u64>,
9278 pos_seq: &CudaSlice<i32>, o: &mut CudaSlice<f32>,
9279 head_dim: usize, n_head: usize, n_head_kv: usize,
9280 b_n: usize, t_kv_max: usize, scale: f32,
9281 split_keys: usize, k_tok_bytes: usize, v_tok_bytes: usize)
9282 -> Result<(), Box<dyn std::error::Error>> {
9283 debug_assert!(head_dim == 256, "seqs twin is v4-stamped (hd256 only)");
9284 let n_splits_max = (t_kv_max + split_keys - 1) / split_keys;
9285 let o_len = b_n * n_head * n_splits_max * head_dim;
9286 let ml_len = b_n * n_head * n_splits_max;
9287 let mut part_guard = self.fa_part_pool.lock().unwrap();
9288 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9289 let old = part_guard.take();
9300 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9301 if let Some(old) = old {
9302 self.fa_part_retired.lock().unwrap().push(old);
9303 }
9304 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9305 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9306 }
9307 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9308 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9309 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9310 }
9311 let pg = part_guard.as_mut().unwrap();
9312 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9313 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9314 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9315 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9316 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9317 let (nspm, spk) = (n_splits_max as i32, split_keys as i32);
9318 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9319 let gqa = (n_head / n_head_kv).max(1) as u32;
9320 let f = self.func("fa_decode_vec_q_seqs_v4");
9321 let shmem = (11520 + 32 * head_dim * 2) as u32;
9323 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9324 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
9325 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, b_n as u32),
9326 block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9327 {
9328 let __s_b = self.gpu.stream();
9329 let mut b = __s_b.launch_builder(&f);
9330 b.arg(q).arg(kv_ptrs).arg(pos_seq).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9331 .arg(&hd).arg(&nh).arg(&nhkv).arg(&scale).arg(&nspm).arg(&spk).arg(&ktb).arg(&vtb);
9332 unsafe { b.launch(cfg)?; }
9333 }
9334 let fc = self.func("fa_decode_combine_seqs");
9335 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, b_n as u32, 1),
9336 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9337 let __s_b2 = self.gpu.stream();
9338 let mut b2 = __s_b2.launch_builder(&fc);
9339 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9340 .arg(pos_seq).arg(&nspm).arg(&spk);
9341 unsafe { b2.launch(cfg2)?; }
9342 Ok(())
9343 }
9344
9345 #[allow(clippy::too_many_arguments)]
9352 pub fn append_kv_quantized_seqs(&self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>,
9353 kv_ptrs: &cudarc::driver::CudaView<u64>,
9354 pos_seq: &CudaSlice<i32>, b_n: usize,
9355 kv_dim_k: usize, kv_dim_v: usize,
9356 k_tok_bytes: usize, v_tok_bytes: usize)
9357 -> Result<(), Box<dyn std::error::Error>> {
9358 let f = self.func("append_quantize_kv_q8_0_q5_1_seqs");
9359 let nblk = (kv_dim_k.max(kv_dim_v) / 32) as u32;
9360 let cfg = LaunchConfig { grid_dim: (nblk, b_n as u32, 1),
9361 block_dim: (32, 1, 1), shared_mem_bytes: 0 };
9362 let (kdk, kdv) = (kv_dim_k as i32, kv_dim_v as i32);
9363 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9364 let __s_b = self.gpu.stream();
9365 let mut b = __s_b.launch_builder(&f);
9366 b.arg(k_rows).arg(v_rows).arg(kv_ptrs).arg(pos_seq)
9367 .arg(&kdk).arg(&kdv).arg(&ktb).arg(&vtb);
9368 unsafe { b.launch(cfg)?; }
9369 Ok(())
9370 }
9371
9372 pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool {
9378 std::env::var("MEMRA_NO_FA_VEC").is_err()
9379 && std::env::var("MEMRA_FA_ROWS_OFF").is_err()
9380 && base_len + 1 >= fa_vec_min_tkv()
9381 && head_dim <= 256 && head_dim % 32 == 0
9382 }
9383
9384 #[allow(clippy::too_many_arguments)]
9393 pub fn fa_decode_rows(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9394 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9395 head_dim: usize, n_head: usize, n_head_kv: usize,
9396 base_len: usize, t: usize, scale: f32,
9397 k_tok_bytes: usize, v_tok_bytes: usize,
9398 base_dev: Option<(&CudaSlice<i32>, i32)>,
9402 kv_shared: bool,
9405 g: bool,
9409 mut q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9412 -> Result<(), Box<dyn std::error::Error>> {
9413 debug_assert!(base_len + 1 >= fa_vec_min_tkv() && head_dim <= 512 && head_dim % 32 == 0);
9414 let t_kv_max = base_len + t; let mut sp = fa_split_keys(t_kv_max, n_head_kv); if head_dim == 512 {
9421 static SP512: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9422 let v = *SP512.get_or_init(|| std::env::var("MEMRA_FA_SP512").ok()
9425 .and_then(|x| x.parse().ok()).unwrap_or(0));
9426 sp = if v >= 8 { v } else { FA_SP512_DEFAULT.load(std::sync::atomic::Ordering::Relaxed) };
9427 }
9428 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9429 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9430 let gqa = (n_head / n_head_kv).max(1) as u32;
9431 let mut groups: Vec<(usize, usize, usize)> = Vec::new(); if head_dim == 512 || fa_split_keys(base_len + 1, n_head_kv) == sp {
9442 groups.push((0, t, sp));
9443 } else {
9444 let mut r0 = 0usize;
9445 while r0 < t {
9446 let sp_g = fa_split_keys(base_len + r0 + 1, n_head_kv);
9447 let mut r1 = r0 + 1;
9448 while r1 < t && fa_split_keys(base_len + r1 + 1, n_head_kv) == sp_g { r1 += 1; }
9449 groups.push((r0, r1 - r0, sp_g));
9450 r0 = r1;
9451 }
9452 }
9453 static SMEM_TKV_R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9457 let smem_tkv = *SMEM_TKV_R.get_or_init(|| {
9458 std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9459 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9460 });
9461 let v4 = fa_v4_at(base_len + t) && head_dim == 256;
9462 let v3 = fa_v3_active(head_dim);
9463 let smem_rows = head_dim <= 256 && !v3 && !fa_v2_on() && smem_tkv > 0 && t_kv_max >= smem_tkv;
9464 let _ = kv_shared;
9469 let i2 = head_dim == 512 && std::env::var("MEMRA_FA_I2").as_deref() != Ok("0");
9472 static TB512: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9486 let tb512 = head_dim == 512 && sp <= 32 && n_head / n_head_kv.max(1) <= 16
9488 && *TB512.get_or_init(|| std::env::var("MEMRA_FA_TB512").as_deref() != Ok("0"));
9489 let fname = if tb512 { "fa_decode_vec_q_rows_v4_512_tb" }
9490 else if i2 { "fa_decode_vec_q_rows_dpl16_i2" }
9491 else if head_dim == 512 { "fa_decode_vec_q_rows_dpl16" } else if v4 { "fa_decode_vec_q_rows_v4" }
9493 else if v3 { "fa_decode_vec_q_rows_v3" }
9494 else if fa_v2_on() { "fa_decode_vec_q_rows_v2" }
9495 else if smem_rows { "fa_decode_vec_q_rows_smem" }
9496 else { "fa_decode_vec_q_rows" };
9497 let f = if head_dim == 512 { self.fa_func(fname, head_dim) }
9498 else if g {
9499 self.func_g(if smem_rows { "fa_decode_vec_q_rows" } else { fname })
9507 }
9508 else { self.func(fname) };
9509 let shmem = if tb512 {
9510 let gk = Self::gkv_on();
9512 let sh = (8192 + 1024 + 32 * 512 + 32 * 64
9513 + 32 * head_dim * if gk { 1 } else { 2 }) as u32;
9514 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9515 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9516 sh
9517 } else if v4 || v3 || smem_rows || fa_v2_on() {
9518 let sh = (if v4 { 11520 + 32 * head_dim * if g { 1 } else { 2 } }
9520 else if v3 { 32 * head_dim * 2 } else { 2 * 32 * head_dim * 2 }) as u32;
9521 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9522 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9523 sh
9524 } else { 0 };
9525 for &(r0, t_g, sp_g) in &groups {
9529 let n_splits_g = (base_len + r0 + t_g).div_ceil(sp_g);
9530 let (nspm, spk) = (n_splits_g as i32, sp_g as i32);
9531 let base_i = (base_len + r0) as i32;
9532 let o_len = t_g * n_head * n_splits_g * head_dim;
9533 let ml_len = t_g * n_head * n_splits_g;
9534 let mut part_guard = self.fa_part_pool.lock().unwrap();
9535 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9536 let old = part_guard.take();
9547 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9548 if let Some(old) = old {
9549 self.fa_part_retired.lock().unwrap().push(old);
9550 }
9551 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9552 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9553 }
9554 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9555 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9556 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9557 }
9558 let pg = part_guard.as_mut().unwrap();
9559 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9560 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9561 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9562 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9563 let (part_o, part_m, part_l) = (&mut *part_o, &mut *part_m, &mut *part_l);
9564 let qv = self.view(q, t * n_head * head_dim);
9565 let q_g = qv.slice(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
9566 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_g as u32, t_g as u32),
9567 block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9568 {
9569 let __s_b = self.gpu.stream();
9570 let mut b = __s_b.launch_builder(&f);
9571 if tb512 {
9572 let (bd, plus) = base_dev.expect("hd512 rows twin requires a device base counter");
9574 let plus_g = plus + r0 as i32;
9575 let nr = t_g as i32;
9576 if Self::pdl_on() && Self::pdl_wb_on() {
9577 use cudarc::driver::{DevicePtr, DevicePtrMut};
9579 let s = &self.gpu.stream();
9580 let (pq, _b0) = q_g.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9581 let (pv, _b2) = v.device_ptr(s);
9582 let (po, _b3) = part_o.device_ptr_mut(s);
9583 let (pm, _b4) = part_m.device_ptr_mut(s);
9584 let (pl, _b5) = part_l.device_ptr_mut(s);
9585 let (pb, _b6) = bd.device_ptr(s);
9586 let mut ps = [
9587 &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9588 &pv as *const _ as *mut _, &po as *const _ as *mut _,
9589 &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9590 &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9591 &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9592 &plus_g as *const _ as *mut _, &scale as *const _ as *mut _,
9593 &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9594 &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9595 &nr as *const _ as *mut _,
9596 ];
9597 unsafe { self.launch_pdl_flash(Self::gkv_on(),
9598 "fa_decode_vec_q_rows_v4_512_tb",
9599 (n_head_kv as u32, n_splits_g as u32, 1), (32, gqa, 1),
9600 shmem, &mut ps)?; }
9601 } else {
9602 let cfg_tb = LaunchConfig {
9603 grid_dim: (n_head_kv as u32, n_splits_g as u32, 1),
9604 block_dim: (32, gqa, 1), shared_mem_bytes: shmem };
9605 b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9606 .arg(&hd).arg(&nh).arg(&nhkv).arg(bd).arg(&plus_g).arg(&scale).arg(&nspm).arg(&spk)
9607 .arg(&ktb).arg(&vtb).arg(&nr);
9608 unsafe { b.launch(cfg_tb)?; }
9609 }
9610 } else if head_dim == 512 {
9611 let (bd, plus) = base_dev.expect("hd512 rows twin requires a device base counter");
9612 let plus_g = plus + r0 as i32;
9613 b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9614 .arg(&hd).arg(&nh).arg(&nhkv).arg(bd).arg(&plus_g).arg(&scale).arg(&nspm).arg(&spk)
9615 .arg(&ktb).arg(&vtb);
9616 unsafe { b.launch(cfg)?; }
9617 } else {
9618 b.arg(&q_g).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9619 .arg(&hd).arg(&nh).arg(&nhkv).arg(&base_i).arg(&scale).arg(&nspm).arg(&spk)
9620 .arg(&ktb).arg(&vtb);
9621 unsafe { b.launch(cfg)?; }
9622 }
9623 }
9624 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t_g as u32, 1),
9625 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9626 let mut o_g = o.slice_mut(r0 * n_head * head_dim..(r0 + t_g) * n_head * head_dim);
9627 if head_dim == 512 {
9628 let (bd, plus) = base_dev.unwrap();
9631 let plus_g = plus + r0 as i32;
9632 if let Some((oq, od)) = q8_out.as_mut() {
9633 debug_assert!(t == 1, "rows q8 emit is a t=1 decode arm");
9635 if Self::pdl_on() && Self::pdl_wb_on() {
9636 use cudarc::driver::{DevicePtr, DevicePtrMut};
9638 let s = &self.gpu.stream();
9639 let (po, _g0) = part_o.device_ptr(s); let (pm, _g1) = part_m.device_ptr(s);
9640 let (pl, _g2) = part_l.device_ptr(s);
9641 let (pq, _g3) = oq.device_ptr_mut(s); let (pd, _g4) = od.device_ptr_mut(s);
9642 let (pb, _g5) = bd.device_ptr(s);
9643 let mut ps = [
9644 &po as *const _ as *mut std::ffi::c_void, &pm as *const _ as *mut _,
9645 &pl as *const _ as *mut _, &pq as *const _ as *mut _,
9646 &pd as *const _ as *mut _, &hd as *const _ as *mut _,
9647 &nh as *const _ as *mut _, &pb as *const _ as *mut _,
9648 &plus_g as *const _ as *mut _, &nspm as *const _ as *mut _,
9649 &spk as *const _ as *mut _,
9650 ];
9651 unsafe { self.launch_pdl_flash(Self::gkv_on(),
9652 "fa_decode_combine_rows_dc_q8_1",
9653 cfg2.grid_dim, cfg2.block_dim, 0, &mut ps)?; }
9654 continue;
9655 }
9656 let fc = self.fa_func("fa_decode_combine_rows_dc_q8_1", head_dim);
9657 let __s_b2 = self.gpu.stream();
9658 let mut b2 = __s_b2.launch_builder(&fc);
9659 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut **oq).arg(&mut **od)
9660 .arg(&hd).arg(&nh).arg(bd).arg(&plus_g).arg(&nspm).arg(&spk);
9661 unsafe { b2.launch(cfg2)?; }
9662 continue;
9663 }
9664 let fc = self.fa_func("fa_decode_combine_rows_dc", head_dim);
9665 let __s_b2 = self.gpu.stream();
9666 let mut b2 = __s_b2.launch_builder(&fc);
9667 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut o_g).arg(&hd).arg(&nh)
9668 .arg(bd).arg(&plus_g).arg(&nspm).arg(&spk);
9669 unsafe { b2.launch(cfg2)?; }
9670 } else {
9671 assert!(q8_out.is_none(), "rows q8 emit requires the hd512 dc combine");
9674 let fc = self.func("fa_decode_combine_rows");
9675 let __s_b2 = self.gpu.stream();
9676 let mut b2 = __s_b2.launch_builder(&fc);
9677 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(&mut o_g).arg(&hd).arg(&nh)
9678 .arg(&base_i).arg(&nspm).arg(&spk);
9679 unsafe { b2.launch(cfg2)?; }
9680 }
9681 }
9682 Ok(())
9683 }
9684
9685 #[allow(clippy::too_many_arguments)]
9689 pub fn fa_decode_rows_w(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9690 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9691 head_dim: usize, n_head: usize, n_head_kv: usize,
9692 base_dev: &CudaSlice<i32>, base_plus: i32, t: usize, scale: f32,
9693 window: usize, k_tok_bytes: usize, v_tok_bytes: usize,
9694 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
9695 -> Result<(), Box<dyn std::error::Error>> {
9696 debug_assert!(head_dim == 256);
9701 let sp = {
9709 static SPW: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9710 let v = *SPW.get_or_init(|| std::env::var("MEMRA_FA_SPW").ok()
9711 .and_then(|x| x.parse().ok()).unwrap_or(0));
9712 if v >= 8 { v } else { FA_SPW_DEFAULT.load(std::sync::atomic::Ordering::Relaxed) }
9713 };
9714 let n_splits_max = (window + sp - 1) / sp;
9715 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9716 let (nspm, spk, wini) = (n_splits_max as i32, sp as i32, window as i32);
9717 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9718 let gqa = (n_head / n_head_kv).max(1) as u32;
9719 let o_len = t * n_head * n_splits_max * head_dim;
9720 let ml_len = t * n_head * n_splits_max;
9721 let mut part_guard = self.fa_part_pool.lock().unwrap();
9722 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9723 let old = part_guard.take();
9734 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9735 if let Some(old) = old {
9736 self.fa_part_retired.lock().unwrap().push(old);
9737 }
9738 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9739 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9740 }
9741 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9742 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9743 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9744 }
9745 let pg = part_guard.as_mut().unwrap();
9746 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9747 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9748 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9749 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9750 static SMEM_TKV_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
9756 let smem_tkv = *SMEM_TKV_W.get_or_init(|| {
9757 std::env::var("MEMRA_FA_SMEM_TKV").ok().and_then(|v| v.parse().ok())
9758 .unwrap_or_else(|| FA_SMEM_TKV_DEFAULT.load(std::sync::atomic::Ordering::Relaxed))
9759 });
9760 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9766 let wg = Self::wkv_on();
9771 let sp2 = gqa <= 4 && fa_v4_at(window)
9774 && std::env::var("MEMRA_FA_SPW2").as_deref() != Ok("0");
9775 if sp2 {
9776 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
9777 if Self::pdl_on() && Self::pdl_wb_on() {
9778 use cudarc::driver::{DevicePtr, DevicePtrMut};
9780 let s = &self.gpu.stream();
9781 let (pq, _b0) = q.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9782 let (pv, _b2) = v.device_ptr(s);
9783 let (po, _b3) = part_o.device_ptr_mut(s);
9784 let (pm, _b4) = part_m.device_ptr_mut(s);
9785 let (pl, _b5) = part_l.device_ptr_mut(s);
9786 let (pb, _b6) = base_dev.device_ptr(s);
9787 let mut ps = [
9788 &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9789 &pv as *const _ as *mut _, &po as *const _ as *mut _,
9790 &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9791 &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9792 &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9793 &base_plus as *const _ as *mut _, &scale as *const _ as *mut _,
9794 &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9795 &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9796 &wini as *const _ as *mut _,
9797 ];
9798 unsafe { self.launch_pdl_flash(wg, "fa_decode_vec_q_rows_v4_w_sp",
9799 (n_head_kv as u32, n_splits_max as u32, t as u32), (32, gqa + 1, 1),
9800 sh, &mut ps)?; }
9801 } else {
9802 let f = if wg { self.func_g("fa_decode_vec_q_rows_v4_w_sp") }
9803 else { self.func("fa_decode_vec_q_rows_v4_w_sp") };
9804 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9805 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9806 block_dim: (32, gqa + 1, 1), shared_mem_bytes: sh };
9807 let __s_b = self.gpu.stream();
9808 let mut b = __s_b.launch_builder(&f);
9809 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9810 .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale).arg(&nspm).arg(&spk)
9811 .arg(&ktb).arg(&vtb).arg(&wini);
9812 unsafe { b.launch(cfg)?; }
9813 }
9814 } else {
9815 if fa_v4_at(window) && Self::pdl_on() && Self::pdl_wb_on() {
9816 let sh = (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32;
9818 use cudarc::driver::{DevicePtr, DevicePtrMut};
9819 let s = &self.gpu.stream();
9820 let (pq, _b0) = q.device_ptr(s); let (pk, _b1) = k.device_ptr(s);
9821 let (pv, _b2) = v.device_ptr(s);
9822 let (po, _b3) = part_o.device_ptr_mut(s);
9823 let (pm, _b4) = part_m.device_ptr_mut(s);
9824 let (pl, _b5) = part_l.device_ptr_mut(s);
9825 let (pb, _b6) = base_dev.device_ptr(s);
9826 let mut ps = [
9827 &pq as *const _ as *mut std::ffi::c_void, &pk as *const _ as *mut _,
9828 &pv as *const _ as *mut _, &po as *const _ as *mut _,
9829 &pm as *const _ as *mut _, &pl as *const _ as *mut _,
9830 &hd as *const _ as *mut _, &nh as *const _ as *mut _,
9831 &nhkv as *const _ as *mut _, &pb as *const _ as *mut _,
9832 &base_plus as *const _ as *mut _, &scale as *const _ as *mut _,
9833 &nspm as *const _ as *mut _, &spk as *const _ as *mut _,
9834 &ktb as *const _ as *mut _, &vtb as *const _ as *mut _,
9835 &wini as *const _ as *mut _,
9836 ];
9837 unsafe { self.launch_pdl_flash(wg, "fa_decode_vec_q_rows_v4_w",
9838 (n_head_kv as u32, n_splits_max as u32, t as u32), (32, gqa, 1),
9839 sh, &mut ps)?; }
9840 } else {
9841 let pick = |name: &str| if wg { self.func_g(name) } else { self.func(name) };
9842 let (f, sh) = if fa_v4_at(window) {
9843 let f = pick("fa_decode_vec_q_rows_v4_w");
9844 (f, (11520 + 32 * head_dim * if wg { 1 } else { 2 }) as u32)
9845 } else if smem_tkv > 0 && window >= smem_tkv {
9846 (pick("fa_decode_vec_q_rows_smem_w"), (2 * 32 * head_dim * 2) as u32)
9849 } else {
9850 (pick("fa_decode_vec_q_rows_reg_w"), 0u32)
9851 };
9852 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9853 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9854 block_dim: (32, gqa, 1), shared_mem_bytes: sh };
9855 let __s_b = self.gpu.stream();
9856 let mut b = __s_b.launch_builder(&f);
9857 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9858 .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale).arg(&nspm).arg(&spk)
9859 .arg(&ktb).arg(&vtb).arg(&wini);
9860 unsafe { b.launch(cfg)?; }
9861 }
9862 }
9863 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
9864 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9865 if let Some((oq, od)) = q8_out {
9866 if Self::pdl_on() && Self::pdl_wb_on() {
9869 use cudarc::driver::{DevicePtr, DevicePtrMut};
9871 let s = &self.gpu.stream();
9872 let (po, _g0) = part_o.device_ptr(s); let (pm, _g1) = part_m.device_ptr(s);
9873 let (pl, _g2) = part_l.device_ptr(s);
9874 let (pq, _g3) = oq.device_ptr_mut(s); let (pd, _g4) = od.device_ptr_mut(s);
9875 let mut ps = [
9876 &po as *const _ as *mut std::ffi::c_void, &pm as *const _ as *mut _,
9877 &pl as *const _ as *mut _, &pq as *const _ as *mut _,
9878 &pd as *const _ as *mut _, &hd as *const _ as *mut _,
9879 &nh as *const _ as *mut _, &nspm as *const _ as *mut _,
9880 &spk as *const _ as *mut _, &wini as *const _ as *mut _,
9881 ];
9882 unsafe { self.launch_pdl_flash(wg, "fa_decode_combine_rows_w_q8_1",
9883 cfg2.grid_dim, cfg2.block_dim, 0, &mut ps)?; }
9884 return Ok(());
9885 }
9886 let fc = if wg { self.func_g("fa_decode_combine_rows_w_q8_1") }
9887 else { self.func("fa_decode_combine_rows_w_q8_1") };
9888 let __s_b2 = self.gpu.stream();
9889 let mut b2 = __s_b2.launch_builder(&fc);
9890 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh)
9891 .arg(&nspm).arg(&spk).arg(&wini);
9892 unsafe { b2.launch(cfg2)?; }
9893 return Ok(());
9894 }
9895 let fc = if wg { self.func_g("fa_decode_combine_rows_w") }
9896 else { self.func("fa_decode_combine_rows_w") };
9897 let __s_b2 = self.gpu.stream();
9898 let mut b2 = __s_b2.launch_builder(&fc);
9899 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9900 .arg(&nspm).arg(&spk).arg(&wini);
9901 unsafe { b2.launch(cfg2)?; }
9902 Ok(())
9903 }
9904
9905 #[allow(clippy::too_many_arguments)]
9911 pub fn fa_decode_rows_dc(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
9912 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
9913 head_dim: usize, n_head: usize, n_head_kv: usize,
9914 base_dev: &CudaSlice<i32>, t_kv_upper: usize, t: usize, scale: f32,
9915 k_tok_bytes: usize, v_tok_bytes: usize, base_plus: i32, g: bool)
9916 -> Result<(), Box<dyn std::error::Error>> {
9917 let v4 = head_dim == 256 && fa_v4_at(t_kv_upper);
9918 assert!(v4 || fa_v3_active(head_dim), "stream fa rows requires the v3 or v4 lane");
9919 assert!(v4 || base_plus == 0, "v3_dc kernel takes no plus arg");
9920 if v4 {
9921 let sp = fa_split_keys(t_kv_upper, n_head_kv);
9922 let n_splits_max = (t_kv_upper + sp - 1) / sp;
9923 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9924 let (nspm, spk) = (n_splits_max as i32, sp as i32);
9925 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9926 let gqa = (n_head / n_head_kv).max(1) as u32;
9927 let o_len = t * n_head * n_splits_max * head_dim;
9928 let ml_len = t * n_head * n_splits_max;
9929 let mut part_guard = self.fa_part_pool.lock().unwrap();
9930 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9931 let old = part_guard.take();
9942 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
9943 if let Some(old) = old {
9944 self.fa_part_retired.lock().unwrap().push(old);
9945 }
9946 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
9947 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
9948 }
9949 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
9950 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
9951 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
9952 }
9953 let pg = part_guard.as_mut().unwrap();
9954 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
9955 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
9956 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
9957 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
9958 let f = if g { self.func_g("fa_decode_vec_q_rows_v4_dc") }
9959 else { self.func("fa_decode_vec_q_rows_v4_dc") };
9960 let sh = (11520 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
9961 use cudarc::driver::sys::CUfunction_attribute_enum as A;
9962 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
9963 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
9964 block_dim: (32, gqa, 1), shared_mem_bytes: sh };
9965 let __s_b = self.gpu.stream();
9966 let mut b = __s_b.launch_builder(&f);
9967 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
9968 .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&base_plus).arg(&scale)
9969 .arg(&nspm).arg(&spk).arg(&ktb).arg(&vtb);
9970 unsafe { b.launch(cfg)?; }
9971 let fc = self.func("fa_decode_combine_rows_dc");
9972 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
9973 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
9974 let __s_b2 = self.gpu.stream();
9975 let mut b2 = __s_b2.launch_builder(&fc);
9976 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
9977 .arg(base_dev).arg(&base_plus).arg(&nspm).arg(&spk);
9978 unsafe { b2.launch(cfg2)?; }
9979 return Ok(());
9980 }
9981 let sp = fa_split_keys(t_kv_upper, n_head_kv);
9982 let n_splits_max = (t_kv_upper + sp - 1) / sp;
9983 let (hd, nh, nhkv) = (head_dim as i32, n_head as i32, n_head_kv as i32);
9984 let (nspm, spk) = (n_splits_max as i32, sp as i32);
9985 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
9986 let gqa = (n_head / n_head_kv).max(1) as u32;
9987 let o_len = t * n_head * n_splits_max * head_dim;
9988 let ml_len = t * n_head * n_splits_max;
9989 let mut part_guard = self.fa_part_pool.lock().unwrap();
9990 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
9991 let old = part_guard.take();
10002 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
10003 if let Some(old) = old {
10004 self.fa_part_retired.lock().unwrap().push(old);
10005 }
10006 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
10007 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
10008 }
10009 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
10010 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
10011 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
10012 }
10013 let pg = part_guard.as_mut().unwrap();
10014 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
10015 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
10016 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
10017 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
10018 let f = self.func("fa_decode_vec_q_rows_v3_dc");
10019 let sh = (32 * head_dim * 2) as u32;
10020 use cudarc::driver::sys::CUfunction_attribute_enum as A;
10021 f.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, sh as i32)?;
10022 let cfg = LaunchConfig { grid_dim: (n_head_kv as u32, n_splits_max as u32, t as u32),
10023 block_dim: (32, gqa, 1), shared_mem_bytes: sh };
10024 let __s_b = self.gpu.stream();
10025 let mut b = __s_b.launch_builder(&f);
10026 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
10027 .arg(&hd).arg(&nh).arg(&nhkv).arg(base_dev).arg(&scale).arg(&nspm).arg(&spk)
10028 .arg(&ktb).arg(&vtb);
10029 unsafe { b.launch(cfg)?; }
10030 let fc = self.func("fa_decode_combine_rows_dc");
10031 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, t as u32, 1),
10032 block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
10033 let plus0 = 0i32;
10034 let __s_b2 = self.gpu.stream();
10035 let mut b2 = __s_b2.launch_builder(&fc);
10036 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh)
10037 .arg(base_dev).arg(&plus0).arg(&nspm).arg(&spk);
10038 unsafe { b2.launch(cfg2)?; }
10039 Ok(())
10040 }
10041
10042 pub fn fa_decode_dc(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
10053 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
10054 head_dim: usize, n_head: usize, n_head_kv: usize,
10055 t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32,
10056 k_tok_bytes: usize, v_tok_bytes: usize, g: bool)
10057 -> Result<(), Box<dyn std::error::Error>> {
10058 self.fa_decode_dc_q8(q, k, v, o, head_dim, n_head, n_head_kv, t_kv_dev, bucket_max,
10059 scale, k_tok_bytes, v_tok_bytes, g, None)
10060 }
10061
10062 #[allow(clippy::too_many_arguments)]
10065 pub fn fa_decode_dc_q8(&self, q: &CudaSlice<f32>, k: &cudarc::driver::CudaView<u8>,
10066 v: &cudarc::driver::CudaView<u8>, o: &mut CudaSlice<f32>,
10067 head_dim: usize, n_head: usize, n_head_kv: usize,
10068 t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32,
10069 k_tok_bytes: usize, v_tok_bytes: usize, g: bool,
10070 q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>)
10071 -> Result<(), Box<dyn std::error::Error>> {
10072 let mut fa_vec = std::env::var("MEMRA_NO_FA_VEC").is_err() && bucket_max >= fa_vec_min_tkv();
10080 if g && head_dim == 256 && !fa_v4_at(bucket_max) { fa_vec = false; } let sp = fa_split_keys(bucket_max, n_head_kv);
10082 let n_splits = if fa_vec { ((bucket_max + sp - 1) / sp).max(1) } else { ((bucket_max + 255) / 256).max(1) };
10083 let o_len = n_head * n_splits * head_dim;
10084 let ml_len = n_head * n_splits;
10085 let mut part_guard = self.fa_part_pool.lock().unwrap();
10086 if part_guard.as_ref().map(|pp| pp.0.len() < o_len || pp.1.len() < ml_len).unwrap_or(true) {
10087 let old = part_guard.take();
10098 let (co, cm) = old.as_ref().map(|pp| (pp.0.len(), pp.1.len())).unwrap_or((0, 0));
10099 if let Some(old) = old {
10100 self.fa_part_retired.lock().unwrap().push(old);
10101 }
10102 if std::env::var("MEMRA_DEBUG_FAPOOL").is_ok() {
10103 eprintln!("[fa-pool] REALLOC o {} -> {} ml {} -> {} (old retired)", co, o_len, cm, ml_len);
10104 }
10105 *part_guard = Some((self.alloc_uninit::<f32>(o_len.max(2 * co))?,
10106 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?,
10107 self.alloc_uninit::<f32>(ml_len.max(2 * cm))?));
10108 }
10109 let pg = part_guard.as_mut().unwrap();
10110 self.gpu.stream().memset_zeros(&mut pg.0.slice_mut(0..o_len))?;
10111 self.gpu.stream().memset_zeros(&mut pg.1.slice_mut(0..ml_len))?;
10112 self.gpu.stream().memset_zeros(&mut pg.2.slice_mut(0..ml_len))?;
10113 let (part_o, part_m, part_l) = (&mut pg.0, &mut pg.1, &mut pg.2);
10114 let (hd, nh, nhkv, nsp) = (head_dim as i32, n_head as i32, n_head_kv as i32, n_splits as i32);
10115 let (ktb, vtb) = (k_tok_bytes as i64, v_tok_bytes as i64);
10116 let fa_vec = fa_vec && head_dim <= 512 && head_dim % 32 == 0;
10117 let deep = fa_vec && head_dim == 256 && fa_v4_at(bucket_max) && !g
10120 && fa_deep_at(bucket_max) && !matches!(fa_v4_mode(), "noB3" | "stage");
10121 let (f, cfg) = if fa_vec && head_dim == 512 && bucket_max >= {
10122 static FA512_MIN_DC: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
10123 *FA512_MIN_DC.get_or_init(|| std::env::var("MEMRA_FA512_MIN").ok()
10124 .and_then(|v| v.parse().ok()).unwrap_or(512))
10125 } {
10126 let gqa = (n_head / n_head_kv).max(1) as u32;
10128 (self.fa_func("fa_decode_vec_q_dpl16_dc", head_dim),
10129 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10130 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
10131 } else if fa_vec && head_dim == 512 {
10132 return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
10135 0, Some(t_kv_dev), scale, n_splits, sp,
10136 k_tok_bytes, v_tok_bytes, g,
10137 &mut *part_o, &mut *part_m, &mut *part_l, q8_out);
10138 } else if fa_vec && head_dim == 256 && fa_v4_at(bucket_max) {
10139 let gqa = (n_head / n_head_kv).max(1) as u32;
10142 let fv = if g { self.func_g("fa_decode_vec_q_v4_dc") }
10143 else if deep { self.func("fa_decode_vec_q_v4_deep_dc") }
10144 else { self.func("fa_decode_vec_q_v4_dc") };
10145 let shmem = (if deep { 12160 } else { 11520 }
10146 + 32 * head_dim * if g { 1 } else { 2 }) as u32;
10147 use cudarc::driver::sys::CUfunction_attribute_enum as A;
10148 fv.set_attribute(A::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, shmem as i32)?;
10149 (fv, LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10150 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10151 } else if fa_vec && fa_v3_active(head_dim) {
10152 let gqa = (n_head / n_head_kv).max(1) as u32;
10155 let fv = if g { self.func_g("fa_decode_vec_q_v3_dc") } else { self.func("fa_decode_vec_q_v3_dc") };
10156 let shmem = (32 * head_dim * 2) as u32; (fv,
10158 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10159 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10160 } else if fa_vec && fa_v2_on() {
10161 let gqa = (n_head / n_head_kv).max(1) as u32;
10165 let fv = if g { self.func_g("fa_decode_vec_q_v2_dc") } else { self.func("fa_decode_vec_q_v2_dc") };
10166 let shmem = (2 * 32 * head_dim * 2) as u32; (fv,
10168 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10169 block_dim: (32, gqa, 1), shared_mem_bytes: shmem })
10170 } else if fa_vec {
10171 let gqa = (n_head / n_head_kv).max(1) as u32;
10172 let fv = if g { self.func_g("fa_decode_vec_q_dc") } else { self.func("fa_decode_vec_q_dc") };
10174 (fv,
10175 LaunchConfig { grid_dim: (n_head_kv as u32, n_splits as u32, 1),
10176 block_dim: (32, gqa, 1), shared_mem_bytes: 0 })
10177 } else {
10178 return self.fa_decode_scalar_unified(q, k, v, o, head_dim, n_head, n_head_kv,
10179 0, Some(t_kv_dev), scale, n_splits,
10180 if fa_vec { sp } else { 256 },
10181 k_tok_bytes, v_tok_bytes, g,
10182 &mut *part_o, &mut *part_m, &mut *part_l, q8_out);
10183 };
10184 let ski = sp as i32; let __s_b = self.gpu.stream();
10186 let mut b = __s_b.launch_builder(&f);
10187 b.arg(q).arg(k).arg(v).arg(&mut *part_o).arg(&mut *part_m).arg(&mut *part_l)
10188 .arg(&hd).arg(&nh).arg(&nhkv).arg(t_kv_dev).arg(&scale).arg(&nsp).arg(&ski)
10189 .arg(&ktb).arg(&vtb);
10190 unsafe { b.launch(cfg)?; }
10191 let cfg2 = LaunchConfig { grid_dim: (n_head as u32, 1, 1), block_dim: (head_dim as u32, 1, 1), shared_mem_bytes: 0 };
10192 if let Some((oq, od)) = q8_out {
10193 let fc = if g { self.func_g("fa_decode_combine_q8_1") }
10194 else { self.fa_func("fa_decode_combine_q8_1", head_dim) };
10195 let __s_b2 = self.gpu.stream();
10196 let mut b2 = __s_b2.launch_builder(&fc);
10197 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(oq).arg(od).arg(&hd).arg(&nh).arg(&nsp);
10198 unsafe { b2.launch(cfg2)?; }
10199 return Ok(());
10200 }
10201 let fc = if g { self.func_g("fa_decode_combine_f32") } else { self.fa_func("fa_decode_combine_f32", head_dim) };
10202 let __s_b2 = self.gpu.stream();
10203 let mut b2 = __s_b2.launch_builder(&fc);
10204 b2.arg(&*part_o).arg(&*part_m).arg(&*part_l).arg(o).arg(&hd).arg(&nh).arg(&nsp);
10205 unsafe { b2.launch(cfg2)?; }
10206 Ok(())
10207 }
10208
10209 pub fn fa_geom_eager(&self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool) -> (bool, usize) {
10215 let fa_ok = std::env::var("MEMRA_NO_FA_VEC").is_err() && t_kv >= fa_vec_min_tkv();
10219 let vec512 = fa_ok && head_dim == 512 && t_kv >= fa512_min_tkv();
10225 let mut fa_vec = vec512 || (fa_ok && head_dim <= 256 && head_dim % 32 == 0);
10226 if g && head_dim == 256 && !fa_v4_at(t_kv) { fa_vec = false; }
10232 let sp = fa_split_keys(t_kv, n_head_kv);
10233 let n_splits = if fa_vec { ((t_kv + sp - 1) / sp).max(1) } else { ((t_kv + 255) / 256).max(1) };
10234 (fa_vec, n_splits)
10235 }
10236
10237 pub fn fa_bucket_key(&self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool) -> (bool, usize) {
10243 self.fa_geom_eager(t_kv, head_dim, n_head_kv, g)
10244 }
10245
10246 pub fn capture_graph_retained<F>(&self, step: F)
10258 -> Result<(cudarc::driver::CudaGraph, Vec<Box<dyn std::any::Any + Send>>), Box<dyn std::error::Error>>
10259 where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10260 {
10261 use cudarc::driver::sys::CUgraphInstantiate_flags;
10262 self.capture_graph_retained_flags(
10263 CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, step)
10264 }
10265
10266 pub fn capture_graph_retained_flags<F>(&self,
10271 flags: cudarc::driver::sys::CUgraphInstantiate_flags, mut step: F)
10272 -> Result<(cudarc::driver::CudaGraph, Vec<Box<dyn std::any::Any + Send>>), Box<dyn std::error::Error>>
10273 where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10274 {
10275 use cudarc::driver::sys::CUstreamCaptureMode;
10276 self.capture_keep.lock().unwrap().clear();
10284 let was_tracking = self.gpu.ctx.is_event_tracking();
10285 if was_tracking { unsafe { self.gpu.ctx.disable_event_tracking(); } }
10286 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
10287 self.capture_keep_on.store(true, std::sync::atomic::Ordering::Relaxed);
10288 let w = (|| { step(self)?; step(self) })();
10289 self.capture_keep_on.store(false, std::sync::atomic::Ordering::Relaxed);
10290 w?;
10291 self.gpu.stream().synchronize()?;
10292 self.gpu.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
10293 let r = step(self);
10294 let g = self.gpu.stream().end_capture(flags);
10295 r?;
10296 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
10297 graph.upload()?;
10298 Ok(graph)
10299 };
10300 let result = run();
10301 self.capture_keep_on.store(false, std::sync::atomic::Ordering::Relaxed);
10302 if was_tracking { unsafe { self.gpu.ctx.enable_event_tracking(); } }
10303 let keeper = std::mem::take(&mut *self.capture_keep.lock().unwrap());
10304 Ok((result?, keeper))
10305 }
10306
10307 pub fn capture_graph<F>(&self, mut step: F) -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>>
10308 where F: FnMut(&Engine) -> Result<(), Box<dyn std::error::Error>>
10309 {
10310 use cudarc::driver::sys::{CUstreamCaptureMode, CUgraphInstantiate_flags};
10311 let was_tracking = self.gpu.ctx.is_event_tracking();
10319 if was_tracking { unsafe { self.gpu.ctx.disable_event_tracking(); } }
10320 let mut run = || -> Result<cudarc::driver::CudaGraph, Box<dyn std::error::Error>> {
10321 step(self)?;
10323 step(self)?;
10324 self.gpu.stream().synchronize()?;
10325 self.gpu.stream().begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)?;
10327 let r = step(self);
10330 let g = self.gpu.stream().end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
10331 r?;
10332 let graph = g?.ok_or("capture produced no graph (stream was not capturing)")?;
10333 graph.upload()?;
10334 Ok(graph)
10335 };
10336 let result = run();
10337 if was_tracking { unsafe { self.gpu.ctx.enable_event_tracking(); } }
10338 result
10339 }
10340
10341 pub fn gdn_scan_s128_view(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10343 g: &CudaSlice<f32>, beta: &CudaSlice<f32>,
10344 state_in: &cudarc::driver::CudaView<f32>,
10345 state_out: &mut cudarc::driver::CudaViewMut<f32>,
10346 o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32)
10347 -> Result<(), Box<dyn std::error::Error>> {
10348 let f = self.func("gdn_scan_s128");
10349 const S_V: u32 = 128; const WARP: u32 = 32; const COLS: u32 = 4;
10350 let cfg = LaunchConfig { grid_dim: (n_head as u32, 1, S_V / COLS), block_dim: (WARP, COLS, 1), shared_mem_bytes: 0 };
10351 let (h, ti) = (n_head as i32, t as i32);
10352 let __s_b = self.gpu.stream();
10353 let mut b = __s_b.launch_builder(&f);
10354 b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o).arg(&h).arg(&ti).arg(&scale);
10355 unsafe { b.launch(cfg)?; }
10356 Ok(())
10357 }
10358
10359 pub fn ssm_conv1d_view(&self, x: &cudarc::driver::CudaView<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10361 conv_dim: usize, t: usize, d_conv: usize, silu: bool)
10362 -> Result<(), Box<dyn std::error::Error>> {
10363 let f = self.func("ssm_conv1d_silu_f32");
10364 let cfg = LaunchConfig { grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
10366 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10367 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
10368 let __s_b = self.gpu.stream();
10369 let mut b = __s_b.launch_builder(&f);
10370 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
10371 unsafe { b.launch(cfg)?; }
10372 Ok(())
10373 }
10374
10375 pub fn ssm_conv1d_tm(&self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10382 conv_dim: usize, t: usize, d_conv: usize)
10383 -> Result<(), Box<dyn std::error::Error>> {
10384 let f = self.func("ssm_conv1d_tm_f32");
10385 let cfg = LaunchConfig {
10386 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10387 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10388 };
10389 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10390 let __s_b = self.gpu.stream();
10391 let mut b = __s_b.launch_builder(&f);
10392 b.arg(qkv_tm).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10393 unsafe { b.launch(cfg)?; }
10394 Ok(())
10395 }
10396
10397 pub fn ssm_conv1d_tm_state(&self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
10405 w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10406 conv_dim: usize, t: usize, d_conv: usize)
10407 -> Result<(), Box<dyn std::error::Error>> {
10408 self.ssm_conv1d_tm_state_pad(qkv_tm, conv_state, w, y, conv_dim, t, d_conv, None)
10409 }
10410
10411 #[allow(clippy::too_many_arguments)]
10414 pub fn ssm_conv1d_tm_state_pad(&self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
10415 w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10416 conv_dim: usize, t: usize, d_conv: usize,
10417 pad_len: Option<&CudaSlice<i32>>)
10418 -> Result<(), Box<dyn std::error::Error>> {
10419 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
10420 let ring_old = if t < d_conv - 1 { Some(self.clone_dtod(conv_state)?) } else { None };
10424 {
10425 let f = self.func("ssm_conv1d_tm_state_f32");
10426 let cfg = LaunchConfig {
10427 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10428 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10429 };
10430 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10431 let __s_b = self.gpu.stream();
10432 let mut b = __s_b.launch_builder(&f);
10433 b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10434 unsafe { b.launch(cfg)?; }
10435 }
10436 match (ring_old, pad_len) {
10437 (None, Some(len_d)) => {
10438 let f = self.func("ssm_conv_ring_update_dev_f32");
10439 let n = conv_dim * (d_conv - 1);
10440 let cfg = LaunchConfig::for_num_elems(n as u32);
10441 let (cd, dc) = (conv_dim as i32, d_conv as i32);
10442 let __s_b = self.gpu.stream();
10443 let mut b = __s_b.launch_builder(&f);
10444 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10445 unsafe { b.launch(cfg)?; }
10446 }
10447 (None, None) => {
10448 let f = self.func("ssm_conv_ring_update_f32");
10449 let n = conv_dim * (d_conv - 1);
10450 let cfg = LaunchConfig::for_num_elems(n as u32);
10451 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10452 let __s_b = self.gpu.stream();
10453 let mut b = __s_b.launch_builder(&f);
10454 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10455 unsafe { b.launch(cfg)?; }
10456 }
10457 (Some(old), _) => self.ssm_conv_ring_rebuild(qkv_tm, &old, conv_state, conv_dim, t, d_conv)?,
10458 }
10459 Ok(())
10460 }
10461
10462 pub fn ssm_conv1d_tm_state_pad_v(&self, qkv_tm: &cudarc::driver::CudaView<f32>, conv_state: &mut CudaSlice<f32>,
10464 w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10465 conv_dim: usize, t: usize, d_conv: usize,
10466 pad_len: Option<&CudaSlice<i32>>)
10467 -> Result<(), Box<dyn std::error::Error>> {
10468 assert!(t >= 1, "ssm_conv1d_tm_state requires T >= 1");
10469 let ring_old = if t < d_conv - 1 { Some(self.clone_dtod(conv_state)?) } else { None };
10473 {
10474 let f = self.func("ssm_conv1d_tm_state_f32");
10475 let cfg = LaunchConfig {
10476 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10477 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10478 };
10479 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10480 let __s_b = self.gpu.stream();
10481 let mut b = __s_b.launch_builder(&f);
10482 b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc);
10483 unsafe { b.launch(cfg)?; }
10484 }
10485 match (ring_old, pad_len) {
10486 (None, Some(len_d)) => {
10487 let f = self.func("ssm_conv_ring_update_dev_f32");
10488 let n = conv_dim * (d_conv - 1);
10489 let cfg = LaunchConfig::for_num_elems(n as u32);
10490 let (cd, dc) = (conv_dim as i32, d_conv as i32);
10491 let __s_b = self.gpu.stream();
10492 let mut b = __s_b.launch_builder(&f);
10493 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10494 unsafe { b.launch(cfg)?; }
10495 }
10496 (None, None) => {
10497 let f = self.func("ssm_conv_ring_update_f32");
10498 let n = conv_dim * (d_conv - 1);
10499 let cfg = LaunchConfig::for_num_elems(n as u32);
10500 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10501 let __s_b = self.gpu.stream();
10502 let mut b = __s_b.launch_builder(&f);
10503 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10504 unsafe { b.launch(cfg)?; }
10505 }
10506 (Some(_), _) => unreachable!(
10507 "ssm_conv1d_tm_state_pad_v: T < d_conv-1 has no view path (PRIME_MIN_T gates it)"),
10508 }
10509 Ok(())
10510 }
10511
10512 pub fn ssm_conv_ring_rebuild(&self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>,
10517 conv_state: &mut CudaSlice<f32>,
10518 conv_dim: usize, tc: usize, d_conv: usize)
10519 -> Result<(), Box<dyn std::error::Error>> {
10520 let f = self.func("ssm_conv_ring_rebuild_f32");
10521 let n = conv_dim * (d_conv - 1);
10522 let cfg = LaunchConfig::for_num_elems(n as u32);
10523 let (cd, ti, dc) = (conv_dim as i32, tc as i32, d_conv as i32);
10524 let __s_b = self.gpu.stream();
10525 let mut b = __s_b.launch_builder(&f);
10526 b.arg(qkv_tm).arg(ring_old).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10527 unsafe { b.launch(cfg)?; }
10528 Ok(())
10529 }
10530
10531 #[allow(clippy::too_many_arguments)]
10536 pub fn gdn_prep_decode(&self, conv_out: &CudaSlice<f32>, beta_raw: &CudaSlice<f32>,
10537 alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10538 q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10539 beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
10540 d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32)
10541 -> Result<(), Box<dyn std::error::Error>> {
10542 let f = self.func("gdn_prep_decode_f32");
10543 let cfg = LaunchConfig { grid_dim: (num_v as u32, 1, 1), block_dim: (32, 4, 1), shared_mem_bytes: 0 };
10544 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
10545 let __s_b = self.gpu.stream();
10546 let mut b = __s_b.launch_builder(&f);
10547 b.arg(conv_out).arg(beta_raw).arg(alpha).arg(dt_bias).arg(a)
10548 .arg(q_l2).arg(k_l2).arg(v_g).arg(beta).arg(g_log)
10549 .arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&eps);
10550 unsafe { b.launch(cfg)?; }
10551 Ok(())
10552 }
10553
10554 #[allow(clippy::too_many_arguments)]
10558 pub fn ssm_conv1d_gdn(&self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>,
10559 q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10560 conv_dim: usize, t: usize, d_conv: usize,
10561 d_state: usize, num_v: usize, num_k: usize, key_dim: usize)
10562 -> Result<(), Box<dyn std::error::Error>> {
10563 let f = self.func("ssm_conv1d_gdn_f32");
10564 let cfg = LaunchConfig {
10565 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10566 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10567 };
10568 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10569 let (ds, nv, nk, kd) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
10570 let __s_b = self.gpu.stream();
10571 let mut b = __s_b.launch_builder(&f);
10572 b.arg(qkv_tm).arg(w).arg(q_g).arg(k_g).arg(v_g)
10573 .arg(&cd).arg(&ti).arg(&dc).arg(&ds).arg(&nv).arg(&nk).arg(&kd);
10574 unsafe { b.launch(cfg)?; }
10575 Ok(())
10576 }
10577
10578 pub fn ssm_conv1d(&self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
10579 conv_dim: usize, t: usize, d_conv: usize, silu: bool)
10580 -> Result<(), Box<dyn std::error::Error>> {
10581 let f = self.func("ssm_conv1d_silu_f32");
10582 let cfg = LaunchConfig { grid_dim: (conv_dim as u32, ((t as u32 + 255) / 256).max(1), 1),
10583 block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10584 let (cd, ti, dc, s) = (conv_dim as i32, t as i32, d_conv as i32, silu as i32);
10585 let __s_b = self.gpu.stream();
10586 let mut b = __s_b.launch_builder(&f);
10587 b.arg(x).arg(w).arg(y).arg(&cd).arg(&ti).arg(&dc).arg(&s);
10588 unsafe { b.launch(cfg)?; }
10589 Ok(())
10590 }
10591
10592 pub fn gdn_scan_s128(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10595 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
10596 state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
10597 n_head: usize, t: usize, scale: f32)
10598 -> Result<(), Box<dyn std::error::Error>> {
10599 let f = self.func("gdn_scan_s128");
10600 const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
10601 let cfg = LaunchConfig {
10602 grid_dim: (n_head as u32, 1, S_V / COLS_PER_BLOCK),
10603 block_dim: (WARP, COLS_PER_BLOCK, 1),
10604 shared_mem_bytes: 0,
10605 };
10606 let (h, ti) = (n_head as i32, t as i32);
10607 let __s_b = self.gpu.stream();
10608 let mut b = __s_b.launch_builder(&f);
10609 b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in).arg(state_out).arg(o).arg(&h).arg(&ti).arg(&scale);
10610 unsafe { b.launch(cfg)?; }
10611 Ok(())
10612 }
10613
10614 #[allow(clippy::too_many_arguments)]
10619 pub fn ssm_conv1d_fused_decode_b(
10620 &self, qkv_cols: &CudaSlice<f32>, conv_state_ptrs: &cudarc::driver::CudaView<u64>,
10621 w: &CudaSlice<f32>, conv_outs: &mut CudaSlice<f32>, conv_dim: usize, d_conv: usize,
10622 b_n: usize) -> Result<(), Box<dyn std::error::Error>> {
10623 let f = self.func("ssm_conv1d_fused_decode_b_f32");
10624 let cfg = LaunchConfig {
10625 grid_dim: (((conv_dim + 255) / 256) as u32, 1, b_n as u32),
10626 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10627 };
10628 let (cd, dc) = (conv_dim as i32, d_conv as i32);
10629 let __s_b = self.gpu.stream();
10630 let mut b = __s_b.launch_builder(&f);
10631 b.arg(qkv_cols).arg(conv_state_ptrs).arg(w).arg(conv_outs).arg(&cd).arg(&dc);
10632 unsafe { b.launch(cfg)?; }
10633 Ok(())
10634 }
10635
10636 #[allow(clippy::too_many_arguments)]
10637 pub fn gdn_prep_decode_b(
10638 &self, conv_outs: &CudaSlice<f32>, beta_raws: &CudaSlice<f32>, alphas: &CudaSlice<f32>,
10639 dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10640 q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
10641 beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>,
10642 d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32,
10643 conv_dim: usize, b_n: usize) -> Result<(), Box<dyn std::error::Error>> {
10644 let f = self.func("gdn_prep_decode_b_f32");
10645 let cfg = LaunchConfig {
10646 grid_dim: (num_v as u32, 1, b_n as u32),
10647 block_dim: (32, 4, 1), shared_mem_bytes: 0,
10648 };
10649 let (ds, nv, nk, kd, cd) =
10650 (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, conv_dim as i32);
10651 let __s_b = self.gpu.stream();
10652 let mut b = __s_b.launch_builder(&f);
10653 b.arg(conv_outs).arg(beta_raws).arg(alphas).arg(dt_bias).arg(a)
10654 .arg(q_l2).arg(k_l2).arg(v_g).arg(beta).arg(g_log)
10655 .arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&eps).arg(&cd);
10656 unsafe { b.launch(cfg)?; }
10657 Ok(())
10658 }
10659
10660 #[allow(clippy::too_many_arguments)]
10661 pub fn gdn_scan_s128_batched(
10662 &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10663 g: &CudaSlice<f32>, beta: &CudaSlice<f32>,
10664 state_in_ptrs: &cudarc::driver::CudaView<u64>,
10665 state_out_ptrs: &cudarc::driver::CudaView<u64>,
10666 o: &mut CudaSlice<f32>, n_head: usize, b_n: usize, scale: f32)
10667 -> Result<(), Box<dyn std::error::Error>> {
10668 let f = self.func("gdn_scan_s128_b");
10669 const S_V: u32 = 128; const WARP: u32 = 32; const COLS_PER_BLOCK: u32 = 4;
10670 let cfg = LaunchConfig {
10671 grid_dim: (n_head as u32, b_n as u32, S_V / COLS_PER_BLOCK),
10672 block_dim: (WARP, COLS_PER_BLOCK, 1), shared_mem_bytes: 0,
10673 };
10674 let h = n_head as i32;
10675 let __s_b = self.gpu.stream();
10676 let mut b = __s_b.launch_builder(&f);
10677 b.arg(q).arg(k).arg(v).arg(g).arg(beta).arg(state_in_ptrs).arg(state_out_ptrs)
10678 .arg(o).arg(&h).arg(&scale);
10679 unsafe { b.launch(cfg)?; }
10680 Ok(())
10681 }
10682
10683 pub fn gdn_chunked_enabled() -> bool {
10692 static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10693 *E.get_or_init(|| std::env::var("MEMRA_GDN_CHUNKED").map(|v| v != "0").unwrap_or(true))
10694 }
10695
10696 pub fn gdn_chunk_size() -> usize {
10701 static C: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
10702 *C.get_or_init(|| {
10703 let c: usize = std::env::var("MEMRA_GDN_CHUNK").ok()
10704 .and_then(|v| v.parse().ok()).unwrap_or(32);
10705 c.clamp(32, 128) / 32 * 32
10706 })
10707 }
10708
10709 #[allow(clippy::too_many_arguments)]
10714 #[allow(clippy::too_many_arguments, clippy::type_complexity)]
10717 #[allow(clippy::too_many_arguments)]
10718 pub fn gdn_chunk_k123(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
10719 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, wb16: Option<&mut CudaSlice<u8>>,
10720 n_head: usize, t: usize, c: usize, hk: usize,
10721 k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>)
10722 -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10723 const D: usize = 128;
10724 let h = n_head;
10725 let nc = (t + c - 1) / c;
10726 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
10727 let mut gcum = self.uninit(t * h)?;
10728 let mut a = self.uninit(nc * h * c * c)?;
10729 let mut p = self.uninit(nc * h * c * c)?;
10730 let mut u = self.uninit(nc * h * c * D)?;
10731 let mut w = self.uninit(nc * h * c * D)?;
10732 { let f = self.func("gdn_chunk_cumgate_f32");
10734 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
10735 let __s_b = self.gpu.stream();
10736 let mut b = __s_b.launch_builder(&f);
10737 b.arg(g).arg(&mut gcum).arg(&hi).arg(&ti).arg(&ci);
10738 unsafe { b.launch(cfg)?; }
10739 }
10740 if let Some((qb, kb, pb)) = k2w {
10741 assert!(c == 32, "gdn_k2_wgmma is a C==32 tile");
10744 let f = self.func("gdn_k2_wgmma");
10745 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
10746 let hki = hk as i32;
10747 let __s_b = self.gpu.stream();
10748 let mut b = __s_b.launch_builder(&f);
10749 b.arg(qb).arg(kb).arg(&gcum).arg(beta).arg(&mut a).arg(&mut *pb).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
10750 unsafe { b.launch(cfg)?; }
10751 } else if c <= 64 && !portable_mma_gated() { let f = self.func("gdn_chunk_attn_f32");
10753 let jt = ((c + 31) / 32) as u32;
10754 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10755 let hki = hk as i32;
10756 let __s_b = self.gpu.stream();
10757 let mut b = __s_b.launch_builder(&f);
10758 b.arg(q).arg(k).arg(&gcum).arg(beta).arg(&mut a).arg(&mut p).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
10759 unsafe { b.launch(cfg)?; }
10760 } else { assert!(hk == h, "generic K2 is broadcast-only (de-broadcast rides C==32)");
10762 let f = self.func("gdn_chunk_attn_g_f32");
10763 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (32, 8, 1), shared_mem_bytes: 0 };
10764 let __s_b = self.gpu.stream();
10765 let mut b = __s_b.launch_builder(&f);
10766 b.arg(q).arg(k).arg(&gcum).arg(beta).arg(&mut a).arg(&mut p).arg(&hi).arg(&ti).arg(&ci);
10767 unsafe { b.launch(cfg)?; }
10768 }
10769 { let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10771 match c {
10772 32 | 64 => {
10773 let f = self.func(if c == 32 { "gdn_chunk_solve32_f32" } else { "gdn_chunk_solve64_f32" });
10774 let wb: u64 = match wb16 { Some(d) => self.addr_u8(d), None => 0 };
10776 let hki = hk as i32;
10777 let __s_b = self.gpu.stream();
10778 let mut b = __s_b.launch_builder(&f);
10779 b.arg(v).arg(k).arg(&a).arg(&gcum).arg(&mut u).arg(&mut w).arg(&wb).arg(&hi).arg(&ti).arg(&hki);
10780 unsafe { b.launch(cfg)?; }
10781 }
10782 _ => {
10783 assert!(hk == h, "generic K3 is broadcast-only");
10784 let f = self.func("gdn_chunk_solve_f32");
10785 let __s_b = self.gpu.stream();
10786 let mut b = __s_b.launch_builder(&f);
10787 b.arg(v).arg(k).arg(&a).arg(&gcum).arg(&mut u).arg(&mut w).arg(&hi).arg(&ti).arg(&ci);
10788 unsafe { b.launch(cfg)?; }
10789 }
10790 }
10791 }
10792 Ok((gcum, p, u, w))
10793 }
10794
10795 pub fn gdn_db_on() -> bool {
10799 std::env::var("MEMRA_GDN_DB").as_deref() != Ok("0")
10800 }
10801
10802 pub fn gdn_mma_enabled(&self, c: usize) -> bool {
10805 !portable_mma_gated() && c == 32
10806 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
10807 Ok("1") => true,
10808 Ok("0") => false,
10809 _ => cfg!(memra_hopper_mma),
10810 }
10811 }
10812
10813 pub fn gdn_wgmma_on(&self, c: usize) -> bool {
10816 self.gdn_mma_enabled(c)
10817 && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
10818 Ok("0") => false,
10819 Ok("1") => true,
10820 _ => cfg!(memra_hopper_mma),
10821 }
10822 }
10823
10824 #[allow(clippy::too_many_arguments)]
10829 pub fn ssm_conv1d_gdn_state_pad(&self, qkv_tm: &cudarc::driver::CudaView<f32>,
10830 conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>,
10831 q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>,
10832 v_g: &mut CudaSlice<f32>,
10833 conv_dim: usize, t: usize, d_conv: usize,
10834 d_state: usize, num_v: usize, num_k: usize, key_dim: usize,
10835 hk: usize,
10836 pad_len: Option<&CudaSlice<i32>>)
10837 -> Result<(), Box<dyn std::error::Error>> {
10838 assert!(t >= d_conv - 1, "fused state conv requires T >= pad (PRIME_MIN_T gates)");
10839 {
10840 let f = self.func("ssm_conv1d_gdn_state_f32");
10841 let cfg = LaunchConfig {
10842 grid_dim: (((conv_dim + 255) / 256) as u32, t as u32, 1),
10843 block_dim: (256, 1, 1), shared_mem_bytes: 0,
10844 };
10845 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10846 let (ds, nv, nk, kd, hki) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, hk as i32);
10847 let __s_b = self.gpu.stream();
10848 let mut b = __s_b.launch_builder(&f);
10849 b.arg(qkv_tm).arg(&*conv_state).arg(w).arg(q_g).arg(k_g).arg(v_g)
10850 .arg(&cd).arg(&ti).arg(&dc).arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&hki);
10851 unsafe { b.launch(cfg)?; }
10852 }
10853 match pad_len {
10854 Some(len_d) => {
10855 let f = self.func("ssm_conv_ring_update_dev_f32");
10856 let n = conv_dim * (d_conv - 1);
10857 let cfg = LaunchConfig::for_num_elems(n as u32);
10858 let (cd, dc) = (conv_dim as i32, d_conv as i32);
10859 let __s_b = self.gpu.stream();
10860 let mut b = __s_b.launch_builder(&f);
10861 b.arg(qkv_tm).arg(conv_state).arg(len_d).arg(&cd).arg(&dc);
10862 unsafe { b.launch(cfg)?; }
10863 }
10864 None => {
10865 let f = self.func("ssm_conv_ring_update_f32");
10866 let n = conv_dim * (d_conv - 1);
10867 let cfg = LaunchConfig::for_num_elems(n as u32);
10868 let (cd, ti, dc) = (conv_dim as i32, t as i32, d_conv as i32);
10869 let __s_b = self.gpu.stream();
10870 let mut b = __s_b.launch_builder(&f);
10871 b.arg(qkv_tm).arg(conv_state).arg(&cd).arg(&ti).arg(&dc);
10872 unsafe { b.launch(cfg)?; }
10873 }
10874 }
10875 Ok(())
10876 }
10877
10878 pub fn gdn_chunk_alloc(&self, n_head: usize, t: usize, c: usize, hk: usize)
10882 -> Result<GdnChunkBufs, Box<dyn std::error::Error>> {
10883 const D: usize = 128;
10884 assert!(c == 32, "gdn_chunk_alloc: varlen chain is the C==32 mma pair");
10885 let h = n_head;
10886 let nc = (t + c - 1) / c;
10887 Ok(GdnChunkBufs {
10888 gcum: self.uninit(t * h)?,
10889 a: self.uninit(nc * h * c * c)?,
10890 p: self.uninit(nc * h * c * c)?,
10891 u: self.uninit(nc * h * c * D)?,
10892 w: self.uninit(nc * h * c * D)?,
10893 kb16: self.alloc_u8_uninit(t * hk * D * 2)?,
10894 wb16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
10895 y16: self.alloc_u8_uninit(nc * h * c * D * 2)?,
10896 ssnap16: self.alloc_u8_uninit(nc * h * D * D * 2)?,
10897 qb16: self.alloc_u8_uninit(t * hk * D * 2)?,
10898 pb16: self.alloc_u8_uninit(nc * h * c * c * 2)?,
10899 o: self.uninit(D * h * t)?,
10900 t, nc,
10901 })
10902 }
10903
10904 pub fn f32_to_bf16_v(&self, x: &cudarc::driver::CudaView<f32>, dst: &mut CudaSlice<u8>, n: usize)
10906 -> Result<(), Box<dyn std::error::Error>> {
10907 let f = self.func("f32_to_bf16_bulk");
10908 let ni = n as i64;
10909 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10910 let __s_b = self.gpu.stream();
10911 let mut b = __s_b.launch_builder(&f);
10912 b.arg(x).arg(dst).arg(&ni);
10913 unsafe { b.launch(cfg)?; }
10914 Ok(())
10915 }
10916
10917 pub fn f32_to_bf16_into(&self, x: &CudaSlice<f32>, dst: &mut CudaSlice<u8>, n: usize)
10919 -> Result<(), Box<dyn std::error::Error>> {
10920 let f = self.func("f32_to_bf16_bulk");
10921 let ni = n as i64;
10922 let cfg = LaunchConfig::for_num_elems((n as u32).div_ceil(4));
10923 let __s_b = self.gpu.stream();
10924 let mut b = __s_b.launch_builder(&f);
10925 b.arg(x).arg(dst).arg(&ni);
10926 unsafe { b.launch(cfg)?; }
10927 Ok(())
10928 }
10929
10930 pub fn gdn_chunk_k123_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, hk: usize,
10933 wq: Option<&GdnWVl8>)
10934 -> Result<(), Box<dyn std::error::Error>> {
10935 let b = seqs.len();
10936 assert!(b >= 1 && b <= 8, "gdn_chunk_k123_vl8: 1..=8 sequences");
10937 let mut packed = [GdnSeqVl::default(); 8];
10938 packed[..b].copy_from_slice(seqs);
10939 let v = GdnVl8(packed);
10940 let (hi, ci) = (n_head as i32, 32i32);
10941 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
10942 {
10943 let f = self.func("gdn_chunk_cumgate_vl");
10944 let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (32, 1, 1), shared_mem_bytes: 0 };
10945 let __s_lb = self.gpu.stream();
10946 let mut lb = __s_lb.launch_builder(&f);
10947 lb.arg(&v).arg(&hi).arg(&ci);
10948 unsafe { lb.launch(cfg)?; }
10949 }
10950 let hki = hk as i32;
10951 if let Some(w) = wq { let f = self.func("gdn_k2_wgmma_vl");
10953 let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
10954 let __s_lb = self.gpu.stream();
10955 let mut lb = __s_lb.launch_builder(&f);
10956 lb.arg(&v).arg(w).arg(&hi).arg(&ci).arg(&hki);
10957 unsafe { lb.launch(cfg)?; }
10958 } else {
10959 let f = self.func("gdn_chunk_attn_vl");
10960 let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10961 let __s_lb = self.gpu.stream();
10962 let mut lb = __s_lb.launch_builder(&f);
10963 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
10964 unsafe { lb.launch(cfg)?; }
10965 }
10966 {
10967 let f = self.func("gdn_chunk_solve32_vl");
10968 let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10969 let __s_lb = self.gpu.stream();
10970 let mut lb = __s_lb.launch_builder(&f);
10971 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
10972 unsafe { lb.launch(cfg)?; }
10973 }
10974 Ok(())
10975 }
10976
10977 #[allow(clippy::too_many_arguments)]
10981 pub fn gdn_prep_vl8(&self, seqs: &[GdnPrepVl], conv_w: &CudaSlice<f32>,
10982 dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
10983 conv_dim: usize, d_conv: usize, d_state: usize,
10984 num_v: usize, num_k: usize, key_dim: usize, hk: usize, eps: f32)
10985 -> Result<(), Box<dyn std::error::Error>> {
10986 let b = seqs.len();
10987 assert!(b >= 1 && b <= 8);
10988 let mut packed = [GdnPrepVl::default(); 8];
10989 packed[..b].copy_from_slice(seqs);
10990 let v = GdnPrepVl8(packed);
10991 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
10992 let (cdi, dci) = (conv_dim as i32, d_conv as i32);
10993 let conv_fuse = std::env::var("MEMRA_CONV_FUSE").as_deref() != Ok("0");
10994 assert!(conv_fuse || hk == num_v, "de-broadcast requires the fused conv");
10995 if conv_fuse {
10996 let f = self.func("ssm_conv1d_gdn_state_vl");
10997 let cfg = LaunchConfig { grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
10998 let (dsi, nvi, nki, kdi, hki) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, hk as i32);
10999 let __s_lb = self.gpu.stream();
11000 let mut lb = __s_lb.launch_builder(&f);
11001 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi).arg(&hki);
11002 unsafe { lb.launch(cfg)?; }
11003 } else {
11004 let f = self.func("ssm_conv1d_tm_state_vl");
11005 let cfg = LaunchConfig { grid_dim: ((conv_dim as u32).div_ceil(256), max_t, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11006 let __s_lb = self.gpu.stream();
11007 let mut lb = __s_lb.launch_builder(&f);
11008 lb.arg(&v).arg(conv_w).arg(&cdi).arg(&dci);
11009 unsafe { lb.launch(cfg)?; }
11010 }
11011 {
11012 let f = self.func("ssm_conv_ring_update_vl");
11013 let n = (conv_dim * (d_conv - 1)) as u32;
11014 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11015 let __s_lb = self.gpu.stream();
11016 let mut lb = __s_lb.launch_builder(&f);
11017 lb.arg(&v).arg(&cdi).arg(&dci);
11018 unsafe { lb.launch(cfg)?; }
11019 }
11020 if !conv_fuse {
11021 let f = self.func("qkv_to_gdn_repack_vl");
11022 let n = max_t * (num_v * d_state) as u32;
11023 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11024 let (dsi, nvi, nki, kdi) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32);
11025 let __s_lb = self.gpu.stream();
11026 let mut lb = __s_lb.launch_builder(&f);
11027 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&nki).arg(&kdi);
11028 unsafe { lb.launch(cfg)?; }
11029 }
11030 if Self::l2_v2_on(d_state) {
11031 let f = self.func("gdn_l2_v2_vl");
11032 let cfg = LaunchConfig { grid_dim: ((max_t * hk as u32).div_ceil(8), 2, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11033 let (dsi, nvi) = (d_state as i32, hk as i32);
11034 let __s_lb = self.gpu.stream();
11035 let mut lb = __s_lb.launch_builder(&f);
11036 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
11037 unsafe { lb.launch(cfg)?; }
11038 } else {
11039 let f = self.func("gdn_l2_vl");
11040 let cfg = LaunchConfig { grid_dim: (max_t * hk as u32, 2, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11041 let (dsi, nvi) = (d_state as i32, hk as i32);
11042 let __s_lb = self.gpu.stream();
11043 let mut lb = __s_lb.launch_builder(&f);
11044 lb.arg(&v).arg(&dsi).arg(&nvi).arg(&eps);
11045 unsafe { lb.launch(cfg)?; }
11046 }
11047 {
11048 let f = self.func("gdn_gate_prep_vl");
11049 let n = max_t * num_v as u32;
11050 let cfg = LaunchConfig { grid_dim: (n.div_ceil(256), 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11051 let nvi = num_v as i32;
11052 let __s_lb = self.gpu.stream();
11053 let mut lb = __s_lb.launch_builder(&f);
11054 lb.arg(&v).arg(dt_bias).arg(a).arg(&nvi);
11055 unsafe { lb.launch(cfg)?; }
11056 }
11057 Ok(())
11058 }
11059
11060 pub fn gdn_mirror_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, which: i32, hk: usize)
11062 -> Result<(), Box<dyn std::error::Error>> {
11063 let b = seqs.len();
11064 assert!(b >= 1 && b <= 8);
11065 let mut packed = [GdnSeqVl::default(); 8];
11066 packed[..b].copy_from_slice(seqs);
11067 let v = GdnVl8(packed);
11068 let ept = (if which == 0 { hk } else { n_head } * 128) as i32;
11069 let max_n = seqs.iter().map(|s| if which == 0 { s.t as i64 * ept as i64 }
11070 else { s.nc as i64 * ept as i64 * 32 }).max().unwrap();
11071 let f = self.func("gdn_mirror_vl");
11072 let blocks = ((max_n as u32).div_ceil(4)).div_ceil(256);
11073 let cfg = LaunchConfig { grid_dim: (blocks, 1, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11074 let __s_lb = self.gpu.stream();
11075 let mut lb = __s_lb.launch_builder(&f);
11076 lb.arg(&v).arg(&ept).arg(&which);
11077 unsafe { lb.launch(cfg)?; }
11078 Ok(())
11079 }
11080
11081 pub fn gdn_tail_vl8(&self, seqs: &[GdnPrepVl], norm_w: &CudaSlice<f32>,
11083 d_state: usize, num_v: usize, eps: f32)
11084 -> Result<(), Box<dyn std::error::Error>> {
11085 let b = seqs.len();
11086 assert!(b >= 1 && b <= 8);
11087 let mut packed = [GdnPrepVl::default(); 8];
11088 packed[..b].copy_from_slice(seqs);
11089 let v = GdnPrepVl8(packed);
11090 let max_t = seqs.iter().map(|s| s.t).max().unwrap() as u32;
11091 let f = self.func("gated_rmsnorm_f16out_vl");
11092 let cfg = LaunchConfig { grid_dim: (max_t * num_v as u32, 1, b as u32), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11094 let (dsi, nvi) = (d_state as i32, num_v as i32);
11095 let __s_lb = self.gpu.stream();
11096 let mut lb = __s_lb.launch_builder(&f);
11097 lb.arg(&v).arg(norm_w).arg(&dsi).arg(&nvi).arg(&eps);
11098 unsafe { lb.launch(cfg)?; }
11099 Ok(())
11100 }
11101
11102 pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64 {
11105 use cudarc::driver::DevicePtr;
11106 let s = self.gpu.stream();
11107 let (p, _g) = x.device_ptr(&s);
11108 p as u64
11109 }
11110 pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64 {
11111 use cudarc::driver::DevicePtrMut;
11112 let s = self.gpu.stream();
11113 let (p, _g) = x.device_ptr_mut(&s);
11114 p as u64
11115 }
11116 pub fn addr_f32v(&self, x: &cudarc::driver::CudaView<f32>) -> u64 {
11117 use cudarc::driver::DevicePtr;
11118 let s = self.gpu.stream();
11119 let (p, _g) = x.device_ptr(&s);
11120 p as u64
11121 }
11122 pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64 {
11123 use cudarc::driver::DevicePtr;
11124 let s = self.gpu.stream();
11125 let (p, _g) = x.device_ptr(&s);
11126 p as u64
11127 }
11128
11129 pub fn gdn_chunk_vl8(&self, seqs: &[GdnSeqVl], n_head: usize, scale: f32, hk: usize,
11133 wq: Option<&GdnWVl8>)
11134 -> Result<(), Box<dyn std::error::Error>> {
11135 const NSPLIT: u32 = 4;
11136 let b = seqs.len();
11137 assert!(b >= 1 && b <= 8, "gdn_chunk_vl8: 1..=8 sequences");
11138 let mut packed = [GdnSeqVl::default(); 8];
11139 packed[..b].copy_from_slice(seqs);
11140 let v = GdnVl8(packed);
11141 let (hi, ci) = (n_head as i32, 32i32);
11142 let max_nc = seqs.iter().map(|a| a.nc).max().unwrap() as u32;
11143 let hki = hk as i32;
11144 if let Some(w) = wq {
11145 let f = self.func("gdn_k45_wgmma_vl");
11147 let cfg = LaunchConfig { grid_dim: (n_head as u32, NSPLIT, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11148 let __s_lb = self.gpu.stream();
11149 let mut lb = __s_lb.launch_builder(&f);
11150 lb.arg(&v).arg(w).arg(&scale).arg(&hi).arg(&ci).arg(&hki);
11151 unsafe { lb.launch(cfg)?; }
11152 let _ = max_nc;
11153 return Ok(());
11154 }
11155 {
11156 let f = self.func("gdn_chunk_state_mma_vl");
11157 let cfg = LaunchConfig { grid_dim: (n_head as u32, NSPLIT, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11158 let __s_lb = self.gpu.stream();
11159 let mut lb = __s_lb.launch_builder(&f);
11160 lb.arg(&v).arg(&hi).arg(&ci).arg(&hki);
11161 unsafe { lb.launch(cfg)?; }
11162 }
11163 {
11164 let f = self.func("gdn_chunk_output_mma_vl");
11165 let cfg = LaunchConfig { grid_dim: (max_nc, n_head as u32, b as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11166 let __s_lb = self.gpu.stream();
11167 let mut lb = __s_lb.launch_builder(&f);
11168 lb.arg(&v).arg(&hi).arg(&ci).arg(&scale).arg(&hki);
11169 unsafe { lb.launch(cfg)?; }
11170 }
11171 Ok(())
11172 }
11173 pub fn gdn_scan_chunked(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11174 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>,
11175 qb16_pre: Option<&CudaSlice<u8>>,
11176 state_in: &CudaSlice<f32>,
11177 state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11178 n_head: usize, t: usize, scale: f32, c: usize, hk: usize)
11179 -> Result<(), Box<dyn std::error::Error>> {
11180 const D: usize = 128;
11181 const NSPLIT: u32 = 4;
11182 assert!(c >= 1 && c <= 128, "gdn_scan_chunked: C must be in 1..=128");
11183 let h = n_head;
11184 let nc = (t + c - 1) / c;
11185 let (hi, ti, ci) = (h as i32, t as i32, c as i32);
11186 let gdn_mma_pre = !portable_mma_gated() && c == 32
11190 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
11191 Ok("1") => true,
11192 Ok("0") => false,
11193 _ => cfg!(memra_hopper_mma),
11194 };
11195 let mut wb16_pre: Option<CudaSlice<u8>> = if gdn_mma_pre {
11196 Some(self.alloc_u8_uninit(nc * h * c * D * 2)?)
11197 } else { None };
11198 let gdn_wgmma_pre = gdn_mma_pre
11202 && match std::env::var("MEMRA_GDN_WGMMA").as_deref() {
11203 Ok("0") => false,
11204 Ok("1") => true,
11205 _ => cfg!(memra_hopper_mma),
11206 };
11207 let nk = t * hk * D;
11208 let mut kb16_local: Option<CudaSlice<u8>> = None;
11209 if gdn_mma_pre && kb16_pre.is_none() {
11210 let mut kb = self.alloc_u8_uninit(nk * 2)?;
11211 let f = self.func("f32_to_bf16_bulk");
11212 let n2 = nk as i64;
11213 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
11214 let __s_b = self.gpu.stream();
11215 let mut b = __s_b.launch_builder(&f);
11216 b.arg(k).arg(&mut kb).arg(&n2);
11217 unsafe { b.launch(cfg2)?; }
11218 kb16_local = Some(kb);
11219 }
11220 let kb16_ref0: Option<&CudaSlice<u8>> = kb16_local.as_ref().or(kb16_pre);
11221 if let Some(kb) = kb16_pre { assert!(kb.len() >= nk * 2, "kb16_pre too small"); }
11222 let mut qb16: Option<CudaSlice<u8>> = None;
11223 let mut pb16: Option<CudaSlice<u8>> = None;
11224 if gdn_wgmma_pre {
11225 if qb16_pre.is_none() {
11228 let mut qb = self.alloc_u8_uninit(nk * 2)?;
11229 let f = self.func("f32_to_bf16_bulk");
11230 let n2 = nk as i64;
11231 let cfg2 = LaunchConfig::for_num_elems((nk as u32).div_ceil(4));
11232 let __s_b = self.gpu.stream();
11233 let mut b = __s_b.launch_builder(&f);
11234 b.arg(q).arg(&mut qb).arg(&n2);
11235 unsafe { b.launch(cfg2)?; }
11236 qb16 = Some(qb);
11237 } else if let Some(qb) = qb16_pre {
11238 assert!(qb.len() >= nk * 2, "qb16_pre too small");
11239 }
11240 pb16 = Some(self.alloc_u8_uninit(nc * h * c * c * 2)?);
11241 }
11242 let qb16_ref0: Option<&CudaSlice<u8>> = qb16.as_ref().or(qb16_pre);
11243 let k2w = if gdn_wgmma_pre {
11244 Some((*qb16_ref0.as_ref().unwrap(),
11245 *kb16_ref0.as_ref().unwrap(),
11246 pb16.as_mut().unwrap()))
11247 } else { None };
11248 let (gcum, p, u, w) = self.gdn_chunk_k123(q, k, v, g, beta, wb16_pre.as_mut(), n_head, t, c, hk, k2w)?;
11249 let _ = &w;
11250 let mut y = self.uninit(nc * h * c * D)?;
11251 let mut ssnap = self.uninit(nc * h * D * D)?; let gdn_mma = !portable_mma_gated() && c == 32
11263 && match std::env::var("MEMRA_GDN_MMA").as_deref() {
11264 Ok("1") => true,
11265 Ok("0") => false,
11266 _ => cfg!(memra_hopper_mma),
11267 };
11268 if gdn_mma {
11269 let wb16 = wb16_pre.take().expect("mma path pre-allocates wb16 (K3 store fold)");
11270 let kb16_ref: &CudaSlice<u8> = kb16_ref0.expect("mma path pre-builds kb16 above K123");
11271 if gdn_wgmma_pre {
11283 let qb16 = qb16_ref0.unwrap();
11285 let pb16 = pb16.as_ref().unwrap();
11286 {
11287 let f = self.func("gdn_k45_wgmma");
11288 let cfg = LaunchConfig { grid_dim: (h as u32, 4, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11289 let hki = hk as i32;
11290 let __s_b = self.gpu.stream();
11291 let mut b = __s_b.launch_builder(&f);
11292 b.arg(kb16_ref).arg(&gcum).arg(beta).arg(&u).arg(&wb16).arg(qb16).arg(pb16)
11293 .arg(o).arg(&scale).arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
11294 unsafe { b.launch(cfg)?; }
11295 }
11296 return Ok(());
11297 }
11298 let mut y16 = self.alloc_u8_uninit(nc * h * c * D * 2)?;
11302 let mut ssnap16 = self.alloc_u8_uninit(nc * h * D * D * 2)?;
11303 {
11304 let f = self.func("gdn_chunk_state_mma");
11305 let cfg = LaunchConfig { grid_dim: (h as u32, NSPLIT, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11306 let hki = hk as i32;
11307 let __s_b = self.gpu.stream();
11308 let mut b = __s_b.launch_builder(&f);
11309 b.arg(kb16_ref).arg(&gcum).arg(beta).arg(&u).arg(&wb16).arg(&mut y16).arg(&mut ssnap16)
11310 .arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci).arg(&hki);
11311 unsafe { b.launch(cfg)?; }
11312 }
11313 { let f = self.func("gdn_chunk_output_mma");
11315 let jt = ((c + 31) / 32) as u32;
11316 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11317 let hki = hk as i32;
11318 let __s_b = self.gpu.stream();
11319 let mut b = __s_b.launch_builder(&f);
11320 b.arg(q).arg(&gcum).arg(&p).arg(&y16).arg(&ssnap16).arg(o).arg(&hi).arg(&ti).arg(&ci).arg(&scale).arg(&hki);
11321 unsafe { b.launch(cfg)?; }
11322 }
11323 return Ok(());
11324 }
11325 { let f = self.func("gdn_chunk_state_f32");
11327 let cfg = LaunchConfig { grid_dim: (h as u32, NSPLIT, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11328 let __s_b = self.gpu.stream();
11329 let mut b = __s_b.launch_builder(&f);
11330 b.arg(k).arg(&gcum).arg(beta).arg(&u).arg(&w).arg(&mut y).arg(&mut ssnap)
11331 .arg(state_in).arg(&mut *state_out).arg(&hi).arg(&ti).arg(&ci);
11332 unsafe { b.launch(cfg)?; }
11333 }
11334 { let f = self.func("gdn_chunk_output_f32");
11336 let jt = ((c + 31) / 32) as u32;
11337 let cfg = LaunchConfig { grid_dim: (nc as u32, h as u32, jt), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
11338 let __s_b = self.gpu.stream();
11339 let mut b = __s_b.launch_builder(&f);
11340 b.arg(q).arg(&gcum).arg(&p).arg(&y).arg(&ssnap).arg(o).arg(&hi).arg(&ti).arg(&ci).arg(&scale);
11341 unsafe { b.launch(cfg)?; }
11342 }
11343 Ok(())
11344 }
11345
11346 #[allow(clippy::too_many_arguments)]
11355 #[allow(clippy::too_many_arguments)]
11356 pub fn gdn_scan_prefill(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11357 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>,
11358 qb16_pre: Option<&CudaSlice<u8>>,
11359 state_in: &CudaSlice<f32>,
11360 state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11361 n_head: usize, t: usize, scale: f32, hk: usize)
11362 -> Result<(), Box<dyn std::error::Error>> {
11363 if std::env::var("MEMRA_GDN_DIFF").is_ok() && t >= 16 {
11364 assert!(hk == n_head, "GDN_DIFF oracle is broadcast-only");
11365 return self.gdn_scan_diff(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale);
11366 }
11367 if Self::gdn_chunked_enabled() && t >= 16 {
11368 self.gdn_scan_chunked(q, k, v, g, beta, kb16_pre, qb16_pre, state_in, state_out, o, n_head, t, scale,
11369 Self::gdn_chunk_size(), hk)
11370 } else {
11371 assert!(hk == n_head, "s128 scan is broadcast-only (prep guarantees by predicate)");
11372 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)
11373 }
11374 }
11375
11376 #[allow(clippy::too_many_arguments)]
11378 fn gdn_scan_diff(&self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>,
11379 g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>,
11380 state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>,
11381 n_head: usize, t: usize, scale: f32)
11382 -> Result<(), Box<dyn std::error::Error>> {
11383 static CALL: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
11384 let call = CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11385 let mut o_c = self.uninit(o.len())?;
11386 let mut st_c = self.uninit(state_out.len())?;
11387 self.gdn_scan_chunked(q, k, v, g, beta, None, None, state_in, &mut st_c, &mut o_c,
11388 n_head, t, scale, Self::gdn_chunk_size(), n_head)?;
11389 self.gdn_scan_s128(q, k, v, g, beta, state_in, state_out, o, n_head, t, scale)?;
11390 let (oh_s, oh_c) = (self.dtoh(o)?, self.dtoh(&o_c)?);
11391 let (sh_s, sh_c) = (self.dtoh(state_out)?, self.dtoh(&st_c)?);
11392 let stats = |a: &[f32], b: &[f32]| -> (f32, f32, f64) {
11393 let mut max_abs = 0f32; let mut max_rel = 0f32; let mut sum_rel = 0f64;
11394 for (x, y) in a.iter().zip(b) {
11395 let ad = (x - y).abs();
11396 let rel = ad / x.abs().max(y.abs()).max(1e-3);
11397 if ad > max_abs { max_abs = ad; }
11398 if rel > max_rel { max_rel = rel; }
11399 sum_rel += rel as f64;
11400 }
11401 (max_abs, max_rel, sum_rel / a.len() as f64)
11402 };
11403 let (o_ma, o_mr, o_mean) = stats(&oh_s, &oh_c);
11404 let (s_ma, s_mr, s_mean) = stats(&sh_s, &sh_c);
11405 println!("[gdn-diff call {call:3} T={t} C={}] out: max_abs={o_ma:.3e} max_rel={o_mr:.3e} mean_rel={o_mean:.3e} | \
11406 state: max_abs={s_ma:.3e} max_rel={s_mr:.3e} mean_rel={s_mean:.3e}",
11407 Self::gdn_chunk_size());
11408 Ok(())
11409 }
11410
11411 pub fn gdn_glog(&self, alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>,
11413 g_log: &mut CudaSlice<f32>, n_head: usize, t: usize)
11414 -> Result<(), Box<dyn std::error::Error>> {
11415 let f = self.func("gdn_glog_f32");
11416 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
11417 let (h, ti) = (n_head as i32, t as i32);
11418 let __s_b = self.gpu.stream();
11419 let mut b = __s_b.launch_builder(&f);
11420 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
11421 unsafe { b.launch(cfg)?; }
11422 Ok(())
11423 }
11424
11425 pub fn sigmoid_v(&self, x: &cudarc::driver::CudaView<f32>, y: &mut CudaSlice<f32>, n: usize)
11428 -> Result<(), Box<dyn std::error::Error>> {
11429 let f = self.func("sigmoid_f32");
11430 let cfg = LaunchConfig::for_num_elems(n as u32);
11431 let ni = n as i32;
11432 let __s_b = self.gpu.stream();
11433 let mut b = __s_b.launch_builder(&f);
11434 b.arg(x).arg(y).arg(&ni);
11435 unsafe { b.launch(cfg)?; }
11436 Ok(())
11437 }
11438
11439 pub fn gdn_glog_v(&self, alpha: &cudarc::driver::CudaView<f32>, dt_bias: &CudaSlice<f32>,
11440 a: &CudaSlice<f32>, g_log: &mut CudaSlice<f32>, n_head: usize, t: usize)
11441 -> Result<(), Box<dyn std::error::Error>> {
11442 let f = self.func("gdn_glog_f32");
11443 let cfg = LaunchConfig::for_num_elems((n_head * t) as u32);
11444 let (h, ti) = (n_head as i32, t as i32);
11445 let __s_b = self.gpu.stream();
11446 let mut b = __s_b.launch_builder(&f);
11447 b.arg(alpha).arg(dt_bias).arg(a).arg(g_log).arg(&h).arg(&ti);
11448 unsafe { b.launch(cfg)?; }
11449 Ok(())
11450 }
11451
11452 pub fn sigmoid(&self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize)
11453 -> Result<(), Box<dyn std::error::Error>> {
11454 let f = self.func("sigmoid_f32");
11455 let cfg = LaunchConfig::for_num_elems(n as u32);
11456 let ni = n as i32;
11457 let __s_b = self.gpu.stream();
11458 let mut b = __s_b.launch_builder(&f);
11459 b.arg(x).arg(y).arg(&ni);
11460 unsafe { b.launch(cfg)?; }
11461 Ok(())
11462 }
11463
11464 pub fn sig_mul_f16out(&self, a: &CudaSlice<f32>, g: &CudaSlice<f32>,
11467 dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize)
11468 -> Result<(), Box<dyn std::error::Error>> {
11469 let f = self.func("sig_mul_f16out_f32");
11470 let cfg = LaunchConfig::for_num_elems(n as u32);
11471 let ni = n as i32;
11472 let __s_b = self.gpu.stream();
11473 let mut b = __s_b.launch_builder(&f);
11474 b.arg(a).arg(g).arg(dst).arg(dst16).arg(&ni);
11475 unsafe { b.launch(cfg)?; }
11476 Ok(())
11477 }
11478
11479 pub fn gated_rmsnorm(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11481 dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
11482 -> Result<(), Box<dyn std::error::Error>> {
11483 let f = self.func("gated_rmsnorm_f32");
11484 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11485 let (nc, e) = (ncols as i32, eps);
11486 let __s_b = self.gpu.stream();
11487 let mut b = __s_b.launch_builder(&f);
11488 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
11489 unsafe { b.launch(cfg)?; }
11490 Ok(())
11491 }
11492
11493 pub fn gated_rmsnorm_f16out(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11496 dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
11497 ncols: usize, nrows: usize, eps: f32)
11498 -> Result<(), Box<dyn std::error::Error>> {
11499 let f = self.func("gated_rmsnorm_f16out_f32");
11500 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11502 let (nc, e) = (ncols as i32, eps);
11503 let __s_b = self.gpu.stream();
11504 let mut b = __s_b.launch_builder(&f);
11505 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
11506 unsafe { b.launch(cfg)?; }
11507 Ok(())
11508 }
11509
11510 #[allow(clippy::too_many_arguments)]
11514 pub fn add_rms_norm_zq8(&self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>,
11515 res: &mut CudaSlice<f32>, z: &mut CudaSlice<f32>,
11516 ncols: usize, nrows: usize, eps: f32)
11517 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11518 assert!(ncols % 32 == 0);
11519 let mut q = self.alloc_uninit::<i8>(nrows * ncols)?;
11520 let mut d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11521 let f = self.func("add_rms_norm_zq8");
11522 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
11523 let (nc, ep) = (ncols as i32, eps);
11524 let __s_b = self.gpu.stream();
11525 let mut b = __s_b.launch_builder(&f);
11526 b.arg(a).arg(b_in).arg(w).arg(res).arg(z).arg(&mut q).arg(&mut d).arg(&nc).arg(&ep);
11527 unsafe { b.launch(cfg)?; }
11528 Ok((q, d))
11529 }
11530
11531 pub fn gated_rmsnorm_zv(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>,
11536 z: &cudarc::driver::CudaView<f32>,
11537 dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32)
11538 -> Result<(), Box<dyn std::error::Error>> {
11539 let f = self.func("gated_rmsnorm_f32");
11540 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11541 let (nc, e) = (ncols as i32, eps);
11542 let __s_b = self.gpu.stream();
11543 let mut b = __s_b.launch_builder(&f);
11544 b.arg(o).arg(w).arg(z).arg(dst).arg(&nc).arg(&e);
11545 unsafe { b.launch(cfg)?; }
11546 Ok(())
11547 }
11548
11549 pub fn gated_rmsnorm_f16out_zv(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>,
11550 z: &cudarc::driver::CudaView<f32>,
11551 dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>,
11552 ncols: usize, nrows: usize, eps: f32)
11553 -> Result<(), Box<dyn std::error::Error>> {
11554 let f = self.func("gated_rmsnorm_f16out_f32");
11555 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11557 let (nc, e) = (ncols as i32, eps);
11558 let __s_b = self.gpu.stream();
11559 let mut b = __s_b.launch_builder(&f);
11560 b.arg(o).arg(w).arg(z).arg(dst).arg(dst16).arg(&nc).arg(&e);
11561 unsafe { b.launch(cfg)?; }
11562 Ok(())
11563 }
11564
11565 pub fn gated_rmsnorm_q8_1(&self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>,
11566 ncols: usize, nrows: usize, eps: f32)
11567 -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn std::error::Error>> {
11568 assert!(ncols % 32 == 0);
11569 let f = self.func("gated_rmsnorm_q8_1");
11570 let mut out_q = self.alloc_uninit::<i8>(nrows * ncols)?;
11571 let mut out_d = self.alloc_uninit::<f32>(nrows * (ncols / 32))?;
11572 let cfg = LaunchConfig { grid_dim: (nrows as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 };
11573 let (nc, ep) = (ncols as i32, eps);
11574 let __s_b = self.gpu.stream();
11575 let mut b = __s_b.launch_builder(&f);
11576 b.arg(o).arg(w).arg(z).arg(&mut out_q).arg(&mut out_d).arg(&nc).arg(&ep);
11577 unsafe { b.launch(cfg)?; }
11578 Ok((out_q, out_d))
11579 }
11580
11581 pub fn transpose(&self, inp: &CudaSlice<f32>, rows: usize, cols: usize)
11583 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11584 let f = self.func("transpose_f32");
11585 let mut out = self.zeros(rows * cols)?;
11586 let cfg = LaunchConfig::for_num_elems((rows * cols) as u32);
11587 let (r, c) = (rows as i32, cols as i32);
11588 let __s_b = self.gpu.stream();
11589 let mut b = __s_b.launch_builder(&f);
11590 b.arg(inp).arg(&mut out).arg(&r).arg(&c);
11591 unsafe { b.launch(cfg)?; }
11592 Ok(out)
11593 }
11594
11595 pub fn repeat_heads(&self, inp: &CudaSlice<f32>, out: &mut CudaSlice<f32>,
11597 head_dim: usize, n_in: usize, n_out: usize, t: usize)
11598 -> Result<(), Box<dyn std::error::Error>> {
11599 let f = self.func("repeat_heads_f32");
11600 let cfg = LaunchConfig::for_num_elems((head_dim * n_out * t) as u32);
11601 let (hd, ni, no, ti) = (head_dim as i32, n_in as i32, n_out as i32, t as i32);
11602 let __s_b = self.gpu.stream();
11603 let mut b = __s_b.launch_builder(&f);
11604 b.arg(inp).arg(out).arg(&hd).arg(&ni).arg(&no).arg(&ti);
11605 unsafe { b.launch(cfg)?; }
11606 Ok(())
11607 }
11608
11609 pub fn q_gate_split(&self, qf: &CudaSlice<f32>, q_out: &mut CudaSlice<f32>,
11612 gate_out: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, t: usize)
11613 -> Result<(), Box<dyn std::error::Error>> {
11614 let f = self.func("q_gate_split_f32");
11615 let cfg = LaunchConfig::for_num_elems((head_dim * n_head * t) as u32);
11616 let (hd, nh, ti) = (head_dim as i32, n_head as i32, t as i32);
11617 let __s_b = self.gpu.stream();
11618 let mut b = __s_b.launch_builder(&f);
11619 b.arg(qf).arg(q_out).arg(gate_out).arg(&hd).arg(&nh).arg(&ti);
11620 unsafe { b.launch(cfg)?; }
11621 Ok(())
11622 }
11623
11624 pub fn qkv_to_gdn_repack(&self, conv_out: &CudaSlice<f32>, q_g: &mut CudaSlice<f32>,
11628 k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>,
11629 d_state: usize, num_v: usize, num_k: usize, key_dim: usize, t: usize)
11630 -> Result<(), Box<dyn std::error::Error>> {
11631 let f = self.func("qkv_to_gdn_repack_f32");
11632 let cfg = LaunchConfig::for_num_elems((d_state * num_v * t) as u32);
11633 let (ds, nv, nk, kd, ti) = (d_state as i32, num_v as i32, num_k as i32, key_dim as i32, t as i32);
11634 let __s_b = self.gpu.stream();
11635 let mut b = __s_b.launch_builder(&f);
11636 b.arg(conv_out).arg(q_g).arg(k_g).arg(v_g).arg(&ds).arg(&nv).arg(&nk).arg(&kd).arg(&ti);
11637 unsafe { b.launch(cfg)?; }
11638 Ok(())
11639 }
11640
11641 pub fn conv_left_pad(&self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>,
11644 conv_dim: usize, t: usize, pad: usize)
11645 -> Result<(), Box<dyn std::error::Error>> {
11646 let f = self.func("conv_left_pad_f32");
11647 let cfg = LaunchConfig::for_num_elems((conv_dim * t) as u32);
11648 let (cd, ti, p) = (conv_dim as i32, t as i32, pad as i32);
11649 let __s_b = self.gpu.stream();
11650 let mut b = __s_b.launch_builder(&f);
11651 b.arg(src).arg(dst).arg(&cd).arg(&ti).arg(&p);
11652 unsafe { b.launch(cfg)?; }
11653 Ok(())
11654 }
11655
11656 pub fn conv_assemble_and_roll(&self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
11660 conv_in: &mut CudaSlice<f32>, conv_dim: usize, pad: usize)
11661 -> Result<(), Box<dyn std::error::Error>> {
11662 let f = self.func("conv_assemble_and_roll_f32");
11663 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
11664 let (cd, p) = (conv_dim as i32, pad as i32);
11665 let __s_b = self.gpu.stream();
11666 let mut b = __s_b.launch_builder(&f);
11667 b.arg(qkv_col).arg(conv_state).arg(conv_in).arg(&cd).arg(&p);
11668 unsafe { b.launch(cfg)?; }
11669 Ok(())
11670 }
11671
11672 pub fn ssm_conv1d_fused_decode(&self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>,
11678 w: &CudaSlice<f32>, conv_out: &mut CudaSlice<f32>,
11679 conv_dim: usize, d_conv: usize)
11680 -> Result<(), Box<dyn std::error::Error>> {
11681 let f = self.func("ssm_conv1d_fused_decode_f32");
11682 let cfg = LaunchConfig::for_num_elems(conv_dim as u32);
11683 let (cd, dc) = (conv_dim as i32, d_conv as i32);
11684 let __s_b = self.gpu.stream();
11685 let mut b = __s_b.launch_builder(&f);
11686 b.arg(qkv_col).arg(conv_state).arg(w).arg(conv_out).arg(&cd).arg(&dc);
11687 unsafe { b.launch(cfg)?; }
11688 Ok(())
11689 }
11690
11691 pub fn slice_range(&self, src: &CudaSlice<f32>, start: usize, len: usize)
11694 -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11695 let host = self.gpu.stream().clone_dtoh(src)?;
11696 self.gpu.stream().synchronize()?;
11697 Ok(self.htod(&host[start..start + len])?)
11698 }
11699}
11700
11701#[cfg(test)]
11702mod target_dispatch_tests {
11703 use super::legacy_quant_gemm_allowed;
11704
11705 #[test]
11706 fn legacy_quant_gemm_arch_policy_honors_the_escape_hatch() {
11707 assert!(legacy_quant_gemm_allowed(false, false, false));
11709 assert!(!legacy_quant_gemm_allowed(false, false, true));
11710 assert!(!legacy_quant_gemm_allowed(true, false, false));
11712 assert!(!legacy_quant_gemm_allowed(true, false, true));
11713 assert!(legacy_quant_gemm_allowed(true, true, false));
11715 assert!(!legacy_quant_gemm_allowed(true, true, true));
11716 }
11717
11718 #[cfg(all(memra_portable_cuda, not(memra_hopper_mma)))]
11719 #[test]
11720 fn portable_build_disables_legacy_quant_gemm_without_an_env_override() {
11721 assert!(!legacy_quant_gemm_allowed(cfg!(memra_portable_cuda), cfg!(memra_hopper_mma), false));
11722 }
11723
11724 #[cfg(memra_hopper_mma)]
11725 #[test]
11726 fn hopper_mma_build_re_admits_legacy_quant_gemm() {
11727 assert!(legacy_quant_gemm_allowed(cfg!(memra_portable_cuda), cfg!(memra_hopper_mma), false));
11728 assert!(super::portable_mma_gated() == false);
11729 }
11730}
11731
11732impl memra_kv::KvDev for Engine {
11735 fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11736 Engine::zeros(self, n)
11737 }
11738 fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11739 Engine::uninit(self, n)
11740 }
11741 fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn std::error::Error>> {
11742 Engine::alloc_u8(self, n)
11743 }
11744 fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
11745 Engine::htod_i32(self, v)
11746 }
11747 fn clone_dtod(&self, src: &CudaSlice<f32>) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
11748 Engine::clone_dtod(self, src)
11749 }
11750 fn copy_into(&self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize)
11751 -> Result<(), Box<dyn std::error::Error>> {
11752 Engine::copy_into(self, dst, off, src, len)
11753 }
11754 fn set_i32_one(&self, d: &mut CudaSlice<i32>, v: i32) -> Result<(), Box<dyn std::error::Error>> {
11755 Engine::set_i32_one(self, d, v)
11756 }
11757}